langchain 1.0.0a12__py3-none-any.whl → 1.0.4__py3-none-any.whl

langchain/agents/middleware/types.py

@@ -2,8 +2,8 @@
 
 from __future__ import annotations
 
-from collections.abc import Callable
-from dataclasses import dataclass, field
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass, field, replace
 from inspect import iscoroutinefunction
 from typing import (
     TYPE_CHECKING,
@@ -16,24 +16,29 @@ from typing import (
     overload,
 )
 
-from langchain_core.runnables import run_in_executor
-
 if TYPE_CHECKING:
     from collections.abc import Awaitable
 
-# needed as top level import for pydantic schema generation on AgentState
-from langchain_core.messages import AnyMessage  # noqa: TC002
+# Needed as top level import for Pydantic schema generation on AgentState
+from typing import TypeAlias
+
+from langchain_core.messages import (  # noqa: TC002
+    AIMessage,
+    AnyMessage,
+    BaseMessage,
+    ToolMessage,
+)
 from langgraph.channels.ephemeral_value import EphemeralValue
-from langgraph.channels.untracked_value import UntrackedValue
 from langgraph.graph.message import add_messages
+from langgraph.prebuilt.tool_node import ToolCallRequest, ToolCallWrapper
+from langgraph.types import Command  # noqa: TC002
 from langgraph.typing import ContextT
-from typing_extensions import NotRequired, Required, TypedDict, TypeVar
+from typing_extensions import NotRequired, Required, TypedDict, TypeVar, Unpack
 
 if TYPE_CHECKING:
     from langchain_core.language_models.chat_models import BaseChatModel
     from langchain_core.tools import BaseTool
     from langgraph.runtime import Runtime
-    from langgraph.types import Command
 
     from langchain.agents.structured_output import ResponseFormat
 
@@ -42,15 +47,19 @@ __all__ = [
     "AgentState",
     "ContextT",
     "ModelRequest",
+    "ModelResponse",
     "OmitFromSchema",
-    "PublicAgentState",
+    "ResponseT",
+    "StateT_co",
+    "ToolCallRequest",
+    "ToolCallWrapper",
     "after_agent",
     "after_model",
     "before_agent",
     "before_model",
     "dynamic_prompt",
     "hook_config",
-    "modify_model_request",
+    "wrap_tool_call",
 ]
 
 JumpTo = Literal["tools", "model", "end"]
@@ -59,6 +68,18 @@ JumpTo = Literal["tools", "model", "end"]
 ResponseT = TypeVar("ResponseT")
 
 
+class _ModelRequestOverrides(TypedDict, total=False):
+    """Possible overrides for ModelRequest.override() method."""
+
+    model: BaseChatModel
+    system_prompt: str | None
+    messages: list[AnyMessage]
+    tool_choice: Any | None
+    tools: list[BaseTool | dict]
+    response_format: ResponseFormat | None
+    model_settings: dict[str, Any]
+
+
 @dataclass
 class ModelRequest:
     """Model request information for the agent."""
@@ -69,8 +90,65 @@ class ModelRequest:
     tool_choice: Any | None
     tools: list[BaseTool | dict]
     response_format: ResponseFormat | None
+    state: AgentState
+    runtime: Runtime[ContextT]  # type: ignore[valid-type]
     model_settings: dict[str, Any] = field(default_factory=dict)
 
+    def override(self, **overrides: Unpack[_ModelRequestOverrides]) -> ModelRequest:
+        """Replace the request with a new request with the given overrides.
+
+        Returns a new `ModelRequest` instance with the specified attributes replaced.
+        This follows an immutable pattern, leaving the original request unchanged.
+
+        Args:
+            **overrides: Keyword arguments for attributes to override. Supported keys:
+                - model: BaseChatModel instance
+                - system_prompt: Optional system prompt string
+                - messages: List of messages
+                - tool_choice: Tool choice configuration
+                - tools: List of available tools
+                - response_format: Response format specification
+                - model_settings: Additional model settings
+
+        Returns:
+            New ModelRequest instance with specified overrides applied.
+
+        Examples:
+            ```python
+            # Create a new request with different model
+            new_request = request.override(model=different_model)
+
+            # Override multiple attributes
+            new_request = request.override(system_prompt="New instructions", tool_choice="auto")
+            ```
+        """
+        return replace(self, **overrides)
+
+
+@dataclass
+class ModelResponse:
+    """Response from model execution including messages and optional structured output.
+
+    The result will usually contain a single AIMessage, but may include
+    an additional ToolMessage if the model used a tool for structured output.
+    """
+
+    result: list[BaseMessage]
+    """List of messages from model execution."""
+
+    structured_response: Any = None
+    """Parsed structured output if response_format was specified, None otherwise."""
+
+
+# Type alias for middleware return type - allows returning either full response or just AIMessage
+ModelCallResult: TypeAlias = "ModelResponse | AIMessage"
+"""Type alias for model call handler return value.
+
+Middleware can return either:
+- ModelResponse: Full response with messages and optional structured output
+- AIMessage: Simplified return for simple use cases
+"""
+
 
 @dataclass
 class OmitFromSchema:
@@ -99,21 +177,23 @@ class AgentState(TypedDict, Generic[ResponseT]):
     messages: Required[Annotated[list[AnyMessage], add_messages]]
     jump_to: NotRequired[Annotated[JumpTo | None, EphemeralValue, PrivateStateAttr]]
     structured_response: NotRequired[Annotated[ResponseT, OmitFromInput]]
-    thread_model_call_count: NotRequired[Annotated[int, PrivateStateAttr]]
-    run_model_call_count: NotRequired[Annotated[int, UntrackedValue, PrivateStateAttr]]
 
 
-class PublicAgentState(TypedDict, Generic[ResponseT]):
-    """Public state schema for the agent.
+class _InputAgentState(TypedDict):  # noqa: PYI049
+    """Input state schema for the agent."""
 
-    Just used for typing purposes.
-    """
+    messages: Required[Annotated[list[AnyMessage | dict], add_messages]]
+
+
+class _OutputAgentState(TypedDict, Generic[ResponseT]):  # noqa: PYI049
+    """Output state schema for the agent."""
 
     messages: Required[Annotated[list[AnyMessage], add_messages]]
     structured_response: NotRequired[ResponseT]
 
 
 StateT = TypeVar("StateT", bound=AgentState, default=AgentState)
+StateT_co = TypeVar("StateT_co", bound=AgentState, default=AgentState, covariant=True)
 StateT_contra = TypeVar("StateT_contra", bound=AgentState, contravariant=True)
 
 
@@ -154,24 +234,6 @@ class AgentMiddleware(Generic[StateT, ContextT]):
     ) -> dict[str, Any] | None:
         """Async logic to run before the model is called."""
 
