Update agent.py

agent.py

@@ -1,203 +1,191 @@
 # /home/user/app/agent.py
+import os
 from langchain_google_genai import ChatGoogleGenerativeAI
-from langchain.agents import AgentExecutor, create_structured_chat_agent # Using a more general agent
-# If you want to try Gemini's native function calling (experimental and might require specific model versions):
-# from langchain_google_genai.chat_models import GChatVertexAI # For Vertex AI
-# from langchain_google_genai import HarmBlockThreshold, HarmCategory # For safety settings
+from langchain.agents import AgentExecutor, create_structured_chat_agent
+# from langchain_google_genai import HarmBlockThreshold, HarmCategory # Optional for safety
 
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
 from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
-# from langchain_community.chat_message_histories import ChatMessageHistory # Not used directly here
 
-from tools import (
-    UMLSLookupTool, BioPortalLookupTool, QuantumTreatmentOptimizerTool
-    # GeminiTool might be redundant if the main LLM is Gemini, unless it's for a different Gemini model/task.
-)
+# --- Import your defined tools ---
+from tools.bioportal_tool import BioPortalLookupTool, BioPortalInput
+from tools.gemini_tool import GeminiTool, GeminiInput # For using Gemini as a specific sub-task tool
+from tools.umls_tool import UMLSLookupTool, UMLSInput
+from tools.quantum_treatment_optimizer_tool import QuantumTreatmentOptimizerTool, QuantumOptimizerInput # Assuming this path and model name
+
 from config.settings import settings
 from services.logger import app_logger
 
 # --- Initialize LLM (Gemini) ---
-# Ensure GOOGLE_API_KEY is set in your environment, or pass it directly:
-# api_key=settings.GEMINI_API_KEY (if settings.GEMINI_API_KEY maps to GOOGLE_API_KEY)
 try:
+    # Ensure GOOGLE_API_KEY is set in your environment (HuggingFace Secrets)
+    # or settings.GEMINI_API_KEY correctly maps to it.
+    if not (settings.GEMINI_API_KEY or os.getenv("GOOGLE_API_KEY")):
+        raise ValueError("GOOGLE_API_KEY (for Gemini) not found in settings or environment.")
+
     llm = ChatGoogleGenerativeAI(
-        model="gemini-pro", # Or "gemini-1.5-pro-latest" if available and preferred
+        model="gemini-1.5-pro-latest", # Using a more capable Gemini model if available
+        # model="gemini-pro", # Fallback if 1.5-pro is not yet available or preferred
         temperature=0.3,
-        # google_api_key=settings.GEMINI_API_KEY, # Explicitly pass if needed
-        # convert_system_message_to_human=True, # Sometimes helpful for models not strictly adhering to system role
-        # safety_settings={ # Optional: configure safety settings
-        #     HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE,
-        #     HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_LOW_AND_ABOVE,
+        # google_api_key=settings.GEMINI_API_KEY, # Explicitly pass if GOOGLE_API_KEY env var isn't set
+        convert_system_message_to_human=True, # Can be helpful for some models
+        # safety_settings={ # Example safety settings
+        #     HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,
+        #     HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
         # }
     )
-    app_logger.info("ChatGoogleGenerativeAI (Gemini) initialized successfully.")
+    app_logger.info(f"ChatGoogleGenerativeAI ({llm.model}) initialized successfully.")
 except Exception as e:
     app_logger.error(f"Failed to initialize ChatGoogleGenerativeAI: {e}", exc_info=True)
-    # Raise an error that can be caught by get_agent_executor to inform the user
-    raise ValueError(f"Gemini LLM initialization failed: {e}. Check API key and configurations.")
+    raise ValueError(f"Gemini LLM initialization failed: {e}. Check API key and configurations in HF Secrets.")
 
 
 # --- Initialize Tools ---
-# Ensure your tools' descriptions are very clear, especially for non-OpenAI function calling agents.
+# Ensure each tool's description is clear and guides the LLM on when and how to use it.
+# Also, ensure their args_schema is correctly defined.
 tools = [
     UMLSLookupTool(),
     BioPortalLookupTool(),
     QuantumTreatmentOptimizerTool(),
+    # GeminiTool(), # Consider if this is needed. The main LLM is already Gemini.
+                   # Useful if this tool performs a very specific, different task with Gemini,
+                   # or uses a different Gemini model (e.g., for vision if main is text).
+                   # If it's just for general queries, the main agent LLM can handle it.
 ]
 app_logger.info(f"Tools initialized: {[tool.name for tool in tools]}")
 
 
