fix: cot bug

Author: jasperan

agentic_rag/README.md

@@ -119,6 +119,7 @@ The API will be available at `http://localhost:8000`. You can then use the API e
 ### 2. Using the Gradio Interface (Recommended)
 
 The system provides a user-friendly web interface using Gradio, which allows you to:
+- Select and pull `ollama` models directly from the interface
 - Upload and process PDF documents
 - Process web content from URLs
 - Chat with your documents using either local or OpenAI models

agentic_rag/local_rag_agent.py

@@ -272,35 +272,76 @@ def _process_query_with_cot(self, query: str) -> Dict[str, Any]:
                 self.agents = create_agents(self.llm, self.vector_store)
 
             # Get planning step
-            planning_result = self.agents["planner"].plan(query, context)
-            logger.info("Planning step completed")
+            try:
+                planning_result = self.agents["planner"].plan(query, context)
+                logger.info("Planning step completed")
+            except Exception as e:
+                logger.error(f"Error in planning step: {str(e)}")
+                logger.info("Falling back to general response")
+                return self._generate_general_response(query)
 
             # Get research step
-            research_result = self.agents["researcher"].research(query, context)
-            logger.info("Research step completed")
+            research_results = []
+            if self.agents.get("researcher") is not None and context:
+                for step in planning_result.split("\n"):
+                    if not step.strip():
+                        continue
+                    try:
+                        step_research = self.agents["researcher"].research(query, step)
+                        # Extract findings from research result
+                        findings = step_research.get("findings", []) if isinstance(step_research, dict) else []
+                        research_results.append({"step": step, "findings": findings})
+                        
+                        # Log which sources were used for this step
+                        try:
+                            source_indices = [context.index(finding) + 1 for finding in findings if finding in context]
+                            logger.info(f"Research for step: {step}\nUsing sources: {source_indices}")
+                        except ValueError as ve:
+                            logger.warning(f"Could not find some findings in initial context: {str(ve)}")
+                    except Exception as e:
+                        logger.error(f"Error during research for step '{step}': {str(e)}")
+                        research_results.append({"step": step, "findings": []})
+            else:
+                # If no researcher or no context, use the steps directly
+                research_results = [{"step": step, "findings": []} for step in planning_result.split("\n") if step.strip()]
+                logger.info("No research performed (no researcher agent or no context available)")
 
             # Get reasoning step
-            reasoning_result = self.agents["reasoner"].reason(query, research_result["context"])
-            logger.info("Reasoning step completed")
+            reasoning_steps = []
+            if not self.agents.get("reasoner"):
+                logger.warning("No reasoner agent available, using direct response")
+                return self._generate_general_response(query)
+            
+            for result in research_results:
+                try:
+                    step_reasoning = self.agents["reasoner"].reason(
+                        query,
+                        result["step"],
+                        result["findings"] if result["findings"] else [{"content": "Using general knowledge", "metadata": {"source": "General Knowledge"}}]
+                    )
+                    reasoning_steps.append(step_reasoning)
+                    logger.info(f"Reasoning for step: {result['step']}\n{step_reasoning}")
+                except Exception as e:
+                    logger.error(f"Error in reasoning for step '{result['step']}': {str(e)}")
+                    reasoning_steps.append(f"Error in reasoning for this step: {str(e)}")
 
             # Get synthesis step
-            synthesis_result = self.agents["synthesizer"].synthesize(
-                query, 
-                planning_result["context"],
-                research_result["context"],
-                reasoning_result["context"]
-            )
-            logger.info("Synthesis step completed")
+            if not self.agents.get("synthesizer"):
+                logger.warning("No synthesizer agent available, using direct response")
+                return self._generate_general_response(query)
+            
+            try:
+                synthesis_result = self.agents["synthesizer"].synthesize(query, reasoning_steps)
+                logger.info("Synthesis step completed")
+            except Exception as e:
+                logger.error(f"Error in synthesis step: {str(e)}")
+                logger.info("Falling back to general response")
+                return self._generate_general_response(query)
 
             return {
                 "answer": synthesis_result["answer"],
-                "reasoning_steps": [
-                    planning_result["answer"],
-                    research_result["answer"],
-                    reasoning_result["answer"],
-                    synthesis_result["answer"]
-                ],
-                "context": synthesis_result["context"]
+                "reasoning_steps": reasoning_steps,
+                "context": context
             }
 
         except Exception as e:

agentic_rag/rag_agent.py

