TestRequestModification Class — langchain Architecture

Source Code

libs/langchain_v1/tests/unit_tests/agents/middleware/core/test_wrap_model_call.py lines 559–601

class TestRequestModification:
    """Test request modification with wrap_model_call."""

    def test_add_system_prompt(self) -> None:
        """Test middleware that adds a system prompt to requests."""
        received_requests = []

        class SystemPromptMiddleware(AgentMiddleware):
            def __init__(self, system_prompt: str):
                super().__init__()
                self.system_prompt = system_prompt

            def wrap_model_call(
                self,
                request: ModelRequest,
                handler: Callable[[ModelRequest], ModelResponse],
            ) -> ModelCallResult:
                # Modify request to add system prompt
                modified_request = ModelRequest(
                    model=request.model,
                    system_prompt=self.system_prompt,
                    messages=request.messages,
                    tools=request.tools,
                    tool_choice=request.tool_choice,
                    response_format=request.response_format,
                    model_settings=request.model_settings,
                    state=request.state,
                    runtime=request.runtime,
                )
                received_requests.append(modified_request)
                return handler(modified_request)

        model = GenericFakeChatModel(messages=iter([AIMessage(content="Response")]))
        agent = create_agent(
            model=model,
            middleware=[SystemPromptMiddleware(system_prompt="You are a helpful assistant.")],
        )

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

        assert len(received_requests) == 1
        assert received_requests[0].system_prompt == "You are a helpful assistant."
        assert result["messages"][1].content == "Response"