""" Multi-Step Reasoning Agent - Performs structured reasoning for complex queries. Breaks down complex medical queries into steps: 1. Problem understanding 2. Evidence gathering 3. Condition comparison 4. Answer generation 5. Validation """ import sys from pathlib import Path from typing import Dict, List, Optional from langchain_openai import ChatOpenAI from langchain_core.messages import SystemMessage, HumanMessage from loguru import logger sys.path.append(str(Path(__file__).parent.parent)) from utils.config import config class ReasoningAgent: """ Performs multi-step reasoning for complex medical queries. Steps: 1. Analyze problem 2. Organize evidence 3. Compare conditions 4. Generate answer 5. Validate answer """ PROBLEM_ANALYSIS_PROMPT = """You are a medical problem analyzer. Break down this query into key components. Identify: 1. Main medical concern 2. Relevant symptoms/conditions mentioned 3. What the user wants to know 4. Any context or constraints Return a structured analysis in 2-3 sentences.""" EVIDENCE_ORGANIZATION_PROMPT = """You are organizing medical evidence. Given these retrieved documents, organize them by relevance and topic. Group evidence into: 1. Directly relevant (answers the question) 2. Supporting information (provides context) 3. Related but tangential Summarize each group in 1-2 sentences.""" CONDITION_COMPARISON_PROMPT = """You are comparing medical conditions or treatments. Given the evidence, compare the relevant options. For each option: 1. Key characteristics 2. Pros/cons or benefits/risks 3. When it's appropriate 4. Evidence strength Be objective and evidence-based.""" ANSWER_GENERATION_PROMPT = """You are generating a final medical answer. Based on the reasoning steps, provide a clear, accurate answer. Structure: 1. Direct answer to the question 2. Key supporting points (2-3) 3. Important caveats or warnings 4. When to seek professional help Use simple language. Be accurate and helpful.""" VALIDATION_PROMPT = """You are validating a medical answer. Check if the answer: 1. Answers the original question 2. Is grounded in the provided evidence 3. Contains appropriate medical disclaimers 4. Is clear and understandable 5. Doesn't make claims beyond the evidence Return: "VALID" or "NEEDS_REVISION" with brief reason.""" def __init__(self): self.llm = ChatOpenAI( api_key=config.OPENAI_API_KEY, model="gpt-4o-mini", temperature=0.1 ) logger.info("[ReasoningAgent] Initialized") async def reason( self, query: str, evidence: List[str], context: Optional[str] = None ) -> Dict: """ Perform multi-step reasoning. Args: query: User's question evidence: Retrieved evidence/documents context: Optional conversation context Returns: Dict with reasoning steps and final answer """ reasoning_steps = [] try: # Step 1: Analyze Problem problem_analysis = await self._analyze_problem(query, context) reasoning_steps.append({ "step": 1, "name": "Problem Analysis", "output": problem_analysis }) # Step 2: Organize Evidence organized_evidence = await self._organize_evidence(query, evidence) reasoning_steps.append({ "step": 2, "name": "Evidence Organization", "output": organized_evidence }) # Step 3: Compare Conditions (if applicable) comparison = await self._compare_conditions(query, organized_evidence) reasoning_steps.append({ "step": 3, "name": "Condition Comparison", "output": comparison }) # Step 4: Generate Answer answer = await self._generate_answer(query, reasoning_steps) reasoning_steps.append({ "step": 4, "name": "Answer Generation", "output": answer }) # Step 5: Validate Answer validation = await self._validate_answer(query, answer, evidence) reasoning_steps.append({ "step": 5, "name": "Validation", "output": validation }) # Calculate confidence based on validation confidence = 0.9 if "VALID" in validation else 0.6 logger.info(f"[ReasoningAgent] Completed 5-step reasoning with confidence {confidence}") return { "answer": answer, "reasoning_steps": reasoning_steps, "confidence": confidence, "validation_status": "VALID" if "VALID" in validation else "NEEDS_REVIEW" } except Exception as e: logger.error(f"[ReasoningAgent] Reasoning failed: {e}") return { "answer": "I encountered an error during reasoning. Please try rephrasing your question.", "reasoning_steps": reasoning_steps, "confidence": 0.3, "validation_status": "ERROR", "error": str(e) } async def _analyze_problem(self, query: str, context: Optional[str]) -> str: """Step 1: Analyze the problem""" messages = [ SystemMessage(content=self.PROBLEM_ANALYSIS_PROMPT), HumanMessage(content=f"Query: {query}\n\nContext: {context or 'None'}") ] response = await self.llm.ainvoke(messages) return response.content.strip() async def _organize_evidence(self, query: str, evidence: List[str]) -> str: """Step 2: Organize evidence""" evidence_text = "\n\n---\n\n".join(evidence[:10]) if evidence else "No evidence provided" messages = [ SystemMessage(content=self.EVIDENCE_ORGANIZATION_PROMPT), HumanMessage(content=f"Query: {query}\n\nEvidence:\n{evidence_text}") ] response = await self.llm.ainvoke(messages) return response.content.strip() async def _compare_conditions(self, query: str, organized_evidence: str) -> str: """Step 3: Compare conditions/options""" messages = [ SystemMessage(content=self.CONDITION_COMPARISON_PROMPT), HumanMessage(content=f"Query: {query}\n\nOrganized Evidence:\n{organized_evidence}") ] response = await self.llm.ainvoke(messages) return response.content.strip() async def _generate_answer(self, query: str, reasoning_steps: List[Dict]) -> str: """Step 4: Generate final answer""" # Combine all reasoning steps reasoning_summary = "\n\n".join([ f"Step {step['step']} - {step['name']}:\n{step['output']}" for step in reasoning_steps ]) messages = [ SystemMessage(content=self.ANSWER_GENERATION_PROMPT), HumanMessage(content=f"Original Query: {query}\n\nReasoning Steps:\n{reasoning_summary}") ] response = await self.llm.ainvoke(messages) return response.content.strip() async def _validate_answer(self, query: str, answer: str, evidence: List[str]) -> str: """Step 5: Validate the answer""" evidence_text = "\n".join(evidence[:5]) if evidence else "No evidence" messages = [ SystemMessage(content=self.VALIDATION_PROMPT), HumanMessage(content=f"Query: {query}\n\nAnswer: {answer}\n\nEvidence: {evidence_text}") ] response = await self.llm.ainvoke(messages) return response.content.strip() # Singleton instance reasoning_agent = ReasoningAgent()