test_decorator_with_custom_state_schema() — langchain Function Reference

Source Code

libs/langchain_v1/tests/unit_tests/agents/middleware/core/test_wrap_model_call.py lines 1191–1207

    def test_decorator_with_custom_state_schema(self) -> None:
        """Test decorator with custom state schema."""

        class CustomState(TypedDict):
            messages: list[Any]
            custom_field: str

        @wrap_model_call(state_schema=CustomState)
        def middleware_with_schema(
            request: ModelRequest,
            handler: Callable[[ModelRequest], ModelResponse],
        ) -> ModelCallResult:
            return handler(request)

        assert isinstance(middleware_with_schema, AgentMiddleware)
        # Custom state schema should be set
        assert middleware_with_schema.state_schema == CustomState