Update agent.py

agent.py

@@ -2,20 +2,27 @@
 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
+# --- Import your defined tools FROM THE 'tools' PACKAGE ---
+from tools import ( # This now relies on tools/__init__.py
+    BioPortalLookupTool,
+    # GeminiTool, # Uncomment if you decide to use it
+    UMLSLookupTool,
+    QuantumTreatmentOptimizerTool,
+    # You might not need to import the Input schemas here if the tools use them internally
+    # and the agent only needs the tool instances.
+    # BioPortalInput, GeminiInput, UMLSInput, QuantumOptimizerInput
+)
+# If QuantumOptimizerInput is needed for type hinting or direct use in agent.py:
+from tools import QuantumOptimizerInput # Or from tools.quantum_treatment_optimizer_tool import QuantumOptimizerInput
 
 from config.settings import settings
 from services.logger import app_logger
 
+# ... (rest of your agent.py file) ...
+
 # --- Initialize LLM (Gemini) ---
 try:
     # Ensure GOOGLE_API_KEY is set in your environment (HuggingFace Secrets)