@@ -102,8 +102,13 @@ def _process_query_with_cot(self, query: str) -> Dict[str, Any]:
                 logger.warning("No planner agent available, using direct response")
                 return self._generate_general_response(query)
 
-            plan = self.agents["planner"].plan(query, initial_context)
-            logger.info(f"Generated plan:\n{plan}")
+            try:
+                plan = self.agents["planner"].plan(query, initial_context)
+                logger.info(f"Generated plan:\n{plan}")
+            except Exception as e:
+                logger.error(f"Error in planning step: {str(e)}")
+                logger.info("Falling back to general response")
+                return self._generate_general_response(query)
 
             # Step 2: Research each step (if researcher is available)
             logger.info("Step 2: Research")
@@ -112,11 +117,21 @@ def _process_query_with_cot(self, query: str) -> Dict[str, Any]:
                 for step in plan.split("\n"):
                     if not step.strip():
                         continue
-                    step_research = self.agents["researcher"].research(query, step)
-                    research_results.append({"step": step, "findings": step_research})
-                    # Log which sources were used for this step
-                    source_indices = [initial_context.index(finding) + 1 for finding in step_research if finding in initial_context]
-                    logger.info(f"Research for step: {step}\nUsing sources: {source_indices}")
+                    try:
+                        step_research = self.agents["researcher"].research(query, step)
+                        # Extract findings from research result
+                        findings = step_research.get("findings", []) if isinstance(step_research, dict) else []
+                        research_results.append({"step": step, "findings": findings})
+                        
+                        # Log which sources were used for this step
+                        try:
+                            source_indices = [initial_context.index(finding) + 1 for finding in findings if finding in initial_context]
+                            logger.info(f"Research for step: {step}\nUsing sources: {source_indices}")
+                        except ValueError as ve:
+                            logger.warning(f"Could not find some findings in initial context: {str(ve)}")
+                    except Exception as e:
+                        logger.error(f"Error during research for step '{step}': {str(e)}")
+                        research_results.append({"step": step, "findings": []})
             else:
                 # If no researcher or no context, use the steps directly
                 research_results = [{"step": step, "findings": []} for step in plan.split("\n") if step.strip()]
@@ -130,30 +145,39 @@ def _process_query_with_cot(self, query: str) -> Dict[str, Any]:
 
             reasoning_steps = []
             for result in research_results:
-                step_reasoning = self.agents["reasoner"].reason(
-                    query,
-                    result["step"],
-                    result["findings"] if result["findings"] else [{"content": "Using general knowledge", "metadata": {"source": "General Knowledge"}}]
-                )
-                reasoning_steps.append(step_reasoning)
-                logger.info(f"Reasoning for step: {result['step']}\n{step_reasoning}")
+                try:
+                    step_reasoning = self.agents["reasoner"].reason(
+                        query,
+                        result["step"],
+                        result["findings"] if result["findings"] else [{"content": "Using general knowledge", "metadata": {"source": "General Knowledge"}}]
+                    )
+                    reasoning_steps.append(step_reasoning)
+                    logger.info(f"Reasoning for step: {result['step']}\n{step_reasoning}")
+                except Exception as e:
+                    logger.error(f"Error in reasoning for step '{result['step']}': {str(e)}")
+                    reasoning_steps.append(f"Error in reasoning for this step: {str(e)}")
 
             # Step 4: Synthesize final answer
             logger.info("Step 4: Synthesis")
             if not self.agents.get("synthesizer"):
                 logger.warning("No synthesizer agent available, using direct response")
                 return self._generate_general_response(query)
 
-            final_answer = self.agents["synthesizer"].synthesize(query, reasoning_steps)
-            logger.info(f"Final synthesized answer:\n{final_answer}")
+            try:
+                final_answer = self.agents["synthesizer"].synthesize(query, reasoning_steps)
+                logger.info(f"Final synthesized answer:\n{final_answer}")
+            except Exception as e:
+                logger.error(f"Error in synthesis step: {str(e)}")
+                logger.info("Falling back to general response")
+                return self._generate_general_response(query)
 
             return {
                 "answer": final_answer,
                 "context": initial_context,
                 "reasoning_steps": reasoning_steps
             }
         except Exception as e:
-            logger.error(f"Error in CoT processing: {str(e)}")
+            logger.error(f"Error in CoT processing: {str(e)}", exc_info=True)
             logger.info("Falling back to general response")
             return self._generate_general_response(query)