Update agent.py

agent.py

@@ -7,6 +7,7 @@ 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
 
 
 class AgentState(TypedDict):
@@ -21,6 +22,7 @@ class BasicAgent():
         self.tools = [extract_text_from_image.__call__, describe_image.__call__, 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()
         print("BasicAgent initialized.")
 
     
@@ -39,6 +41,17 @@ class BasicAgent():
         # Compile the graph
         self.agent = builder.compile()
 
+    def _initialize_telemetry(self):
+        LANGFUSE_PUBLIC_KEY = os.getenv("LANGFUSE_PUBLIC_KEY")
+        LANGFUSE_SECRET_KEY = os.getenv("LANGFUSE_SECRET_KEY")
+        LANGFUSE_HOST = "https://cloud.langfuse.com"
+
+        self.langfuse_handler = CallbackHandler(
+            public_key=LANGFUSE_PUBLIC_KEY,
+            secret_key=LANGFUSE_SECRET_KEY,
+            host=LANGFUSE_HOST
+        )
+
     
     def __call__(self, question: str, file_name : str) -> str:
         print(f"Agent received question: {question}. /nProvided file: {file_name}.")
@@ -46,7 +59,7 @@ class BasicAgent():
             messages=[HumanMessage(content=f"{question}. The filename you have access to is {file_name}.")]
         else:
             messages=[HumanMessage(content=question)]
-        response = self.agent.invoke({"messages":messages})
+        response = self.agent.invoke({"messages":messages}, config={"callbacks": [self.langfuse_handler]})
         answer = response['messages'][-1].content
         print(f"Agent returning answer: {answer}")
         return answer