feat: adicionando resumo do cursor

_utils/resumo_completo_cursor.py

@@ -0,0 +1,221 @@
+import os
+from typing import List, Dict, Tuple
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.document_loaders import PyPDFLoader
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.vectorstores import Chroma
+from langchain.chat_models import ChatOpenAI
+from langchain.chains import create_extraction_chain
+from langchain.prompts import PromptTemplate
+from dataclasses import dataclass
+import uuid
+import json
+from langchain_huggingface import HuggingFaceEndpoint
+from setup.environment import default_model
+
+os.environ["LANGCHAIN_TRACING_V2"]="true"
+os.environ["LANGCHAIN_ENDPOINT"]="https://api.smith.langchain.com"
+os.environ.get("LANGCHAIN_API_KEY")
+os.environ["LANGCHAIN_PROJECT"]="VELLA"
+
+@dataclass
+class DocumentChunk:
+  content: str
+  page_number: int
+  chunk_id: str
+  start_char: int
+  end_char: int
+
+class DocumentSummarizer:
+  def __init__(self, openai_api_key: str, model, embedding, chunk_config, system_prompt):
+      self.model = model
+      self.system_prompt = system_prompt
+      self.openai_api_key = openai_api_key
+      self.embeddings = HuggingFaceEmbeddings(
+          model_name=embedding
+      )
+      self.text_splitter = RecursiveCharacterTextSplitter(
+          chunk_size=chunk_config["size"],
+          chunk_overlap=chunk_config["overlap"]
+      )
+      self.chunk_metadata = {}  # Store chunk metadata for tracing
+  
+  def load_and_split_document(self, pdf_path: str) -> List[DocumentChunk]:
+      """Load PDF and split into chunks with metadata"""
+      loader = PyPDFLoader(pdf_path)
+      pages = loader.load()
+      chunks = []
+      char_count = 0
+      
+      for page in pages:
+          text = page.page_content
+          # Split the page content
+          page_chunks = self.text_splitter.split_text(text)
+          
+          for chunk in page_chunks:
+              chunk_id = str(uuid.uuid4())
+              start_char = text.find(chunk)
+              end_char = start_char + len(chunk)
+              
+              doc_chunk = DocumentChunk(
+                  content=chunk,
+                  page_number=page.metadata.get('page') + 1,  # 1-based page numbering
+                  chunk_id=chunk_id,
+                  start_char=char_count + start_char,
+                  end_char=char_count + end_char
+              )
+              chunks.append(doc_chunk)
+              
+              # Store metadata for later retrieval
+              self.chunk_metadata[chunk_id] = {
+                  'page': doc_chunk.page_number,
+                  'start_char': doc_chunk.start_char,
+                  'end_char': doc_chunk.end_char
+              }
+          
+          char_count += len(text)
+      
+      return chunks
+
+  def create_vector_store(self, chunks: List[DocumentChunk]) -> Chroma:
+      """Create vector store with metadata"""
+      texts = [chunk.content for chunk in chunks]
+      metadatas = [{
+          'chunk_id': chunk.chunk_id,
+          'page': chunk.page_number,
+          'start_char': chunk.start_char,
+          'end_char': chunk.end_char
+      } for chunk in chunks]
+      
+      vector_store = Chroma.from_texts(
+          texts=texts,
+          metadatas=metadatas,
+          embedding=self.embeddings
+      )
+      return vector_store
+
+  def generate_summary_with_sources(
+      self, 
+      vector_store: Chroma, 
+      query: str = "Summarize the main points of this document"
+  ) -> List[Dict]:
+      """Generate summary with source citations, returning structured JSON data"""
+      # Retrieve relevant chunks with metadata
+      relevant_docs = vector_store.similarity_search_with_score(query, k=5)
+      
+      # Prepare context and track sources
+      contexts = []
+      sources = []
+      
+      for doc, score in relevant_docs:
+          chunk_id = doc.metadata['chunk_id']
+          context = doc.page_content
+          contexts.append(context)
+          
+          sources.append({
+              'content': context,
+              'page': doc.metadata['page'],
+              'chunk_id': chunk_id,
+              'relevance_score': score
+          })
+      
+      prompt = PromptTemplate(
+          template=self.system_prompt,
+          input_variables=["context"]
+      )
+      llm = ""
+
+      if (self.model == default_model):
+        llm = ChatOpenAI(
+            temperature=0,
+            model_name="gpt-4o-mini",
+            api_key=self.openai_api_key
+        )
+      else:
+        llm = HuggingFaceEndpoint(
+          repo_id=self.model,
+          task="text-generation",
+          max_new_tokens=1100,
+          do_sample=False,
+          huggingfacehub_api_token=os.environ.get("HUGGINGFACEHUB_API_TOKEN")
+        )
+
+      
+      response = llm.predict(prompt.format(context="\n\n".join(contexts)))
+      
+      # Split the response into paragraphs
+      summaries = [p.strip() for p in response.split('\n\n') if p.strip()]
+      
+      # Create structured output
+      structured_output = []
+      for idx, summary in enumerate(summaries):
+          # Associate each summary with the most relevant source
+          structured_output.append({
+              "content": summary,
+              "source": {
+                  "page": sources[min(idx, len(sources)-1)]['page'],
+                  "text": sources[min(idx, len(sources)-1)]['content'][:200] + "...",
+                  "relevance_score": sources[min(idx, len(sources)-1)]['relevance_score']
+              }
+          })
+      
+      return structured_output
+
+  def get_source_context(self, chunk_id: str, window: int = 100) -> Dict:
+      """Get extended context around a specific chunk"""
+      metadata = self.chunk_metadata.get(chunk_id)
+      if not metadata:
+          return None
+      
+      return {
+          'page': metadata['page'],
+          'start_char': metadata['start_char'],
+          'end_char': metadata['end_char']
+      }
+
+def get_llm_summary_answer_by_cursor(serializer, listaPDFs):
+  # By Luan
+  allPdfsChunks = []
+
+  # Initialize summarizer
+  summarizer = DocumentSummarizer(
+    openai_api_key=os.environ.get("OPENAI_API_KEY"),
+    embedding=serializer["hf_embedding"],
+    chunk_config={"size": serializer["chunk_size"], "overlap": serializer["chunk_overlap"]},
+    system_prompt=serializer["system_prompt"],
+    model=serializer["model"]
+  )
+  
+  # Load and process document
+  for pdf in listaPDFs:
+    pdf_path = pdf
+    chunks = summarizer.load_and_split_document(pdf_path)
+    allPdfsChunks = allPdfsChunks + chunks
+
+  vector_store = summarizer.create_vector_store(allPdfsChunks)
+  
+  # Generate structured summary
+  structured_summaries = summarizer.generate_summary_with_sources(vector_store)
+  
+  # Print or return the structured data
+  # print(structured_summaries)
+  json_data = json.dumps(structured_summaries)
+  print("\n\n")
+  print(json_data)
+  return structured_summaries
+  # If you need to send to frontend, you can just return structured_summaries
+  # It will be in the format:
+  # [
+  #     {
+  #         "content": "Summary point 1...",
+  #         "source": {
+  #             "page": 1,
+  #             "text": "Source text...",
+  #             "relevance_score": 0.95
+  #         }
+  #     },
+  #     ...
+  # ]
+
+if __name__ == "__main__":
+    get_llm_summary_answer_by_cursor()

