Fix output keys format with wrapper function

app/core/agent.py

@@ -50,23 +50,13 @@ Assistant:"""
             
             # Use the RAG chain to get an answer
             response = self.rag_chain({"question": question})
+            logger.info(f"Raw response keys: {list(response.keys())}")
             
-            # Extract the answer and source documents
-            logger.debug(f"RAG chain response keys: {response.keys()}")
+            # Extract the answer (should now be normalized by our wrapper)
+            answer = response.get("answer", "I couldn't generate a proper response.")
             
-            if "answer" not in response:
-                logger.warning(f"Missing 'answer' key in response. Available keys: {response.keys()}")
-                # Create a fallback answer if the expected key is missing
-                answer = "I'm sorry, I encountered an issue processing your request. Let me try a simpler response."
-            else:
-                answer = response["answer"]
-            
-            # Handle different variations of source document keys
-            source_docs = []
-            if "source_documents" in response:
-                source_docs = response["source_documents"]
-            elif "sources" in response:
-                source_docs = response["sources"]
+            # Extract sources (should now be normalized by our wrapper)
+            source_docs = response.get("sources", [])
             
             # Format source documents for display
             sources = []

app/core/memory.py

@@ -126,7 +126,29 @@ class MemoryManager:
                 return_source_documents=True,
                 return_generated_question=False,
             )
-            return chain
+            
+            # Create a wrapper function that normalizes the chain output format
+            def normalized_chain(inputs):
+                logger.info("Executing RAG chain with normalizer")
+                try:
+                    # Execute the original chain
+                    response = chain(inputs)
+                    logger.info(f"Original chain output keys: {list(response.keys())}")
+                    
+                    # Create a normalized response
+                    normalized = {
+                        "answer": response.get("answer", "No answer generated"),
+                        "sources": response.get("source_documents", [])
+                    }
+                    return normalized
+                except Exception as e:
+                    logger.error(f"Error in normalized chain: {e}")
+                    return {
+                        "answer": f"Error processing your query: {str(e)}",
+                        "sources": []
+                    }
+            
+            return normalized_chain
         except Exception as e:
             logger.error(f"Error creating RAG chain: {e}")
             
@@ -138,7 +160,7 @@ class MemoryManager:
                 logger.info(f"Mock chain received query: {inputs.get('question', '')}")
                 return {
                     "answer": "I'm having trouble accessing the knowledge base. I can only answer general questions right now.",
-                    "source_documents": []
+                    "sources": []
                 }
             
             return mock_chain