test_create_agent_jump() — langchain Function Reference

Source Code

libs/langchain_v1/tests/unit_tests/agents/middleware/core/test_framework.py lines 140–199

def test_create_agent_jump(
    snapshot: SnapshotAssertion,
    sync_checkpointer: BaseCheckpointSaver[str],
) -> None:
    calls = []

    class NoopSeven(AgentMiddleware):
        def before_model(self, state: AgentState[Any], runtime: Runtime) -> None:
            calls.append("NoopSeven.before_model")

        def wrap_model_call(
            self,
            request: ModelRequest,
            handler: Callable[[ModelRequest], ModelResponse],
        ) -> ModelCallResult:
            calls.append("NoopSeven.wrap_model_call")
            return handler(request)

        def after_model(self, state: AgentState[Any], runtime: Runtime) -> None:
            calls.append("NoopSeven.after_model")

    class NoopEight(AgentMiddleware):
        @hook_config(can_jump_to=["end"])
        def before_model(self, state: AgentState[Any], runtime: Runtime) -> dict[str, Any]:
            calls.append("NoopEight.before_model")
            return {"jump_to": "end"}

        def wrap_model_call(
            self,
            request: ModelRequest,
            handler: Callable[[ModelRequest], ModelResponse],
        ) -> ModelCallResult:
            calls.append("NoopEight.wrap_model_call")
            return handler(request)

        def after_model(self, state: AgentState[Any], runtime: Runtime) -> None:
            calls.append("NoopEight.after_model")

    @tool
    def my_tool(value: str) -> str:
        """A great tool."""
        calls.append("my_tool")
        return value.upper()

    agent_one = create_agent(
        model=FakeToolCallingModel(
            tool_calls=[[ToolCall(id="1", name="my_tool", args={"value": "yo"})]],
        ),
        tools=[my_tool],
        system_prompt="You are a helpful assistant.",
        middleware=[NoopSeven(), NoopEight()],
        checkpointer=sync_checkpointer,
    )

    if isinstance(sync_checkpointer, InMemorySaver):
        assert agent_one.get_graph().draw_mermaid() == snapshot

    thread1 = {"configurable": {"thread_id": "1"}}
    assert agent_one.invoke({"messages": []}, thread1) == {"messages": []}
    assert calls == ["NoopSeven.before_model", "NoopEight.before_model"]