Update agent.py

agent.py

@@ -6,34 +6,44 @@ from langchain.agents import AgentExecutor, create_openai_functions_agent
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
 from langchain_core.messages import AIMessage, HumanMessage
 
+# --- Import your defined tools FROM THE 'tools' PACKAGE ---
+# This relies on tools/__init__.py correctly exporting these names.
 from tools import (
     BioPortalLookupTool,
     UMLSLookupTool,
     QuantumTreatmentOptimizerTool,
+    # QuantumOptimizerInput, # Only if needed for type hints directly in this file
 )
 
 from config.settings import settings
 from services.logger import app_logger
 
+# --- Initialize LLM (OpenAI) ---
 llm = None
 try:
     if not settings.OPENAI_API_KEY:
         app_logger.error("CRITICAL: OPENAI_API_KEY not found in settings. Agent cannot initialize.")
         raise ValueError("OpenAI API Key not configured. Please set it in Hugging Face Space secrets as OPENAI_API_KEY.")
+
     llm = ChatOpenAI(
-        model_name="gpt-4-turbo-preview",
-        temperature=0.1,
+        model_name="gpt-4-turbo-preview", # More capable for function calling & instruction following
+        temperature=0.1, # Low for more deterministic tool use
         openai_api_key=settings.OPENAI_API_KEY
     )
     app_logger.info(f"ChatOpenAI ({llm.model_name}) initialized successfully for agent.")
+
 except Exception as e:
     detailed_error_message = str(e)
-    user_facing_error = f"OpenAI LLM initialization failed: {detailed_error_message}."
+    user_facing_error = f"OpenAI LLM initialization failed: {detailed_error_message}. Check API key and model name."
     if "api_key" in detailed_error_message.lower() or "authenticate" in detailed_error_message.lower():
-        user_facing_error = "OpenAI LLM initialization failed: API key issue. Check HF Secrets."
-    app_logger.error(user_facing_error, exc_info=True)
+        user_facing_error = "OpenAI LLM initialization failed: API key issue. Ensure OPENAI_API_KEY is correctly set in Hugging Face Secrets and is valid."
+        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(),
@@ -41,67 +51,92 @@ tools_list = [
 ]
 app_logger.info(f"Agent tools initialized: {[tool.name for tool in tools_list]}")
 