-# --- Agent Prompt (Adapted for a general structured chat agent) ---
-# This prompt needs to guide the LLM to decide on tool use and format tool calls.
-# It's more complex than relying on native function calling.
-# We might need to instruct it to output a specific JSON structure for tool calls.
-# For now, let's try a ReAct-style approach with create_structured_chat_agent.
-# LangChain Hub has prompts for this: hub.pull("hwchase17/structured-chat-agent")
-# Or we define a custom one:
-
-# This is a simplified prompt. For robust tool use with Gemini (without native function calling),
-# you'd often use a ReAct prompt or a prompt that guides it to output JSON for tool calls.
-# `create_structured_chat_agent` is designed to work with models that can follow instructions
-# to produce a structured output, often for tool usage.
-
-# For Gemini, which now supports tool calling via its API directly (though LangChain integration might vary),
-# we *could* try to structure the prompt for that if `ChatGoogleGenerativeAI` has good support.
-# Let's assume a more general structured chat approach first if direct tool calling isn't as smooth
-# as with OpenAI's function calling agents.
-
-# If trying Gemini's newer tool calling features with LangChain (ensure your langchain-google-genai is up to date):
-# You might be able to bind tools directly to the LLM and use a simpler agent structure.
-# llm_with_tools = llm.bind_tools(tools) # This is the newer LangChain syntax
-# Then create an agent that leverages this.
-
-# For now, let's use `create_structured_chat_agent` which is more general.
-# This prompt is similar to hwchase17/structured-chat-agent on Langsmith Hub
-SYSTEM_PROMPT_TEXT = (
-    "You are 'Quantum Health Navigator', a helpful AI assistant for healthcare professionals. "
-    "Respond to the human as helpfully and accurately as possible. "
-    "You have access to the following tools:\n\n"
-    "{tools}\n\n" # This will be filled by the agent with tool names and descriptions
-    "To use a tool, you can use the following format:\n\n"
-    "```json\n"
-    "{{\n"
-    '  "action": "tool_name",\n'
-    '  "action_input": "input_to_tool"\n' # For tools with single string input
-    # Or for tools with structured input (like QuantumTreatmentOptimizerTool):
-    # '  "action_input": {{"arg1": "value1", "arg2": "value2"}}\n'
-    "}}\n"
-    "```\n\n"
-    "If you use a tool, the system will give you the observation from the tool. "
-    "Then you must respond to the human based on this observation and your knowledge. "
-    "If the human asks a question that doesn't require a tool, answer directly. "
-    "When asked about treatment optimization for a specific patient based on provided context, "
-    "you MUST use the 'quantum_treatment_optimizer' tool. "
-    "For general medical knowledge, you can answer directly or use UMLS/BioPortal. "
-    "Always cite the tool you used if its output is part of your final response. "
-    "Do not provide medical advice directly for specific patient cases without using the 'quantum_treatment_optimizer' tool. "
-    "Patient Context for this session (if provided by the user earlier): {patient_context}\n" # Added patient_context
+# --- Agent Prompt (Adapted for Structured Chat with Gemini and your tools) ---
+# This prompt guides the LLM to:
+# 1. Understand its role and capabilities.
+# 2. Know which tools are available and their purpose (from {tools}).
+# 3. Format tool invocations as a JSON blob with "action" and "action_input".
+#    - "action_input" should be a string for simple tools (UMLSInput, GeminiInput).
+#    - "action_input" should be a dictionary for tools with multiple args (BioPortalInput, QuantumOptimizerInput).
+# 4. Use the provided {patient_context}.
+# 5. Refer to {chat_history}.
+# 6. Process the new {input}.
+# 7. Use {agent_scratchpad} for its internal monologue/tool results.
+
+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. "
+    "You must adhere to the following guidelines:\n"
+    "1.  Disclaimers: Always remind the user that you are an AI, not a human medical professional, and your information "
+        "is for support, not a substitute for clinical judgment. Do not provide direct medical advice for specific patient cases "
+        "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"
+    "   To use a tool, respond with a JSON markdown code block like this:\n"
+    "   ```json\n"
+    "   {{\n"
+    '     "action": "tool_name",\n'
+    '     "action_input": "query string for the tool" OR {{"arg1": "value1", "arg2": "value2", ...}} \n'
+    "   }}\n"
+    "   ```\n"
+    "   - For `umls_lookup` and `google_gemini_chat`, `action_input` is a single string (the 'term' or 'query').\n"
+    "   - For `bioportal_lookup`, `action_input` is a dictionary like `{{\"term\": \"search_term\", \"ontology\": \"ONTOLOGY_CODE\"}}`. If ontology is not specified by user, you can default to SNOMEDCT or ask.\n"
+    "   - For `quantum_treatment_optimizer`, `action_input` is a dictionary like `{{\"patient_data\": {{...patient details...}}, \"current_treatments\": [\"med1\"], \"conditions\": [\"cond1\"]}}`. You MUST populate 'patient_data' using the overall {patient_context} if available and relevant.\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"
+    "5.  Specific Tool Guidance:\n"
+    "   - If asked about treatment optimization for a specific patient (especially if 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) or answering directly if confident.\n"
+    "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}"
 )
 
 prompt = ChatPromptTemplate.from_messages([
-    ("system", SYSTEM_PROMPT_TEXT),
-    # MessagesPlaceholder(variable_name="chat_history"), # chat_history is now part of the system prompt
-    # ("human", "{input}"), # input is now part of the system prompt
-    MessagesPlaceholder(variable_name="agent_scratchpad"), # For structured chat agent, this is important
+    ("system", SYSTEM_PROMPT_TEMPLATE),
+    # For structured chat agent, HumanMessage/AIMessage sequence is often handled by MessagesPlaceholder("agent_scratchpad")
+    # or by how the agent formats history into the main prompt.
+    # The key is that the {chat_history} and {input} placeholders are in the system prompt.
+    MessagesPlaceholder(variable_name="agent_scratchpad"),
 ])
