LangGraph

Run LangGraph graphs on Flyte. This adapter is shaped differently from the others: LangGraph is a framework for building your own control flow, so rather than hiding the graph behind run_agent, it gives you durable node factories and expects you to wire them up yourself.

You build the StateGraph. ai_node and tool_node are the pieces Flyte makes durable and observable.

Installation

pip install flyteplugins-agents-langgraph

Requires Python 3.10 or later. Install the provider integration you need alongside it, for example langchain-openai.

Quick start

Build the graph with the two node factories, compile it, and hand the compiled graph to run_agent.

import flyte
from flyteplugins.agents.langgraph import ai_node, run_agent, tool, tool_node

env = flyte.TaskEnvironment(
    "langgraph-agent",
    secrets=[flyte.Secret(key="openai_api_key", as_env_var="OPENAI_API_KEY")],
    image=flyte.Image.from_debian_base().with_pip_packages(
        "flyteplugins-agents-langgraph", "langchain-openai",
    ),
)


@tool
@env.task(cache="auto", retries=3)
async def get_weather(city: str) -> str:
    """Get the current weather for a city."""
    return f"The weather in {city} is sunny, 22C."


def build_city_graph():
    """A standard tool-calling loop: ai, then tools, then back to ai."""
    from langchain_openai import ChatOpenAI
    from langgraph.graph import START, MessagesState, StateGraph
    from langgraph.prebuilt import tools_condition

    tools = [get_weather]
    builder = StateGraph(MessagesState)
    builder.add_node("ai", ai_node(ChatOpenAI(model="gpt-4o"), tools))
    builder.add_node("tools", tool_node(tools))
    builder.add_edge(START, "ai")
    builder.add_conditional_edges("ai", tools_condition)
    builder.add_edge("tools", "ai")
    return builder.compile()


@env.task(report=True, retries=3)
async def city_agent(city: str) -> str:
    return await run_agent(
        f"What's the weather in {city}?",
        agent=build_city_graph(),
    )

Build the graph inside the task, where the provider API key is available.

The node factories

ai_node(model, tools, *, name="ai", durable=True, observability=True)

The model-calling node. It binds the tools to your chat model and runs one turn over state["messages"], appending the response. With durable=True, each turn is recorded as a flyte.trace leaf keyed by a fingerprint of the message list, so a retry replays the recorded response instead of calling the model again.

Returns an async node with the signature state -> {"messages": [ai_message]}.

tool_node(tools, *, name="tools", observability=True)

The tool-executing node. It reads the tool calls off the last message and runs each one, appending a ToolMessage per call. Tools wrapped with tool run as durable Flyte child actions; anything else runs as the tool defines. Tool errors are caught and surfaced back to the model as the tool result rather than failing the node.

Returns an async node with the signature state -> {"messages": [tool_message, ...]}.

Both render their activity into the task report.

How it maps to Flyte

Tools: tool turns an @env.task into a LangChain StructuredTool. It is a first-class LangGraph tool, so it works with model.bind_tools(...), with tool_node, and with LangGraph’s own ToolNode. The args schema comes from the task’s typed signature.

Model turns: Durability lives in ai_node, not in a model wrapper. That means any chat model works, including a provider:model string, and durability applies uniformly.

Skipping the graph

If you do not need a custom topology, pass tools and a model instead of agent, and run_agent assembles the same tool-calling loop from the same two factories.

@env.task(report=True, retries=3)
async def quick_city_agent(city: str) -> str:
    from langchain_openai import ChatOpenAI

    return await run_agent(
        f"What's the weather in {city}?",
        tools=[get_weather],
        model=ChatOpenAI(model="gpt-4o"),
        instructions="You are a concise assistant. Use the tools to answer.",
    )

model accepts a chat model instance or a provider:model string. The string form is resolved through init_chat_model, which requires the langchain package.

agent and tools are mutually exclusive.

Custom state

input accepts a full graph input state as a dict, not just a prompt string, so a graph with a state schema beyond MessagesState works.

await run_agent({"messages": [...], "budget": 3}, agent=graph)

When memory is in play, prior messages are merged into the messages key of whatever state you pass.

Memory

await run_agent(question, agent=graph, memory_key="user-alice")

The conversation transcript is persisted to a durable, keyed MemoryStore and prepended to the graph’s messages on the next run with the same key. On a resumed run the system prompt is not re-added, since it already lives in the prior transcript.

run_agent parameters

Returns the final assistant message as a string. Use run_agent_sync with the same signature from a sync task.

Examples

Full runnable examples live in the SDK repository under plugins/agents/langgraph/examples:

• langgraph_custom_agent.py: building a StateGraph from ai_node and tool_node, plus the default-graph shortcut.
• langgraph_durable_agent.py: a single durable agent with traced model turns.
• langgraph_multi_agent.py: a planner, parallel researchers and an editor, each its own durable action.
• langgraph_crash_resume.py: the task crashes on its first attempt and replays completed turns on retry.
• langgraph_memory.py: two separate runs sharing a memory_key.