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:

# Function to 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]}
 
# Function to 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
 
# Function to handle tool authorization
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"]}

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.