Update agent.py

agent.py

@@ -1,250 +1,477 @@
-"""
-SynapseAI Clinical Decision Support System
-Expert-Level Implementation with Safety-Centric Architecture
-"""
-
 import os
 import re
 import json
 import logging
-from typing import (Any, Dict, List, Optional, TypedDict, 
-                    Callable, Sequence, Tuple, Union)
+import traceback
 from functools import lru_cache
-from enum import Enum
+from typing import List, Dict, Any, Optional, TypedDict
 
 import requests
-from pydantic import BaseModel, Field, ValidationError
 from langchain_groq import ChatGroq
-from langchain_core.messages import (HumanMessage, SystemMessage, 
-                                   AIMessage, ToolMessage)
-from langchain_core.tools import BaseTool
-from langgraph.graph import StateGraph, END
+from langchain_community.tools.tavily_search import TavilySearchResults
+from langchain_core.messages import HumanMessage, SystemMessage, AIMessage, ToolMessage
+from langchain_core.pydantic_v1 import BaseModel, Field
+from langchain_core.tools import tool
 from langgraph.prebuilt import ToolExecutor
+from langgraph.graph import StateGraph, END
+
+# ── Logging Configuration ──────────────────────────────────────────────
+logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO)
+
+# ── Environment Variables ──────────────────────────────────────────────
+UMLS_API_KEY   = os.getenv("UMLS_API_KEY")
+GROQ_API_KEY   = os.getenv("GROQ_API_KEY")
+TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")
+
+if not all([UMLS_API_KEY, GROQ_API_KEY, TAVILY_API_KEY]):
+    logger.error("Missing one or more required API keys: UMLS_API_KEY, GROQ_API_KEY, TAVILY_API_KEY")
+    raise RuntimeError("Missing required API keys")
+
+# ── Agent Configuration ──────────────────────────────────────────────
+AGENT_MODEL_NAME   = "llama3-70b-8192"
+AGENT_TEMPERATURE  = 0.1
+MAX_SEARCH_RESULTS = 3
+
+class ClinicalPrompts:
+    SYSTEM_PROMPT = """
+    You are SynapseAI, an expert AI clinical assistant engaged in an interactive consultation...
+    [SYSTEM PROMPT CONTENT HERE]
+    """
+
+# ── Helper Functions ─────────────────────────────────────────────────────
+UMLS_AUTH_ENDPOINT  = "https://utslogin.nlm.nih.gov/cas/v1/api-key"
+RXNORM_API_BASE     = "https://rxnav.nlm.nih.gov/REST"
+OPENFDA_API_BASE    = "https://api.fda.gov/drug/label.json"
+
+@lru_cache(maxsize=256)
+def get_rxcui(drug_name: str) -> Optional[str]:
+    """Lookup RxNorm CUI for a given drug name."""
+    drug_name = (drug_name or "").strip()
+    if not drug_name:
+        return None
+    logger.info(f"Looking up RxCUI for '{drug_name}'")
+    try:
+        params = {"name": drug_name, "search": 1}
+        r = requests.get(f"{RXNORM_API_BASE}/rxcui.json", params=params, timeout=10)
+        r.raise_for_status()
+        ids = r.json().get("idGroup", {}).get("rxnormId")
+        if ids:
+            logger.info(f"Found RxCUI {ids[0]} for '{drug_name}'")
+            return ids[0]
+        r = requests.get(f"{RXNORM_API_BASE}/drugs.json", params={"name": drug_name}, timeout=10)
+        r.raise_for_status()
+        for grp in r.json().get("drugGroup", {}).get("conceptGroup", []):
+            props = grp.get("conceptProperties")
+            if props:
+                logger.info(f"Found RxCUI {props[0]['rxcui']} via /drugs for '{drug_name}'")
+                return props[0]["rxcui"]
+    except Exception:
+        logger.exception(f"Error fetching RxCUI for '{drug_name}'")
+    return None
 