resumos/serializer.py

@@ -1,25 +1,29 @@
 from rest_framework import serializers
+from setup.environment import default_model
+# from _utils.utils import DEFAULT_SYSTEM_PROMPT
 
-# Exemplo de retorno que devo enviar ao frontend:
-# {
-#   "nome_do_memorial": "[b]Memorial de Defesa em Ação de Indenização por Danos Morais[/b]",
-#   "argumentos": "[b]Argumentos:[/b]\n[i]• Responsabilidade Civil do Réu:[/i] [i]O réu agiu de forma negligente ao causar o dano. Há dever de indenizar baseado no artigo 186 do Código Civil Brasileiro.[/i]\n[i]• Dano Moral Comprovado:[/i] [i]O dano sofrido pela parte autora é evidente, gerando sofrimento e abalo psicológico. O nexo causal entre a ação do réu e o dano sofrido é claro.[/i]",
-#   "jurisprudencia": "[b]Jurisprudência:[/b]\n[i]• STJ, REsp 123456/DF -[/i] [i]O Superior Tribunal de Justiça entendeu que a indenização por danos morais deve ser fixada de acordo com a gravidade do ato ilícito.[/i]\n[i]• STJ, REsp 654321/SP -[/i] [i]A jurisprudência confirma que o réu tem o dever de reparar integralmente o dano causado.[/i]",
-#   "doutrina": "[b]Doutrina:[/b]\n[i]• Carlos Roberto Gonçalves, Responsabilidade Civil -[/i] [i]A responsabilidade civil é objetiva quando há risco para os direitos da personalidade da vítima.[/i]\n[i]• Maria Helena Diniz, Curso de Direito Civil -[/i] [i]O dano moral é configurado pela violação dos direitos da personalidade.[/i]",
-#   "argumentos_faltantes": "[b]Argumentos Faltantes:[/b]\n[i]• Prova pericial de impacto psicológico.[/i]\n[i]• Estudo comparativo com casos análogos para quantificação do valor da indenização.[/i]",
-#   "palavras_chave": [
-#     "[i]Responsabilidade civil[/i]",
-#     "[i]Dano moral[/i]",
-#     "[i]Nexo causal[/i]",
-#     "[i]Indenização[/i]"
-#   ]
-# }
-
-# pecam para a AI formatar em BBcode
+prompt_template = """
+    Based on the following context, provide multiple key points from the document.
+    For each point, create a new paragraph.
+    Each paragraph should be a complete, self-contained insight.
+    
+    Context: {context}
+    
+    Key points:
+    """
 
 class ResumoPDFSerializer(serializers.Serializer):
     files = serializers.ListField(child=serializers.FileField(), required=True)
     system_prompt = serializers.CharField(required=False)
     user_message = serializers.CharField(required=False, default="")
     model = serializers.CharField(required=False)
