TestStateAndRuntime Class — langchain Architecture

Source Code

libs/langchain_v1/tests/unit_tests/agents/middleware/core/test_wrap_model_call.py lines 604–675

class TestStateAndRuntime:
    """Test state and runtime access in wrap_model_call."""

    def test_access_state_in_middleware(self) -> None:
        """Test middleware can read and use state."""
        state_values = []

        class StateAwareMiddleware(AgentMiddleware):
            def wrap_model_call(
                self,
                request: ModelRequest,
                handler: Callable[[ModelRequest], ModelResponse],
            ) -> ModelCallResult:
                # Access state from request
                state_values.append(
                    {
                        "messages_count": len(request.state.get("messages", [])),
                    }
                )
                return handler(request)

        model = GenericFakeChatModel(messages=iter([AIMessage(content="Response")]))
        agent = create_agent(model=model, middleware=[StateAwareMiddleware()])

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

        assert len(state_values) == 1
        assert state_values[0]["messages_count"] == 1  # Just The HumanMessage
        assert result["messages"][1].content == "Response"

    def test_retry_with_state_tracking(self) -> None:
        """Test middleware that tracks retry count in state."""

        class StateTrackingRetryMiddleware(AgentMiddleware):
            def wrap_model_call(
                self,
                request: ModelRequest,
                handler: Callable[[ModelRequest], ModelResponse],
            ) -> ModelCallResult:
                max_retries = 2
                for attempt in range(max_retries):
                    try:
                        return handler(request)
                    except Exception:
                        if attempt == max_retries - 1:
                            raise
                pytest.fail("Should have raised an exception")

        call_count = {"value": 0}

        class FailOnceThenSucceed(GenericFakeChatModel):
            @override
            def _generate(
                self,
                messages: list[BaseMessage],
                stop: list[str] | None = None,
                run_manager: CallbackManagerForLLMRun | None = None,
                **kwargs: Any,
            ) -> ChatResult:
                call_count["value"] += 1
                if call_count["value"] == 1:
                    msg = "First fails"
                    raise ValueError(msg)
                return super()._generate(messages, **kwargs)

        model = FailOnceThenSucceed(messages=iter([AIMessage(content="Success")]))
        agent = create_agent(model=model, middleware=[StateTrackingRetryMiddleware()])

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

        assert call_count["value"] == 2  # Failed once, succeeded second time
        assert result["messages"][1].content == "Success"