Update agent.py

agent.py

@@ -23,7 +23,6 @@ logging.basicConfig(level=logging.INFO)
 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")
@@ -43,8 +42,8 @@ class ClinicalPrompts:
 def wrap_message(msg: Any) -> AIMessage:
     """
     Ensures 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 it is a dict, extracts the 'content' field (or serializes the dict).
+    Otherwise, converts the message to a string.
     """
     if isinstance(msg, AIMessage):
         return msg
@@ -358,6 +357,7 @@ def tool_node(state: AgentState) -> Dict[str, Any]:
         new_state = {"messages": []}
         return propagate_state(new_state, state)
     last = wrap_message(messages_list[-1])
+    # Safely retrieve pending tool_calls from the message's __dict__
     tool_calls = last.__dict__.get("tool_calls")
     if not (isinstance(last, AIMessage) and tool_calls):
         logger.warning("tool_node invoked without pending tool_calls")
@@ -463,6 +463,7 @@ def should_continue(state: AgentState) -> str:
         state["done"] = True
         return "end_conversation_turn"
     state["done"] = False
+    # Return "start" to loop back.
     return "start"
 
 def after_tools_router(state: AgentState) -> str:
@@ -479,15 +480,17 @@ class ClinicalAgent:
         wf.add_node("tools", tool_node)
         wf.add_node("reflection", reflection_node)
         wf.set_entry_point("start")
+        # Note: Added a "start" branch in the conditional edges.
         wf.add_conditional_edges("start", should_continue, {
             "continue_tools": "tools",
+            "start": "start",
             "end_conversation_turn": END
         })
         wf.add_conditional_edges("tools", after_tools_router, {
             "reflection": "reflection",
             "end_conversation_turn": END
         })
-        # Removed edge from reflection back to start.
+        # Removed edge from reflection back to start to break the cycle.
         self.graph_app = wf.compile()
         logger.info("ClinicalAgent ready")