Update agent.py

agent.py

@@ -39,6 +39,20 @@ class ClinicalPrompts:
     [SYSTEM PROMPT CONTENT HERE]
     """
 
+# ── Helper: Message Wrapper ─────────────────────────────────────────────
+def wrap_message(msg: Any) -> AIMessage:
+    """
+    Ensures that the given message is an AIMessage.
+    If it is a dict, it extracts the 'content' field (or serializes the dict).
+    Otherwise, it converts the message to a string.
+    """
+    if isinstance(msg, AIMessage):
+        return msg
+    elif isinstance(msg, dict):
+        return AIMessage(content=msg.get("content", json.dumps(msg)))
+    else:
+        return AIMessage(content=str(msg))
+
 # ── Helper Functions ─────────────────────────────────────────────────────
 UMLS_AUTH_ENDPOINT  = "https://utslogin.nlm.nih.gov/cas/v1/api-key"
 RXNORM_API_BASE     = "https://rxnav.nlm.nih.gov/REST"
@@ -318,22 +332,30 @@ def agent_node(state: AgentState) -> Dict[str, Any]:
     logger.info(f"Invoking LLM with {len(msgs)} messages")
     try:
         response = model_with_tools.invoke(msgs)
+        response = wrap_message(response)
         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}")]}
+        new_state = {"messages": [wrap_message(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):
+    messages_list = state.get("messages", [])
+    if not messages_list:
+        logger.warning("tool_node invoked with no messages")
+        new_state = {"messages": []}
+        return propagate_state(new_state, state)
+    last = wrap_message(messages_list[-1])
+    # Check for pending tool_calls using dict.get if necessary.
+    tool_calls = last.tool_calls if hasattr(last, "tool_calls") else last.__dict__.get("tool_calls")
+    if not (isinstance(last, AIMessage) and tool_calls):
         logger.warning("tool_node invoked without pending tool_calls")
         new_state = {"messages": []}
         return propagate_state(new_state, state)
-    calls = last.tool_calls
+    calls = tool_calls
     blocked_ids = set()
     for call in calls:
         if call["name"] == "prescribe_medication":
@@ -387,8 +409,9 @@ def reflection_node(state: AgentState) -> Dict[str, Any]:
         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
+        wrapped = wrap_message(msg)
+        if isinstance(wrapped, AIMessage) and wrapped.__dict__.get("tool_calls"):
+            triggering = wrapped
             break
     if not triggering:
         new_state = {"messages": [AIMessage(content="Internal Error: reflection context missing.")]}
@@ -401,7 +424,7 @@ def reflection_node(state: AgentState) -> Dict[str, Any]:
     )
     try:
         resp = llm.invoke([SystemMessage(content="Safety reflection"), HumanMessage(content=prompt)])
-        new_state = {"messages": [AIMessage(content=resp.content)]}
+        new_state = {"messages": [wrap_message(resp)]}
         return propagate_state(new_state, state)
     except Exception as e:
         logger.exception("Error during reflection")
@@ -413,7 +436,6 @@ 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
@@ -421,11 +443,11 @@ def should_continue(state: AgentState) -> str:
     if not state.get("messages"):
         state["done"] = True
         return "end_conversation_turn"
-    last = state["messages"][-1]
+    last = wrap_message(state["messages"][-1])
     if not isinstance(last, AIMessage):
         state["done"] = True
         return "end_conversation_turn"
-    if getattr(last, "tool_calls", None):
+    if last.__dict__.get("tool_calls"):
         return "continue_tools"
     if "consultation complete" in last.content.lower():
         state["done"] = True
@@ -434,7 +456,6 @@ def should_continue(state: AgentState) -> str:
     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"
@@ -456,13 +477,12 @@ class ClinicalAgent:
             "reflection": "reflection",
             "end_conversation_turn": END
         })
-        # Removed the edge from reflection back to agent to break the cycle.
+        # Removed edge from reflection to agent to break 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)