Update agent.py

agent.py

@@ -4,9 +4,8 @@ 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 # SystemMessage can be implicitly created
+from langchain_core.messages import AIMessage, HumanMessage
 
-# --- Import your defined tools FROM THE 'tools' PACKAGE ---
 from tools import (
     BioPortalLookupTool,
     UMLSLookupTool,
@@ -16,32 +15,25 @@ from tools import (
 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", # More capable for complex instruction following
-        temperature=0.1, # Low temperature for more deterministic and structured output for tool calls
+        model_name="gpt-4-turbo-preview",
+        temperature=0.1,
         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}. Check API key and model name."
+    user_facing_error = f"OpenAI LLM initialization failed: {detailed_error_message}."
     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. 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)
+        user_facing_error = "OpenAI LLM initialization failed: API key issue. Check HF Secrets."
+    app_logger.error(user_facing_error, exc_info=True)
     raise ValueError(user_facing_error)
 
-
-# --- Initialize Tools List ---
 tools_list = [
     UMLSLookupTool(),
     BioPortalLookupTool(),
@@ -49,20 +41,19 @@ tools_list = [
 ]
 app_logger.info(f"Agent tools initialized: {[tool.name for tool in tools_list]}")
 
-
-# --- Agent Prompt (for OpenAI Functions Agent - Enhanced System Prompt) ---
-OPENAI_SYSTEM_PROMPT_TEXT_ENHANCED = (
-    "You are 'Quantum Health Navigator', an AI assistant for healthcare professionals. "
+# --- 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. "
     "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"
+    "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}
     "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
+    "Patient Context for this session (if provided by the user earlier): {patient_context}\n"
 
     "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\"]}}. " # Added symptoms to example
+    "     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"
     "   - '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"
@@ -74,18 +65,16 @@ OPENAI_SYSTEM_PROMPT_TEXT_ENHANCED = (
 )
 
 prompt = ChatPromptTemplate.from_messages([
-    ("system", OPENAI_SYSTEM_PROMPT_TEXT_ENHANCED),
+    ("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")
 ])
-app_logger.info("Agent prompt template (enhanced for OpenAI Functions) created.")
+app_logger.info("Agent prompt template (with explicit tools/tool_names in system message) created.")
 
-# --- Create Agent ---
 if llm is None:
-    app_logger.critical("LLM object is None at agent creation (OpenAI). Application cannot proceed.")
+    app_logger.critical("LLM object is None at agent creation (OpenAI). Cannot proceed.")
     raise SystemExit("Agent LLM failed to initialize.")
-
 try:
     agent = create_openai_functions_agent(llm=llm, tools=tools_list, prompt=prompt)
     app_logger.info("OpenAI Functions agent created successfully.")
@@ -93,83 +82,51 @@ except Exception as e:
     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,
     verbose=True,
-    handle_parsing_errors=True, # Important for robust function calling
-    max_iterations=7, # Good to have a limit
-    # return_intermediate_steps=True, # Enable this for deeper debugging of agent steps
+    handle_parsing_errors=True,
+    max_iterations=7,
 )
 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). Initialization likely failed.")
-        raise RuntimeError("Agent executor (OpenAI) was not properly initialized. Check application startup logs for errors.")
-    if not settings.OPENAI_API_KEY: # Final check
+        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:
         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 file or as an environment variable to run the test.")
+        print("🚨 Please set your OPENAI_API_KEY in .env or environment.")
     else:
         print("\n🚀 Quantum Health Navigator (OpenAI 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)
-
-        try:
-            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 and prompt variables are set.")
-            exit()
-
-        current_chat_history_for_test_run = []
-        test_patient_context_summary_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 Patient Context for this test run: {test_patient_context_summary_str}\n")
-
+        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")
         while True:
-            user_input_str = input("👤 You: ").strip()
-            if user_input_str.lower() in ["exit", "quit"]:
-                print("👋 Exiting test console.")
-                break
-            if not user_input_str:
-                continue
-
+            usr_in = input("👤 You: ").strip()
+            if usr_in.lower() in ["exit", "quit"]: print("👋 Exiting."); break
+            if not usr_in: continue
             try:
-                app_logger.info(f"__main__ test (OpenAI): Invoking with input: '{user_input_str}'")
-                response_dict = test_executor.invoke({
-                    "input": user_input_str,
-                    "chat_history": current_chat_history_for_test_run,
-                    "patient_context": test_patient_context_summary_str
+                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
                 })
-                
-                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_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)
+                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}")