-# Note: The `create_structured_chat_agent` expects `input` and `chat_history` to be implicitly handled
-# or passed through `agent_scratchpad` based on how it formats things.
-# The prompt structure might need adjustment based on the exact agent behavior.
-# Often, for these agents, you pass "input" and "chat_history" to invoke, and the prompt template variables
-# are {input}, {chat_history}, {agent_scratchpad}, {tools}, {tool_names}.
-
-# For create_structured_chat_agent, the prompt should guide the LLM to produce
-# either a final answer or a JSON blob for a tool call.
-# The input variables for the prompt are typically 'input', 'chat_history', 'agent_scratchpad', 'tools', 'tool_names'.
-# Our SYSTEM_PROMPT_TEXT includes these implicitly or explicitly.
-
+app_logger.info("Agent prompt template created.")
 
 # --- Create Agent ---
 try:
-    # `create_structured_chat_agent` is designed for LLMs that can follow complex instructions
-    # and output structured data (like JSON for tool calls) when prompted to do so.
     agent = create_structured_chat_agent(llm=llm, tools=tools, prompt=prompt)
-    app_logger.info("Structured chat agent created successfully with Gemini.")
+    app_logger.info("Structured chat agent created successfully with Gemini LLM.")
 except Exception as e:
     app_logger.error(f"Failed to create structured chat agent: {e}", exc_info=True)
-    raise ValueError(f"Agent creation failed: {e}")
+    raise ValueError(f"Gemini agent creation failed: {e}")
 
 
 # --- Create Agent Executor ---
 agent_executor = AgentExecutor(
     agent=agent,
     tools=tools,
-    verbose=True,
-    handle_parsing_errors=True, # Important for agents that parse LLM output for tool calls
-    # Example: "Could not parse LLM output: `...`" - then it can retry or return this error.
-    # max_iterations=7, # Good to prevent overly long chains
-    # return_intermediate_steps=True # Useful for debugging to see thought process
+    verbose=True, # Set to True for debugging, False for production
+    handle_parsing_errors=True, # Crucial for LLM-generated JSON for tool calls
+    max_iterations=10, # Increased slightly for potentially complex tool interactions
+    # return_intermediate_steps=True, # Enable if you need to see thoughts/tool calls in the response object
+    early_stopping_method="generate", # Stop if LLM generates a stop token or a final answer
 )
-app_logger.info("AgentExecutor created successfully.")
+app_logger.info("AgentExecutor with Gemini agent created successfully.")
 
 
 # --- Getter Function for Streamlit App ---
 def get_agent_executor():
     """Returns the configured agent executor for Gemini."""
