Update agent.py

agent.py

@@ -1,35 +1,30 @@
-# /home/user/app/agent.py
 import os
 import sys
-from typing import List
+from typing import List, Union
 
-from langchain.agents import AgentExecutor, create_structured_chat_agent
 from langchain.prompts import ChatPromptTemplate
 from langchain.prompts.chat import MessagesPlaceholder
-from langchain.messages import AIMessage, HumanMessage, SystemMessage
+from langchain.schema import BaseMessage, AIMessage, HumanMessage, SystemMessage
+from langchain.agents import AgentExecutor, create_structured_chat_agent
 from langchain_google_genai import ChatGoogleGenerativeAI
 
 from config.settings import settings
 from services.logger import app_logger
-from tools import (
-    BioPortalLookupTool,
-    UMLSLookupTool,
-    QuantumTreatmentOptimizerTool,
-)
+from tools import BioPortalLookupTool, UMLSLookupTool, QuantumTreatmentOptimizerTool
 
 # -----------------------------------------------------------------------------
-# 1. Initialize Gemini LLM
+# 1. Initialize the Gemini LLM
 # -----------------------------------------------------------------------------
 def _init_llm() -> ChatGoogleGenerativeAI:
     """
-    Initialize the Google Gemini LLM client with the configured API key.
+    Initialize the Google Gemini LLM with the configured API key.
     Raises ValueError if no key is found or initialization fails.
     """
     api_key = settings.GEMINI_API_KEY or os.getenv("GOOGLE_API_KEY")
     if not api_key:
-        msg = "Gemini API key missing: set GEMINI_API_KEY in settings or GOOGLE_API_KEY env var"
-        app_logger.error(msg)
-        raise ValueError(msg)
+        err = "Gemini API key not found: set GEMINI_API_KEY in settings or GOOGLE_API_KEY in env"
+        app_logger.error(err)
+        raise ValueError(err)
 
     try:
         llm = ChatGoogleGenerativeAI(
@@ -38,54 +33,59 @@ def _init_llm() -> ChatGoogleGenerativeAI:
             google_api_key=api_key,
             convert_system_message_to_human=True,
         )
-        app_logger.info(f"Initialized Gemini LLM ({llm.model}) successfully.")
+        app_logger.info(f"Gemini LLM initialized ({llm.model})")
         return llm
     except Exception as e:
-        msg = f"Failed to initialize Gemini LLM: {e}"
-        app_logger.error(msg, exc_info=True)
-        raise ValueError(msg)
+        err = f"Failed to initialize Gemini LLM: {e}"
+        app_logger.error(err, exc_info=True)
+        raise ValueError(err)
 
 # -----------------------------------------------------------------------------
-# 2. Build Prompt Template
+# 2. Build the structured chat prompt
 # -----------------------------------------------------------------------------
 def _build_prompt_template(tool_names: List[str], tools) -> ChatPromptTemplate:
     """
-    Build a structured chat prompt template including system instructions,
-    chat history, user input, and an agent scratchpad for intermediate tool calls.
+    Construct a ChatPromptTemplate that includes:
+    - a system instruction block,
+    - a placeholder for chat_history (List[BaseMessage]),
+    - the current human input,
+    - a placeholder for agent_scratchpad (List[BaseMessage]) to manage tool calls.
     """
     system_text = """
 You are Quantum Health Navigator, an AI assistant for healthcare professionals.
-• Always disclose you are an AI and not a substitute for clinical judgment.
-• Patient context: {patient_context}
-• Tools available: {tool_names}
+
+• Disclaim: you are an AI, not a substitute for clinical judgment.
+• Patient context is provided as: {patient_context}
+• Available tools: {tool_names}
 {tools}
 
-To invoke a tool, respond only with a JSON code block containing:
+To call a tool, reply *only* with a JSON block:
 {"action": "<tool_name>", "action_input": <input>}
 
-After the tool observation, craft a user-facing answer, citing the tool when used.
+After receiving the tool’s output, craft a full answer for the user, citing any tools used.
 """.strip()
 
     return ChatPromptTemplate.from_messages([
         ("system", system_text),
-        MessagesPlaceholder(variable_name="chat_history"),      # List[BaseMessage]
-        ("human", "{input}"),                                   # Current user input
-        MessagesPlaceholder(variable_name="agent_scratchpad"),   # List[BaseMessage]
+        MessagesPlaceholder(variable_name="chat_history"),
+        ("human", "{input}"),
+        MessagesPlaceholder(variable_name="agent_scratchpad"),
     ])
 
 # -----------------------------------------------------------------------------