-    def modify_model_request(
-        self,
-        request: ModelRequest,
-        state: StateT,  # noqa: ARG002
-        runtime: Runtime[ContextT],  # noqa: ARG002
-    ) -> ModelRequest:
-        """Logic to modify request kwargs before the model is called."""
-        return request
-
-    async def amodify_model_request(
-        self,
-        request: ModelRequest,
-        state: StateT,
-        runtime: Runtime[ContextT],
-    ) -> ModelRequest:
-        """Async logic to modify request kwargs before the model is called."""
-        return await run_in_executor(None, self.modify_model_request, request, state, runtime)
-
     def after_model(self, state: StateT, runtime: Runtime[ContextT]) -> dict[str, Any] | None:
         """Logic to run after the model is called."""
 
@@ -180,53 +242,133 @@ class AgentMiddleware(Generic[StateT, ContextT]):
     ) -> dict[str, Any] | None:
         """Async logic to run after the model is called."""
 
-    def retry_model_request(
+    def wrap_model_call(
         self,
-        error: Exception,  # noqa: ARG002
-        request: ModelRequest,  # noqa: ARG002
-        state: StateT,  # noqa: ARG002
-        runtime: Runtime[ContextT],  # noqa: ARG002
-        attempt: int,  # noqa: ARG002
-    ) -> ModelRequest | None:
-        """Logic to handle model invocation errors and optionally retry.
+        request: ModelRequest,
+        handler: Callable[[ModelRequest], ModelResponse],
+    ) -> ModelCallResult:
+        """Intercept and control model execution via handler callback.
+
+        The handler callback executes the model request and returns a `ModelResponse`.
+        Middleware can call the handler multiple times for retry logic, skip calling
+        it to short-circuit, or modify the request/response. Multiple middleware
+        compose with first in list as outermost layer.
 
         Args:
-            error: The exception that occurred during model invocation.
-            request: The original model request that failed.
-            state: The current agent state.
-            runtime: The langgraph runtime.
-            attempt: The current attempt number (1-indexed).
+            request: Model request to execute (includes state and runtime).
+            handler: Callback that executes the model request and returns
+                `ModelResponse`. Call this to execute the model. Can be called multiple
+                times for retry logic. Can skip calling it to short-circuit.
 
         Returns:
-            ModelRequest: Modified request to retry with.
-            None: Propagate the error (re-raise).
+            `ModelCallResult`
+
+        Examples:
+            Retry on error:
+            ```python
+            def wrap_model_call(self, request, handler):
+                for attempt in range(3):
+                    try:
+                        return handler(request)
+                    except Exception:
+                        if attempt == 2:
+                            raise
+            ```
+
+            Rewrite response:
+            ```python
+            def wrap_model_call(self, request, handler):
+                response = handler(request)
+                ai_msg = response.result[0]
+                return ModelResponse(
+                    result=[AIMessage(content=f"[{ai_msg.content}]")],
+                    structured_response=response.structured_response,
+                )
+            ```
+
+            Error to fallback:
+            ```python
+            def wrap_model_call(self, request, handler):
+                try:
+                    return handler(request)
+                except Exception:
+                    return ModelResponse(result=[AIMessage(content="Service unavailable")])
+            ```
+
+            Cache/short-circuit:
+            ```python
+            def wrap_model_call(self, request, handler):
+                if cached := get_cache(request):
+                    return cached  # Short-circuit with cached result
+                response = handler(request)
+                save_cache(request, response)
+                return response
+            ```
+
+            Simple AIMessage return (converted automatically):
+            ```python
+            def wrap_model_call(self, request, handler):
+                response = handler(request)
+                # Can return AIMessage directly for simple cases
+                return AIMessage(content="Simplified response")
+            ```
         """
-        return None
+        msg = (
+            "Synchronous implementation of wrap_model_call is not available. "
+            "You are likely encountering this error because you defined only the async version "
+            "(awrap_model_call) and invoked your agent in a synchronous context "
+            "(e.g., using `stream()` or `invoke()`). "
+            "To resolve this, either: "
+            "(1) subclass AgentMiddleware and implement the synchronous wrap_model_call method, "
+            "(2) use the @wrap_model_call decorator on a standalone sync function, or "
+            "(3) invoke your agent asynchronously using `astream()` or `ainvoke()`."
+        )
+        raise NotImplementedError(msg)
 
-    async def aretry_model_request(
+    async def awrap_model_call(
         self,
-        error: Exception,
         request: ModelRequest,
-        state: StateT,
-        runtime: Runtime[ContextT],
-        attempt: int,
-    ) -> ModelRequest | None:
-        """Async logic to handle model invocation errors and optionally retry.
+        handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
+    ) -> ModelCallResult:
+        """Intercept and control async model execution via handler callback.
+
+        The handler callback executes the model request and returns a `ModelResponse`.
+        Middleware can call the handler multiple times for retry logic, skip calling
+        it to short-circuit, or modify the request/response. Multiple middleware
+        compose with first in list as outermost layer.
 
         Args:
-            error: The exception that occurred during model invocation.
-            request: The original model request that failed.
-            state: The current agent state.
-            runtime: The langgraph runtime.
-            attempt: The current attempt number (1-indexed).
+            request: Model request to execute (includes state and runtime).
+            handler: Async callback that executes the model request and returns
+                `ModelResponse`. Call this to execute the model. Can be called multiple
+                times for retry logic. Can skip calling it to short-circuit.
 
         Returns:
-            ModelRequest: Modified request to retry with.
-            None: Propagate the error (re-raise).
+            ModelCallResult
+
+        Examples:
+            Retry on error:
+            ```python
+            async def awrap_model_call(self, request, handler):
+                for attempt in range(3):
+                    try:
+                        return await handler(request)
+                    except Exception:
+                        if attempt == 2:
+                            raise
+            ```
         """
