File size: 5,587 Bytes
7a40d3a a741f9e 7a40d3a a741f9e 7a40d3a a741f9e 7a40d3a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 | import os
from typing import Annotated, TypedDict
from langchain.tools import tool
from langchain_core.messages import AnyMessage, HumanMessage, SystemMessage
from langchain_openai import ChatOpenAI
from langfuse import Langfuse
from langfuse.langchain import CallbackHandler
from langgraph.graph import START, StateGraph
from langgraph.graph.message import add_messages
from langgraph.prebuilt import ToolNode, tools_condition
try:
from tools import (
DescribeImage,
ExtractTextFromImage,
arxiv_search,
download_youtube_video,
extract_audio_from_video,
read_excel,
read_python,
transcribe_audio,
web_search,
wiki_search,
add,
divide,
multiply,
)
except:
from .tools import (
DescribeImage,
ExtractTextFromImage,
arxiv_search,
download_youtube_video,
extract_audio_from_video,
read_excel,
read_python,
transcribe_audio,
web_search,
wiki_search,
add,
divide,
multiply,
)
class AgentState(TypedDict):
"""Class representing the state for agent graph."""
messages: Annotated[list[AnyMessage], add_messages]
class SmartAgent:
def __init__(self, chat):
"""Initialize agent, multimodal model and tools."""
self.multimodal_model = ChatOpenAI(model="gpt-4o")
extract_text_from_image = tool(
ExtractTextFromImage(self.multimodal_model).__call_extract_text_from_image__
)
describe_image = tool(
DescribeImage(self.multimodal_model).__call_describe_image__
)
self.tools = [
extract_text_from_image,
describe_image,
transcribe_audio,
read_excel,
read_python,
wiki_search,
web_search,
arxiv_search,
download_youtube_video,
extract_audio_from_video,
add,
divide,
multiply,
]
self.chat_with_tools = chat.bind_tools(self.tools)
self._initialize_graph()
self._initialize_telemetry()
def _initialize_graph(self):
"""Initialize and compile the agent graph."""
builder = StateGraph(AgentState)
# Define nodes
builder.add_node("assistant", self.assistant)
builder.add_node("tools", ToolNode(self.tools))
# Define edges
builder.add_edge(START, "assistant")
builder.add_conditional_edges("assistant", tools_condition)
builder.add_edge("tools", "assistant")
# Compile the graph
self.agent = builder.compile()
print("Agent initialized.")
def _initialize_telemetry(self):
"""Initialize langfuse telemetry using CallbackHandler."""
LANGFUSE_PUBLIC_KEY = os.getenv("LANGFUSE_PUBLIC_KEY")
LANGFUSE_SECRET_KEY = os.getenv("LANGFUSE_SECRET_KEY")
LANGFUSE_HOST = "https://cloud.langfuse.com"
langfuse = Langfuse(
public_key=LANGFUSE_PUBLIC_KEY,
secret_key=LANGFUSE_SECRET_KEY,
host=LANGFUSE_HOST, # or your custom host if applicable
)
# Create a Langchain callback handler using the initialized client
self.langfuse_handler = CallbackHandler()
print("Telemetry initialized.")
def __call__(self, question: str, file_name: str | None = None) -> str:
"""Call the agent, passing system prompt and eventual file name."""
sys_msg = SystemMessage(
content="""You are a general AI assistant. You will be asked a factual question.
1. Reason step by step and search for the information using available tools if needed.
2. Finish your response with this exact format:
FINAL ANSWER: [YOUR FINAL ANSWER]
IMPORTANT RULES for [YOUR FINAL ANSWER]:
- If the answer is a number, provide only the number, with no commas, units, or symbols, do not write it as a string.
- If the answer is a string, provide only the core noun phrase with no articles or abbreviations.
- If the answer is a list, return a comma-separated list applying the above rules per item.
- DO NOT include any other text before or after the final answer.
- DO NOT explain or justify the answer after it is given.
- DO NOT repeat the question.
- DO NOT include the words 'FINAL ANSWER: '.
Strictly follow these formatting rules.
"""
)
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):
"""Assistant node which calls the model initialized with tools."""
response = self.chat_with_tools.invoke(state["messages"])
return {
"messages": state["messages"] + [response],
}
|