Update agent.py

agent.py

@@ -2,78 +2,58 @@
 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
+
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
 from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
 
 # --- Import your defined tools FROM THE 'tools' PACKAGE ---
-from tools import ( # This now relies on tools/__init__.py
+# This relies on tools/__init__.py correctly exporting these names.
+from tools import (
     BioPortalLookupTool,
-    # GeminiTool, # Uncomment if you decide to use it
     UMLSLookupTool,
     QuantumTreatmentOptimizerTool,
-    # You might not need to import the Input schemas here if the tools use them internally
-    # and the agent only needs the tool instances.
-    # BioPortalInput, GeminiInput, UMLSInput, QuantumOptimizerInput
+    # QuantumOptimizerInput, # Only if needed for type hints directly in this file
+    # GeminiTool, # Uncomment and add to __all__ in tools/__init__.py if you decide to use it
 )
-# If QuantumOptimizerInput is needed for type hinting or direct use in agent.py:
-from tools import QuantumOptimizerInput # Or from tools.quantum_treatment_optimizer_tool import QuantumOptimizerInput
 
 from config.settings import settings
 from services.logger import app_logger
 
-# ... (rest of your agent.py file) ...
-
 # --- Initialize LLM (Gemini) ---
 try:
-    # Ensure GOOGLE_API_KEY is set in your environment (HuggingFace Secrets)
-    # or settings.GEMINI_API_KEY correctly maps to it.
     if not (settings.GEMINI_API_KEY or os.getenv("GOOGLE_API_KEY")):
-        raise ValueError("GOOGLE_API_KEY (for Gemini) not found in settings or environment.")
+        # This check is crucial. If no key, LLM init will fail.
+        app_logger.error("CRITICAL: GOOGLE_API_KEY (for Gemini) not found in settings or environment. Agent cannot initialize.")
+        raise ValueError("GOOGLE_API_KEY (for Gemini) not configured.")
 
     llm = ChatGoogleGenerativeAI(
-        model="gemini-1.5-pro-latest", # Using a more capable Gemini model if available
-        # model="gemini-pro", # Fallback if 1.5-pro is not yet available or preferred
-        temperature=0.3,
-        # google_api_key=settings.GEMINI_API_KEY, # Explicitly pass if GOOGLE_API_KEY env var isn't set
-        convert_system_message_to_human=True, # Can be helpful for some models
-        # safety_settings={ # Example safety settings
-        #     HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,
-        #     HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
-        # }
+        model="gemini-1.5-pro-latest", # Or "gemini-pro"
+        temperature=0.2, # Lower temperature for more deterministic tool use
+        # google_api_key=settings.GEMINI_API_KEY, # Explicitly pass if GOOGLE_API_KEY env var might not be picked up
+        convert_system_message_to_human=True, # Can help with models that don't strictly follow system role
+        # safety_settings={...} # Optional safety settings
     )
-    app_logger.info(f"ChatGoogleGenerativeAI ({llm.model}) initialized successfully.")
+    app_logger.info(f"ChatGoogleGenerativeAI ({llm.model}) initialized successfully for agent.")
 except Exception as e:
-    app_logger.error(f"Failed to initialize ChatGoogleGenerativeAI: {e}", exc_info=True)
+    app_logger.error(f"Failed to initialize ChatGoogleGenerativeAI for agent: {e}", exc_info=True)
+    # This error needs to be propagated so get_agent_executor fails clearly
     raise ValueError(f"Gemini LLM initialization failed: {e}. Check API key and configurations in HF Secrets.")
 
 
-# --- Initialize Tools ---
-# Ensure each tool's description is clear and guides the LLM on when and how to use it.
-# Also, ensure their args_schema is correctly defined.
-tools = [
+# --- Initialize Tools List ---
+# The tool instances are created here. Their internal logic (like API calls)
+# will be executed when the agent calls their .run() or ._run() method.
+tools_list = [
     UMLSLookupTool(),
     BioPortalLookupTool(),
     QuantumTreatmentOptimizerTool(),
-    # GeminiTool(), # Consider if this is needed. The main LLM is already Gemini.
-                   # Useful if this tool performs a very specific, different task with Gemini,
-                   # or uses a different Gemini model (e.g., for vision if main is text).
-                   # If it's just for general queries, the main agent LLM can handle it.
+    # GeminiTool(), # Add if using
 ]
-app_logger.info(f"Tools initialized: {[tool.name for tool in tools]}")
+app_logger.info(f"Agent tools initialized: {[tool.name for tool in tools_list]}")
 
 
 # --- Agent Prompt (Adapted for Structured Chat with Gemini and your tools) ---
-# This prompt guides the LLM to:
-# 1. Understand its role and capabilities.
-# 2. Know which tools are available and their purpose (from {tools}).
-# 3. Format tool invocations as a JSON blob with "action" and "action_input".
-#    - "action_input" should be a string for simple tools (UMLSInput, GeminiInput).
-#    - "action_input" should be a dictionary for tools with multiple args (BioPortalInput, QuantumOptimizerInput).
-# 4. Use the provided {patient_context}.
-# 5. Refer to {chat_history}.
-# 6. Process the new {input}.
-# 7. Use {agent_scratchpad} for its internal monologue/tool results.
-
 SYSTEM_PROMPT_TEMPLATE = (
     "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. "
@@ -83,43 +63,42 @@ SYSTEM_PROMPT_TEMPLATE = (
         "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:\n{tools}\n"
-    "   To use a tool, respond with a JSON markdown code block like this:\n"
-    "   ```json\n"
-    "   {{\n"
-    '     "action": "tool_name",\n'
-    '     "action_input": "query string for the tool" OR {{"arg1": "value1", "arg2": "value2", ...}} \n'
-    "   }}\n"
-    "   ```\n"
-    "   - For `umls_lookup` and `google_gemini_chat`, `action_input` is a single string (the 'term' or 'query').\n"
-    "   - For `bioportal_lookup`, `action_input` is a dictionary like `{{\"term\": \"search_term\", \"ontology\": \"ONTOLOGY_CODE\"}}`. If ontology is not specified by user, you can default to SNOMEDCT or ask.\n"
-    "   - For `quantum_treatment_optimizer`, `action_input` is a dictionary like `{{\"patient_data\": {{...patient details...}}, \"current_treatments\": [\"med1\"], \"conditions\": [\"cond1\"]}}`. You MUST populate 'patient_data' using the overall {patient_context} if available and relevant.\n"
+    "3.  Tool Usage: You have access to the following tools:\n{tools}\n" # {tools} is filled by the agent with tool names and descriptions
+    "   To use a tool, respond with a JSON markdown code block with the 'action' and 'action_input' keys. "
+    "   The 'action_input' should 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\"}}, \"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 answer. Cite the tool if you used one (e.g., 'According to UMLS Lookup...').\n"
     "5.  Specific Tool Guidance:\n"
-    "   - If asked about treatment optimization for a specific patient (especially if context is provided), you MUST use the `quantum_treatment_optimizer` tool.\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"
-    "   - If the query is very general, complex, or creative beyond simple lookups, you might consider using `google_gemini_chat` (if enabled) or answering directly if confident.\n"
+    # "   - If the query is very general, complex, or creative beyond simple lookups, you might consider using `google_gemini_chat` (if enabled as a tool) or answering directly if confident.\n" # If GeminiTool is used
     "6.  Conversation Flow: Refer to the `Previous conversation history` to maintain context.\n\n"
     "Begin!\n\n"
     "Previous conversation history:\n"
     "{chat_history}\n\n"
     "New human question: {input}\n"
-    "{agent_scratchpad}"
+    "{agent_scratchpad}" # Placeholder for agent's thoughts and tool outputs
 )
 
+# Create the prompt template
+# The input_variables are what agent_executor.invoke expects, plus what create_structured_chat_agent adds.
+# create_structured_chat_agent uses 'tools' and 'tool_names' internally when formatting the prompt for the LLM.
+# The primary inputs we pass to invoke are 'input', 'chat_history', and 'patient_context'.
 prompt = ChatPromptTemplate.from_messages([
     ("system", SYSTEM_PROMPT_TEMPLATE),
-    # For structured chat agent, HumanMessage/AIMessage sequence is often handled by MessagesPlaceholder("agent_scratchpad")
-    # or by how the agent formats history into the main prompt.
-    # The key is that the {chat_history} and {input} placeholders are in the system prompt.
     MessagesPlaceholder(variable_name="agent_scratchpad"),
 ])
-app_logger.info("Agent prompt template created.")
+app_logger.info("Agent prompt template created for Gemini structured chat agent.")
 
 # --- Create Agent ---
 try:
-    agent = create_structured_chat_agent(llm=llm, tools=tools, prompt=prompt)
-    app_logger.info("Structured chat agent created successfully with Gemini LLM.")
+    # create_structured_chat_agent is suitable for LLMs that can follow instructions
+    # to produce structured output (like JSON for tool calls) when prompted.
+    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)
     raise ValueError(f"Gemini agent creation failed: {e}")
@@ -128,70 +107,80 @@ except Exception as e:
 # --- Create Agent Executor ---
 agent_executor = AgentExecutor(
     agent=agent,
-    tools=tools,
-    verbose=True, # Set to True for debugging, False for production
-    handle_parsing_errors=True, # Crucial for LLM-generated JSON for tool calls
-    max_iterations=10, # Increased slightly for potentially complex tool interactions
-    # return_intermediate_steps=True, # Enable if you need to see thoughts/tool calls in the response object
-    early_stopping_method="generate", # Stop if LLM generates a stop token or a final answer
+    tools=tools_list,
+    verbose=True, # Essential for debugging tool usage
+    handle_parsing_errors=True, # Gracefully handle if LLM output for tool call isn't perfect JSON
+    max_iterations=10, # Prevents overly long or runaway chains
+    # return_intermediate_steps=True, # Set to True to get thoughts/actions in the response dict
+    early_stopping_method="generate", # Sensible default
 )
 app_logger.info("AgentExecutor with Gemini agent created successfully.")
 
 
 # --- Getter Function for Streamlit App ---
 def get_agent_executor():
-    """Returns the configured agent executor for Gemini."""
-    # Initialization happens above when the module is loaded.
-    # This function just returns the already created executor.
-    # A check for API key is good practice, though it would have failed earlier if not set.
+    """
+    Returns the configured agent executor for Gemini.
+    Initialization of LLM, tools, agent, and executor happens when this module is imported.
+    """
+    # A final check for API key availability, though LLM initialization should have caught it.
     if not (settings.GEMINI_API_KEY or os.getenv("GOOGLE_API_KEY")):
-        # This log might be redundant if LLM init failed, but good as a sanity check here.
-        app_logger.error("CRITICAL: GOOGLE_API_KEY (for Gemini) is not available at get_agent_executor call.")
+        app_logger.critical("CRITICAL: GOOGLE_API_KEY (for Gemini) is not available when get_agent_executor is called. This indicates an earlier init failure or misconfiguration.")
         raise ValueError("Google API Key for Gemini not configured. Agent cannot function.")
     return agent_executor
 
 # --- Example Usage (for local testing of this agent.py file) ---
 if __name__ == "__main__":
     if not (settings.GEMINI_API_KEY or os.getenv("GOOGLE_API_KEY")):
-        print("Please set your GOOGLE_API_KEY in .env file or as an environment variable.")
+        print("🚨 Please set your GOOGLE_API_KEY in .env file or as an environment variable to run the test.")
     else:
-        print("\nQuantum Health Navigator (Gemini Agent Test Console)")
+        print("\n🚀 Quantum Health Navigator (Gemini Agent Test Console) 🚀")
+        print("-----------------------------------------------------------")
         print("Type 'exit' or 'quit' to stop.")
-        print("Example queries:")
-        print("  - What is hypertension?")
-        print("  - Lookup 'myocardial infarction' in UMLS.")
-        print("  - Search for 'diabetes mellitus type 2' in BioPortal using SNOMEDCT ontology.")
-        print("  - Optimize treatment for a patient (context will be simulated).")
-        print("-" * 30)
+        print("Example topics: medical definitions, treatment optimization (will use simulated patient context).")
+        print("-" * 59)
 
-        executor = get_agent_executor()
-        current_chat_history_for_test = [] # List of HumanMessage, AIMessage
+        test_executor = get_agent_executor() # Get the globally defined executor
+        current_chat_history_for_test_run = [] # List of HumanMessage, AIMessage
 
         # Simulated patient context for testing the {patient_context} variable
-        test_patient_context_summary = (
-            "Age: 45; Gender: Male; Chief Complaint: Intermittent chest pain; "
-            "Key Medical History: Hyperlipidemia; Current Medications: Atorvastatin 20mg."
+        test_patient_context_summary_str = (
+            "Age: 62; Gender: Female; Chief Complaint: Fatigue and increased thirst; "
+            "Key Medical History: Obesity, family history of diabetes; "
+            "Current Medications: None reported; Allergies: Sulfa drugs."
         )
+        print(f"ℹ️  Simulated Patient Context for this test run: {test_patient_context_summary_str}\n")
+
 
         while True:
-            user_input_str = input("\n👤 You: ")
+            user_input_str = input("👤 You: ")
             if user_input_str.lower() in ["exit", "quit"]:
-                print("Exiting test console.")
+                print("👋 Exiting test console.")
                 break
 
+            if not user_input_str.strip():
+                continue
+
             try:
                 app_logger.info(f"__main__ test: Invoking agent with input: '{user_input_str}'")
-                response = executor.invoke({
+                # These are the keys expected by the prompt template
+                # and processed by create_structured_chat_agent
+                response_dict = test_executor.invoke({
                     "input": user_input_str,
-                    "chat_history": current_chat_history_for_test,
-                    "patient_context": test_patient_context_summary # Passing the context
+                    "chat_history": current_chat_history_for_test_run,
+                    "patient_context": test_patient_context_summary_str
                 })
                 
-                ai_output_str = response.get('output', "Agent did not produce an output.")
+                ai_output_str = response_dict.get('output', "Agent did not produce an 'output' key.")
                 print(f"🤖 Agent: {ai_output_str}")
                 
-                current_chat_history_for_test.append(HumanMessage(content=user_input_str))
-                current_chat_history_for_test.append(AIMessage(content=ai_output_str))
+                # Update history for the next turn
+                current_chat_history_for_test_run.append(HumanMessage(content=user_input_str))
+                current_chat_history_for_test_run.append(AIMessage(content=ai_output_str))
+                
+                # Optional: Limit history length
+                if len(current_chat_history_for_test_run) > 10: # Keep last 5 pairs
+                    current_chat_history_for_test_run = current_chat_history_for_test_run[-10:]
 
             except Exception as e:
                 print(f"⚠️ Error during agent invocation: {e}")