Update agent.py

agent.py

@@ -2,7 +2,7 @@
 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_google_genai import HarmBlockThreshold, HarmCategory # Optional for safety settings
 
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
 from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
@@ -13,47 +13,80 @@ from tools import (
     BioPortalLookupTool,
     UMLSLookupTool,
     QuantumTreatmentOptimizerTool,
-    # 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
+    # QuantumOptimizerInput, # Only if needed for type hints directly in this file for some reason
+    # GeminiTool, # Assuming not used for now as main LLM is Gemini
 )
 
-from config.settings import settings
+from config.settings import settings # This loads your HF secrets into the settings object
 from services.logger import app_logger
 
 # --- Initialize LLM (Gemini) ---
+# This block is critical for ensuring the API key is used.
+llm = None # Initialize to None in case of failure
 try:
-    if not (settings.GEMINI_API_KEY or os.getenv("GOOGLE_API_KEY")):
-        # 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.")
+    # Prioritize the API key from settings (loaded from HF Secrets)
+    # settings.GEMINI_API_KEY should be populated by Pydantic BaseSettings from the HF Secret
+    gemini_api_key_from_settings = settings.GEMINI_API_KEY
+    
+    # Fallback to environment variable GOOGLE_API_KEY if settings.GEMINI_API_KEY is not found/set
+    # (though ideally, settings.GEMINI_API_KEY should be the primary source via HF Secrets)
+    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. "
+            "Ensure GEMINI_API_KEY is set in Hugging Face Space secrets and loaded into settings, "
+            "or GOOGLE_API_KEY is set as an environment variable."
+        )
+        raise ValueError(
+            "Gemini API Key not configured. Please set it in Hugging Face Space secrets "
+            "as GEMINI_API_KEY or ensure GOOGLE_API_KEY environment variable is available."
+        )
 
     llm = ChatGoogleGenerativeAI(
-        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
+        model="gemini-1.5-pro-latest", # Using a more capable model
+        # model="gemini-pro", # Fallback if 1.5-pro is not available or for cost reasons
+        temperature=0.2,
+        google_api_key=api_key_to_use, # *** EXPLICITLY PASS THE KEY HERE ***
+        convert_system_message_to_human=True, # Often useful for non-OpenAI models
+        # Example safety settings (optional, adjust as needed)
+        # safety_settings={
+        #     HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,
+        #     HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
+        # }
     )
-    app_logger.info(f"ChatGoogleGenerativeAI ({llm.model}) initialized successfully for agent.")
+    app_logger.info(f"ChatGoogleGenerativeAI ({llm.model}) initialized successfully using provided API key.")
+
 except Exception as e:
-    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.")
+    # This broad exception catch is to provide a clear error message if LLM init fails for any reason.
+    detailed_error_message = str(e)
+    user_facing_error = f"Gemini LLM initialization failed: {detailed_error_message}. " \
+                        "Check API key validity, model name, and configurations in Hugging Face Secrets."
+    
+    if "default credentials were not found" in detailed_error_message.lower() or \
+       "could not find default credentials" in detailed_error_message.lower() or \
+       "api_key" in detailed_error_message.lower(): # Catch common API key related messages
+        user_facing_error = "Gemini LLM initialization failed: API key issue or missing credentials. " \
+                            "Ensure GEMINI_API_KEY is correctly set in Hugging Face Secrets and is valid."
+        app_logger.error(user_facing_error + f" Original error details: {detailed_error_message}", exc_info=False)
+    else:
+        app_logger.error(user_facing_error, exc_info=True) # Log full traceback for other errors
+    
+    # Re-raise to stop agent setup if LLM fails. This will be caught in get_agent_executor.
+    raise ValueError(user_facing_error)
 
 
 # --- 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(), # Add if using
+    # GeminiTool(), # Add if you have a specific reason to use Gemini as a sub-tool
 ]
 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) ---
+# --- Agent Prompt (for Structured Chat with Gemini) ---
 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. "
@@ -63,30 +96,25 @@ 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" # {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"
+    "3.  Tool Usage: You have access to the following tools:\n{tools}\n" # {tools} is filled by the agent
+    "   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\"}}, \"current_treatments\": [\"metformin\"], \"conditions\": [\"Type 2 Diabetes\"]}}}}\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 answer. Cite the tool if you used one (e.g., 'According to UMLS Lookup...').\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"
-    # "   - 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}" # Placeholder for agent's thoughts and tool outputs
+    "{agent_scratchpad}" # Placeholder for agent's internal thoughts, tool calls, and tool observations
 )
 