-    iterative_refinement = serializers.BooleanField(required=False, default=False)
+    iterative_refinement = serializers.BooleanField(required=False, default=False)
+
+class ResumoCursorSerializer(serializers.Serializer):
+    files = serializers.ListField(child=serializers.FileField(), required=True)
+    system_prompt = serializers.CharField(required=False, default=prompt_template)
+    user_message = serializers.CharField(required=False, default="")
+    model = serializers.CharField(required=False, default=default_model)
+    hf_embedding = serializers.CharField(required=False, default="all-MiniLM-L6-v2")
+    chunk_size = serializers.IntegerField(required=False, default=1000)
+    chunk_overlap = serializers.IntegerField(required=False, default=200)

resumos/views.py

@@ -2,8 +2,9 @@ from rest_framework.views import APIView
 import tempfile, os
 from rest_framework.response import Response
 
+from _utils.resumo_completo_cursor import get_llm_summary_answer_by_cursor
 from _utils.utils import DEFAULT_SYSTEM_PROMPT
-from .serializer import ResumoPDFSerializer
+from .serializer import ResumoPDFSerializer, ResumoCursorSerializer
 from _utils.main import get_llm_answer_summary, get_llm_answer_summary_with_embedding
 from setup.environment import default_model
 from rest_framework.parsers import MultiPartParser
@@ -68,6 +69,35 @@ class ResumoEmbeddingView(APIView):
         system_prompt = data.get("system_prompt", DEFAULT_SYSTEM_PROMPT)
         resposta_llm = get_llm_answer_summary_with_embedding(system_prompt, data["user_message"], listaPDFs, model=model, isIterativeRefinement=data["iterative_refinement"])
 
+        for file in listaPDFs:
+            os.remove(file)
+
+        return Response({"resposta": resposta_llm})
+
+class ResumoCompletoCursorView(APIView):
+  parser_classes = [MultiPartParser]
+
+  @extend_schema(
+    request=ResumoCursorSerializer,
+  )
+  def post(self, request):
+    serializer = ResumoCursorSerializer(data=request.data)
+    if serializer.is_valid(raise_exception=True):
+        listaPDFs = []
+        data = serializer.validated_data
+        print('\nserializer.validated_data: ', serializer.validated_data)
+
+        for file in serializer.validated_data['files']:
+            file.seek(0)
+            with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as temp_file: # Create a temporary file to save the uploaded PDF
+                for chunk in file.chunks(): # Write the uploaded file content to the temporary file
+                    temp_file.write(chunk)
+                temp_file_path = temp_file.name  # Get the path of the temporary file
+                listaPDFs.append(temp_file_path)
+        print('listaPDFs: ', listaPDFs)
+
+        resposta_llm = get_llm_summary_answer_by_cursor(data, listaPDFs)
+
         for file in listaPDFs:
             os.remove(file)
 

setup/urls.py

@@ -5,7 +5,7 @@ from drf_spectacular.views import SpectacularSwaggerView, SpectacularAPIView
 
 
 from pdfs.views import getPDF
-from resumos.views import ResumoView
+from resumos.views import ResumoView, ResumoCompletoCursorView
 from modelos_usuarios.views import ListCreateModeloUsuarioView, CreateUpdateDeleteModeloUsuarioView, ListModelosPorUsuarioView
 
 router = routers.DefaultRouter()
@@ -16,9 +16,9 @@ urlpatterns = [
     path('swagger/', SpectacularSwaggerView.as_view(url_name='schema'), name='swagger-ui'),
     path("admin/", admin.site.urls),
     path('', include(router.urls)),
-
     path('pdf', getPDF, name='upload-pdf'),
     path('resumo', ResumoView.as_view(), name='summary-pdf'),
+    path('resumo/cursor', ResumoCompletoCursorView.as_view(), name='summary-cursor-pdf'),
     path("modelo", ListCreateModeloUsuarioView.as_view()),
     path("modelo/<int:pk>", CreateUpdateDeleteModeloUsuarioView.as_view()),
     path("usuario/<int:user_id>/modelos", ListModelosPorUsuarioView.as_view())