-        return await run_in_executor(
-            None, self.retry_model_request, error, request, state, runtime, attempt
+        msg = (
+            "Asynchronous implementation of awrap_model_call is not available. "
+            "You are likely encountering this error because you defined only the sync version "
+            "(wrap_model_call) and invoked your agent in an asynchronous context "
+            "(e.g., using `astream()` or `ainvoke()`). "
+            "To resolve this, either: "
+            "(1) subclass AgentMiddleware and implement the asynchronous awrap_model_call method, "
+            "(2) use the @wrap_model_call decorator on a standalone async function, or "
+            "(3) invoke your agent synchronously using `stream()` or `invoke()`."
         )
+        raise NotImplementedError(msg)
 
     def after_agent(self, state: StateT, runtime: Runtime[ContextT]) -> dict[str, Any] | None:
         """Logic to run after the agent execution completes."""
@@ -236,9 +378,140 @@ class AgentMiddleware(Generic[StateT, ContextT]):
     ) -> dict[str, Any] | None:
         """Async logic to run after the agent execution completes."""
 
+    def wrap_tool_call(
+        self,
+        request: ToolCallRequest,
+        handler: Callable[[ToolCallRequest], ToolMessage | Command],
+    ) -> ToolMessage | Command:
+        """Intercept tool execution for retries, monitoring, or modification.
+
+        Multiple middleware compose automatically (first defined = outermost).
+        Exceptions propagate unless `handle_tool_errors` is configured on `ToolNode`.
+
+        Args:
+            request: Tool call request with call `dict`, `BaseTool`, state, and runtime.
+                Access state via `request.state` and runtime via `request.runtime`.
+            handler: Callable to execute the tool (can be called multiple times).
+
+        Returns:
+            `ToolMessage` or `Command` (the final result).
+
+        The handler callable can be invoked multiple times for retry logic.
+        Each call to handler is independent and stateless.
+
+        Examples:
+            Modify request before execution:
+
+            ```python
+            def wrap_tool_call(self, request, handler):
+                request.tool_call["args"]["value"] *= 2
+                return handler(request)
+            ```
+
+            Retry on error (call handler multiple times):
+
+            ```python
+            def wrap_tool_call(self, request, handler):
+                for attempt in range(3):
+                    try:
+                        result = handler(request)
+                        if is_valid(result):
+                            return result
+                    except Exception:
+                        if attempt == 2:
+                            raise
+                return result
+            ```
+
+            Conditional retry based on response:
+
+            ```python
+            def wrap_tool_call(self, request, handler):
+                for attempt in range(3):
+                    result = handler(request)
+                    if isinstance(result, ToolMessage) and result.status != "error":
+                        return result
+                    if attempt < 2:
+                        continue
+                    return result
+            ```
+        """
+        msg = (
+            "Synchronous implementation of wrap_tool_call is not available. "
+            "You are likely encountering this error because you defined only the async version "
+            "(awrap_tool_call) and invoked your agent in a synchronous context "
+            "(e.g., using `stream()` or `invoke()`). "
+            "To resolve this, either: "
+            "(1) subclass AgentMiddleware and implement the synchronous wrap_tool_call method, "
+            "(2) use the @wrap_tool_call decorator on a standalone sync function, or "
+            "(3) invoke your agent asynchronously using `astream()` or `ainvoke()`."
+        )
+        raise NotImplementedError(msg)
+
+    async def awrap_tool_call(
+        self,
+        request: ToolCallRequest,
+        handler: Callable[[ToolCallRequest], Awaitable[ToolMessage | Command]],
+    ) -> ToolMessage | Command:
+        """Intercept and control async tool execution via handler callback.
+
+        The handler callback executes the tool call and returns a `ToolMessage` or
+        `Command`. Middleware can call the handler multiple times for retry logic, skip
+        calling it to short-circuit, or modify the request/response. Multiple middleware
+        compose with first in list as outermost layer.
+
+        Args:
+            request: Tool call request with call `dict`, `BaseTool`, state, and runtime.
+                Access state via `request.state` and runtime via `request.runtime`.
+            handler: Async callable to execute the tool and returns `ToolMessage` or
+                `Command`. Call this to execute the tool. Can be called multiple times
+                for retry logic. Can skip calling it to short-circuit.
+
+        Returns:
+            `ToolMessage` or `Command` (the final result).
+
+        The handler callable can be invoked multiple times for retry logic.
+        Each call to handler is independent and stateless.
+
+        Examples:
+            Async retry on error:
+            ```python
+            async def awrap_tool_call(self, request, handler):
+                for attempt in range(3):
+                    try:
+                        result = await handler(request)
+                        if is_valid(result):
+                            return result
+                    except Exception:
+                        if attempt == 2:
+                            raise
+                return result
+            ```
+
+            ```python
+            async def awrap_tool_call(self, request, handler):
+                if cached := await get_cache_async(request):
+                    return ToolMessage(content=cached, tool_call_id=request.tool_call["id"])
+                result = await handler(request)
+                await save_cache_async(request, result)
+                return result
+            ```
+        """
+        msg = (
+            "Asynchronous implementation of awrap_tool_call is not available. "
+            "You are likely encountering this error because you defined only the sync version "
+            "(wrap_tool_call) and invoked your agent in an asynchronous context "
+            "(e.g., using `astream()` or `ainvoke()`). "
+            "To resolve this, either: "
+            "(1) subclass AgentMiddleware and implement the asynchronous awrap_tool_call method, "
+            "(2) use the @wrap_tool_call decorator on a standalone async function, or "
+            "(3) invoke your agent synchronously using `stream()` or `invoke()`."
+        )
+        raise NotImplementedError(msg)
+
 
 class _CallableWithStateAndRuntime(Protocol[StateT_contra, ContextT]):
