Update agent.py

agent.py

@@ -1,15 +1,17 @@
 import os
+from typing import TypedDict, Annotated, Optional
 from langgraph.graph.message import add_messages
 from langchain_core.messages import AnyMessage, HumanMessage, AIMessage, SystemMessage
+from langchain_openai import ChatOpenAI
+from langchain.tools import tool
 from langgraph.prebuilt import ToolNode
 from langgraph.graph import START, StateGraph
 from langgraph.prebuilt import tools_condition
 from langchain_huggingface import HuggingFaceEndpoint, ChatHuggingFace
-from tools import ExtractTextFromImage, DescribeImage, read_excel, read_python, wiki_search, web_search, arxiv_search
-from langchain_openai import ChatOpenAI
-from typing import TypedDict, Annotated, Optional
 from langfuse.callback import CallbackHandler
 
+from tools import ExtractTextFromImage, DescribeImage, read_excel, read_python, wiki_search, web_search, arxiv_search
+
 
 class AgentState(TypedDict):
     messages: Annotated[list[AnyMessage], add_messages]
@@ -17,10 +19,10 @@ class AgentState(TypedDict):
 
 class BasicAgent():
     def __init__(self, chat, vision_llm):
-        extract_text_from_image = ExtractTextFromImage(vision_llm)
-        describe_image = DescribeImage(vision_llm)
+        extract_text_from_image = tool(ExtractTextFromImage(vision_llm).__call__)
+        describe_image = tool(DescribeImage(vision_llm).__call__)
 
-        self.tools = [extract_text_from_image.__call__, describe_image.__call__, read_excel, read_python, wiki_search, web_search, arxiv_search]
+        self.tools = [extract_text_from_image, describe_image, read_excel, read_python, wiki_search, web_search, arxiv_search]
         self.chat_with_tools = chat.bind_tools(self.tools)
         self._initialize_graph()
         self._initialize_telemetry()