Update agent.py

agent.py

@@ -21,15 +21,9 @@ from config.settings import settings # This loads your HF secrets into the setti
 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:
-    # 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:
@@ -44,35 +38,25 @@ try:
         )
 
     llm = ChatGoogleGenerativeAI(
-        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
+        model="gemini-1.5-pro-latest",
         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,
-        # }
+        google_api_key=api_key_to_use,
+        convert_system_message_to_human=True,
     )
     app_logger.info(f"ChatGoogleGenerativeAI ({llm.model}) initialized successfully using provided API key.")
 
 except Exception as e:
-    # 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
+       "api_key" in detailed_error_message.lower():
         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.
+        app_logger.error(user_facing_error, exc_info=True)
     raise ValueError(user_facing_error)
 
 
@@ -81,12 +65,11 @@ tools_list = [
     UMLSLookupTool(),
     BioPortalLookupTool(),
     QuantumTreatmentOptimizerTool(),
-    # 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 (for Structured Chat with Gemini) ---
+# --- Agent Prompt (for Structured Chat with Gemini and your tools) ---
 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. "
@@ -96,7 +79,7 @@ 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
+    "3.  Tool Usage: You have access to the following tools (names: {tool_names}):\n{tools}\n"  # <--- {tool_names} ADDED HERE
     "   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"
@@ -112,7 +95,7 @@ SYSTEM_PROMPT_TEMPLATE = (
     "Previous conversation history:\n"
     "{chat_history}\n\n"
     "New human question: {input}\n"
-    "{agent_scratchpad}" # Placeholder for agent's internal thoughts, tool calls, and tool observations
+    "{agent_scratchpad}"
 )
 
 prompt = ChatPromptTemplate.from_messages([
@@ -122,22 +105,18 @@ 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:
     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:
+    # The error "Prompt missing required variables: {'tool_names'}" would be caught here
+    # if the placeholder wasn't correctly handled or if others are missing.
     app_logger.error(f"Failed to create structured chat agent: {e}", exc_info=True)
-    raise ValueError(f"Gemini agent creation failed: {e}")
-
+    raise ValueError(f"Gemini agent creation failed: {e}") # This is what you saw in the UI
 
 # --- Create Agent Executor ---
 agent_executor = AgentExecutor(
@@ -147,55 +126,40 @@ agent_executor = AgentExecutor(
     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
+_agent_executor_instance = agent_executor
 
 def get_agent_executor():
-    """
-    Returns the configured agent executor for Gemini.
-    The executor is initialized when this module is first imported.
-    """
     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.
+        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 of this agent.py file) ---
+# --- Example Usage (for local testing) ---
 if __name__ == "__main__":
-    # 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("-----------------------------------------------------------")
-        print("Type 'exit' or 'quit' to stop.")
-        print("Example topics: medical definitions, treatment optimization (will use simulated patient context).")
-        print("-" * 59)
-
+        # ... (rest of __main__ block from previous full agent.py) ...
         try:
-            test_executor = get_agent_executor() # Get the executor
+            test_executor = get_agent_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
+            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; "
+            "Age: 62; Gender: Female; Chief Complaint: Recent onset of blurry vision and fatigue; "
             "Key Medical History: Prediabetes, Mild dyslipidemia; "
-            "Current Medications: None; Allergies: None known."
+            "Current Medications: None reported; Allergies: Sulfa drugs."
         )
         print(f"ℹ️  Simulated Patient Context for this test run: {test_patient_context_summary_str}\n")
 
@@ -206,7 +170,6 @@ if __name__ == "__main__":
                 break
             if not user_input_str:
                 continue
-
             try:
                 app_logger.info(f"__main__ test: Invoking agent with input: '{user_input_str}'")
                 response_dict = test_executor.invoke({
@@ -214,16 +177,12 @@ if __name__ == "__main__":
                     "chat_history": current_chat_history_for_test_run,
                     "patient_context": test_patient_context_summary_str
                 })
-                
                 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_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:
                     current_chat_history_for_test_run = current_chat_history_for_test_run[-10:]
-
             except Exception as e:
                 print(f"⚠️ Error during agent invocation: {e}")
                 app_logger.error(f"Error in __main__ agent test invocation: {e}", exc_info=True)