test_overrides.py — langchain Source File

Source Code

"""Unit tests for override() methods on ModelRequest and ToolCallRequest."""

from typing import Any
from unittest.mock import Mock

import pytest
from langchain_core.language_models.fake_chat_models import GenericFakeChatModel
from langchain_core.messages import (
    AIMessage,
    AnyMessage,
    HumanMessage,
    SystemMessage,
    ToolCall,
)
from langchain_core.tools import tool

from langchain.agents import AgentState
from langchain.agents.middleware.types import ModelRequest, ToolCallRequest


class TestModelRequestOverride:
    """Test the ModelRequest.override() method."""

    def test_override_single_attribute(self) -> None:
        """Test overriding a single attribute."""
        model = GenericFakeChatModel(messages=iter([AIMessage(content="Hello")]))
        original_request = ModelRequest(
            model=model,
            system_message=SystemMessage("Original prompt"),
            messages=[HumanMessage("Hi")],
            tool_choice=None,
            tools=[],
            response_format=None,
            state=AgentState(messages=[]),
            runtime=None,
        )

        new_request = original_request.override(system_message=SystemMessage("New prompt"))

        # New request should have the overridden value
        assert new_request.system_prompt == "New prompt"
        # Original request should be unchanged (immutability)
        assert original_request.system_prompt == "Original prompt"
        # Other attributes should be the same
        assert new_request.model == original_request.model
        assert new_request.messages == original_request.messages

    def test_override_multiple_attributes(self) -> None:
        """Test overriding multiple attributes at once."""

        class CustomState(AgentState[Any]):
            count: int

        model = GenericFakeChatModel(messages=iter([AIMessage(content="Hello")]))
        original_request = ModelRequest(
            model=model,
            system_message=SystemMessage("Original prompt"),
            messages=[HumanMessage("Hi")],
            tool_choice=None,
            tools=[],
// ... (383 more lines)