Update backend/main.py

backend/main.py

@@ -35,8 +35,7 @@ app.add_middleware(
 # --- 2. Configurações de IA ---
 HF_EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
 
-# MUDANÇA CRÍTICA: Removi o os.getenv para ignorar configurações antigas do ambiente
-# e forçar o uso do modelo novo e funcional.
+# Configuração do Modelo (Ignorando variáveis antigas para garantir funcionamento)
 model = ChatGroq(
     model="llama-3.3-70b-versatile",
     temperature=0.3 
@@ -47,7 +46,6 @@ embeddings = HuggingFaceEmbeddings(
     model_kwargs={'device': 'cpu'} 
 ) 
 
-# System Prompt Otimizado
 rag_system_prompt = (
     "Você é um assistente experiente e prestativo. "
     "Sua tarefa é fornecer respostas detalhadas e ricas em contexto com base nas informações fornecidas. "
@@ -60,7 +58,9 @@ rag_prompt = ChatPromptTemplate.from_messages(
     [("system", rag_system_prompt), ("human", "{input}")]
 )
 
+# Variáveis globais para armazenar a "inteligência" do RAG
 rag_chain = None
+global_retriever = None  # MUDANÇA: Retriever global para acessarmos no chat
 
 # --- 3. Modelo de Dados ---
 class ChatRequest(BaseModel):
@@ -69,7 +69,7 @@ class ChatRequest(BaseModel):
 # --- 4. Helpers ---
 def format_docs(docs):
     return "\n\n---\n\n".join(
-        f"Conteúdo: {doc.page_content}\n(Página {doc.metadata.get('page', 'N/A')})" 
+        f"📄 Conteúdo: {doc.page_content}\n(🔖 Fonte: Página {doc.metadata.get('page', 'N/A')})" 
         for doc in docs
     )
 
@@ -77,7 +77,7 @@ def format_docs(docs):
 
 @app.post("/upload-document")
 async def upload_document(file: UploadFile = File(...)):
-    global rag_chain
+    global rag_chain, global_retriever
 
     if file.content_type != "application/pdf":
         raise HTTPException(status_code=400, detail="Apenas arquivos PDF são suportados.")
@@ -91,24 +91,16 @@ async def upload_document(file: UploadFile = File(...)):
         loader = PyPDFLoader(temp_path)
         docs = loader.load()
 
-        # Chunking otimizado
         text_splitter = RecursiveCharacterTextSplitter(chunk_size=1200, chunk_overlap=200)
         splits = text_splitter.split_documents(docs)
 
         vectorstore = FAISS.from_documents(documents=splits, embedding=embeddings)
         
-        # Recuperação otimizada (k=6)
-        retriever = vectorstore.as_retriever(search_kwargs={"k": 6}) 
+        # MUDANÇA: Salvamos o retriever globalmente
+        global_retriever = vectorstore.as_retriever(search_kwargs={"k": 6}) 
         
-        # Chain
-        rag_chain = (
-            RunnablePassthrough.assign(
-                context=(lambda x: x["input"]) | retriever | format_docs
-            )
-            | rag_prompt
-            | model
-            | StrOutputParser()
-        )
+        # A chain agora é simples, pois faremos a recuperação manual no endpoint chat
+        rag_chain = rag_prompt | model | StrOutputParser()
 
         return {
             "message": "Processamento concluído!",
@@ -126,19 +118,48 @@ async def upload_document(file: UploadFile = File(...)):
 
 @app.post("/chat")
 async def chat(request: ChatRequest):
+    """
+    Endpoint de chat com Auditoria (envia contexto ao final)
+    """
     current_chain = rag_chain
+    context_str = ""
+    docs_source = []
 
-    if current_chain is None:
-        base_prompt = ChatPromptTemplate.from_messages(
-            [("system", "Você é um assistente útil. Nenhum documento foi carregado ainda."), ("human", "{input}")]
+    # 1. Recuperação Manual de Contexto (Se o RAG estiver ativo)
+    if global_retriever:
+        try:
+            # Busca os documentos relevantes
+            docs_source = global_retriever.invoke(request.content)
+            context_str = format_docs(docs_source)
+        except Exception as e:
+            print(f"Erro na recuperação: {e}")
+            context_str = "Erro ao recuperar contexto."
+    else:
+        # Fallback se não houver PDF
+        current_chain = (
+            ChatPromptTemplate.from_messages([("system", "Você é um assistente útil."), ("human", "{input}")])
+            | model 
+            | StrOutputParser()
         )
-        current_chain = base_prompt | model | StrOutputParser()
 
+    # 2. Gerador de Streaming com "Payload Oculto"
     async def stream_generator():
         try:
-            async for chunk in current_chain.astream({"input": request.content}):
+            # Passa o contexto manualmente para o prompt
+            input_data = {"input": request.content}
+            if context_str:
+                input_data["context"] = context_str
+
+            # Stream da resposta da IA
+            async for chunk in current_chain.astream(input_data):
                 if chunk:
                     yield chunk
+            
+            # MUDANÇA: Ao final, enviamos um separador e o contexto para auditoria
+            if context_str:
+                debug_data = f"\n\n###__DEBUG__###\n**Auditoria de Contexto (RAG):**\n\n{context_str}"
+                yield debug_data
+
         except Exception as e:
              print(f"Erro stream: {e}")
              yield f"Erro no serviço de IA: {e}"