-# 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),
     MessagesPlaceholder(variable_name="agent_scratchpad"),
@@ -94,9 +122,16 @@ prompt = ChatPromptTemplate.from_messages([
 app_logger.info("Agent prompt template created for Gemini structured chat agent.")
 
 # --- Create Agent ---
+# This assumes `llm` was successfully initialized above.
+if llm is None:
+    # This case should ideally not be reached if the ValueError was raised during LLM init,
+    # but as a defensive measure:
+    app_logger.critical("LLM object is None at agent creation stage. Cannot proceed.")
+    # The ValueError from LLM init should have already stopped the module loading.
+    # If somehow execution reaches here with llm=None, something is very wrong.
+    raise SystemExit("Agent LLM failed to initialize. Application cannot start.")
+
 try:
-    # 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:
@@ -108,31 +143,39 @@ except Exception as e:
 agent_executor = AgentExecutor(
     agent=agent,
     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
+    verbose=True,
+    handle_parsing_errors=True,
+    max_iterations=10,
+    early_stopping_method="generate",
+    # return_intermediate_steps=True, # Good for debugging, makes response a dict with 'intermediate_steps'
 )
 app_logger.info("AgentExecutor with Gemini agent created successfully.")
 
 
 # --- Getter Function for Streamlit App ---
+_agent_executor_instance = agent_executor # Store the initialized executor
+
 def get_agent_executor():
     """
     Returns the configured agent executor for Gemini.
-    Initialization of LLM, tools, agent, and executor happens when this module is imported.
+    The executor is initialized when this module is first 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")):
-        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
+    global _agent_executor_instance
+    if _agent_executor_instance is None:
+        # This should not happen if module initialization was successful.
+        # It might indicate an issue where the module is reloaded or init failed silently.
+        app_logger.critical("CRITICAL: Agent executor is None when get_agent_executor is called. Re-initialization attempt or fundamental error.")
+        # You could try to re-initialize here, but it's better to ensure init works on first load.
+        # For now, raise an error to make it obvious.
+        raise RuntimeError("Agent executor was not properly initialized. Check application startup logs.")
+    return _agent_executor_instance
 
 # --- 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 to run the test.")
+    # Check if the API key is available for the test
+    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 file or as an environment variable to run the test.")
     else:
         print("\n🚀 Quantum Health Navigator (Gemini Agent Test Console) 🚀")
         print("-----------------------------------------------------------")
@@ -140,31 +183,32 @@ if __name__ == "__main__":
         print("Example topics: medical definitions, treatment optimization (will use simulated patient context).")
         print("-" * 59)
 
-        test_executor = get_agent_executor() # Get the globally defined executor
-        current_chat_history_for_test_run = [] # List of HumanMessage, AIMessage
+        try:
+            test_executor = get_agent_executor() # Get the executor
+        except ValueError as e_init:
+            print(f"⚠️ Agent initialization failed during test startup: {e_init}")
+            print("Ensure your API key is correctly configured.")
+            exit() # Exit if agent can't be initialized
+
+        current_chat_history_for_test_run = []
 
-        # Simulated patient context for testing the {patient_context} variable
         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."
+            "Age: 58; Gender: Female; Chief Complaint: Recent onset of blurry vision and fatigue; "
+            "Key Medical History: Prediabetes, Mild dyslipidemia; "
+            "Current Medications: None; Allergies: None known."
         )
         print(f"ℹ️  Simulated Patient Context for this test run: {test_patient_context_summary_str}\n")
 
-
         while True:
-            user_input_str = input("👤 You: ")
+            user_input_str = input("👤 You: ").strip()
             if user_input_str.lower() in ["exit", "quit"]:
                 print("👋 Exiting test console.")
                 break
-
-            if not user_input_str.strip():
+            if not user_input_str:
                 continue
 
             try:
                 app_logger.info(f"__main__ test: Invoking agent with input: '{user_input_str}'")
-                # 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_run,
@@ -174,12 +218,10 @@ if __name__ == "__main__":
                 ai_output_str = response_dict.get('output', "Agent did not produce an 'output' key.")
                 print(f"🤖 Agent: {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
+                if len(current_chat_history_for_test_run) > 10:
                     current_chat_history_for_test_run = current_chat_history_for_test_run[-10:]
 
             except Exception as e: