TestAgentMiddlewareHooks Class — langchain Architecture

Source Code

libs/langchain_v1/tests/unit_tests/agents/middleware/core/test_framework.py lines 769–875

class TestAgentMiddlewareHooks:
    """Test before_agent and after_agent middleware hooks."""

    @pytest.mark.parametrize("is_async", [False, True])
    @pytest.mark.parametrize("hook_type", ["before", "after"])
    async def test_hook_execution(self, *, is_async: bool, hook_type: str) -> None:
        """Test that agent hooks are called in both sync and async modes."""
        execution_log: list[str] = []

        if is_async:
            if hook_type == "before":

                @before_agent
                async def log_hook(
                    state: AgentState[Any], *_args: Any, **_kwargs: Any
                ) -> dict[str, Any] | None:
                    execution_log.append(f"{hook_type}_agent_called")
                    execution_log.append(f"message_count: {len(state['messages'])}")
                    return None

            else:

                @after_agent
                async def log_hook(
                    state: AgentState[Any], *_args: Any, **_kwargs: Any
                ) -> dict[str, Any] | None:
                    execution_log.append(f"{hook_type}_agent_called")
                    execution_log.append(f"message_count: {len(state['messages'])}")
                    return None

        elif hook_type == "before":

            @before_agent
            def log_hook(
                state: AgentState[Any], *_args: Any, **_kwargs: Any
            ) -> dict[str, Any] | None:
                execution_log.append(f"{hook_type}_agent_called")
                execution_log.append(f"message_count: {len(state['messages'])}")
                return None

        else:

            @after_agent
            def log_hook(
                state: AgentState[Any], *_args: Any, **_kwargs: Any
            ) -> dict[str, Any] | None:
                execution_log.append(f"{hook_type}_agent_called")
                execution_log.append(f"message_count: {len(state['messages'])}")
                return None

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

        if is_async:
            await agent.ainvoke({"messages": [HumanMessage("Hi")]})
        else:
            agent.invoke({"messages": [HumanMessage("Hi")]})

        assert f"{hook_type}_agent_called" in execution_log
        assert any("message_count:" in log for log in execution_log)

    @pytest.mark.parametrize("is_async", [False, True])
    @pytest.mark.parametrize("hook_type", ["before", "after"])
    async def test_hook_with_class_inheritance(self, *, is_async: bool, hook_type: str) -> None:
        """Test agent hooks using class inheritance in both sync and async modes."""
        execution_log: list[str] = []

        class AsyncCustomMiddleware(AgentMiddleware):
            async def abefore_agent(
                self, state: AgentState[Any], runtime: Runtime
            ) -> dict[str, Any] | None:
                if hook_type == "before":
                    execution_log.append("hook_called")
                return None

            async def aafter_agent(
                self, state: AgentState[Any], runtime: Runtime
            ) -> dict[str, Any] | None:
                if hook_type == "after":
                    execution_log.append("hook_called")
                return None