Update agent.py

agent.py

@@ -1,187 +1,168 @@
 # /home/user/app/agent.py
 import os
-from langchain_google_genai import ChatGoogleGenerativeAI
-from langchain.agents import AgentExecutor, create_structured_chat_agent
-# from langchain_google_genai import HarmBlockThreshold, HarmCategory # Optional for safety settings
+import sys
+from typing import List
 
-from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
-from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
+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_google_genai import ChatGoogleGenerativeAI
 
-# --- Import your defined tools FROM THE 'tools' PACKAGE ---
+from config.settings import settings
+from services.logger import app_logger
 from tools import (
     BioPortalLookupTool,
     UMLSLookupTool,
     QuantumTreatmentOptimizerTool,
-    # QuantumOptimizerInput, # Import if needed for type hints directly in this file
-    # GeminiTool, # Uncomment if using
-)
-
-from config.settings import settings
-from services.logger import app_logger
-
-# --- Initialize LLM (Gemini) ---
-llm = None
-try:
-    gemini_api_key_from_settings = settings.GEMINI_API_KEY
-    api_key_to_use = gemini_api_key_from_settings or os.getenv("GOOGLE_API_KEY")
-
-    if not api_key_to_use:
-        app_logger.error("CRITICAL: Gemini API Key not found in settings or environment.")
-        raise ValueError("Gemini API Key not configured.")
-
-    llm = ChatGoogleGenerativeAI(
-        model="gemini-1.5-pro-latest",
-        temperature=0.2,
-        google_api_key=api_key_to_use,
-        convert_system_message_to_human=True,
-    )
-    app_logger.info(f"ChatGoogleGenerativeAI ({llm.model}) initialized successfully for agent.")
-except Exception as e:
-    detailed_error_message = str(e)
-    user_facing_error = f"Gemini LLM initialization failed: {detailed_error_message}."
-    if "credential" in detailed_error_message.lower() or "api_key" in detailed_error_message.lower():
-        user_facing_error = "Gemini LLM initialization failed: API key/credential issue. Check HF Secrets."
-        app_logger.error(user_facing_error + f" Original: {detailed_error_message}", exc_info=False)
-    else:
-        app_logger.error(user_facing_error, exc_info=True)
-    raise ValueError(user_facing_error)
-
-
-# --- Initialize Tools List ---
-tools_list = [
-    UMLSLookupTool(),
-    BioPortalLookupTool(),
-    QuantumTreatmentOptimizerTool(),
-]
-app_logger.info(f"Agent tools initialized: {[tool.name for tool in tools_list]}")
-
-
-# --- Agent Prompt (Revised Structure) ---
-# System prompt contains general instructions and placeholders for context/tools.
-# Chat history, current input, and agent scratchpad are handled as separate message sequence parts.
-REVISED_SYSTEM_PROMPT_TEXT = (
-    "You are 'Quantum Health Navigator', an advanced AI assistant for healthcare professionals. "
-    "Your primary goal is to provide accurate information and insights based on user queries and available tools. "
-    "You must adhere to the following guidelines:\n"
-    "1.  Disclaimers: Always remind the user that you are an AI, not a human medical professional, and your information "
-        "is for support, not a substitute for clinical judgment. Do not provide direct medical advice for specific patient cases "
-        "unless it's the direct output of a specialized tool like 'quantum_treatment_optimizer'.\n"
-    "2.  Patient Context: The user may provide patient context at the start of the session. This context is available as: {patient_context}. "
-        "You MUST consider this context when it's relevant to the query, especially for the 'quantum_treatment_optimizer' tool.\n"
-    "3.  Tool Usage: You have access to the following tools (names: {tool_names}):\n{tools}\n"
-    "   To use a tool, respond *only* with a JSON markdown code block with 'action' and 'action_input' keys. "
-    "   The 'action_input' must match the schema for the specified tool. Examples:\n"
-    "   For `umls_lookup`: ```json\n{{\"action\": \"umls_lookup\", \"action_input\": \"myocardial infarction\"}}\n```\n"
-    "   For `bioportal_lookup`: ```json\n{{\"action\": \"bioportal_lookup\", \"action_input\": {{\"term\": \"diabetes mellitus\", \"ontology\": \"SNOMEDCT\"}}}}\n```\n"
-    "   For `quantum_treatment_optimizer`: ```json\n{{\"action\": \"quantum_treatment_optimizer\", \"action_input\": {{\"patient_data\": {{\"age\": 55, \"gender\": \"Male\", \"symptoms\": [\"chest pain\"]}}, \"current_treatments\": [\"metformin\"], \"conditions\": [\"Type 2 Diabetes\"]}}}}\n```\n"
-    "   Ensure the `action_input` for `quantum_treatment_optimizer` includes a `patient_data` dictionary populated from the overall {patient_context}.\n"
-    "4.  Responding to User: After using a tool, you will receive an observation. Use this observation and your knowledge to formulate a comprehensive final answer to the human. Cite the tool if you used one (e.g., 'According to UMLS Lookup...'). Do not output a tool call again unless necessary for a multi-step process.\n"
-    "5.  Specific Tool Guidance:\n"
-    "   - If asked about treatment optimization for a specific patient (especially if patient context is provided), you MUST use the `quantum_treatment_optimizer` tool.\n"
-    "   - For definitions, codes, or general medical concepts, `umls_lookup` or `bioportal_lookup` are appropriate.\n"
-    "6.  Conversation Flow: Maintain context from the chat history.\n\n"
-    "Begin!"
-    # Note: {chat_history}, {input}, and {agent_scratchpad} are NOT in this string anymore.
-    # They are handled by the ChatPromptTemplate.from_messages structure.
-)
-
-# Create the prompt template
-# Input variables for this prompt will be:
-# - patient_context (from invoke call)
-# - tool_names (provided by create_structured_chat_agent)
-# - tools (provided by create_structured_chat_agent)
-# - chat_history (from invoke call, via MessagesPlaceholder)
-# - input (from invoke call, via ("human", "{input}"))
-# - agent_scratchpad (managed by agent, via MessagesPlaceholder)
-prompt = ChatPromptTemplate.from_messages([
-    ("system", REVISED_SYSTEM_PROMPT_TEXT),
-    MessagesPlaceholder(variable_name="chat_history"),
-    ("human", "{input}"), # The current human input
-    MessagesPlaceholder(variable_name="agent_scratchpad") # For agent's intermediate work (must be list of BaseMessages)
-])
-app_logger.info("Agent prompt template created for Gemini structured chat agent.")
-
-# --- Create Agent ---
-if llm is None: # Should have been caught by now, but defensive check
-    app_logger.critical("LLM object is None at agent creation stage. Application cannot proceed.")
-    raise SystemExit("Agent LLM failed to initialize. Application cannot start.")
-
-try:
-    agent = create_structured_chat_agent(llm=llm, tools=tools_list, prompt=prompt)
-    app_logger.info("Structured chat agent created successfully with Gemini LLM and tools.")
-except Exception as e:
-    app_logger.error(f"Failed to create structured chat agent: {e}", exc_info=True)
-    # This is where "Prompt missing required variables: {'tool_names'}" was caught previously.
-    # Or "variable agent_scratchpad should be a list of base messages" if structure is wrong.
-    raise ValueError(f"Gemini agent creation failed: {e}")
-
-
-# --- Create Agent Executor ---
-agent_executor = AgentExecutor(
-    agent=agent,
-    tools=tools_list,
-    verbose=True,
-    handle_parsing_errors=True, # Important for structured output parsing
-    max_iterations=10,
-    early_stopping_method="generate",
 )