-    """Callable with AgentState and Runtime as arguments."""
+    """Callable with `AgentState` and `Runtime` as arguments."""
 
     def __call__(
         self, state: StateT_contra, runtime: Runtime[ContextT]
@@ -247,23 +520,43 @@ class _CallableWithStateAndRuntime(Protocol[StateT_contra, ContextT]):
         ...
 
 
-class _CallableWithModelRequestAndStateAndRuntime(Protocol[StateT_contra, ContextT]):
-    """Callable with ModelRequest, AgentState, and Runtime as arguments."""
+class _CallableReturningPromptString(Protocol[StateT_contra, ContextT]):  # type: ignore[misc]
+    """Callable that returns a prompt string given `ModelRequest` (contains state and runtime)."""
+
+    def __call__(self, request: ModelRequest) -> str | Awaitable[str]:
+        """Generate a system prompt string based on the request."""
+        ...
+
+
+class _CallableReturningModelResponse(Protocol[StateT_contra, ContextT]):  # type: ignore[misc]
+    """Callable for model call interception with handler callback.
+
+    Receives handler callback to execute model and returns `ModelResponse` or
+    `AIMessage`.
+    """
 
     def __call__(
-        self, request: ModelRequest, state: StateT_contra, runtime: Runtime[ContextT]
-    ) -> ModelRequest | Awaitable[ModelRequest]:
-        """Perform some logic with the model request, state, and runtime."""
+        self,
+        request: ModelRequest,
+        handler: Callable[[ModelRequest], ModelResponse],
+    ) -> ModelCallResult:
+        """Intercept model execution via handler callback."""
         ...
 
 
-class _CallableReturningPromptString(Protocol[StateT_contra, ContextT]):
-    """Callable that returns a prompt string given ModelRequest, AgentState, and Runtime."""
+class _CallableReturningToolResponse(Protocol):
+    """Callable for tool call interception with handler callback.
+
+    Receives handler callback to execute tool and returns final `ToolMessage` or
+    `Command`.
+    """
 
     def __call__(
-        self, request: ModelRequest, state: StateT_contra, runtime: Runtime[ContextT]
-    ) -> str | Awaitable[str]:
-        """Generate a system prompt string based on the request, state, and runtime."""
+        self,
+        request: ToolCallRequest,
+        handler: Callable[[ToolCallRequest], ToolMessage | Command],
+    ) -> ToolMessage | Command:
+        """Intercept tool execution via handler callback."""
         ...
 
 
@@ -348,22 +641,22 @@ def before_model(
     Callable[[_CallableWithStateAndRuntime[StateT, ContextT]], AgentMiddleware[StateT, ContextT]]
     | AgentMiddleware[StateT, ContextT]
 ):
-    """Decorator used to dynamically create a middleware with the before_model hook.
+    """Decorator used to dynamically create a middleware with the `before_model` hook.
 
     Args:
         func: The function to be decorated. Must accept:
             `state: StateT, runtime: Runtime[ContextT]` - State and runtime context
         state_schema: Optional custom state schema type. If not provided, uses the default
-            AgentState schema.
+            `AgentState` schema.
         tools: Optional list of additional tools to register with this middleware.
         can_jump_to: Optional list of valid jump destinations for conditional edges.
