Behaviour of checkpointer with create_agent and subagents

Please take a look at this example (example taken from another thread from this forum).

  • Has a Supervisor Agent
  • which has 2 tools , and each tool has an agent

Scenario 1:

InMemorySaver checkpointer defined on Supervisor agent, Meal Agent & Weekend agent

Behavior - >

  • Only the Supervisor agent has values in the state
  • Meal Agent and Weekend agent have no values
  • StateSnapshot with subgraphs=True is same as without subgraphs

Scenario 2:

InMemorySaver Checkpointer defined only on Supervisor agent

Behavior - >

  • Only the Supervisor agent has values in the state
  • Meal Agent and Weekend check throws exception about missing checkpointer

Scenario 3:

InMemorySaver Checkpointer defined on Supervisor agent, Meal Agent and Weekend Agent. All three have a different thread_id

Behavior - >

  • All agent has values in the StateSnapshot

Is this how the Checkpointers supposed to work?

  • The thread ids needs to be unique across subagents
  • checkpointer needs to be defined for each subagent
  • subgraphs=True will not return the values

Code:

import logging
import random
from datetime import datetime

from langchain.agents import create_agent
from langchain_core.messages import HumanMessage
from langchain_core.tools import tool
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.prebuilt import ToolRuntime

from config import llm_config

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)


# ----------------------------------------------------------------------------------
# SUB-AGENT 1: Activity planning agent
# ----------------------------------------------------------------------------------
@tool
def get_weather(city: str, date: str) -> dict:
    """Returns weather data for a given city and date."""
    logger.info(f"Getting weather for {city} on {date}")
    if random.random() < 0.05:
        return {
            "temperature": 72,
            "description": "Sunny",
        }
    else:
        return {
            "temperature": 60,
            "description": "Rainy",
        }


@tool
def get_activities(city: str, date: str) -> list:
    """Returns a list of activities for a given city and date."""
    logger.info(f"Getting activities for {city} on {date}")
    return [
        {"name": "Hiking", "location": city},
        {"name": "Beach", "location": city},
        {"name": "Museum", "location": city},
    ]


@tool
def get_current_date() -> str:
    """Gets the current date from the system and returns as a string in format YYYY-MM-DD."""
    logger.info("Getting current date")
    return datetime.now().strftime("%Y-%m-%d")


weekend_agent = create_agent(
    model=llm_config.get_default_model(),
    system_prompt=(
        "You help users plan their weekends and choose the best activities for the given weather."
        "If an activity would be unpleasant in the weather, don't suggest it."
        "Include the date of the weekend in your response."
    ),
    tools=[get_weather, get_activities, get_current_date],
    checkpointer=InMemorySaver(),
)


@tool
def weekend_agent_tool(query: str, runtime: ToolRuntime) -> str:
    """Invoke the activity planning agent and return its final response as plain text."""
    logger.info("Tool:weekend_agent invoked")
    response = weekend_agent.invoke(
        {"messages": [HumanMessage(content=query)]}, config=runtime.config
    )
    final = response["messages"][-1].content
    return final


# ----------------------------------------------------------------------------------
# SUB-AGENT 2: Recipe planning agent
# ----------------------------------------------------------------------------------


@tool
def find_recipes(query: str) -> list[dict]:
    """Returns recipes based on a query."""
    logger.info(f"Finding recipes for '{query}'")
    if "pasta" in query.lower():
        return [
            {
                "title": "Pasta Primavera",
                "ingredients": ["pasta", "vegetables", "olive oil"],
                "steps": ["Cook pasta.", "Sauté vegetables."],
            }
        ]
    elif "tofu" in query.lower():
        return [
            {
                "title": "Tofu Stir Fry",
                "ingredients": ["tofu", "soy sauce", "vegetables"],
                "steps": ["Cube tofu.", "Stir fry veggies."],
            }
        ]
    else:
        return [
            {
                "title": "Grilled Cheese Sandwich",
                "ingredients": ["bread", "cheese", "butter"],
                "steps": [
                    "Butter bread.",
                    "Place cheese between slices.",
                    "Grill until golden brown.",
                ],
            }
        ]


@tool
def check_fridge() -> list[str]:
    """Returns a list of ingredients currently in the fridge."""
    logger.info("Checking fridge for current ingredients")
    if random.random() < 0.5:
        return ["pasta", "tomato sauce", "bell peppers", "olive oil"]
    else:
        return ["tofu", "soy sauce", "broccoli", "carrots"]


meal_agent = create_agent(
    model=llm_config.get_default_model(),
    system_prompt=(
        "You help users plan meals and choose the best recipes."
        "Include the ingredients and cooking instructions in your response."
        "Indicate what user needs to buy from store when their fridge is missing ingredients."
    ),
    checkpointer=InMemorySaver(),
    tools=[find_recipes, check_fridge],
)


@tool
def meal_agent_tool(query: str, runtime: ToolRuntime) -> str:
    """Invoke the recipe planning agent and return its final response as plain text."""
    logger.info("Tool:meal_agent invoked")
    response = meal_agent.invoke(
        {"messages": [HumanMessage(content=query)]}, config=runtime.config
    )
    final = response["messages"][-1].content
    return final


# ----------------------------------------------------------------------------------
# SUPERVISOR AGENT: Manages the sub-agents
# ----------------------------------------------------------------------------------
supervisor_agent = create_agent(
    model=llm_config.get_default_model(),
    system_prompt=(
        "You are a supervisor, managing an activity planning agent and recipe planning agent."
        "Assign work to them as needed in order to answer user's question."
    ),
    tools=[weekend_agent_tool, meal_agent_tool],
    checkpointer=InMemorySaver(),
)

config = {"configurable": {"thread_id": "123"}}


def test_supervisor_agent_checkpointer():
    supervisor_agent.invoke(
        {"messages": [{"role": "user", "content": "my kids want pasta for dinner"}]},
        config=config,
    )

    supervisor_agent.invoke(
        {
            "messages": [
                {
                    "role": "user",
                    "content": "I want to plan something for the weekend in Boston",
                }
            ]
        },
        config=config,
    )

    supervisor_agent.invoke(
        {
            "messages": [
                {
                    "role": "user",
                    "content": "I want to plan something for the weekend in New your",
                }
            ]
        },
        config=config,
    )

    supervisor_agent.invoke(
        {
            "messages": [
                {
                    "role": "user",
                    "content": "I don't like Pasta. What about Chicken Briyani?",
                }
            ]
        },
        config=config,
    )
    supervisor_agent_state = supervisor_agent.get_state(config=config)
    logger.info(supervisor_agent_state.values)

    supervisor_agent_state_with_subgraphs = supervisor_agent.get_state(
        config=config, subgraphs=True
    )
    logger.info(supervisor_agent_state_with_subgraphs.values)

    meal_agent_state = meal_agent.get_state(config=config)
    logger.info(meal_agent_state.values)

    weekend_agent_state = weekend_agent.get_state(config=config)
    logger.info(weekend_agent_state.values)
```