-    # The llm and agent_executor are already initialized.
-    # We check for the API key here as a safeguard, though initialization would have failed earlier.
-    if not (settings.GEMINI_API_KEY or os.environ.get("GOOGLE_API_KEY")): # Check both setting and env var
-        app_logger.error("GOOGLE_API_KEY (for Gemini) not set. Agent will not function.")
-        raise ValueError("Google API Key for Gemini not configured. Agent cannot be initialized.")
+    # Initialization happens above when the module is loaded.
+    # This function just returns the already created executor.
+    # A check for API key is good practice, though it would have failed earlier if not set.
+    if not (settings.GEMINI_API_KEY or os.getenv("GOOGLE_API_KEY")):
+        # This log might be redundant if LLM init failed, but good as a sanity check here.
+        app_logger.error("CRITICAL: GOOGLE_API_KEY (for Gemini) is not available at get_agent_executor call.")
+        raise ValueError("Google API Key for Gemini not configured. Agent cannot function.")
     return agent_executor
 
-# --- Example Usage (for local testing) ---
-import os # For checking GOOGLE_API_KEY from environment
-
+# --- Example Usage (for local testing of this agent.py file) ---
 if __name__ == "__main__":
-    if not (settings.GEMINI_API_KEY or os.environ.get("GOOGLE_API_KEY")):
-        print("Please set your GOOGLE_API_KEY (for Gemini) in .env file or environment.")
+    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.")
     else:
-        print("Gemini Agent Test Console (type 'exit' or 'quit' to stop)")
+        print("\nQuantum Health Navigator (Gemini Agent Test Console)")
+        print("Type 'exit' or 'quit' to stop.")
+        print("Example queries:")
+        print("  - What is hypertension?")
+        print("  - Lookup 'myocardial infarction' in UMLS.")
+        print("  - Search for 'diabetes mellitus type 2' in BioPortal using SNOMEDCT ontology.")
+        print("  - Optimize treatment for a patient (context will be simulated).")
+        print("-" * 30)
+
         executor = get_agent_executor()
-        
-        # For structured chat agents, chat_history is often passed in the invoke call.
-        # The agent prompt includes {chat_history}.
-        current_chat_history = [] # List of HumanMessage, AIMessage
-
-        # Initial patient context (simulated for testing)
-        patient_context_for_test = {
-            "age": 35,
-            "gender": "Male",
-            "key_medical_history": "Type 2 Diabetes, Hypertension",
-            "current_medications": "Metformin, Lisinopril"
-        }
-        context_summary_parts_test = [f"{k.replace('_', ' ').title()}: {v}" for k, v in patient_context_for_test.items() if v]
-        patient_context_str_test = "; ".join(context_summary_parts_test) if context_summary_parts_test else "None provided."
+        current_chat_history_for_test = [] # List of HumanMessage, AIMessage
 
+        # Simulated patient context for testing the {patient_context} variable
+        test_patient_context_summary = (
+            "Age: 45; Gender: Male; Chief Complaint: Intermittent chest pain; "
+            "Key Medical History: Hyperlipidemia; Current Medications: Atorvastatin 20mg."
+        )
 
         while True:
-            user_input = input("You: ")
-            if user_input.lower() in ["exit", "quit"]:
+            user_input_str = input("\n👤 You: ")
+            if user_input_str.lower() in ["exit", "quit"]:
+                print("Exiting test console.")
                 break
 
             try:
+                app_logger.info(f"__main__ test: Invoking agent with input: '{user_input_str}'")
                 response = executor.invoke({
-                    "input": user_input,
-                    "chat_history": current_chat_history,
-                    "patient_context": patient_context_str_test # Passing patient context
-                    # `tools` and `tool_names` are usually handled by the agent constructor
+                    "input": user_input_str,
+                    "chat_history": current_chat_history_for_test,
+                    "patient_context": test_patient_context_summary # Passing the context
                 })
                 
-                ai_output = response.get('output', "No output from agent.")
-                print(f"Agent: {ai_output}")
+                ai_output_str = response.get('output', "Agent did not produce an output.")
+                print(f"🤖 Agent: {ai_output_str}")
                 
-                current_chat_history.append(HumanMessage(content=user_input))
-                current_chat_history.append(AIMessage(content=ai_output))
+                current_chat_history_for_test.append(HumanMessage(content=user_input_str))
+                current_chat_history_for_test.append(AIMessage(content=ai_output_str))
 
             except Exception as e:
-                print(f"Error invoking agent: {e}")
-                app_logger.error(f"Error in __main__ agent test: {e}", exc_info=True)
+                print(f"⚠️ Error during agent invocation: {e}")
+                app_logger.error(f"Error in __main__ agent test invocation: {e}", exc_info=True)