Update rag_query.py

Author: reichaves

tools/rag_query.py

@@ -17,7 +17,7 @@ def __init__(self, *args, **kwargs):
         super().__init__()
         self.is_initialized = True
 
-    def forward(self, question: str, vectorstore: Any, llm_api_key: str) -> str:
+    def forward(self, question: str, vectorstore: Any, llm_api_key: str, use_general_knowledge: bool = True) -> str:
         """Performs a RAG query against the video transcript"""
         try:
             if vectorstore is None:
@@ -36,22 +36,28 @@ def forward(self, question: str, vectorstore: Any, llm_api_key: str) -> str:
             context = "\n\n".join([doc.page_content for doc in docs])
 
             # Generate response with Groq
-            system_content = """Você é um assistente especializado em vídeos que responde perguntas com base 
-            em transcrições. Responda apenas com informações encontradas no contexto fornecido. 
-            Se a resposta não estiver no contexto, admita que não pode responder com base na 
-            transcrição disponível."""
-            
+            system_content = (
+                "Você é um assistente especialista em análise de vídeos e política brasileira.\n\n"
+                "Responda à pergunta com base:\n"
+                "1. No contexto da transcrição do vídeo\n"
+                "2. No seu próprio conhecimento geral\n"
+                "3. (Opcional) Em fatos recentes, se aplicável\n\n"
+                "Deixe claro quando uma parte da resposta vem da transcrição e quando vem do seu conhecimento geral."
+                if use_general_knowledge else
+                "Você é um assistente especializado em vídeos que responde perguntas com base apenas na transcrição fornecida.\n"
+                "Responda exclusivamente com informações contidas no contexto.\n"
+                "Se a informação não estiver no contexto, diga que não pode responder com base no que foi fornecido."
+            )
+
             user_content = f"""
-            Responda à seguinte pergunta com base na transcrição do vídeo:
-            
-            Contexto da transcrição:
-            {context}
-            
-            Pergunta: {question}
-            """
-            
+Pergunta: {question}
+
+Transcrição do vídeo (para análise):
+{context}
+"""
+
             chat_response = client.chat.completions.create(
-                model="deepseek-coder-33b-instruct",  # Modelo DeepSeek disponível no Groq
+                model="deepseek-coder-33b-instruct",
                 messages=[
                     {"role": "system", "content": system_content},
                     {"role": "user", "content": user_content}