107 lines
3.8 KiB
Python
107 lines
3.8 KiB
Python
import os
|
|
from datetime import datetime
|
|
|
|
from configuration import AgentConfigurable
|
|
from langchain_arcade import ArcadeToolManager
|
|
from langchain_core.messages import HumanMessage, ToolMessage
|
|
from langchain_core.prompts import ChatPromptTemplate
|
|
from langchain_openai import ChatOpenAI
|
|
from langgraph.graph import END, START, MessagesState, StateGraph
|
|
from langgraph.prebuilt import ToolNode
|
|
|
|
# Initialize the Arcade Tool Manager with your API key
|
|
arcade_api_key = os.getenv("ARCADE_API_KEY")
|
|
openai_api_key = os.getenv("OPENAI_API_KEY")
|
|
|
|
toolkit = ArcadeToolManager(api_key=arcade_api_key)
|
|
# Retrieve tools compatible with LangGraph
|
|
tools = toolkit.get_tools(langgraph=True)
|
|
tool_node = ToolNode(tools)
|
|
|
|
PROMPT_TEMPLATE = f"""
|
|
You are a helpful assistant who can use tools to help users with tasks
|
|
Today's date is {datetime.now().strftime("%Y-%m-%d")}
|
|
|
|
ALL RESPONSES should be in plain text and not markdown.
|
|
"""
|
|
# prompt for the main agent
|
|
prompt = ChatPromptTemplate.from_messages([
|
|
("system", PROMPT_TEMPLATE),
|
|
("placeholder", "{messages}"),
|
|
])
|
|
# Initialize the language model with your OpenAI API key
|
|
model = ChatOpenAI(model="gpt-4o", api_key=openai_api_key).bind_tools(tools)
|
|
prompted_model = prompt | model
|
|
|
|
|
|
def call_agent(state):
|
|
"""Define the agent function that invokes the model"""
|
|
messages = state["messages"]
|
|
# replace placeholder with messages from state
|
|
response = prompted_model.invoke({"messages": messages})
|
|
return {"messages": [response]}
|
|
|
|
|
|
def should_continue(state: MessagesState, config: dict):
|
|
"""Function to determine the next step based on the model's response"""
|
|
last_message = state["messages"][-1]
|
|
if last_message.tool_calls:
|
|
user_id = config["configurable"].get("user_id")
|
|
tool_name = state["messages"][-1].tool_calls[0]["name"]
|
|
auth_response = toolkit.authorize(tool_name, user_id)
|
|
if auth_response.status == "completed":
|
|
return "tools"
|
|
else:
|
|
# If the tool requires authorization, proceed to the authorization step
|
|
return "authorization"
|
|
# If no tool calls are present, end the workflow
|
|
return END
|
|
|
|
|
|
def wait_for_auth(state: MessagesState):
|
|
last_message = state["messages"][-1]
|
|
if isinstance(last_message, HumanMessage):
|
|
return "agent"
|
|
return "tools"
|
|
|
|
|
|
def authorize(state: MessagesState, config: dict):
|
|
"""Function to handle tool authorization"""
|
|
user_id = config["configurable"].get("user_id")
|
|
tool_name = state["messages"][-1].tool_calls[0]["name"]
|
|
auth_response = toolkit.authorize(tool_name, user_id)
|
|
|
|
auth_message = (
|
|
f"Please authorize the application in your browser:\n\n {auth_response.authorization_url}"
|
|
)
|
|
tool_call_id = state["messages"][-1].tool_calls[0]["id"]
|
|
response = ToolMessage(
|
|
content=auth_message,
|
|
tool_call_id=tool_call_id,
|
|
)
|
|
# Add the new message to the message history and add a new human message
|
|
# saying that the agent should try again
|
|
try_message = HumanMessage(
|
|
content="Please try the previous tool call again now that you are authorized."
|
|
)
|
|
return {"messages": [response, try_message]}
|
|
|
|
|
|
# Build the workflow graph
|
|
workflow = StateGraph(MessagesState, AgentConfigurable)
|
|
|
|
# Add nodes to the graph
|
|
workflow.add_node("agent", call_agent)
|
|
workflow.add_node("tools", tool_node)
|
|
workflow.add_node("authorization", authorize)
|
|
# workflow.add_node("wait_for_auth", wait_for_auth)
|
|
|
|
# Define the edges and control flow
|
|
workflow.add_edge(START, "agent")
|
|
workflow.add_conditional_edges("agent", should_continue, ["authorization", "tools", END])
|
|
workflow.add_edge("authorization", "agent")
|
|
workflow.add_edge("tools", "agent")
|
|
|
|
# Compile the graph with an interrupt after the authorization node
|
|
# so that we can prompt the user to authorize the application
|
|
graph = workflow.compile(interrupt_after=["authorization"])
|