-app_logger.info("AgentExecutor with Gemini agent created successfully.")
 
+# -----------------------------------------------------------------------------
+# 1. Initialize Gemini LLM
+# -----------------------------------------------------------------------------
+def _init_llm() -> ChatGoogleGenerativeAI:
+    """
+    Initialize the Google Gemini LLM client 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)
+
+    try:
+        llm = ChatGoogleGenerativeAI(
+            model="gemini-1.5-pro-latest",
+            temperature=0.2,
+            google_api_key=api_key,
+            convert_system_message_to_human=True,
+        )
+        app_logger.info(f"Initialized Gemini LLM ({llm.model}) successfully.")
+        return llm
+    except Exception as e:
+        msg = f"Failed to initialize Gemini LLM: {e}"
+        app_logger.error(msg, exc_info=True)
+        raise ValueError(msg)
+
+# -----------------------------------------------------------------------------
+# 2. Build Prompt Template
+# -----------------------------------------------------------------------------
+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.
+    """
+    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}
+{tools}
+
+To invoke a tool, respond only with a JSON code block containing:
+{"action": "<tool_name>", "action_input": <input>}
+
+After the tool observation, craft a user-facing answer, citing the tool when 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]
+    ])
+
+# -----------------------------------------------------------------------------
+# 3. Assemble Agent and Executor
+# -----------------------------------------------------------------------------
+def get_agent_executor() -> AgentExecutor:
+    """
+    Lazily initialize and return the singleton AgentExecutor.
+    """
+    global _agent_executor_instance
+    if '_agent_executor_instance' not in globals():
+        # Initialize LLM
+        llm = _init_llm()
+
+        # Prepare tools
+        tools_list = [
+            UMLSLookupTool(),
+            BioPortalLookupTool(),
+            QuantumTreatmentOptimizerTool(),
+        ]
+        app_logger.info(f"Loaded tools: {[t.name for t in tools_list]}")
+
+        # 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.")
 
