TestShortCircuit Class — langchain Architecture

Source Code

libs/langchain_v1/tests/unit_tests/agents/middleware/core/test_wrap_model_call.py lines 496–556

class TestShortCircuit:
    """Test short-circuit patterns with wrap_model_call."""

    def test_cache_short_circuit(self) -> None:
        """Test middleware that short-circuits with cached response."""
        cache: dict[str, ModelResponse] = {}
        model_calls = []

        class CachingMiddleware(AgentMiddleware):
            def wrap_model_call(
                self,
                request: ModelRequest,
                handler: Callable[[ModelRequest], ModelResponse],
            ) -> ModelCallResult:
                # Simple cache key based on last message
                cache_key = str(request.messages[-1].content) if request.messages else ""

                if cache_key in cache:
                    # Short-circuit with cached result
                    return cache[cache_key]
                # Execute and cache
                result = handler(request)
                cache[cache_key] = result
                return result

        class TrackingModel(GenericFakeChatModel):
            @override
            def _generate(
                self,
                messages: list[BaseMessage],
                stop: list[str] | None = None,
                run_manager: CallbackManagerForLLMRun | None = None,
                **kwargs: Any,
            ) -> ChatResult:
                model_calls.append(len(messages))
                return super()._generate(messages, **kwargs)

        model = TrackingModel(
            messages=iter(
                [
                    AIMessage(content="Response 1"),
                    AIMessage(content="Response 2"),
                ]
            )
        )
        agent = create_agent(model=model, middleware=[CachingMiddleware()])

        # First call - cache miss, calls model
        result1 = agent.invoke({"messages": [HumanMessage("Hello")]})
        assert result1["messages"][1].content == "Response 1"
        assert len(model_calls) == 1

        # Second call with same message - cache hit, doesn't call model
        result2 = agent.invoke({"messages": [HumanMessage("Hello")]})
        assert result2["messages"][1].content == "Response 1"
        assert len(model_calls) == 1  # Still 1, no new call

        # Third call with different message - cache miss, calls model
        result3 = agent.invoke({"messages": [HumanMessage("Goodbye")]})
        assert result3["messages"][1].content == "Response 2"
        assert len(model_calls) == 2  # New call