-# --- Agent Prompt (Reverting to version that EXPLICITLY includes {tools} and {tool_names} in the system message string) ---
-OPENAI_SYSTEM_PROMPT_TEXT_WITH_TOOLS_EXPLICIT = (
-    "You are 'Quantum Health Navigator', an advanced AI assistant for healthcare professionals. "
+
+# --- Agent Prompt (for OpenAI Functions Agent - Explicitly including {tools} and {tool_names}) ---
+# The KeyError indicated that ChatPromptTemplate was expecting 'tools' and 'tool_names' as input variables.
+# create_openai_functions_agent should populate these if these placeholders are in the system message.
+OPENAI_SYSTEM_PROMPT_WITH_EXPLICIT_TOOLS_VARS = (
+    "You are 'Quantum Health Navigator', an AI assistant for healthcare professionals. "
     "Your primary goal is to assist with medical information lookup, treatment optimization queries, and general medical Q&A. "
-    "You have access to a set of specialized tools (their names are: {tool_names}). Their detailed descriptions are available to you: {tools}. Use them when a user's query can be best answered by one of them.\n" # Explicit {tools} and {tool_names}
+    "You have access to a set of specialized tools. Their names are: {tool_names}. Their detailed descriptions are: {tools}. Use them when a user's query can be best answered by one of them.\n"
     "Disclaimers: Always state that you are for informational support and not a substitute for clinical judgment. Do not provide direct medical advice for specific patient cases without using the 'quantum_treatment_optimizer' tool if relevant.\n"
-    "Patient Context for this session (if provided by the user earlier): {patient_context}\n"
+    "Patient Context for this session (if provided by the user earlier): {patient_context}\n" # This variable is passed from invoke
 
     "Tool Usage Guidelines:\n"
     "1. When using the 'quantum_treatment_optimizer' tool, its 'action_input' argument requires three main keys: 'patient_data', 'current_treatments', and 'conditions'.\n"
     "   - The 'patient_data' key MUST be a dictionary. Populate this dictionary by extracting relevant details from the {patient_context}. "
     "     For example, if {patient_context} is 'Age: 50; Gender: Male; Key Medical History: Hypertension; Chief Complaint: headache', "
     "     then 'patient_data' could be {{\"age\": 50, \"gender\": \"Male\", \"relevant_history\": [\"Hypertension\"], \"symptoms\": [\"headache\"]}}. "
-    "     Include details like age, gender, chief complaint, key medical history, and current medications from {patient_context} within this 'patient_data' dictionary. If a value is not present in context, omit the key or use null/None if appropriate for the tool, but prioritize providing what is available.\n"
+    "     Include details like age, gender, chief complaint, key medical history, and current medications from {patient_context} within this 'patient_data' dictionary.\n"
     "   - 'current_treatments' should be a list of strings derived from the 'Current Medications' part of {patient_context}.\n"
     "   - 'conditions' should be a list of strings, including primary conditions from the 'Key Medical History' or 'Chief Complaint' parts of {patient_context}, and any conditions explicitly mentioned or implied by the current user query.\n"
     "2. For `bioportal_lookup`, the 'action_input' should be a dictionary like {{\"term\": \"search_term\", \"ontology\": \"ONTOLOGY_ACRONYM\"}}. If the user doesn't specify an ontology, you may ask for clarification or default to 'SNOMEDCT_US'.\n"
     "3. For `umls_lookup`, the 'action_input' is a single string: the medical term to search.\n"
-    "4. After using a tool, you will receive an observation. Use this observation and your general knowledge to formulate a comprehensive final answer to the human. Clearly cite the tool if its output forms a key part of your answer (e.g., 'According to UMLS Lookup...').\n"
+    "4. After using a tool, you will receive an observation. Use this observation and your general knowledge to formulate a comprehensive final answer to the human. Clearly cite the tool if its output forms a key part of your answer.\n"
     "5. If a user's query seems to ask for treatment advice or medication suggestions for a specific scenario (especially if patient context is available), you MUST prioritize using the 'quantum_treatment_optimizer' tool.\n"
     "6. For general medical knowledge questions not requiring patient-specific optimization or specific ontology/CUI lookups, you may answer directly from your training data, but always include the standard disclaimer."
 )
 
+# ChatPromptTemplate defines the sequence of messages.
+# Variables here are what the agent_executor.invoke will ultimately need to provide or what the agent manages.
 prompt = ChatPromptTemplate.from_messages([
-    ("system", OPENAI_SYSTEM_PROMPT_TEXT_WITH_TOOLS_EXPLICIT), # Using the version with explicit {tools} and {tool_names}
-    MessagesPlaceholder(variable_name="chat_history"),
-    ("human", "{input}"),
-    MessagesPlaceholder(variable_name="agent_scratchpad")
+    ("system", OPENAI_SYSTEM_PROMPT_WITH_EXPLICIT_TOOLS_VARS), # System instructions, expects {patient_context}, {tools}, {tool_names}
+    MessagesPlaceholder(variable_name="chat_history"),      # For past Human/AI messages
+    ("human", "{input}"),                                  # For the current user query
+    MessagesPlaceholder(variable_name="agent_scratchpad")   # For agent's internal work (function calls/responses)
 ])
 app_logger.info("Agent prompt template (with explicit tools/tool_names in system message) created.")
+# Log the input variables that this prompt structure will expect.
+# `create_openai_functions_agent` should provide 'tools' and 'tool_names' to this prompt.
+# The user (via invoke) provides 'input', 'chat_history', 'patient_context'.
+# 'agent_scratchpad' is managed by the AgentExecutor.
+app_logger.debug(f"Prompt expected input variables: {prompt.input_variables}")
+
 
+# --- Create Agent ---
 if llm is None:
-    app_logger.critical("LLM object is None at agent creation (OpenAI). Cannot proceed.")
+    app_logger.critical("LLM object is None at agent creation (OpenAI). Application cannot proceed.")
     raise SystemExit("Agent LLM failed to initialize.")
+
 try:
+    # `create_openai_functions_agent` is given the llm, the raw tools_list, and the prompt.
+    # It should process `tools_list` to make them available as OpenAI functions AND
+    # populate the `{tools}` and `{tool_names}` placeholders in the prompt.
     agent = create_openai_functions_agent(llm=llm, tools=tools_list, prompt=prompt)
     app_logger.info("OpenAI Functions agent created successfully.")
 except Exception as e:
+    # This is where the KeyError "Input to ChatPromptTemplate is missing variables {'tools', 'tool_names'}"
+    # was occurring.
     app_logger.error(f"Failed to create OpenAI Functions agent: {e}", exc_info=True)
     raise ValueError(f"OpenAI agent creation failed: {e}")
 
+
+# --- Create Agent Executor ---
 agent_executor = AgentExecutor(
     agent=agent,
-    tools=tools_list,
+    tools=tools_list, # Tools are also provided to the executor
     verbose=True,
     handle_parsing_errors=True,
     max_iterations=7,
+    # return_intermediate_steps=True, # Good for debugging
 )
 app_logger.info("AgentExecutor with OpenAI agent created successfully.")
 
+
+# --- Getter Function for Streamlit App ---
 _agent_executor_instance = agent_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 (OpenAI).")
-        raise RuntimeError("Agent executor (OpenAI) was not properly initialized.")
-    if not settings.OPENAI_API_KEY:
+        raise RuntimeError("Agent executor (OpenAI) was not properly initialized. Check application startup logs.")
+    if not settings.OPENAI_API_KEY: # Final check
         app_logger.error("OpenAI API Key is missing at get_agent_executor call. Agent will fail.")
         raise ValueError("OpenAI API Key not configured.")
     return _agent_executor_instance
 
+# --- Example Usage (for local testing) ---
 if __name__ == "__main__":
     if not settings.OPENAI_API_KEY:
         print("🚨 Please set your OPENAI_API_KEY in .env or environment.")
@@ -109,24 +144,36 @@ if __name__ == "__main__":
         print("\n🚀 Quantum Health Navigator (OpenAI Agent Test Console) 🚀")
         try: test_executor = get_agent_executor()
         except ValueError as e_init: print(f"⚠️ Agent init failed: {e_init}"); exit()
+
         history = []
-        context = ("Age: 60; Gender: Male; Chief Complaint: general fatigue and occasional dizziness; "
-                   "Key Medical History: Type 2 Diabetes, Hypertension; "
-                   "Current Medications: Metformin 1000mg daily, Lisinopril 20mg daily; Allergies: None.")
-        print(f"ℹ️  Simulated Context: {context}\n")
+        context_str = ("Age: 60; Gender: Male; Chief Complaint: general fatigue and occasional dizziness; "
+                       "Key Medical History: Type 2 Diabetes, Hypertension; "
+                       "Current Medications: Metformin 1000mg daily, Lisinopril 20mg daily; Allergies: None.")
+        print(f"ℹ️  Simulated Context: {context_str}\n")
+
         while True:
             usr_in = input("👤 You: ").strip()
             if usr_in.lower() in ["exit", "quit"]: print("👋 Exiting."); break
             if not usr_in: continue
             try:
-                res = test_executor.invoke({
-                    "input": usr_in, 
-                    "chat_history": history, 
-                    "patient_context": context,
-                    # DO NOT PASS "tools" or "tool_names" here; the agent constructor does that
-                })
+                # The keys here ('input', 'chat_history', 'patient_context') must match
+                # what the ChatPromptTemplate ultimately expects after create_openai_functions_agent
+                # has done its work with 'tools' and 'tool_names'.
+                payload = {
+                    "input": usr_in,
+                    "chat_history": history,
+                    "patient_context": context_str,
+                    # Note: We do NOT explicitly pass 'tools' or 'tool_names' in invoke.
+                    # The `create_openai_functions_agent` is responsible for making these available
+                    # to the `prompt` object during its formatting process.
+                }
+                app_logger.info(f"__main__ test (OpenAI): Invoking with payload keys: {list(payload.keys())}")
+                res = test_executor.invoke(payload)
+                
                 ai_out = res.get('output', "No output.")
                 print(f"🤖 Agent: {ai_out}")
                 history.extend([HumanMessage(content=usr_in), AIMessage(content=ai_out)])
                 if len(history) > 8: history = history[-8:]
-            except Exception as e_invoke: print(f"⚠️ Invoke Error: {type(e_invoke).__name__} - {e_invoke}")
+            except Exception as e_invoke:
+                print(f"⚠️ Invoke Error: {type(e_invoke).__name__} - {e_invoke}")
+                app_logger.error(f"Error in __main__ OpenAI agent test invocation: {e_invoke}", exc_info=True)