Update agent.py

agent.py

@@ -15,11 +15,11 @@ from langchain_core.tools import tool
 from langgraph.prebuilt import ToolExecutor
 from langgraph.graph import StateGraph, END
 
-# ── Logging Configuration ─────────────────────────────────────────────────────
+# ── Logging Configuration ──────────────────────────────────────────────
 logger = logging.getLogger(__name__)
 logging.basicConfig(level=logging.INFO)
 
-# ── Environment Variables ─────────────────────────────────────────────────────
+# ── 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")
@@ -28,7 +28,7 @@ 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 Configuration ──────────────────────────────────────────────
 AGENT_MODEL_NAME   = "llama3-70b-8192"
 AGENT_TEMPERATURE  = 0.1
 MAX_SEARCH_RESULTS = 3
@@ -39,7 +39,7 @@ class ClinicalPrompts:
     [SYSTEM PROMPT CONTENT HERE]
     """
 
-# ── Helper Functions ──────────────────────────────────────────────────────────
+# ── 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"
@@ -174,7 +174,7 @@ def format_patient_data_for_prompt(data: Dict[str, Any]) -> str:
             lines.append(f"**{title}:** {value}")
     return "\n".join(lines)
 
-# ── Tool Input Schemas ────────────────────────────────────────────────────────
+# ── Tool Input Schemas ─────────────────────────────────────────────────────
 class LabOrderInput(BaseModel):
     test_name: str = Field(...)
     reason: str    = Field(...)
@@ -197,7 +197,7 @@ class FlagRiskInput(BaseModel):
     risk_description: str = Field(...)
     urgency: str          = Field("High")
 
-# ── Tool Implementations ──────────────────────────────────────────────────────
+# ── Tool Implementations ───────────────────────────────────────────────────
 @tool("order_lab_test", args_schema=LabOrderInput)
 def order_lab_test(test_name: str, reason: str, priority: str = "Routine") -> str:
     """
@@ -287,12 +287,12 @@ def flag_risk(risk_description: str, urgency: str = "High") -> str:
 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 & 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 ──────────────────────────────────────────────────────────
+# ── State Definition ─────────────────────────────────────────────────────
 class AgentState(TypedDict):
     messages: List[Any]
     patient_data: Optional[Dict[str, Any]]
@@ -301,19 +301,18 @@ class AgentState(TypedDict):
     done: Optional[bool]
     iterations: Optional[int]
 
-# Helper to propagate state fields
+# 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 ───────────────────────────────────────────────────────────────
+# ── Graph Nodes ─────────────────────────────────────────────────────────
 def agent_node(state: AgentState) -> Dict[str, Any]:
-    # Check for termination
     if state.get("done", False):
         return state
-    msgs = state["messages"]
+    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")
@@ -329,7 +328,7 @@ def agent_node(state: AgentState) -> Dict[str, Any]:
 def tool_node(state: AgentState) -> Dict[str, Any]:
     if state.get("done", False):
         return state
-    last = state["messages"][-1]
+    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": []}
@@ -387,7 +386,7 @@ def reflection_node(state: AgentState) -> Dict[str, Any]:
         new_state = {"messages": []}
         return propagate_state(new_state, state)
     triggering = None
-    for msg in reversed(state["messages"]):
+    for msg in reversed(state.get("messages", [])):
         if isinstance(msg, AIMessage) and getattr(msg, "tool_calls", None):
             triggering = msg
             break
@@ -409,27 +408,35 @@ def reflection_node(state: AgentState) -> Dict[str, Any]:
         new_state = {"messages": [AIMessage(content=f"Error during reflection: {e}")]}
         return propagate_state(new_state, state)
 
-# ── Routing Functions ────────────────────────────────────────────────────────
+# ── 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 a final message and mark done.
     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 state["messages"] else None
+    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"
-    state["done"] = True
-    return "end_conversation_turn"
+    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:
     return "reflection" if state.get("interaction_warnings") else "agent"
 
-# ── ClinicalAgent ────────────────────────────────────────────────────────────
+# ── ClinicalAgent ─────────────────────────────────────────────────────────
 class ClinicalAgent:
     def __init__(self):
         logger.info("Building ClinicalAgent workflow")
@@ -452,7 +459,8 @@ class ClinicalAgent:
 
     def invoke_turn(self, state: Dict[str, Any]) -> Dict[str, Any]:
         try:
-            result = self.graph_app.invoke(state, {"recursion_limit": 15})
+            # Increase the recursion_limit as a temporary workaround if needed.
+            result = self.graph_app.invoke(state, {"recursion_limit": 100})
             result.setdefault("summary", state.get("summary"))
             result.setdefault("interaction_warnings", None)
             return result