Update agent.py

agent.py

@@ -4,30 +4,31 @@ from langchain_openai import ChatOpenAI
 from langchain.agents import AgentExecutor, create_openai_functions_agent
 
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
-from langchain_core.messages import AIMessage, HumanMessage
+from langchain_core.messages import AIMessage, HumanMessage # SystemMessage can be implicitly created or explicit
 
 # --- 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
+    # QuantumOptimizerInput, # Schemas are primarily used by tools themselves
+    # GeminiTool, # Not using in this OpenAI-centric agent
 )
 
 from config.settings import settings
 from services.logger import app_logger
 
 # --- Initialize LLM (OpenAI) ---
-llm = None
+llm = None # Initialize to None for robust error handling if init fails
 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", # More capable for function calling & instruction following
-        temperature=0.1, # Low for more deterministic tool use
+        model_name="gpt-4-turbo-preview", # More capable model for function calling
+        # model_name="gpt-3.5-turbo-0125", # More cost-effective alternative
+        temperature=0.1, # Lower for more deterministic tool use and function calls
         openai_api_key=settings.OPENAI_API_KEY
     )
     app_logger.info(f"ChatOpenAI ({llm.model_name}) initialized successfully for agent.")
@@ -40,7 +41,7 @@ except Exception as e:
         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)
+    raise ValueError(user_facing_error) # Propagate error to stop further agent setup
 
 
 # --- Initialize Tools List ---