-# ── Type Definitions ──────────────────────────────────────────────────
-class ClinicalPriority(str, Enum):
-    STAT = "STAT"
-    URGENT = "Urgent"
-    ROUTINE = "Routine"
-
-class ClinicalState(TypedDict):
-    messages: List[Union[HumanMessage, SystemMessage, AIMessage, ToolMessage]]
-    patient_data: Dict[str, Any]
-    safety_warnings: List[Dict[str, str]]
-    workflow_metadata: Dict[str, Union[int, float, bool]]
-    execution_log: List[Dict[str, str]]
-
-# ── Configuration ─────────────────────────────────────────────────────
-class ClinicalConfig:
-    MAX_ITERATIONS = 6  # Evidence-based conversation turn limit
-    RECURSION_BUFFER = 2  # Safety margin for LangGraph execution
-    DRUG_CHECK_REQUIRED = True  # Hard enforcement for interaction checks
-    
-    SAFETY_PARAMETERS = {
-        'max_bp_systolic': 180,
-        'min_bp_systolic': 90,
-        'max_hr': 120,
-        'min_spo2': 92
+@lru_cache(maxsize=128)
+def get_openfda_label(rxcui: Optional[str] = None, drug_name: Optional[str] = None) -> Optional[Dict[str, Any]]:
+    """Fetch the OpenFDA label for a drug by RxCUI or name."""
+    if not (rxcui or drug_name):
+        return None
+    terms = []
+    if rxcui:
+        terms.append(f'spl_rxnorm_code:"{rxcui}" OR openfda.rxcui:"{rxcui}"')
+    if drug_name:
+        dn = drug_name.lower()
+        terms.append(f'(openfda.brand_name:"{dn}" OR openfda.generic_name:"{dn}")')
+    query = " OR ".join(terms)
+    logger.info(f"Looking up OpenFDA label with query: {query}")
+    try:
+        r = requests.get(OPENFDA_API_BASE, params={"search": query, "limit": 1}, timeout=15)
+        r.raise_for_status()
+        results = r.json().get("results", [])
+        if results:
+            return results[0]
+    except Exception:
+        logger.exception("Error fetching OpenFDA label")
+    return None
+
+def search_text_list(texts: List[str], terms: List[str]) -> List[str]:
+    """Return highlighted snippets from a list of texts containing any of the search terms."""
+    snippets = []
+    lowers = [t.lower() for t in terms if t]
+    for text in texts or []:
+        tl = text.lower()
+        for term in lowers:
+            if term in tl:
+                i = tl.find(term)
+                start = max(0, i - 50)
+                end   = min(len(text), i + len(term) + 100)
+                snippet = text[start:end]
+                snippet = re.sub(f"({re.escape(term)})", r"**\1**", snippet, flags=re.IGNORECASE)
+                snippets.append(f"...{snippet}...")
+                break
+    return snippets
+
+def parse_bp(bp: str) -> Optional[tuple[int, int]]:
+    """Parse 'SYS/DIA' blood pressure string into a (sys, dia) tuple."""
+    if m := re.match(r"(\d{1,3})\s*/\s*(\d{1,3})", (bp or "").strip()):
+        return int(m.group(1)), int(m.group(2))
+    return None
+
+def check_red_flags(patient_data: Dict[str, Any]) -> List[str]:
+    """Identify immediate red flags from patient_data."""
+    flags: List[str] = []
+    hpi    = patient_data.get("hpi", {})
+    vitals = patient_data.get("vitals", {})
+    syms   = [s.lower() for s in hpi.get("symptoms", []) if isinstance(s, str)]
+    mapping = {
+        "chest pain": "Chest pain reported",
+        "shortness of breath": "Shortness of breath reported",
+        "severe headache": "Severe headache reported",
+        "syncope": "Syncope reported",
+        "hemoptysis": "Hemoptysis reported"
     }
+    for term, desc in mapping.items():
+        if term in syms:
+            flags.append(f"Red Flag: {desc}.")
+    temp = vitals.get("temp_c")
+    hr   = vitals.get("hr_bpm")
+    rr   = vitals.get("rr_rpm")
+    spo2 = vitals.get("spo2_percent")
+    bp   = parse_bp(vitals.get("bp_mmhg", ""))
+    if temp is not None and temp >= 38.5:
+        flags.append(f"Red Flag: Fever ({temp}°C).")
+    if hr is not None:
+        if hr >= 120:
+            flags.append(f"Red Flag: Tachycardia ({hr} bpm).")
+        if hr <= 50:
+            flags.append(f"Red Flag: Bradycardia ({hr} bpm).")
+    if rr is not None and rr >= 24:
+        flags.append(f"Red Flag: Tachypnea ({rr} rpm).")
+    if spo2 is not None and spo2 <= 92:
+        flags.append(f"Red Flag: Hypoxia ({spo2}%).")
+    if bp:
+        sys, dia = bp
+        if sys >= 180 or dia >= 110:
+            flags.append(f"Red Flag: Hypertensive urgency/emergency ({sys}/{dia} mmHg).")
+        if sys <= 90 or dia <= 60:
+            flags.append(f"Red Flag: Hypotension ({sys}/{dia} mmHg).")
+    return list(dict.fromkeys(flags))
 
-# ── Core Clinical Tools ──────────────────────────────────────────────
-class ClinicalToolkit:
-    @staticmethod
-    def get_essential_tools() -> List[BaseTool]:
-        """Return validated clinical tools with safety wrappers"""
-        return [
-            ClinicalToolkit.order_lab_test,
-            ClinicalToolkit.prescribe_medication,
-            ClinicalToolkit.check_drug_interactions,
-            ClinicalToolkit.flag_clinical_risk
-        ]
-
-    class LabOrderInput(BaseModel):
-        test_name: str = Field(..., pattern=r"^[A-Za-z0-9\s-]+$")
-        rationale: str = Field(..., min_length=20)
-        priority: ClinicalPriority = ClinicalPriority.ROUTINE
-
-    @tool("order_lab_test", args_schema=LabOrderInput)
-    def order_lab_test(test_name: str, rationale: str, 
-                      priority: ClinicalPriority) -> Dict[str, Any]:
-        """Standardized lab ordering with clinical validation"""
-        # Implementation details...
-        return {"status": "ordered", "details": {...}}
-
-    class PrescriptionSafetyCheck(BaseModel):
-        medication: str
-        rxcui: Optional[str]
-        contraindications: List[str]
-        # Additional safety fields...
-
-    @classmethod
-    def validate_prescription(cls, rx_data: Dict) -> PrescriptionSafetyCheck:
-        """Pharmaceutical safety validation pipeline"""
-        # Comprehensive validation logic...
-        return PrescriptionSafetyCheck(...)
-
-# ── State Management Engine ─────────────��────────────────────────────
-class ClinicalStateManager:
-    @staticmethod
-    def initialize_state(patient_data: Dict) -> ClinicalState:
-        """Create validated initial state with clinical context"""
-        return {
-            "messages": [
-                SystemMessage(content=ClinicalPrompts.SYSTEM_PROMPT),
-                HumanMessage(content="Initiate clinical consultation")
-            ],
-            "patient_data": ClinicalValidator.sanitize_patient_data(patient_data),
-            "safety_warnings": [],
-            "workflow_metadata": {
-                "iterations": 0,
-                "active_alerts": 0,
-                "safety_override": False
-            },
-            "execution_log": []
-        }
-
-    @staticmethod
-    def propagate_state(previous: ClinicalState, 
-                       updates: Dict) -> ClinicalState:
-        """State transition with clinical context preservation"""
-        preserved_fields = {
-            'patient_data': previous['patient_data'],
-            'workflow_metadata': {
-                **previous['workflow_metadata'],
-                **updates.get('workflow_metadata', {})
-            }
-        }
-        return ClinicalValidator.validate_state({
-            **preserved_fields,
-            **updates
-        })
+def format_patient_data_for_prompt(data: Dict[str, Any]) -> str:
+    """Format patient_data dict into a markdown-like prompt section."""
+    if not data:
+        return "No patient data provided."
+    lines: List[str] = []
+    for section, value in data.items():
+        title = section.replace("_", " ").title()
+        if isinstance(value, dict) and any(value.values()):
+            lines.append(f"**{title}:**")
+            for k, v in value.items():
+                if v:
+                    lines.append(f"- {k.replace('_',' ').title()}: {v}")
+        elif isinstance(value, list) and value:
+            lines.append(f"**{title}:** {', '.join(map(str, value))}")
+        elif value:
+            lines.append(f"**{title}:** {value}")
+    return "\n".join(lines)
 
-# ── Clinical Workflow Nodes ─────────────────────────────────────────
-class ClinicalWorkflowNodes:
-    @staticmethod
-    def agent_node(state: ClinicalState) -> ClinicalState:
-        """FDA-compliant clinical reasoning engine"""
-        ClinicalValidator.check_iteration_limit(state)
-        
-        try:
-            response = ClinicalLLM.invoke(state)
-            new_state = ClinicalStateManager.propagate_state(state, {
-                "messages": [response],
-                "workflow_metadata": {
-                    "iterations": state["workflow_metadata"]["iterations"] + 1
-                }
-            })
-            
-            if ClinicalTerminationCriteria.should_terminate(new_state):
-                return ClinicalWorkflowNodes.apply_termination_protocol(new_state)
-                
-            return new_state
-        except CriticalClinicalError as e:
-            return ClinicalErrorHandler.handle_critical_error(state, e)
-
-    @staticmethod
-    def tool_node(state: ClinicalState) -> ClinicalState:
-        """HIPAA-compliant tool execution with safety audit"""
-        ClinicalSafetyEngine.pre_execution_checks(state)
-        
-        tool_results = []
-        for tool_call in state["messages"][-1].tool_calls:
-            result = ClinicalToolExecutor.execute_with_audit(tool_call)
-            tool_results.append(result)
-            
-            if result['category'] == "DRUG_ORDER":
-                ClinicalSafetyEngine.post_drug_order_checks(result)
-        
-        return ClinicalStateManager.propagate_state(state, {
-            "messages": [ToolMessage(...)],
-            "safety_warnings": ClinicalSafetyEngine.aggregate_warnings(tool_results)
-        })
+# ── Tool Input Schemas ─────────────────────────────────────────────────────
+class LabOrderInput(BaseModel):
+    test_name: str = Field(...)
+    reason: str    = Field(...)
+    priority: str  = Field("Routine")
+
+class PrescriptionInput(BaseModel):
+    medication_name: str          = Field(...)
+    dosage: str                   = Field(...)
+    route: str                    = Field(...)
+    frequency: str                = Field(...)
+    duration: str                 = Field("As directed")
+    reason: str                   = Field(...)
+
+class InteractionCheckInput(BaseModel):
+    potential_prescription: str
+    current_medications: Optional[List[str]] = Field(None)
+    allergies: Optional[List[str]]             = Field(None)
 
-# ── Safety Subsystems ───────────────────────────────────────────────
-class ClinicalSafetyEngine:
-    @staticmethod
-    def enforce_prescription_rules(tool_calls: List) -> None:
-        """Hard requirements for medication orders"""
-        drug_orders = [tc for tc in tool_calls if tc.name == "prescribe_medication"]
-        interaction_checks = [tc for tc in tool_calls 
-                            if tc.name == "check_drug_interactions"]
-        
-        if ClinicalConfig.DRUG_CHECK_REQUIRED:
-            for rx in drug_orders:
-                if not any(ic.args['medication'] == rx.args['medication'] 
-                          for ic in interaction_checks):
-                    raise CriticalSafetyViolation(
-                        f"Missing interaction check for {rx.args['medication']}"
-                    )
-
-class ClinicalTerminationCriteria:
-    @staticmethod
-    def should_terminate(state: ClinicalState) -> bool:
-        """Multi-factor clinical conversation termination"""
-        metadata = state["workflow_metadata"]
-        return any([
-            metadata["iterations"] >= ClinicalConfig.MAX_ITERATIONS,
-            metadata["active_alerts"] > 3,
-            "terminate consultation" in state["messages"][-1].content.lower()
-        ])
-
-# ── Execution Framework ─────────────────────────────────────────────
-class ClinicalWorkflow:
+class FlagRiskInput(BaseModel):
+    risk_description: str = Field(...)
+    urgency: str          = Field("High")
+
+# ── Tool Implementations ───────────────────────────────────────────────────
+@tool("order_lab_test", args_schema=LabOrderInput)
+def order_lab_test(test_name: str, reason: str, priority: str = "Routine") -> str:
+    """
+    Place an order for a laboratory test.
+    """
+    logger.info(f"Ordering lab test: {test_name}, reason: {reason}, priority: {priority}")
+    return json.dumps({
+        "status": "success",
+        "message": f"Lab Ordered: {test_name} ({priority})",
+        "details": f"Reason: {reason}"
+    })
+
+@tool("prescribe_medication", args_schema=PrescriptionInput)
+def prescribe_medication(
+    medication_name: str,
+    dosage: str,
+    route: str,
+    frequency: str,
+    duration: str,
+    reason: str
+) -> str:
+    """
+    Prepare a medication prescription.
+    """
+    logger.info(f"Preparing prescription: {medication_name} {dosage}, route: {route}, freq: {frequency}")
+    return json.dumps({
+        "status": "success",
+        "message": f"Prescription Prepared: {medication_name} {dosage} {route} {frequency}",
+        "details": f"Duration: {duration}. Reason: {reason}"
+    })
+
+@tool("check_drug_interactions", args_schema=InteractionCheckInput)
+def check_drug_interactions(
+    potential_prescription: str,
+    current_medications: Optional[List[str]] = None,
+    allergies: Optional[List[str]] = None
+) -> str:
+    """
+    Check for drug–drug interactions and allergy risks.
+    """
+    logger.info(f"Checking interactions for: {potential_prescription}")
+    warnings: List[str] = []
+    pm = [m.lower().strip() for m in (current_medications or []) if m]
+    al = [a.lower().strip() for a in (allergies or []) if a]
+    if potential_prescription.lower().strip() in al:
+        warnings.append(f"CRITICAL ALLERGY: Patient allergic to '{potential_prescription}'.")
+    rxcui = get_rxcui(potential_prescription)
+    label = get_openfda_label(rxcui=rxcui, drug_name=potential_prescription)
+    if not (rxcui or label):
+        warnings.append(f"INFO: Could not identify '{potential_prescription}'. Checks may be incomplete.")
+    for section in ("contraindications", "warnings_and_cautions", "warnings"):
+        items = label.get(section) if label else None
+        if isinstance(items, list):
+            snippets = search_text_list(items, al)
+            if snippets:
+                warnings.append(f"ALLERGY RISK ({section}): {'; '.join(snippets)}")
+    for med in pm:
+        mrxcui = get_rxcui(med)
+        mlabel = get_openfda_label(rxcui=mrxcui, drug_name=med)
+        for sec in ("drug_interactions",):
+            for src_label, src_name in ((label, potential_prescription), (mlabel, med)):
+                items = src_label.get(sec) if src_label else None
+                if isinstance(items, list):
+                    snippets = search_text_list(items, [med if src_name == potential_prescription else potential_prescription])
+                    if snippets:
+                        warnings.append(f"Interaction ({src_name} label): {'; '.join(snippets)}")
+    status  = "warning" if warnings else "clear"
+    message = (
+        f"{len(warnings)} issue(s) found for '{potential_prescription}'."
+        if warnings else
+        f"No major interactions or allergy issues identified for '{potential_prescription}'."
+    )
+    return json.dumps({"status": status, "message": message, "warnings": warnings})
+
+@tool("flag_risk", args_schema=FlagRiskInput)
+def flag_risk(risk_description: str, urgency: str = "High") -> str:
+    """
+    Flag a clinical risk with given urgency.
+    """
+    logger.info(f"Flagging risk: {risk_description} (urgency={urgency})")
+    return json.dumps({
+        "status": "flagged",
+        "message": f"Risk '{risk_description}' flagged with {urgency} urgency."
+    })
+
+# Include the Tavily search tool
+search_tool = TavilySearchResults(max_results=MAX_SEARCH_RESULTS, name="tavily_search_results")
+all_tools   = [order_lab_test, prescribe_medication, check_drug_interactions, flag_risk, search_tool]
+
+# ── LLM & Tool Executor ───────────────────────────────────────────────────
+llm              = ChatGroq(temperature=AGENT_TEMPERATURE, model=AGENT_MODEL_NAME)
+model_with_tools = llm.bind_tools(all_tools)
+tool_executor    = ToolExecutor(all_tools)
+
+# ── State Definition ─────────────────────────────────────────────────────
+class AgentState(TypedDict):
+    messages: List[Any]
+    patient_data: Optional[Dict[str, Any]]
+    summary: Optional[str]
+    interaction_warnings: Optional[List[str]]
+    done: Optional[bool]
+    iterations: Optional[int]
+
+# Helper to propagate state fields between nodes
+def propagate_state(new: Dict[str, Any], old: Dict[str, Any]) -> Dict[str, Any]:
+    for key in ["iterations", "done", "patient_data", "summary", "interaction_warnings"]:
+        if key in old and key not in new:
+            new[key] = old[key]
+    return new
+
+# ── Graph Nodes ─────────────────────────────────────────────────────────
+def agent_node(state: AgentState) -> Dict[str, Any]:
+    if state.get("done", False):
+        return state
+    msgs = state.get("messages", [])
+    if not msgs or not isinstance(msgs[0], SystemMessage):
+        msgs = [SystemMessage(content=ClinicalPrompts.SYSTEM_PROMPT)] + msgs
+    logger.info(f"Invoking LLM with {len(msgs)} messages")
+    try:
+        response = model_with_tools.invoke(msgs)
+        new_state = {"messages": [response]}
+        return propagate_state(new_state, state)
+    except Exception as e:
+        logger.exception("Error in agent_node")
+        new_state = {"messages": [AIMessage(content=f"Error: {e}")]}
+        return propagate_state(new_state, state)
+
+def tool_node(state: AgentState) -> Dict[str, Any]:
+    if state.get("done", False):
+        return state
+    last = state.get("messages", [])[-1]
+    if not isinstance(last, AIMessage) or not getattr(last, "tool_calls", None):
+        logger.warning("tool_node invoked without pending tool_calls")
+        new_state = {"messages": []}
+        return propagate_state(new_state, state)
+    calls = last.tool_calls
+    blocked_ids = set()
+    for call in calls:
+        if call["name"] == "prescribe_medication":
+            med = call["args"].get("medication_name", "").lower()
+            if not any(
+                c["name"] == "check_drug_interactions" and
+                c["args"].get("potential_prescription", "").lower() == med
+                for c in calls
+            ):
+                logger.warning(f"Blocking prescribe_medication for '{med}' without interaction check")
+                blocked_ids.add(call["id"])
+    to_execute = [c for c in calls if c["id"] not in blocked_ids]
+    pd = state.get("patient_data", {})
+    for call in to_execute:
+        if call["name"] == "check_drug_interactions":
+            call["args"].setdefault("current_medications", pd.get("medications", {}).get("current", []))
+            call["args"].setdefault("allergies", pd.get("allergies", []))
+    messages: List[ToolMessage] = []
+    warnings: List[str] = []
+    try:
+        responses = tool_executor.batch(to_execute, return_exceptions=True)
+        for call, resp in zip(to_execute, responses):
+            if isinstance(resp, Exception):
+                logger.exception(f"Error executing tool {call['name']}")
+                content = json.dumps({"status": "error", "message": str(resp)})
+            else:
+                content = str(resp)
+                if call["name"] == "check_drug_interactions":
+                    data = json.loads(content)
+                    if data.get("status") == "warning":
+                        warnings.extend(data.get("warnings", []))
+            messages.append(ToolMessage(content=content, tool_call_id=call["id"], name=call["name"]))
+    except Exception as e:
+        logger.exception("Critical error in tool_node")
+        for call in to_execute:
+            messages.append(ToolMessage(
+                content=json.dumps({"status": "error", "message": str(e)}),
+                tool_call_id=call["id"],
+                name=call["name"]
+            ))
+    new_state = {"messages": messages, "interaction_warnings": warnings or None}
+    return propagate_state(new_state, state)
+
+def reflection_node(state: AgentState) -> Dict[str, Any]:
+    if state.get("done", False):
+        return state
+    warns = state.get("interaction_warnings")
+    if not warns:
+        logger.warning("reflection_node called without warnings")
+        new_state = {"messages": []}
+        return propagate_state(new_state, state)
+    triggering = None
+    for msg in reversed(state.get("messages", [])):
+        if isinstance(msg, AIMessage) and getattr(msg, "tool_calls", None):
+            triggering = msg
+            break
+    if not triggering:
+        new_state = {"messages": [AIMessage(content="Internal Error: reflection context missing.")]}
+        return propagate_state(new_state, state)
+    prompt = (
+        "You are SynapseAI, performing a focused safety review of the following plan:\n\n"
+        f"{triggering.content}\n\n"
+        "Highlight any issues based on these warnings:\n" +
+        "\n".join(f"- {w}" for w in warns)
+    )
+    try:
+        resp = llm.invoke([SystemMessage(content="Safety reflection"), HumanMessage(content=prompt)])
+        new_state = {"messages": [AIMessage(content=resp.content)]}
+        return propagate_state(new_state, state)
+    except Exception as e:
+        logger.exception("Error during reflection")
+        new_state = {"messages": [AIMessage(content=f"Error during reflection: {e}")]}
+        return propagate_state(new_state, state)
+
+# ── Routing Functions ────────────────────────────────────────────────────
+def should_continue(state: AgentState) -> str:
+    state.setdefault("iterations", 0)
+    state["iterations"] += 1
+    logger.info(f"Iteration count: {state['iterations']}")
+    # When iterations exceed threshold, force final output and terminate.
+    if state["iterations"] >= 4:
+        state.setdefault("messages", []).append(AIMessage(content="Final output: consultation complete."))
+        state["done"] = True
+        return "end_conversation_turn"
+    if not state.get("messages"):
+        state["done"] = True
+        return "end_conversation_turn"
+    last = state["messages"][-1]
+    if not isinstance(last, AIMessage):
+        state["done"] = True
+        return "end_conversation_turn"
+    if getattr(last, "tool_calls", None):
+        return "continue_tools"
+    if "consultation complete" in last.content.lower():
+        state["done"] = True
+        return "end_conversation_turn"
+    state["done"] = False
+    return "agent"
+
+def after_tools_router(state: AgentState) -> str:
+    # Instead of routing back to agent, route reflection to END to break the cycle.
+    if state.get("interaction_warnings"):
+        return "reflection"
+    return "end_conversation_turn"
+
+# ── ClinicalAgent ─────────────────────────────────────────────────────────
+class ClinicalAgent:
     def __init__(self):
-        self.workflow = self._build_workflow()
-        self.toolkit = ClinicalToolkit.get_essential_tools()
-        self.llm = ChatGroq(model_name="llama3-70b-8192", temperature=0.1)
-        
-    def _build_workflow(self) -> StateGraph:
-        """Construct ISO 13485-compliant clinical workflow"""
-        workflow = StateGraph(ClinicalState)
-        
-        workflow.add_node("clinical_reasoning", ClinicalWorkflowNodes.agent_node)
-        workflow.add_node("tool_execution", ClinicalWorkflowNodes.tool_node)
-        workflow.add_node("safety_review", ClinicalSafetyProtocols.review_node)
-        
-        workflow.set_entry_point("clinical_reasoning")
-        
-        workflow.add_conditional_edges(
-            "clinical_reasoning",
-            ClinicalDecisionRouter.route_agent_output,
-            {
-                "require_tools": "tool_execution",
-                "need_safety_review": "safety_review",
-                "final_output": END
+        logger.info("Building ClinicalAgent workflow")
+        wf = StateGraph(AgentState)
+        wf.add_node("agent", agent_node)
+        wf.add_node("tools", tool_node)
+        wf.add_node("reflection", reflection_node)
+        wf.set_entry_point("agent")
+        wf.add_conditional_edges("agent", should_continue, {
+            "continue_tools": "tools",
+            "end_conversation_turn": END
+        })
+        wf.add_conditional_edges("tools", after_tools_router, {
+            "reflection": "reflection",
+            "end_conversation_turn": END
+        })
+        # Removed the edge from reflection back to agent to break the cycle.
+        self.graph_app = wf.compile()
+        logger.info("ClinicalAgent ready")
+
+    def invoke_turn(self, state: Dict[str, Any]) -> Dict[str, Any]:
+        try:
+            # Increase recursion limit if needed.
+            result = self.graph_app.invoke(state, {"recursion_limit": 100})
+            result.setdefault("summary", state.get("summary"))
+            result.setdefault("interaction_warnings", None)
+            return result
+        except Exception as e:
+            logger.exception("Error during graph invocation")
+            return {
+                "messages": state.get("messages", []) + [AIMessage(content=f"Error: {e}")],
+                "patient_data": state.get("patient_data"),
+                "summary": state.get("summary"),
+                "interaction_warnings": None
             }
-        )
-        
-        workflow.add_edge("tool_execution", "clinical_reasoning")
-        workflow.add_edge("safety_review", "clinical_reasoning")
-        
-        return workflow.compile()
-
-    def execute_consultation(self, patient_data: Dict) -> ClinicalState:
-        """Execute full clinical workflow with safety audits"""
-        initial_state = ClinicalStateManager.initialize_state(patient_data)
-        return self.workflow.invoke(
-            initial_state,
-            config={"recursion_limit": ClinicalConfig.MAX_ITERATIONS + ClinicalConfig.RECURSION_BUFFER}
-        )
-
-# ── Usage Example ───────────────────────────────────────────────────
-if __name__ == "__main__":
-    # Initialize clinical environment
-    ClinicalValidator.validate_environment()
-    
-    # Sample patient scenario
-    complex_case = {
-        "demographics": {"age": 68, "sex": "F", "weight_kg": 82},
-        "presenting_complaint": "Chest pain radiating to left arm",
-        "medical_history": ["HTN", "Type 2 DM", "HLD"],
-        "current_meds": ["Atenolol 50mg daily", "Simvastatin 40mg HS"]
-    }
-    
-    # Execute clinical workflow
-    workflow = ClinicalWorkflow()
-    result = workflow.execute_consultation(complex_case)
-    
-    # Generate clinical summary
-    final_report = ClinicalDocumentation.generate_report(result)
-    print(json.dumps(final_report, indent=2))