Update backend/main.py

backend/main.py

@@ -12,7 +12,8 @@ from langchain_groq import ChatGroq
 from fastapi.responses import StreamingResponse
 
 # RAG Imports
-from langchain_community.document_loaders import PyPDFLoader
+# MUDANÇA: Adicionado TextLoader para arquivos de texto
+from langchain_community.document_loaders import PyPDFLoader, TextLoader
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_community.vectorstores import FAISS
@@ -35,7 +36,7 @@ app.add_middleware(
 # --- 2. Configurações de IA ---
 HF_EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
 
-# Configuração do Modelo (Ignorando variáveis antigas para garantir funcionamento)
+# Modelo Groq
 model = ChatGroq(
     model="llama-3.3-70b-versatile",
     temperature=0.3 
@@ -50,6 +51,7 @@ rag_system_prompt = (
     "Você é um assistente experiente e prestativo. "
     "Sua tarefa é fornecer respostas detalhadas e ricas em contexto com base nas informações fornecidas. "
     "Ao usar o contexto abaixo, sintetize os pontos principais e explique como eles se relacionam. "
+    "Se o contexto for código (Python/MD), explique o funcionamento ou use como referência."
     "Se a resposta não estiver no contexto, diga honestamente que não sabe, não invente informações."
     "\n\nCONTEXTO DO DOCUMENTO:\n{context}"
 )
@@ -58,9 +60,8 @@ 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
+global_retriever = None
 
 # --- 3. Modelo de Dados ---
 class ChatRequest(BaseModel):
@@ -69,7 +70,7 @@ class ChatRequest(BaseModel):
 # --- 4. Helpers ---
 def format_docs(docs):
     return "\n\n---\n\n".join(
-        f"📄 Conteúdo: {doc.page_content}\n(🔖 Fonte: Página {doc.metadata.get('page', 'N/A')})" 
+        f"📄 Conteúdo: {doc.page_content}\n(🔖 Fonte: {doc.metadata.get('source', 'Desconhecida')})" 
         for doc in docs
     )
 
@@ -79,34 +80,46 @@ def format_docs(docs):
 async def upload_document(file: UploadFile = File(...)):
     global rag_chain, global_retriever
 
-    if file.content_type != "application/pdf":
-        raise HTTPException(status_code=400, detail="Apenas arquivos PDF são suportados.")
+    # 1. Validação de extensão
+    filename = file.filename.lower()
+    allowed_extensions = [".pdf", ".txt", ".md", ".py"]
+    
+    if not any(filename.endswith(ext) for ext in allowed_extensions):
+        raise HTTPException(status_code=400, detail=f"Extensão não suportada. Use: {allowed_extensions}")
 
     try:
-        with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_file:
+        # Salva arquivo temporário com a extensão correta (importante para o Loader)
+        file_ext = os.path.splitext(filename)[1]
+        with tempfile.NamedTemporaryFile(delete=False, suffix=file_ext) as tmp_file:
             content = await file.read()
             tmp_file.write(content)
             temp_path = tmp_file.name
         
-        loader = PyPDFLoader(temp_path)
+        # 2. Seleção do Loader baseado na extensão
+        if filename.endswith(".pdf"):
+            loader = PyPDFLoader(temp_path)
+        else:
+            # Para .txt, .md, .py usamos o TextLoader com UTF-8
+            loader = TextLoader(temp_path, encoding="utf-8")
+
         docs = loader.load()
 
+        # 3. Chunking
+        # Se for código (.py), talvez chunks menores sejam melhores, mas manteremos o padrão por enquanto
         text_splitter = RecursiveCharacterTextSplitter(chunk_size=1200, chunk_overlap=200)
         splits = text_splitter.split_documents(docs)
 
+        # 4. Vetorização
         vectorstore = FAISS.from_documents(documents=splits, embedding=embeddings)
-        
-        # MUDANÇA: Salvamos o retriever globalmente
         global_retriever = vectorstore.as_retriever(search_kwargs={"k": 6}) 
         
-        # A chain agora é simples, pois faremos a recuperação manual no endpoint chat
         rag_chain = rag_prompt | model | StrOutputParser()
 
         return {
             "message": "Processamento concluído!",
             "filename": file.filename,
-            "total_pages": len(docs),
-            "total_chunks": len(splits)
+            "total_chunks": len(splits),
+            "type": file_ext
         }
     
     except Exception as e:
@@ -118,44 +131,33 @@ 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 = []
 
-    # 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()
         )
 
-    # 2. Gerador de Streaming com "Payload Oculto"
     async def stream_generator():
         try:
-            # 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