-# --- Getter Function for Streamlit App ---
-_agent_executor_instance = agent_executor
+        # Create structured agent
+        agent = create_structured_chat_agent(
+            llm=llm,
+            tools=tools_list,
+            prompt=prompt
+        )
+        app_logger.info("Structured chat agent created.")
+
+        # Create executor
+        executor = AgentExecutor(
+            agent=agent,
+            tools=tools_list,
+            verbose=True,
+            handle_parsing_errors=True,
+            max_iterations=10,
+            early_stopping_method="generate",
+        )
+        app_logger.info("AgentExecutor initialized.")
+        _agent_executor_instance = executor
 
-def get_agent_executor():
-    global _agent_executor_instance
-    if _agent_executor_instance is None:
-        app_logger.critical("CRITICAL: Agent executor is None when get_agent_executor is called.")
-        raise RuntimeError("Agent executor was not properly initialized. Check application startup logs.")
     return _agent_executor_instance
 
-# --- Example Usage (for local testing) ---
+# -----------------------------------------------------------------------------
+# 4. If run as script, provide a simple REPL for testing
+# -----------------------------------------------------------------------------
 if __name__ == "__main__":
-    main_test_api_key = settings.GEMINI_API_KEY or os.getenv("GOOGLE_API_KEY")
-    if not main_test_api_key:
-        print("🚨 Please set your GOOGLE_API_KEY (for Gemini) in .env or environment to run test.")
-    else:
-        print("\n🚀 Quantum Health Navigator (Gemini Agent Test Console) 🚀")
-        # ... (rest of the __main__ block from the previous full agent.py, it should work with this prompt structure) ...
+    try:
+        executor = get_agent_executor()
+    except Exception as e:
+        print(f"❌ Initialization error: {e}")
+        sys.exit(1)
+
+    # Example patient context
+    patient_context = (
+        "Age: 58; Gender: Female; Chief Complaint: Blurry vision & fatigue; "
+        "Key History: Prediabetes, mild dyslipidemia; Medications: None."
+    )
+    chat_history: List[SystemMessage | HumanMessage | AIMessage] = []
+
+    print("🚀 Quantum Health Navigator Console (type 'exit' to quit)")
+    while True:
+        user_input = input("👤 You: ").strip()
+        if user_input.lower() in {"exit", "quit"}:
+            print("👋 Goodbye!")
+            break
+        if not user_input:
+            continue
+
+        # Invoke the agent
         try:
-            test_executor = get_agent_executor()
-        except ValueError as e_init:
-            print(f"⚠️ Agent initialization failed: {e_init}")
-            exit()
-
-        current_chat_history_for_test_run = []
-        test_patient_context_summary_str = (
-            "Age: 58; Gender: Female; Chief Complaint: Recent onset of blurry vision and fatigue; "
-            "Key Medical History: Prediabetes, Mild dyslipidemia; "
-            "Current Medications: None reported; Allergies: Sulfa drugs."
-        )
-        print(f"ℹ️  Simulated Patient Context: {test_patient_context_summary_str}\n")
-
-        while True:
-            user_input_str = input("👤 You: ").strip()
-            if user_input_str.lower() in ["exit", "quit"]:
-                print("👋 Exiting.")
-                break
-            if not user_input_str:
-                continue
-            try:
-                app_logger.info(f"__main__ test: Invoking with: '{user_input_str}'")
-                # The keys here must match the input variables expected by the combined prompt and agent
-                response_dict = test_executor.invoke({
-                    "input": user_input_str,
-                    "chat_history": current_chat_history_for_test_run, # List of BaseMessage
-                    "patient_context": test_patient_context_summary_str,
-                    # `tools`, `tool_names`, `agent_scratchpad` are handled internally by the agent executor
-                })
-                ai_output_str = response_dict.get('output', "No 'output' key in response.")
-                print(f"🤖 Agent: {ai_output_str}")
-                current_chat_history_for_test_run.append(HumanMessage(content=user_input_str))
-                current_chat_history_for_test_run.append(AIMessage(content=ai_output_str))
-                if len(current_chat_history_for_test_run) > 10: # Limit history
-                    current_chat_history_for_test_run = current_chat_history_for_test_run[-10:]
-            except Exception as e:
-                print(f"⚠️ Error: {e}")
-                app_logger.error(f"Error in __main__ test invocation: {e}", exc_info=True)
+            result = executor.invoke({
+                "input": user_input,
+                "chat_history": chat_history,
+                "patient_context": patient_context
+            })
+            reply = result.get("output", "")
+            print(f"🤖 Agent: {reply}\n")
+            # Update local history
+            chat_history.append(HumanMessage(content=user_input))
+            chat_history.append(AIMessage(content=reply))
+            # Keep history manageable
+            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)