-# 3. Assemble Agent and Executor
+# 3. Lazily build and return the AgentExecutor singleton
 # -----------------------------------------------------------------------------
 def get_agent_executor() -> AgentExecutor:
     """
-    Lazily initialize and return the singleton AgentExecutor.
+    Returns a singleton AgentExecutor, creating it on first call.
+    Sets up LLM, tools, prompt, and executor params.
     """
     global _agent_executor_instance
-    if '_agent_executor_instance' not in globals():
-        # Initialize LLM
+    if "_agent_executor_instance" not in globals():
+        # 3.1 Initialize LLM
         llm = _init_llm()
 
-        # Prepare tools
+        # 3.2 Prepare tools
         tools_list = [
             UMLSLookupTool(),
             BioPortalLookupTool(),
@@ -93,22 +93,22 @@ def get_agent_executor() -> AgentExecutor:
         ]
         app_logger.info(f"Loaded tools: {[t.name for t in tools_list]}")
 
-        # Build prompt
+        # 3.3 Build prompt
         prompt = _build_prompt_template(
             tool_names=[t.name for t in tools_list],
             tools=tools_list
         )
-        app_logger.info("Prompt template built successfully.")
+        app_logger.info("Prompt template built")
 
-        # Create structured agent
+        # 3.4 Create the structured agent
         agent = create_structured_chat_agent(
             llm=llm,
             tools=tools_list,
             prompt=prompt
         )
-        app_logger.info("Structured chat agent created.")
+        app_logger.info("Structured chat agent created")
 
-        # Create executor
+        # 3.5 Create the executor
         executor = AgentExecutor(
             agent=agent,
             tools=tools_list,
@@ -117,27 +117,27 @@ def get_agent_executor() -> AgentExecutor:
             max_iterations=10,
             early_stopping_method="generate",
         )
-        app_logger.info("AgentExecutor initialized.")
+        app_logger.info("AgentExecutor initialized")
         _agent_executor_instance = executor
 
     return _agent_executor_instance
 
 # -----------------------------------------------------------------------------
-# 4. If run as script, provide a simple REPL for testing
+# 4. Optional REPL for local testing
 # -----------------------------------------------------------------------------
 if __name__ == "__main__":
     try:
         executor = get_agent_executor()
     except Exception as e:
-        print(f"❌ Initialization error: {e}")
+        print(f"❌ Initialization failed: {e}")
         sys.exit(1)
 
-    # Example patient context
+    # Sample patient context for testing
     patient_context = (
         "Age: 58; Gender: Female; Chief Complaint: Blurry vision & fatigue; "
-        "Key History: Prediabetes, mild dyslipidemia; Medications: None."
+        "History: Prediabetes, mild dyslipidemia; Medications: None."
     )
-    chat_history: List[SystemMessage | HumanMessage | AIMessage] = []
+    chat_history: List[Union[SystemMessage, HumanMessage, AIMessage]] = []
 
     print("🚀 Quantum Health Navigator Console (type 'exit' to quit)")
     while True:
@@ -148,7 +148,6 @@ if __name__ == "__main__":
         if not user_input:
             continue
 
-        # Invoke the agent
         try:
             result = executor.invoke({
                 "input": user_input,
@@ -157,12 +156,13 @@ if __name__ == "__main__":
             })
             reply = result.get("output", "")
             print(f"🤖 Agent: {reply}\n")
-            # Update local history
+
+            # Update history
             chat_history.append(HumanMessage(content=user_input))
             chat_history.append(AIMessage(content=reply))
-            # Keep history manageable
+            # Trim to last 20 messages
             if len(chat_history) > 20:
                 chat_history = chat_history[-20:]
         except Exception as err:
-            print(f"⚠️ Error during inference: {err}")
-            app_logger.error("Inference error", exc_info=True)
+            print(f"⚠️ Inference error: {err}")
+            app_logger.error("Runtime error in REPL", exc_info=True)