| import os |
| from langgraph.graph.message import add_messages |
| from langchain_core.messages import AnyMessage, HumanMessage, AIMessage, SystemMessage |
| 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 |
|
|
|
|
| class AgentState(TypedDict): |
| messages: Annotated[list[AnyMessage], add_messages] |
|
|
|
|
| class BasicAgent(): |
| def __init__(self, chat, vision_llm): |
| extract_text_from_image = ExtractTextFromImage(vision_llm) |
| describe_image = DescribeImage(vision_llm) |
|
|
| 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() |
|
|
| def _initialize_graph(self): |
| builder = StateGraph(AgentState) |
|
|
| |
| builder.add_node("assistant", self.assistant) |
| builder.add_node("tools", ToolNode(self.tools)) |
|
|
| |
| builder.add_edge(START, "assistant") |
| builder.add_conditional_edges("assistant",tools_condition) |
| builder.add_edge("tools", "assistant") |
|
|
| |
| self.agent = builder.compile() |
| print("Agent initialized.") |
|
|
| 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 |
| ) |
| print("Telemetry initialized.") |
|
|
| def __call__(self, question: str, file_name : str) -> str: |
|
|
| sys_msg = SystemMessage(content=f""" |
| You are a general AI assistant. I will ask you a question. Reason step by step and search for the information you need using available tools. |
| Finish your answer with the following template: FINAL ANSWER: [YOUR FINAL ANSWER]. |
| YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string. |
| When providing the final answer, ONLY give [YOUR FINAL ANSWER]. Do not add anything else, no additional motivation or explanation, and do not return 'FINAL ANSWER:'. |
| """) |
| |
| print(f"Agent received question: {question}.") |
| |
| if file_name is not None and file_name!='': |
| print(f"Provided file: {file_name}.") |
| messages=[[sys_msg] + HumanMessage(content=f"{question}. The file you have access to is {file_name}.")] |
| else: |
| messages=[[sys_msg] + HumanMessage(content=question)] |
| |
| response = self.agent.invoke({"messages":messages}, config={"callbacks": [self.langfuse_handler]}) |
| answer = response['messages'][-1].content |
| print(f"Agent returning answer: {answer}") |
| return answer |
|
|
| def assistant(self, state: AgentState): |
| response = self.chat_with_tools.invoke(state["messages"]) |
| return { |
| "messages": state["messages"] + [response], |
| } |
|
|
