Update agent.py

agent.py

@@ -278,22 +278,22 @@ class BasicAgent:
                 If Task ID is included in the question, remember to call the relevant read tools [ie. read_file, excel_read, csv_read, mp3_listen, image_caption]
         """
         self.tools = [duckduck_websearch, serper_websearch, visit_webpage, wiki_search, text_splitter, youtube_transcript, read_file, excel_read, csv_read, mp3_listen, image_caption, python_tool]
-        self.binded_model = self.model.bind_tools(self.tools)
+        self.model_with_tools = self.model.bind_tools(self.tools)
         self.sys_msg = SystemMessage(content=self.sys_prompt)
         
         self.prompt = ChatPromptTemplate.from_messages([
             ("system", self.sys_prompt),
             ("human", "{input}")
         ])
-        self.agent = initialize_agent(
-            tools=self.tools,
-            llm=self.model,
-            agent="zero-shot-react-description",  # ReAct agent type
-            verbose=True,
-            system_prompt=self.prompt,
-            handle_parsing_errors="Check your output and make sure it conforms, use the Action/Action Input syntax"
-        )
-        self.app = self.__graph_compile__()
+        # self.agent = initialize_agent(
+        #     tools=self.tools,
+        #     llm=self.model,
+        #     agent="zero-shot-react-description",  # ReAct agent type
+        #     verbose=True,
+        #     system_prompt=self.prompt,
+        #     handle_parsing_errors="Check your output and make sure it conforms, use the Action/Action Input syntax"
+        # )
+        self.graph = self.__graph_compile__()
         print("BasicAgent initialized.")
     
     def __call__(self, task: dict) -> str:
@@ -310,15 +310,19 @@ class BasicAgent:
         # fixed_answer = "This is a default answer."
         # fixed_answer = self.agent.run(question)
         human_message = [HumanMessage(content=question)]
-        messages = self.app.invoke({"message": human_message})
+        messages = self.graph.invoke({"messages": human_message})
         fixed_answer = messages['messages'][-1].content
         print(f"Agent returning fixed answer: {fixed_answer}")
         time.sleep(60)
         return fixed_answer
 
     def __graph_compile__(self):
+        def assistant(self, state: MessagesState):
+            """Assistant Node"""
+            return {"message": [self.model_with_tools.invoke(state["messages"])]}
+        
         builder = StateGraph(MessagesState)
-        builder.add_node("assistant", self.assistant)
+        builder.add_node("assistant", assistant)
         builder.add_node("tools", ToolNode(self.tools))
         builder.add_edge(START, "assistant")
         builder.add_conditional_edges(
@@ -328,8 +332,7 @@ class BasicAgent:
             tools_condition,
         )
         builder.add_edge("tools", "assistant")
-        return builder.compile() # Compile graph
+        # Compile graph
+        return builder.compile()
 
-    def assistant(self, state: MessagesState):
-        """Assistant Node"""
-        return {"message": [self.binded_model.invoke([self.sys_msg] + state["message"])]}
+