Home
LangGraph
LangGraph with authorization

Create a custom LangGraph with authorization

In this guide, we'll walk through how to create a custom LangGraph that prompts the user for authorization of tool calls using Arcade AI tools.

Prerequisites

  • Set up Arcade AI

  • Install the required packages:

    pip install arcade-ai langgraph langchain-openai langchain-arcade

Import the necessary packages

Begin by importing the required libraries:

import os
import time
from langchain_arcade import ArcadeToolManager
from langchain_core.messages import HumanMessage
from langchain_openai import ChatOpenAI
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import END, START, MessagesState, StateGraph
from langgraph.prebuilt import ToolNode

Set up API keys

Ensure your environment variables ARCADE_API_KEY and OPENAI_API_KEY are set with your actual API keys:

arcade_api_key = os.environ["ARCADE_API_KEY"]
openai_api_key = os.environ["OPENAI_API_KEY"]

Initialize the Arcade Tool Manager

Initialize the ArcadeToolManager to fetch and manage tools from Arcade AI:

tool_manager = ArcadeToolManager(api_key=arcade_api_key)

Retrieve tools requiring authorization

Fetch tools from the Github toolkit and wrap them as LangGraph tools:

tools = tool_manager.get_tools(
    toolkits=["Github"],
    langgraph=True,  # Use LangGraph-specific behavior
)
tool_node = ToolNode(tools)

Create the language model

Create an instance of the AI language model and bind it with the tools:

model = ChatOpenAI(model="gpt-4o", api_key=openai_api_key)
model_with_tools = model.bind_tools(tools)

Define functions for the graph nodes

Set up functions that will act as nodes in your LangGraph:

# Invoke the model and get a response
def call_agent(state):
    messages = state["messages"]
    response = model_with_tools.invoke(messages)
    # Return the updated message history
    return {"messages": [*messages, response]}
 
# Determine the next step based on the last message
def should_continue(state: MessagesState):
    last_message = state["messages"][-1]
    if last_message.tool_calls:
        tool_name = last_message.tool_calls[0]["name"]
        if tool_manager.requires_auth(tool_name):
            return "authorization"  # Proceed to authorization if required
        else:
            return "tools"  # Proceed to tool execution if no authorization is needed
    return END  # End the workflow if no tool calls are present
 
# Handle tool authorization and prompt for authorization only when necessary
def authorize(state: MessagesState, config: dict):
    user_id = config["configurable"].get("user_id")
    tool_name = state["messages"][-1].tool_calls[0]["name"]
    auth_response = tool_manager.authorize(tool_name, user_id)
    if auth_response.status == "completed":
        # Authorization completed successfully; continue
        return {"messages": state["messages"]}
    else:
        # Prompt the user to visit the authorization URL
        print(f"Visit the following URL to authorize: {auth_response.authorization_url}")
        # Wait until authorization is completed
        while not tool_manager.is_authorized(auth_response.authorization_id):
            time.sleep(1)
        return {"messages": state["messages"]}

Structuring the nodes in this way allows for users who have already authorized the tool to skip the authorization step so they won't have to login every time they run the tool.

Build the workflow graph

Construct the LangGraph by adding nodes and defining the control flow:

# Build the workflow graph using StateGraph
workflow = StateGraph(MessagesState)
 
# Add nodes (steps) to the graph
workflow.add_node("agent", call_agent)
workflow.add_node("tools", tool_node)
workflow.add_node("authorization", authorize)
 
# Define the edges and control flow between nodes
workflow.add_edge(START, "agent")
workflow.add_conditional_edges("agent", should_continue, ["authorization", "tools", END])
workflow.add_edge("authorization", "tools")
workflow.add_edge("tools", "agent")

Set up memory for checkpointing

Initialize memory for checkpointing the state of the graph:

# Set up memory for checkpointing the state
memory = MemorySaver()
 
# Compile the graph with the checkpointer
graph = workflow.compile(checkpointer=memory)

Define input messages and configuration

Set up the initial user message and configuration parameters:

# Define the input messages from the user
inputs = {
    "messages": [HumanMessage(content="Star arcadeai/arcade-ai on GitHub!")],
}
 
# Configuration with thread and user IDs for authorization purposes
config = {
    "configurable": {
        "thread_id": "4",
        "user_id": "[email protected]",
    }
}

Run the LangGraph and handle authorization

Execute the graph and stream the outputs, handling authorization as needed:

# Run the graph and stream the outputs
for chunk in graph.stream(inputs, config=config, stream_mode="values"):
    # Access the latest message from the conversation
    last_message = chunk["messages"][-1]
    # Print the assistant's message content
    print(last_message.content)

If a tool requires authorization, the program will prompt you to authorize it by visiting a URL. Once authorized, the graph continues execution.

Tips for handling authorization

  • User IDs: Ensure that user_id is unique for each user to manage permissions accurately.