@@ -52,128 +53,139 @@ tools_list = [
 app_logger.info(f"Agent tools initialized: {[tool.name for tool in tools_list]}")
 
 
-# --- 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 = (
+# --- Agent Prompt (for OpenAI Functions Agent - Simplified System Prompt) ---
+# The create_openai_functions_agent implicitly makes tool descriptions available to the LLM
+# via the function-calling mechanism. Explicitly listing {tools} and {tool_names} in the
+# system prompt string might be redundant or conflict with how this agent type works.
+# We will still provide overall instructions and patient_context placeholder.
+OPENAI_SYSTEM_PROMPT_TEXT_SIMPLIFIED = (
     "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: {tools}. Use them when a user's query can be best answered by one of them.\n"
+    "Your 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. Use them when a user's query can be best answered by one of them, based on their descriptions.\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" # 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.\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.\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."
+    "When using the 'quantum_treatment_optimizer' tool, ensure you populate its 'patient_data' argument using the available {patient_context}.\n"
+    "For `bioportal_lookup`, if the user doesn't specify an ontology, you may ask or default to 'SNOMEDCT_US'.\n"
+    "Always be clear and concise. Cite tools if their output forms a key part of your answer."
 )
 
-# ChatPromptTemplate defines the sequence of messages.
-# Variables here are what the agent_executor.invoke will ultimately need to provide or what the agent manages.
+# The ChatPromptTemplate defines the sequence of messages sent to the LLM.
+# `create_openai_functions_agent` expects specific placeholders.
 prompt = ChatPromptTemplate.from_messages([
-    ("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)
+    ("system", OPENAI_SYSTEM_PROMPT_TEXT_SIMPLIFIED), # System instructions, expects {patient_context}
+    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}")
-
+app_logger.info("Agent prompt template (simplified for OpenAI Functions) created.")
 
 # --- Create Agent ---
-if llm is None:
+if llm is None: # Defensive check, should have been caught by earlier raise
     app_logger.critical("LLM object is None at agent creation (OpenAI). Application cannot proceed.")
-    raise SystemExit("Agent LLM failed to initialize.")
+    raise SystemExit("Agent LLM failed to initialize. Application cannot start.")
 
 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.
+    # `create_openai_functions_agent` will use the tools' Pydantic schemas to define
+    # the "functions" that the OpenAI model can call.
     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)
+    # This is where the "Input to ChatPromptTemplate is missing variables" error would occur
+    # if the prompt object was still expecting variables not provided by the agent constructor or invoke.
     raise ValueError(f"OpenAI agent creation failed: {e}")
 
 
 # --- Create Agent Executor ---
 agent_executor = AgentExecutor(
     agent=agent,
-    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
+    tools=tools_list,
+    verbose=True,               # Essential for debugging tool usage and agent thoughts
+    handle_parsing_errors=True, # Tries to gracefully handle LLM output parsing issues
+    max_iterations=7,           # Prevent runaway agent loops
+    # return_intermediate_steps=True, # Set to True to get detailed thought/action steps in the response
 )
 app_logger.info("AgentExecutor with OpenAI agent created successfully.")
 
 
 # --- Getter Function for Streamlit App ---
-_agent_executor_instance = agent_executor
+_agent_executor_instance = agent_executor # Store the successfully initialized executor
 
 def get_agent_executor():
+    """
+    Returns the configured agent executor for OpenAI.
+    The executor is initialized when this module is first imported.
+    """
     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. Check application startup logs.")
-    if not settings.OPENAI_API_KEY: # Final check
+        # This indicates a failure during the initial module load (LLM or agent creation).
+        app_logger.critical("CRITICAL: Agent executor is None when get_agent_executor is called (OpenAI). Initialization likely failed.")
+        raise RuntimeError("Agent executor (OpenAI) was not properly initialized. Check application startup logs for errors (e.g., API key issues, prompt errors).")
+    
+    # Final check for API key, though LLM initialization should be the primary guard.
+    if not settings.OPENAI_API_KEY:
         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) ---
+# --- Example Usage (for local testing of this agent.py file) ---
 if __name__ == "__main__":
     if not settings.OPENAI_API_KEY:
-        print("🚨 Please set your OPENAI_API_KEY in .env or environment.")
+        print("🚨 Please set your OPENAI_API_KEY in .env file or as an environment variable to run the test.")
     else:
         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_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")
+        print("-----------------------------------------------------------")
+        print("Type 'exit' or 'quit' to stop.")
+        print("Example topics: medical definitions, treatment optimization (will use simulated patient context).")
+        print("-" * 59)
+
+        try:
+            test_executor = get_agent_executor() # Get the executor
+        except ValueError as e_init: # Catch errors from get_agent_executor or LLM/agent init
+            print(f"⚠️ Agent initialization failed during test startup: {e_init}")
+            print("Ensure your API key is correctly configured and prompt variables are set.")
+            exit() # Exit if agent can't be initialized
+
+        current_chat_history_for_test_run = [] # List of HumanMessage, AIMessage
+
+        # Simulated patient context for testing the {patient_context} variable
+        test_patient_context_summary_str = (
+            "Age: 70; Gender: Male; Chief Complaint: Shortness of breath on exertion; "
+            "Key Medical History: COPD, Atrial Fibrillation; "
+            "Current Medications: Tiotropium inhaler, Apixaban 5mg BID; Allergies: Penicillin."
+        )
+        print(f"ℹ️  Simulated Patient Context for this test run: {test_patient_context_summary_str}\n")
 
         while True:
-            usr_in = input("👤 You: ").strip()
-            if usr_in.lower() in ["exit", "quit"]: print("👋 Exiting."); break
-            if not usr_in: continue
+            user_input_str = input("👤 You: ").strip()
+            if user_input_str.lower() in ["exit", "quit"]:
+                print("👋 Exiting test console.")
+                break
+            if not user_input_str: # Skip empty input
+                continue
+
             try:
-                # 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)
+                app_logger.info(f"__main__ test (OpenAI): Invoking with input: '{user_input_str}'")
+                # These are the keys expected by the ChatPromptTemplate and agent:
+                # "input", "chat_history", and "patient_context" (because it's in our system prompt)
+                response_dict = test_executor.invoke({
+                    "input": user_input_str,
+                    "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}")
                 
-                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:]
+                # 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 to prevent overly long contexts
+                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_invoke:
-                print(f"⚠️ Invoke Error: {type(e_invoke).__name__} - {e_invoke}")
+                print(f"⚠️ Error during agent invocation: {type(e_invoke).__name__} - {e_invoke}")
                 app_logger.error(f"Error in __main__ OpenAI agent test invocation: {e_invoke}", exc_info=True)