test_composition.py — langchain Source File

Source Code

"""Unit tests for _chain_model_call_handlers handler composition."""

from collections.abc import Callable
from typing import Any, TypedDict, cast

from langchain_core.messages import AIMessage
from langgraph.runtime import Runtime
from langgraph.types import Command

from langchain.agents import AgentState
from langchain.agents.factory import _chain_model_call_handlers, _ComposedExtendedModelResponse
from langchain.agents.middleware.types import ExtendedModelResponse, ModelRequest, ModelResponse


def create_test_request(**kwargs: Any) -> ModelRequest:
    """Helper to create a `ModelRequest` with sensible defaults."""
    defaults: dict[str, Any] = {
        "messages": [],
        "model": None,
        "system_prompt": None,
        "tool_choice": None,
        "tools": [],
        "response_format": None,
        "state": {},
        "runtime": cast("Runtime", object()),
    }
    defaults.update(kwargs)
    return ModelRequest(**defaults)


def create_mock_base_handler(content: str = "test") -> Callable[[ModelRequest], ModelResponse]:
    """Helper to create a base handler that returns `ModelResponse`."""

    def mock_base_handler(req: ModelRequest) -> ModelResponse:
        return ModelResponse(result=[AIMessage(content=content)], structured_response=None)

    return mock_base_handler


class TestChainModelCallHandlers:
    """Test the `_chain_model_call_handlers` composition function."""

    def test_empty_handlers_returns_none(self) -> None:
        """Test that empty handlers list returns None."""
        result = _chain_model_call_handlers([])
        assert result is None

    def test_single_handler_returns_unchanged(self) -> None:
        """Test that single handler is wrapped to normalize output."""

        def handler(
            request: ModelRequest, base_handler: Callable[[ModelRequest], ModelResponse]
        ) -> ModelResponse:
            return base_handler(request)

        result = _chain_model_call_handlers([handler])
        # Result is wrapped to normalize, so it won't be identical
        assert result is not None
        assert callable(result)

// ... (290 more lines)