feat: added tests for cot interface

Author: jasperan

agentic_rag/local_rag_agent.py

@@ -358,8 +358,17 @@ def _process_query_with_cot(self, query: str) -> Dict[str, Any]:
                 logger.info("Falling back to general response")
                 return self._generate_general_response(query)
 
+            # Handle string response from synthesis
+            if isinstance(synthesis_result, str):
+                return {
+                    "answer": synthesis_result,
+                    "reasoning_steps": reasoning_steps,
+                    "context": context
+                }
+            
+            # Handle dictionary response
             return {
-                "answer": synthesis_result["answer"],
+                "answer": synthesis_result.get("answer", synthesis_result) if isinstance(synthesis_result, dict) else synthesis_result,
                 "reasoning_steps": reasoning_steps,
                 "context": context
             }

agentic_rag/tests/test_cot_chat.py

@@ -67,6 +67,17 @@ def test_cot_chat():
             logger.info("Raw response received")
             debug_response_structure(raw_response, "Raw response: ")
 
+            # Verify response structure
+            if not isinstance(raw_response, dict):
+                logger.error(f"Unexpected response type: {type(raw_response)}")
+                raise TypeError(f"Expected dict response, got {type(raw_response)}")
+            
+            required_keys = ["answer", "reasoning_steps", "context"]
+            missing_keys = [key for key in required_keys if key not in raw_response]
+            if missing_keys:
+                logger.error(f"Missing required keys in response: {missing_keys}")
+                raise KeyError(f"Response missing required keys: {missing_keys}")
+            
             # Process through chat function
             logger.info("Processing through chat function...")
             result = chat(
@@ -91,9 +102,21 @@ def test_cot_chat():
         # Save debug information to file
         debug_info = {
             "test_message": test_message,
-            "raw_response": str(raw_response),
-            "final_result": str(result),
-            "history": str(history)
+            "raw_response": {
+                "type": str(type(raw_response)),
+                "keys": list(raw_response.keys()) if isinstance(raw_response, dict) else None,
+                "content": str(raw_response)
+            },
+            "final_result": {
+                "type": str(type(result)),
+                "length": len(result) if isinstance(result, list) else None,
+                "content": str(result)
+            },
+            "history": {
+                "type": str(type(history)),
+                "length": len(history),
+                "content": str(history)
+            }
         }
 
         with open("cot_chat_debug.json", "w") as f: