test_state_schema.py — langchain Source File

Source Code

"""Test state_schema parameter in create_agent.

This module tests that the state_schema parameter allows users to extend
AgentState without needing to create custom middleware.
"""

from __future__ import annotations

from typing import TYPE_CHECKING, Annotated, Any

from langchain_core.messages import HumanMessage
from langchain_core.tools import tool

from langchain.agents import create_agent
from langchain.agents.middleware.types import (
    AgentMiddleware,
    AgentState,
    PrivateStateAttr,
)

# Cannot move ToolRuntime to TYPE_CHECKING as parameters of @tool annotated functions
# are inspected at runtime.
from langchain.tools import ToolRuntime  # noqa: TC001
from tests.unit_tests.agents.model import FakeToolCallingModel

if TYPE_CHECKING:
    from langgraph.runtime import Runtime


@tool
def simple_tool(x: int) -> str:
    """Simple tool for basic tests."""
    return f"Result: {x}"


def test_state_schema_single_custom_field() -> None:
    """Test that a single custom state field is preserved through agent execution."""

    class CustomState(AgentState[Any]):
        custom_field: str

    agent = create_agent(
        model=FakeToolCallingModel(
            tool_calls=[[{"args": {"x": 1}, "id": "call_1", "name": "simple_tool"}], []]
        ),
        tools=[simple_tool],
        state_schema=CustomState,
    )

    result = agent.invoke({"messages": [HumanMessage("Test")], "custom_field": "test_value"})

    assert result["custom_field"] == "test_value"
    assert len(result["messages"]) == 4


def test_state_schema_multiple_custom_fields() -> None:
    """Test that multiple custom state fields are preserved through agent execution."""

    class CustomState(AgentState[Any]):
        user_id: str
// ... (196 more lines)