Update agent.py

agent.py

@@ -12,11 +12,15 @@ from prompts import FORMAT_INSTRUCTIONS, QUESTION_PROMPT, QUESTION_PROMPT1, SUFF
 from tools import make_tools
  
 from rmrkl import ChatZeroShotAgent, RetryAgentExecutor
- 
+from langchain_ollama import OllamaLLM
 import base64
 from io import BytesIO
 from PIL import Image
+ 
 from langchain_openai import ChatOpenAI , OpenAI
+from langchain.agents import load_tools, initialize_agent, AgentType
+from langchain.llms import OpenAI
+
 
 def convert_to_base64(pil_image):
     buffered = BytesIO()
@@ -25,7 +29,7 @@ def convert_to_base64(pil_image):
     return img_str
 
 def _make_llm(model, temp, api_key, streaming: bool = False):
-    if model.startswith("claude") or model.startswith("gpt-3"):
+    if model.startswith("claude")  :
         llm = OpenAI(
        temperature=temp,
        model_name=model,
@@ -33,7 +37,7 @@ max_tokens = 5000,
        openai_api_key=api_key,
        base_url="https://www.dmxapi.com/v1"
    )
-    elif model.startswith("gpt-4o-2024-11-20") or model.startswith("deepseek"):
+    elif model.startswith("gpt") or model.startswith("deepseek"):
         llm = ChatOpenAI(model=model,
             temperature = 0.1,
             
@@ -42,7 +46,10 @@ max_tokens = 5000,
             callbacks=[StreamingStdOutCallbackHandler()],
             openai_api_key=api_key,base_url="https://www.dmxapi.com/v1"
             )
-   
+    elif model.startswith("llama") :
+            llm = OllamaLLM(model=model,
+                temperature = 0.1,    
+                )
     else:
         raise ValueError(f"Invalid model name: {model}")
     return llm
@@ -53,7 +60,7 @@ class TeLLAgent:
         self,
         tools=None,
         model1: str = "deepseek-ai/DeepSeek-R1",
-        model2: str = "deepseek-ai/DeepSeek-V3",
+        model2: str = "gpt-4o-2024-11-20",
         tools_model="gpt-4o-2024-11-20",
         temp=0.1,
         max_iterations=50,
@@ -99,15 +106,15 @@ class TeLLAgent:
                 tools,
                 suffix=SUFFIX,
                 format_instructions=FORMAT_INSTRUCTIONS,
-                question_prompt=QUESTION_PROMPT,return_intermediate_steps=True ,handle_parsing_errors=True 
+                question_prompt=QUESTION_PROMPT  ,handle_parsing_errors=True 
             ),
             verbose=True,
-            max_iterations=max_iterations , return_intermediate_steps=True  ,handle_parsing_errors=True  
+            max_iterations=max_iterations    ,handle_parsing_errors=True  
         )
  
 
     def run(self, prompt):
-         
+        prompt = prompt + ' ' + str(self.file_path) + ' ' + str(self.image_path)
         outputs = self.agent_executor1.invoke( {"input": prompt})
         if outputs["intermediate_steps"] ==[]:
             prompt = str(' ' +outputs["input"]+ ' ' + outputs["output"].split('Action:')[0] )
@@ -115,11 +122,12 @@ class TeLLAgent:
         else: 
             prompt = str(' ' + outputs["input"] + ' ' + outputs["intermediate_steps"][0][0].log.split('Action:')[0])
             outputs = self.agent_executor2.invoke( {"input": prompt})
-        return outputs 
-
+        return outputs['output']
+    
 if __name__ == '__main__':
-        chem_model = TeLLAgent( temp=0.1, streaming=False,
+        chem_model = TeLLAgent( temp=0.1, streaming=  True,
                            openai_api_key =r'sk-itPrztYm9F6XZZpsBMJB9O7Vq0pYUABVVBSoThuBxEGTnDik', 
-                           image_path= r"C:\Users\BM109X32G-10GPU-02\Pictures\1735356359936.jpg"
-                           )
-        chem_model.run(r"""what is Y20""")
+                           image_path= r"C:\Users\BM109X32G-10GPU-02\Desktop\acceptor\1.png" ,file_path = r"..." )
+       # A = chem_model.run(r"""who are you""")
+       # A = chem_model.run(r"""I want to know some basic chemical properties, HOMO /LUMO and PCE values of molecules  CC(C)CCCC(C)CCC1=C(/C=C\2/C(=C(C#N)C#N)C3=C(C=C(C(=C3)F)F)C2=O)SC4=C1N(CCC(C)CCCC(C)C)C5=C4C6=NSN=C6C7=C5N(CCC(C)CCCC(C)C)C8=C7SC(=C8CCC(C)CCCC(C)C)/C=C\9/C(=C(C#N)C#N)C%10=C(C=C(C(=C%10)F)F)C9=O""")
+        a  = chem_model.run(r""" Compare the PCE and similarity of acceptor Y6 and acceptor in image""")