# /home/user/app/agent.py import os from langchain_google_genai import ChatGoogleGenerativeAI 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 # --- 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) --- 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-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 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(f"ChatGoogleGenerativeAI ({llm.model}) initialized successfully.") except Exception as e: app_logger.error(f"Failed to initialize ChatGoogleGenerativeAI: {e}", exc_info=True) raise ValueError(f"Gemini LLM initialization failed: {e}. Check API key and configurations in HF Secrets.") # --- Initialize Tools --- # 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 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}" ) prompt = ChatPromptTemplate.from_messages([ ("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"), ]) app_logger.info("Agent prompt template created.") # --- Create Agent --- try: agent = create_structured_chat_agent(llm=llm, tools=tools, prompt=prompt) 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"Gemini agent creation failed: {e}") # --- Create Agent Executor --- agent_executor = AgentExecutor( agent=agent, tools=tools, 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 with Gemini agent created successfully.") # --- Getter Function for Streamlit App --- def get_agent_executor(): """Returns the configured agent executor for Gemini.""" # 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 of this agent.py file) --- if __name__ == "__main__": 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("\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() 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_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_str, "chat_history": current_chat_history_for_test, "patient_context": test_patient_context_summary # Passing the context }) ai_output_str = response.get('output', "Agent did not produce an output.") print(f"🤖 Agent: {ai_output_str}") 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 during agent invocation: {e}") app_logger.error(f"Error in __main__ agent test invocation: {e}", exc_info=True)