Update agent.py

agent.py

@@ -9,7 +9,7 @@ from pydantic import ValidationError
 from langchain.agents import AgentExecutor
 from langchain.agents.mrkl.base import ZeroShotAgent
 from prompts import FORMAT_INSTRUCTIONS, QUESTION_PROMPT, QUESTION_PROMPT1, SUFFIX
-from tools import make_tools
+from tools import make_tools , drug_tools
 import os
 from rmrkl import ChatZeroShotAgent, RetryAgentExecutor
 from langchain_ollama import OllamaLLM
@@ -96,7 +96,10 @@ class TeLLAgent:
             api_keys["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY")
             tools_llm = _make_llm(tools_model, temp, openai_api_key, streaming)
             tools = make_tools(tools_llm, api_keys=api_keys, verbose=verbose, image_path = image_path, file_path = file_path)
- 
+        if tools == 'drug':
+            api_keys["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY")
+            tools_llm = _make_llm(tools_model, temp, openai_api_key, streaming)
+            tools = drug_tools(tools_llm, api_keys=api_keys, verbose=verbose, image_path = image_path, file_path = file_path)
         # Initialize agent
         self.agent_executor1 = RetryAgentExecutor.from_agent_and_tools(
             tools=tools,
@@ -133,6 +136,4 @@ class TeLLAgent:
         else: 
             prompt = str(' ' + outputs["input"] + ' ' + outputs["intermediate_steps"][0][0].log.split('Action')[0].replace("*", ""))
             outputs = self.agent_executor2.invoke( {"input": prompt})
-        return outputs['output']
-    
- 
+        return outputs['output']