-            Valid values are: "tools", "model", "end"
+            Valid values are: `"tools"`, `"model"`, `"end"`
         name: Optional name for the generated middleware class. If not provided,
             uses the decorated function's name.
 
     Returns:
-        Either an AgentMiddleware instance (if func is provided directly) or a decorator function
-        that can be applied to a function its wrapping.
+        Either an `AgentMiddleware` instance (if func is provided directly) or a
+        decorator function that can be applied to a function it is wrapping.
 
     The decorated function should return:
         - `dict[str, Any]` - State updates to merge into the agent state
@@ -460,143 +753,6 @@ def before_model(
     return decorator
 
 
-@overload
-def modify_model_request(
-    func: _CallableWithModelRequestAndStateAndRuntime[StateT, ContextT],
-) -> AgentMiddleware[StateT, ContextT]: ...
-
-
-@overload
-def modify_model_request(
-    func: None = None,
-    *,
-    state_schema: type[StateT] | None = None,
-    tools: list[BaseTool] | None = None,
-    name: str | None = None,
-) -> Callable[
-    [_CallableWithModelRequestAndStateAndRuntime[StateT, ContextT]],
-    AgentMiddleware[StateT, ContextT],
-]: ...
-
-
-def modify_model_request(
-    func: _CallableWithModelRequestAndStateAndRuntime[StateT, ContextT] | None = None,
-    *,
-    state_schema: type[StateT] | None = None,
-    tools: list[BaseTool] | None = None,
-    name: str | None = None,
-) -> (
-    Callable[
-        [_CallableWithModelRequestAndStateAndRuntime[StateT, ContextT]],
-        AgentMiddleware[StateT, ContextT],
-    ]
-    | AgentMiddleware[StateT, ContextT]
-):
-    r"""Decorator used to dynamically create a middleware with the modify_model_request hook.
-
-    Args:
-        func: The function to be decorated. Must accept:
-            `request: ModelRequest, state: StateT, runtime: Runtime[ContextT]` -
-            Model request, state, and runtime context
-        state_schema: Optional custom state schema type. If not provided, uses the default
-            AgentState schema.
-        tools: Optional list of additional tools to register with this middleware.
-        name: Optional name for the generated middleware class. If not provided,
-            uses the decorated function's name.
-
-    Returns:
-        Either an AgentMiddleware instance (if func is provided) or a decorator function
-        that can be applied to a function.
-
-    The decorated function should return:
-        - `ModelRequest` - The modified model request to be sent to the language model
-
-    Examples:
-        Basic usage to modify system prompt:
-        ```python
-        @modify_model_request
-        def add_context_to_prompt(
-            request: ModelRequest, state: AgentState, runtime: Runtime
-        ) -> ModelRequest:
-            if request.system_prompt:
-                request.system_prompt += "\n\nAdditional context: ..."
-            else:
-                request.system_prompt = "Additional context: ..."
-            return request
-        ```
-
-        Usage with runtime and custom model settings:
-        ```python
-        @modify_model_request
-        def dynamic_model_settings(
-            request: ModelRequest, state: AgentState, runtime: Runtime
-        ) -> ModelRequest:
-            # Use a different model based on user subscription tier
-            if runtime.context.get("subscription_tier") == "premium":
-                request.model = "gpt-4o"
-            else:
-                request.model = "gpt-4o-mini"
-
-            return request
-        ```
-    """
-
-    def decorator(
-        func: _CallableWithModelRequestAndStateAndRuntime[StateT, ContextT],
-    ) -> AgentMiddleware[StateT, ContextT]:
-        is_async = iscoroutinefunction(func)
-
-        if is_async:
-
-            async def async_wrapped(
-                self: AgentMiddleware[StateT, ContextT],  # noqa: ARG001
-                request: ModelRequest,
-                state: StateT,
-                runtime: Runtime[ContextT],
-            ) -> ModelRequest:
-                return await func(request, state, runtime)  # type: ignore[misc]
-
-            middleware_name = name or cast(
-                "str", getattr(func, "__name__", "ModifyModelRequestMiddleware")
-            )
-
-            return type(
-                middleware_name,
-                (AgentMiddleware,),
-                {
-                    "state_schema": state_schema or AgentState,
-                    "tools": tools or [],
-                    "amodify_model_request": async_wrapped,
-                },
-            )()
-
-        def wrapped(
-            self: AgentMiddleware[StateT, ContextT],  # noqa: ARG001
-            request: ModelRequest,
-            state: StateT,
-            runtime: Runtime[ContextT],
-        ) -> ModelRequest:
-            return func(request, state, runtime)  # type: ignore[return-value]
-
-        middleware_name = name or cast(
-            "str", getattr(func, "__name__", "ModifyModelRequestMiddleware")
-        )
-
-        return type(
-            middleware_name,
-            (AgentMiddleware,),
-            {
-                "state_schema": state_schema or AgentState,
-                "tools": tools or [],
-                "modify_model_request": wrapped,
-            },
-        )()
-
-    if func is not None:
-        return decorator(func)
-    return decorator
-
-
 @overload
 def after_model(
     func: _CallableWithStateAndRuntime[StateT, ContextT],
@@ -627,22 +783,22 @@ def after_model(
     Callable[[_CallableWithStateAndRuntime[StateT, ContextT]], AgentMiddleware[StateT, ContextT]]
     | AgentMiddleware[StateT, ContextT]
 ):
-    """Decorator used to dynamically create a middleware with the after_model hook.
+    """Decorator used to dynamically create a middleware with the `after_model` hook.
 
     Args:
         func: The function to be decorated. Must accept:
             `state: StateT, runtime: Runtime[ContextT]` - State and runtime context
-        state_schema: Optional custom state schema type. If not provided, uses the default
-            AgentState schema.
+        state_schema: Optional custom state schema type. If not provided, uses the
+            default `AgentState` schema.
         tools: Optional list of additional tools to register with this middleware.
         can_jump_to: Optional list of valid jump destinations for conditional edges.
-            Valid values are: "tools", "model", "end"
+            Valid values are: `"tools"`, `"model"`, `"end"`
         name: Optional name for the generated middleware class. If not provided,
             uses the decorated function's name.
 
     Returns:
-        Either an AgentMiddleware instance (if func is provided) or a decorator function
-        that can be applied to a function.
+        Either an `AgentMiddleware` instance (if func is provided) or a decorator
+        function that can be applied to a function.
 
     The decorated function should return:
         - `dict[str, Any]` - State updates to merge into the agent state
@@ -758,22 +914,22 @@ def before_agent(
     Callable[[_CallableWithStateAndRuntime[StateT, ContextT]], AgentMiddleware[StateT, ContextT]]
     | AgentMiddleware[StateT, ContextT]
 ):
-    """Decorator used to dynamically create a middleware with the before_agent hook.
+    """Decorator used to dynamically create a middleware with the `before_agent` hook.
 
     Args:
         func: The function to be decorated. Must accept:
             `state: StateT, runtime: Runtime[ContextT]` - State and runtime context
-        state_schema: Optional custom state schema type. If not provided, uses the default
-            AgentState schema.
+        state_schema: Optional custom state schema type. If not provided, uses the
+            default `AgentState` schema.
         tools: Optional list of additional tools to register with this middleware.
         can_jump_to: Optional list of valid jump destinations for conditional edges.
-            Valid values are: "tools", "model", "end"
+            Valid values are: `"tools"`, `"model"`, `"end"`
         name: Optional name for the generated middleware class. If not provided,
             uses the decorated function's name.
 
     Returns:
-        Either an AgentMiddleware instance (if func is provided directly) or a decorator function
-        that can be applied to a function its wrapping.
+        Either an `AgentMiddleware` instance (if func is provided directly) or a
+        decorator function that can be applied to a function it is wrapping.
 
     The decorated function should return:
         - `dict[str, Any]` - State updates to merge into the agent state
@@ -900,22 +1056,22 @@ def after_agent(
     Callable[[_CallableWithStateAndRuntime[StateT, ContextT]], AgentMiddleware[StateT, ContextT]]
     | AgentMiddleware[StateT, ContextT]
 ):
-    """Decorator used to dynamically create a middleware with the after_agent hook.
+    """Decorator used to dynamically create a middleware with the `after_agent` hook.
 
     Args:
         func: The function to be decorated. Must accept:
             `state: StateT, runtime: Runtime[ContextT]` - State and runtime context
-        state_schema: Optional custom state schema type. If not provided, uses the default
-            AgentState schema.
+        state_schema: Optional custom state schema type. If not provided, uses the
+            default `AgentState` schema.
         tools: Optional list of additional tools to register with this middleware.
         can_jump_to: Optional list of valid jump destinations for conditional edges.
-            Valid values are: "tools", "model", "end"
+            Valid values are: `"tools"`, `"model"`, `"end"`
         name: Optional name for the generated middleware class. If not provided,
             uses the decorated function's name.
 
     Returns:
-        Either an AgentMiddleware instance (if func is provided) or a decorator function
-        that can be applied to a function.
+        Either an `AgentMiddleware` instance (if func is provided) or a decorator
+        function that can be applied to a function.
 
     The decorated function should return:
         - `dict[str, Any]` - State updates to merge into the agent state
@@ -1027,14 +1183,13 @@ def dynamic_prompt(
 ):
     """Decorator used to dynamically generate system prompts for the model.
 
-    This is a convenience decorator that creates middleware using `modify_model_request`
+    This is a convenience decorator that creates middleware using `wrap_model_call`
     specifically for dynamic prompt generation. The decorated function should return
     a string that will be set as the system prompt for the model request.
 
     Args:
         func: The function to be decorated. Must accept:
-            `request: ModelRequest, state: StateT, runtime: Runtime[ContextT]` -
-            Model request, state, and runtime context
+            `request: ModelRequest` - Model request (contains state and runtime)
 
     Returns:
         Either an AgentMiddleware instance (if func is provided) or a decorator function
@@ -1047,16 +1202,16 @@ def dynamic_prompt(
         Basic usage with dynamic content:
         ```python
         @dynamic_prompt
-        def my_prompt(request: ModelRequest, state: AgentState, runtime: Runtime) -> str:
-            user_name = runtime.context.get("user_name", "User")
+        def my_prompt(request: ModelRequest) -> str:
+            user_name = request.runtime.context.get("user_name", "User")
             return f"You are a helpful assistant helping {user_name}."
         ```
 
         Using state to customize the prompt:
         ```python
         @dynamic_prompt
-        def context_aware_prompt(request: ModelRequest, state: AgentState, runtime: Runtime) -> str:
-            msg_count = len(state["messages"])
+        def context_aware_prompt(request: ModelRequest) -> str:
+            msg_count = len(request.state["messages"])
             if msg_count > 10:
                 return "You are in a long conversation. Be concise."
             return "You are a helpful assistant."
@@ -1078,12 +1233,11 @@ def dynamic_prompt(
             async def async_wrapped(
                 self: AgentMiddleware[StateT, ContextT],  # noqa: ARG001
                 request: ModelRequest,
-                state: StateT,
-                runtime: Runtime[ContextT],
-            ) -> ModelRequest:
-                prompt = await func(request, state, runtime)  # type: ignore[misc]
+                handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
+            ) -> ModelCallResult:
+                prompt = await func(request)  # type: ignore[misc]
                 request.system_prompt = prompt
-                return request
+                return await handler(request)
 
             middleware_name = cast("str", getattr(func, "__name__", "DynamicPromptMiddleware"))
 
@@ -1093,19 +1247,28 @@ def dynamic_prompt(
                 {
                     "state_schema": AgentState,
                     "tools": [],
-                    "amodify_model_request": async_wrapped,
+                    "awrap_model_call": async_wrapped,
                 },
             )()
 
         def wrapped(
             self: AgentMiddleware[StateT, ContextT],  # noqa: ARG001
             request: ModelRequest,
-            state: StateT,
-            runtime: Runtime[ContextT],
-        ) -> ModelRequest:
-            prompt = cast("str", func(request, state, runtime))
+            handler: Callable[[ModelRequest], ModelResponse],
+        ) -> ModelCallResult:
+            prompt = cast("str", func(request))
             request.system_prompt = prompt
-            return request
+            return handler(request)
+
+        async def async_wrapped_from_sync(
+            self: AgentMiddleware[StateT, ContextT],  # noqa: ARG001
+            request: ModelRequest,
+            handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
+        ) -> ModelCallResult:
+            # Delegate to sync function
+            prompt = cast("str", func(request))
+            request.system_prompt = prompt
+            return await handler(request)
 
         middleware_name = cast("str", getattr(func, "__name__", "DynamicPromptMiddleware"))
 
@@ -1115,7 +1278,301 @@ def dynamic_prompt(
             {
                 "state_schema": AgentState,
                 "tools": [],
-                "modify_model_request": wrapped,
+                "wrap_model_call": wrapped,
+                "awrap_model_call": async_wrapped_from_sync,
+            },
+        )()
+
+    if func is not None:
+        return decorator(func)
+    return decorator
+
+
+@overload
+def wrap_model_call(
+    func: _CallableReturningModelResponse[StateT, ContextT],
+) -> AgentMiddleware[StateT, ContextT]: ...
+
+
+@overload
+def wrap_model_call(
+    func: None = None,
+    *,
+    state_schema: type[StateT] | None = None,
+    tools: list[BaseTool] | None = None,
+    name: str | None = None,
+) -> Callable[
+    [_CallableReturningModelResponse[StateT, ContextT]],
+    AgentMiddleware[StateT, ContextT],
+]: ...
+
+
+def wrap_model_call(
+    func: _CallableReturningModelResponse[StateT, ContextT] | None = None,
+    *,
+    state_schema: type[StateT] | None = None,
+    tools: list[BaseTool] | None = None,
+    name: str | None = None,
+) -> (
+    Callable[
+        [_CallableReturningModelResponse[StateT, ContextT]],
+        AgentMiddleware[StateT, ContextT],
+    ]
+    | AgentMiddleware[StateT, ContextT]
+):
+    """Create middleware with `wrap_model_call` hook from a function.
+
+    Converts a function with handler callback into middleware that can intercept
+    model calls, implement retry logic, handle errors, and rewrite responses.
+
+    Args:
+        func: Function accepting (request, handler) that calls handler(request)
+            to execute the model and returns `ModelResponse` or `AIMessage`.
+            Request contains state and runtime.
+        state_schema: Custom state schema. Defaults to `AgentState`.
+        tools: Additional tools to register with this middleware.
+        name: Middleware class name. Defaults to function name.
+
+    Returns:
+        `AgentMiddleware` instance if func provided, otherwise a decorator.
+
+    Examples:
+        Basic retry logic:
+        ```python
+        @wrap_model_call
+        def retry_on_error(request, handler):
+            max_retries = 3
+            for attempt in range(max_retries):
+                try:
+                    return handler(request)
+                except Exception:
+                    if attempt == max_retries - 1:
+                        raise
+        ```
+
+        Model fallback:
+        ```python
+        @wrap_model_call
+        def fallback_model(request, handler):
+            # Try primary model
+            try:
+                return handler(request)
+            except Exception:
+                pass
+
+            # Try fallback model
+            request.model = fallback_model_instance
+            return handler(request)
+        ```
+
+        Rewrite response content (full ModelResponse):
+        ```python
+        @wrap_model_call
+        def uppercase_responses(request, handler):
+            response = handler(request)
+            ai_msg = response.result[0]
+            return ModelResponse(
+                result=[AIMessage(content=ai_msg.content.upper())],
+                structured_response=response.structured_response,
+            )
+        ```
+
+        Simple AIMessage return (converted automatically):
+        ```python
+        @wrap_model_call
+        def simple_response(request, handler):
+            # AIMessage is automatically converted to ModelResponse
+            return AIMessage(content="Simple response")
+        ```
+    """
+
+    def decorator(
+        func: _CallableReturningModelResponse[StateT, ContextT],
+    ) -> AgentMiddleware[StateT, ContextT]:
+        is_async = iscoroutinefunction(func)
+
+        if is_async:
+
+            async def async_wrapped(
+                self: AgentMiddleware[StateT, ContextT],  # noqa: ARG001
+                request: ModelRequest,
+                handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
+            ) -> ModelCallResult:
+                return await func(request, handler)  # type: ignore[misc, arg-type]
+
+            middleware_name = name or cast(
+                "str", getattr(func, "__name__", "WrapModelCallMiddleware")
+            )
+
+            return type(
+                middleware_name,
+                (AgentMiddleware,),
+                {
+                    "state_schema": state_schema or AgentState,
+                    "tools": tools or [],
+                    "awrap_model_call": async_wrapped,
+                },
+            )()
+
+        def wrapped(
+            self: AgentMiddleware[StateT, ContextT],  # noqa: ARG001
+            request: ModelRequest,
+            handler: Callable[[ModelRequest], ModelResponse],
+        ) -> ModelCallResult:
+            return func(request, handler)
+
+        middleware_name = name or cast("str", getattr(func, "__name__", "WrapModelCallMiddleware"))
+
+        return type(
+            middleware_name,
+            (AgentMiddleware,),
+            {
+                "state_schema": state_schema or AgentState,
+                "tools": tools or [],
+                "wrap_model_call": wrapped,
+            },
+        )()
+
+    if func is not None:
+        return decorator(func)
+    return decorator
+
+
+@overload
+def wrap_tool_call(
+    func: _CallableReturningToolResponse,
+) -> AgentMiddleware: ...
+
+
+@overload
+def wrap_tool_call(
+    func: None = None,
+    *,
+    tools: list[BaseTool] | None = None,
+    name: str | None = None,
+) -> Callable[
+    [_CallableReturningToolResponse],
+    AgentMiddleware,
+]: ...
+
+
+def wrap_tool_call(
+    func: _CallableReturningToolResponse | None = None,
+    *,
+    tools: list[BaseTool] | None = None,
+    name: str | None = None,
+) -> (
+    Callable[
+        [_CallableReturningToolResponse],
+        AgentMiddleware,
+    ]
+    | AgentMiddleware
+):
+    """Create middleware with `wrap_tool_call` hook from a function.
+
+    Converts a function with handler callback into middleware that can intercept
+    tool calls, implement retry logic, monitor execution, and modify responses.
+
+    Args:
+        func: Function accepting (request, handler) that calls
+            handler(request) to execute the tool and returns final `ToolMessage` or
+            `Command`. Can be sync or async.
+        tools: Additional tools to register with this middleware.
+        name: Middleware class name. Defaults to function name.
+
+    Returns:
+        `AgentMiddleware` instance if func provided, otherwise a decorator.
+
+    Examples:
+        Retry logic:
+        ```python
+        @wrap_tool_call
+        def retry_on_error(request, handler):
+            max_retries = 3
+            for attempt in range(max_retries):
+                try:
+                    return handler(request)
+                except Exception:
+                    if attempt == max_retries - 1:
+                        raise
+        ```
+
+        Async retry logic:
+        ```python
+        @wrap_tool_call
+        async def async_retry(request, handler):
+            for attempt in range(3):
+                try:
+                    return await handler(request)
+                except Exception:
+                    if attempt == 2:
+                        raise
+        ```
+
+        Modify request:
+        ```python
+        @wrap_tool_call
+        def modify_args(request, handler):
+            request.tool_call["args"]["value"] *= 2
+            return handler(request)
+        ```
+
+        Short-circuit with cached result:
+        ```python
+        @wrap_tool_call
+        def with_cache(request, handler):
+            if cached := get_cache(request):
+                return ToolMessage(content=cached, tool_call_id=request.tool_call["id"])
+            result = handler(request)
+            save_cache(request, result)
+            return result
+        ```
+    """
+
+    def decorator(
+        func: _CallableReturningToolResponse,
+    ) -> AgentMiddleware:
+        is_async = iscoroutinefunction(func)
+
+        if is_async:
+
+            async def async_wrapped(
+                self: AgentMiddleware,  # noqa: ARG001
+                request: ToolCallRequest,
+                handler: Callable[[ToolCallRequest], Awaitable[ToolMessage | Command]],
+            ) -> ToolMessage | Command:
+                return await func(request, handler)  # type: ignore[arg-type,misc]
+
+            middleware_name = name or cast(
+                "str", getattr(func, "__name__", "WrapToolCallMiddleware")
+            )
+
+            return type(
+                middleware_name,
+                (AgentMiddleware,),
+                {
+                    "state_schema": AgentState,
+                    "tools": tools or [],
+                    "awrap_tool_call": async_wrapped,
+                },
+            )()
+
+        def wrapped(
+            self: AgentMiddleware,  # noqa: ARG001
+            request: ToolCallRequest,
+            handler: Callable[[ToolCallRequest], ToolMessage | Command],
+        ) -> ToolMessage | Command:
+            return func(request, handler)
+
+        middleware_name = name or cast("str", getattr(func, "__name__", "WrapToolCallMiddleware"))
+
+        return type(
+            middleware_name,
+            (AgentMiddleware,),
+            {
+                "state_schema": AgentState,
+                "tools": tools or [],
+                "wrap_tool_call": wrapped,
             },
         )()