langchain 1.0.0a10__py3-none-any.whl → 1.0.0a11__py3-none-any.whl

langchain/agents/middleware/types.py

@@ -2,33 +2,34 @@
 
 from __future__ import annotations
 
+from collections.abc import Callable
 from dataclasses import dataclass, field
-from inspect import signature
+from inspect import iscoroutinefunction
 from typing import (
     TYPE_CHECKING,
     Annotated,
     Any,
-    ClassVar,
     Generic,
     Literal,
     Protocol,
-    TypeAlias,
-    TypeGuard,
     cast,
     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
 from langgraph.channels.ephemeral_value import EphemeralValue
+from langgraph.channels.untracked_value import UntrackedValue
 from langgraph.graph.message import add_messages
-from langgraph.runtime import Runtime
 from langgraph.typing import ContextT
 from typing_extensions import NotRequired, Required, TypedDict, TypeVar
 
 if TYPE_CHECKING:
-    from collections.abc import Callable
-
     from langchain_core.language_models.chat_models import BaseChatModel
     from langchain_core.tools import BaseTool
     from langgraph.runtime import Runtime
@@ -43,6 +44,13 @@ __all__ = [
     "ModelRequest",
     "OmitFromSchema",
     "PublicAgentState",
+    "after_agent",
+    "after_model",
+    "before_agent",
+    "before_model",
+    "dynamic_prompt",
+    "hook_config",
+    "modify_model_request",
 ]
 
 JumpTo = Literal["tools", "model", "end"]
@@ -59,7 +67,7 @@ class ModelRequest:
     system_prompt: str | None
     messages: list[AnyMessage]  # excluding system prompt
     tool_choice: Any | None
-    tools: list[BaseTool]
+    tools: list[BaseTool | dict]
     response_format: ResponseFormat | None
     model_settings: dict[str, Any] = field(default_factory=dict)
 
@@ -90,7 +98,9 @@ class AgentState(TypedDict, Generic[ResponseT]):
 
     messages: Required[Annotated[list[AnyMessage], add_messages]]
     jump_to: NotRequired[Annotated[JumpTo | None, EphemeralValue, PrivateStateAttr]]
-    response: NotRequired[ResponseT]
+    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]):
@@ -100,7 +110,7 @@ class PublicAgentState(TypedDict, Generic[ResponseT]):
     """
 
     messages: Required[Annotated[list[AnyMessage], add_messages]]
-    response: NotRequired[ResponseT]
+    structured_response: NotRequired[ResponseT]
 
 
 StateT = TypeVar("StateT", bound=AgentState, default=AgentState)
@@ -120,15 +130,30 @@ class AgentMiddleware(Generic[StateT, ContextT]):
     tools: list[BaseTool]
     """Additional tools registered by the middleware."""
 
-    before_model_jump_to: ClassVar[list[JumpTo]] = []
-    """Valid jump destinations for before_model hook. Used to establish conditional edges."""
+    @property
+    def name(self) -> str:
+        """The name of the middleware instance.
 
-    after_model_jump_to: ClassVar[list[JumpTo]] = []
-    """Valid jump destinations for after_model hook. Used to establish conditional edges."""
+        Defaults to the class name, but can be overridden for custom naming.
+        """
+        return self.__class__.__name__
+
+    def before_agent(self, state: StateT, runtime: Runtime[ContextT]) -> dict[str, Any] | None:
+        """Logic to run before the agent execution starts."""
+
+    async def abefore_agent(
+        self, state: StateT, runtime: Runtime[ContextT]
+    ) -> dict[str, Any] | None:
+        """Async logic to run before the agent execution starts."""
 
     def before_model(self, state: StateT, runtime: Runtime[ContextT]) -> dict[str, Any] | None:
         """Logic to run before the model is called."""
 
+    async def abefore_model(
+        self, state: StateT, runtime: Runtime[ContextT]
+    ) -> dict[str, Any] | None:
+        """Async logic to run before the model is called."""
+
     def modify_model_request(
         self,
         request: ModelRequest,
@@ -138,16 +163,78 @@ class AgentMiddleware(Generic[StateT, ContextT]):
         """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."""
 
+    async def aafter_model(
+        self, state: StateT, runtime: Runtime[ContextT]
+    ) -> dict[str, Any] | None:
+        """Async logic to run after the model is called."""
+
+    def retry_model_request(
+        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.
+
+        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).
+
+        Returns:
+            ModelRequest: Modified request to retry with.
+            None: Propagate the error (re-raise).
+        """
+        return None
+
+    async def aretry_model_request(
+        self,
+        error: Exception,
+        request: ModelRequest,
+        state: StateT,
+        runtime: Runtime[ContextT],
+        attempt: int,
+    ) -> ModelRequest | None:
+        """Async logic to handle model invocation errors and optionally retry.
+
+        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).
+
+        Returns:
+            ModelRequest: Modified request to retry with.
+            None: Propagate the error (re-raise).
+        """
+        return await run_in_executor(
+            None, self.retry_model_request, error, request, state, runtime, attempt
+        )
 
-class _CallableWithState(Protocol[StateT_contra]):
-    """Callable with AgentState as argument."""
+    def after_agent(self, state: StateT, runtime: Runtime[ContextT]) -> dict[str, Any] | None:
+        """Logic to run after the agent execution completes."""
 
-    def __call__(self, state: StateT_contra) -> dict[str, Any] | Command | None:
-        """Perform some logic with the state."""
-        ...
+    async def aafter_agent(
+        self, state: StateT, runtime: Runtime[ContextT]
+    ) -> dict[str, Any] | None:
+        """Async logic to run after the agent execution completes."""
 
 
 class _CallableWithStateAndRuntime(Protocol[StateT_contra, ContextT]):
@@ -155,53 +242,85 @@ class _CallableWithStateAndRuntime(Protocol[StateT_contra, ContextT]):
 
     def __call__(
         self, state: StateT_contra, runtime: Runtime[ContextT]
-    ) -> dict[str, Any] | Command | None:
+    ) -> dict[str, Any] | Command | None | Awaitable[dict[str, Any] | Command | None]:
         """Perform some logic with the state and runtime."""
         ...
 
 
-class _CallableWithModelRequestAndState(Protocol[StateT_contra]):
-    """Callable with ModelRequest and AgentState as arguments."""
+class _CallableWithModelRequestAndStateAndRuntime(Protocol[StateT_contra, ContextT]):
+    """Callable with ModelRequest, AgentState, and Runtime as arguments."""
 
-    def __call__(self, request: ModelRequest, state: StateT_contra) -> ModelRequest:
-        """Perform some logic with the model request and state."""
+    def __call__(
+        self, request: ModelRequest, state: StateT_contra, runtime: Runtime[ContextT]
+    ) -> ModelRequest | Awaitable[ModelRequest]:
+        """Perform some logic with the model request, state, and runtime."""
         ...
 
 
-class _CallableWithModelRequestAndStateAndRuntime(Protocol[StateT_contra, ContextT]):
-    """Callable with ModelRequest, AgentState, and Runtime as arguments."""
+class _CallableReturningPromptString(Protocol[StateT_contra, ContextT]):
+    """Callable that returns a prompt string given ModelRequest, AgentState, and Runtime."""
 
     def __call__(
         self, request: ModelRequest, state: StateT_contra, runtime: Runtime[ContextT]
-    ) -> ModelRequest:
-        """Perform some logic with the model request, state, and runtime."""
+    ) -> str | Awaitable[str]:
+        """Generate a system prompt string based on the request, state, and runtime."""
         ...
 
 
-_NodeSignature: TypeAlias = (
-    _CallableWithState[StateT] | _CallableWithStateAndRuntime[StateT, ContextT]
-)
-_ModelRequestSignature: TypeAlias = (
-    _CallableWithModelRequestAndState[StateT]
-    | _CallableWithModelRequestAndStateAndRuntime[StateT, ContextT]
-)
+CallableT = TypeVar("CallableT", bound=Callable[..., Any])
 
 
-def is_callable_with_runtime(
-    func: _NodeSignature[StateT, ContextT],
-) -> TypeGuard[_CallableWithStateAndRuntime[StateT, ContextT]]:
-    return "runtime" in signature(func).parameters
+def hook_config(
+    *,
+    can_jump_to: list[JumpTo] | None = None,
+) -> Callable[[CallableT], CallableT]:
+    """Decorator to configure hook behavior in middleware methods.
 
+    Use this decorator on `before_model` or `after_model` methods in middleware classes
+    to configure their behavior. Currently supports specifying which destinations they
+    can jump to, which establishes conditional edges in the agent graph.
 
-def is_callable_with_runtime_and_request(
-    func: _ModelRequestSignature[StateT, ContextT],
-) -> TypeGuard[_CallableWithModelRequestAndStateAndRuntime[StateT, ContextT]]:
-    return "runtime" in signature(func).parameters
+    Args:
+        can_jump_to: Optional list of valid jump destinations. Can be:
+            - "tools": Jump to the tools node
+            - "model": Jump back to the model node
+            - "end": Jump to the end of the graph
+
+    Returns:
+        Decorator function that marks the method with configuration metadata.
+
+    Examples:
+        Using decorator on a class method:
+        ```python
+        class MyMiddleware(AgentMiddleware):
+            @hook_config(can_jump_to=["end", "model"])
+            def before_model(self, state: AgentState) -> dict[str, Any] | None:
+                if some_condition(state):
+                    return {"jump_to": "end"}
+                return None
+        ```
+
+        Alternative: Use the `can_jump_to` parameter in `before_model`/`after_model` decorators:
+        ```python
+        @before_model(can_jump_to=["end"])
+        def conditional_middleware(state: AgentState) -> dict[str, Any] | None:
+            if should_exit(state):
+                return {"jump_to": "end"}
+            return None
+        ```
+    """
+
+    def decorator(func: CallableT) -> CallableT:
+        if can_jump_to is not None:
+            func.__can_jump_to__ = can_jump_to  # type: ignore[attr-defined]
+        return func
+
+    return decorator
 
 
 @overload
 def before_model(
-    func: _NodeSignature[StateT, ContextT],
+    func: _CallableWithStateAndRuntime[StateT, ContextT],
 ) -> AgentMiddleware[StateT, ContextT]: ...
 
 
@@ -211,32 +330,33 @@ def before_model(
     *,
     state_schema: type[StateT] | None = None,
     tools: list[BaseTool] | None = None,
-    jump_to: list[JumpTo] | None = None,
+    can_jump_to: list[JumpTo] | None = None,
     name: str | None = None,
-) -> Callable[[_NodeSignature[StateT, ContextT]], AgentMiddleware[StateT, ContextT]]: ...
+) -> Callable[
+    [_CallableWithStateAndRuntime[StateT, ContextT]], AgentMiddleware[StateT, ContextT]
+]: ...
 
 
 def before_model(
-    func: _NodeSignature[StateT, ContextT] | None = None,
+    func: _CallableWithStateAndRuntime[StateT, ContextT] | None = None,
     *,
     state_schema: type[StateT] | None = None,
     tools: list[BaseTool] | None = None,
-    jump_to: list[JumpTo] | None = None,
+    can_jump_to: list[JumpTo] | None = None,
     name: str | None = None,
 ) -> (
-    Callable[[_NodeSignature[StateT, ContextT]], AgentMiddleware[StateT, ContextT]]
+    Callable[[_CallableWithStateAndRuntime[StateT, ContextT]], AgentMiddleware[StateT, ContextT]]
     | AgentMiddleware[StateT, ContextT]
 ):
     """Decorator used to dynamically create a middleware with the before_model hook.
 
     Args:
-        func: The function to be decorated. Can accept either:
-            - `state: StateT` - Just the agent state
-            - `state: StateT, runtime: Runtime[ContextT]` - State and runtime context
+        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.
         tools: Optional list of additional tools to register with this middleware.
-        jump_to: Optional list of valid jump destinations for conditional edges.
+        can_jump_to: Optional list of valid jump destinations for conditional edges.
             Valid values are: "tools", "model", "end"
         name: Optional name for the generated middleware class. If not provided,
             uses the decorated function's name.
@@ -251,16 +371,16 @@ def before_model(
         - `None` - No state updates or flow control
 
     Examples:
-        Basic usage with state only:
+        Basic usage:
         ```python
         @before_model
-        def log_before_model(state: AgentState) -> None:
+        def log_before_model(state: AgentState, runtime: Runtime) -> None:
             print(f"About to call model with {len(state['messages'])} messages")
         ```
 
-        Advanced usage with runtime and conditional jumping:
+        With conditional jumping:
         ```python
-        @before_model(jump_to=["end"])
+        @before_model(can_jump_to=["end"])
         def conditional_before_model(state: AgentState, runtime: Runtime) -> dict[str, Any] | None:
             if some_condition(state):
                 return {"jump_to": "end"}
@@ -269,34 +389,58 @@ def before_model(
 
         With custom state schema:
         ```python
-        @before_model(
-            state_schema=MyCustomState,
-        )
-        def custom_before_model(state: MyCustomState) -> dict[str, Any]:
+        @before_model(state_schema=MyCustomState)
+        def custom_before_model(state: MyCustomState, runtime: Runtime) -> dict[str, Any]:
             return {"custom_field": "updated_value"}
         ```
     """
 
-    def decorator(func: _NodeSignature[StateT, ContextT]) -> AgentMiddleware[StateT, ContextT]:
-        if is_callable_with_runtime(func):
+    def decorator(
+        func: _CallableWithStateAndRuntime[StateT, ContextT],
+    ) -> AgentMiddleware[StateT, ContextT]:
+        is_async = iscoroutinefunction(func)
 
-            def wrapped_with_runtime(
-                self: AgentMiddleware[StateT, ContextT],  # noqa: ARG001
-                state: StateT,
-                runtime: Runtime[ContextT],
-            ) -> dict[str, Any] | Command | None:
-                return func(state, runtime)
+        func_can_jump_to = (
+            can_jump_to if can_jump_to is not None else getattr(func, "__can_jump_to__", [])
+        )
 
-            wrapped = wrapped_with_runtime
-        else:
+        if is_async:
 
-            def wrapped_without_runtime(
+            async def async_wrapped(
                 self: AgentMiddleware[StateT, ContextT],  # noqa: ARG001
                 state: StateT,
+                runtime: Runtime[ContextT],
             ) -> dict[str, Any] | Command | None:
-                return func(state)  # type: ignore[call-arg]
-
-            wrapped = wrapped_without_runtime  # type: ignore[assignment]
+                return await func(state, runtime)  # type: ignore[misc]
+
+            # Preserve can_jump_to metadata on the wrapped function
+            if func_can_jump_to:
+                async_wrapped.__can_jump_to__ = func_can_jump_to  # type: ignore[attr-defined]
+
+            middleware_name = name or cast(
+                "str", getattr(func, "__name__", "BeforeModelMiddleware")
+            )
+
+            return type(
+                middleware_name,
+                (AgentMiddleware,),
+                {
+                    "state_schema": state_schema or AgentState,
+                    "tools": tools or [],
+                    "abefore_model": async_wrapped,
+                },
+            )()
+
+        def wrapped(
+            self: AgentMiddleware[StateT, ContextT],  # noqa: ARG001
+            state: StateT,
+            runtime: Runtime[ContextT],
+        ) -> dict[str, Any] | Command | None:
+            return func(state, runtime)  # type: ignore[return-value]
+
+        # Preserve can_jump_to metadata on the wrapped function
+        if func_can_jump_to:
+            wrapped.__can_jump_to__ = func_can_jump_to  # type: ignore[attr-defined]
 
         # Use function name as default if no name provided
         middleware_name = name or cast("str", getattr(func, "__name__", "BeforeModelMiddleware"))
@@ -307,7 +451,6 @@ def before_model(
             {
                 "state_schema": state_schema or AgentState,
                 "tools": tools or [],
-                "before_model_jump_to": jump_to or [],
                 "before_model": wrapped,
             },
         )()
@@ -319,7 +462,7 @@ def before_model(
 
 @overload
 def modify_model_request(
-    func: _ModelRequestSignature[StateT, ContextT],
+    func: _CallableWithModelRequestAndStateAndRuntime[StateT, ContextT],
 ) -> AgentMiddleware[StateT, ContextT]: ...
 
 
@@ -330,26 +473,31 @@ def modify_model_request(
     state_schema: type[StateT] | None = None,
     tools: list[BaseTool] | None = None,
     name: str | None = None,
-) -> Callable[[_ModelRequestSignature[StateT, ContextT]], AgentMiddleware[StateT, ContextT]]: ...
+) -> Callable[
+    [_CallableWithModelRequestAndStateAndRuntime[StateT, ContextT]],
+    AgentMiddleware[StateT, ContextT],
+]: ...
 
 
 def modify_model_request(
-    func: _ModelRequestSignature[StateT, ContextT] | None = None,
+    func: _CallableWithModelRequestAndStateAndRuntime[StateT, ContextT] | None = None,
     *,
     state_schema: type[StateT] | None = None,
     tools: list[BaseTool] | None = None,
     name: str | None = None,
 ) -> (
-    Callable[[_ModelRequestSignature[StateT, ContextT]], AgentMiddleware[StateT, ContextT]]
+    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. Can accept either:
-            - `request: ModelRequest, state: StateT` - Model request and agent state
-            - `request: ModelRequest, state: StateT, runtime: Runtime[ContextT]` -
-              Model request, state, and runtime context
+        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.
@@ -367,7 +515,9 @@ def modify_model_request(
         Basic usage to modify system prompt:
         ```python
         @modify_model_request
-        def add_context_to_prompt(request: ModelRequest, state: AgentState) -> ModelRequest:
+        def add_context_to_prompt(
+            request: ModelRequest, state: AgentState, runtime: Runtime
+        ) -> ModelRequest:
             if request.system_prompt:
                 request.system_prompt += "\n\nAdditional context: ..."
             else:
@@ -375,7 +525,7 @@ def modify_model_request(
             return request
         ```
 
-        Advanced usage with runtime and custom model settings:
+        Usage with runtime and custom model settings:
         ```python
         @modify_model_request
         def dynamic_model_settings(
@@ -392,31 +542,42 @@ def modify_model_request(
     """
 
     def decorator(
-        func: _ModelRequestSignature[StateT, ContextT],
+        func: _CallableWithModelRequestAndStateAndRuntime[StateT, ContextT],
     ) -> AgentMiddleware[StateT, ContextT]:
-        if is_callable_with_runtime_and_request(func):
+        is_async = iscoroutinefunction(func)
 
-            def wrapped_with_runtime(
-                self: AgentMiddleware[StateT, ContextT],  # noqa: ARG001
-                request: ModelRequest,
-                state: StateT,
-                runtime: Runtime[ContextT],
-            ) -> ModelRequest:
-                return func(request, state, runtime)
+        if is_async:
 
-            wrapped = wrapped_with_runtime
-        else:
-
-            def wrapped_without_runtime(
+            async def async_wrapped(
                 self: AgentMiddleware[StateT, ContextT],  # noqa: ARG001
                 request: ModelRequest,
                 state: StateT,
+                runtime: Runtime[ContextT],
             ) -> ModelRequest:
-                return func(request, state)  # type: ignore[call-arg]
-
-            wrapped = wrapped_without_runtime  # type: ignore[assignment]
+                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]
 
-        # Use function name as default if no name provided
         middleware_name = name or cast(
             "str", getattr(func, "__name__", "ModifyModelRequestMiddleware")
         )
@@ -438,7 +599,7 @@ def modify_model_request(
 
 @overload
 def after_model(
-    func: _NodeSignature[StateT, ContextT],
+    func: _CallableWithStateAndRuntime[StateT, ContextT],
 ) -> AgentMiddleware[StateT, ContextT]: ...
 
 
@@ -448,32 +609,33 @@ def after_model(
     *,
     state_schema: type[StateT] | None = None,
     tools: list[BaseTool] | None = None,
-    jump_to: list[JumpTo] | None = None,
+    can_jump_to: list[JumpTo] | None = None,
     name: str | None = None,
-) -> Callable[[_NodeSignature[StateT, ContextT]], AgentMiddleware[StateT, ContextT]]: ...
+) -> Callable[
+    [_CallableWithStateAndRuntime[StateT, ContextT]], AgentMiddleware[StateT, ContextT]
+]: ...
 
 
 def after_model(
-    func: _NodeSignature[StateT, ContextT] | None = None,
+    func: _CallableWithStateAndRuntime[StateT, ContextT] | None = None,
     *,
     state_schema: type[StateT] | None = None,
     tools: list[BaseTool] | None = None,
-    jump_to: list[JumpTo] | None = None,
+    can_jump_to: list[JumpTo] | None = None,
     name: str | None = None,
 ) -> (
-    Callable[[_NodeSignature[StateT, ContextT]], AgentMiddleware[StateT, ContextT]]
+    Callable[[_CallableWithStateAndRuntime[StateT, ContextT]], AgentMiddleware[StateT, ContextT]]
     | AgentMiddleware[StateT, ContextT]
 ):
     """Decorator used to dynamically create a middleware with the after_model hook.
 
     Args:
-        func: The function to be decorated. Can accept either:
-            - `state: StateT` - Just the agent state (includes model response)
-            - `state: StateT, runtime: Runtime[ContextT]` - State and runtime context
+        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.
         tools: Optional list of additional tools to register with this middleware.
-        jump_to: Optional list of valid jump destinations for conditional edges.
+        can_jump_to: Optional list of valid jump destinations for conditional edges.
             Valid values are: "tools", "model", "end"
         name: Optional name for the generated middleware class. If not provided,
             uses the decorated function's name.
@@ -491,41 +653,338 @@ def after_model(
         Basic usage for logging model responses:
         ```python
         @after_model
-        def log_latest_message(state: AgentState) -> None:
+        def log_latest_message(state: AgentState, runtime: Runtime) -> None:
             print(state["messages"][-1].content)
         ```
 
         With custom state schema:
         ```python
         @after_model(state_schema=MyCustomState, name="MyAfterModelMiddleware")
-        def custom_after_model(state: MyCustomState) -> dict[str, Any]:
+        def custom_after_model(state: MyCustomState, runtime: Runtime) -> dict[str, Any]:
             return {"custom_field": "updated_after_model"}
         ```
     """
 
-    def decorator(func: _NodeSignature[StateT, ContextT]) -> AgentMiddleware[StateT, ContextT]:
-        if is_callable_with_runtime(func):
+    def decorator(
+        func: _CallableWithStateAndRuntime[StateT, ContextT],
+    ) -> AgentMiddleware[StateT, ContextT]:
+        is_async = iscoroutinefunction(func)
+        # Extract can_jump_to from decorator parameter or from function metadata
+        func_can_jump_to = (
+            can_jump_to if can_jump_to is not None else getattr(func, "__can_jump_to__", [])
+        )
+
+        if is_async:
 
-            def wrapped_with_runtime(
+            async def async_wrapped(
                 self: AgentMiddleware[StateT, ContextT],  # noqa: ARG001
                 state: StateT,
                 runtime: Runtime[ContextT],
             ) -> dict[str, Any] | Command | None:
-                return func(state, runtime)
+                return await func(state, runtime)  # type: ignore[misc]
+
+            # Preserve can_jump_to metadata on the wrapped function
+            if func_can_jump_to:
+                async_wrapped.__can_jump_to__ = func_can_jump_to  # type: ignore[attr-defined]
+
+            middleware_name = name or cast("str", getattr(func, "__name__", "AfterModelMiddleware"))
+
+            return type(
+                middleware_name,
+                (AgentMiddleware,),
+                {
+                    "state_schema": state_schema or AgentState,
+                    "tools": tools or [],
+                    "aafter_model": async_wrapped,
+                },
+            )()
+
+        def wrapped(
+            self: AgentMiddleware[StateT, ContextT],  # noqa: ARG001
+            state: StateT,
+            runtime: Runtime[ContextT],
+        ) -> dict[str, Any] | Command | None:
+            return func(state, runtime)  # type: ignore[return-value]
+
+        # Preserve can_jump_to metadata on the wrapped function
+        if func_can_jump_to:
+            wrapped.__can_jump_to__ = func_can_jump_to  # type: ignore[attr-defined]
+
+        # Use function name as default if no name provided
+        middleware_name = name or cast("str", getattr(func, "__name__", "AfterModelMiddleware"))
+
+        return type(
+            middleware_name,
+            (AgentMiddleware,),
+            {
+                "state_schema": state_schema or AgentState,
+                "tools": tools or [],
+                "after_model": wrapped,
+            },
+        )()
+
+    if func is not None:
+        return decorator(func)
+    return decorator
+
+
+@overload
+def before_agent(
+    func: _CallableWithStateAndRuntime[StateT, ContextT],
+) -> AgentMiddleware[StateT, ContextT]: ...
+
+
+@overload
+def before_agent(
+    func: None = None,
+    *,
+    state_schema: type[StateT] | None = None,
+    tools: list[BaseTool] | None = None,
+    can_jump_to: list[JumpTo] | None = None,
+    name: str | None = None,
+) -> Callable[
+    [_CallableWithStateAndRuntime[StateT, ContextT]], AgentMiddleware[StateT, ContextT]
+]: ...
+
+
+def before_agent(
+    func: _CallableWithStateAndRuntime[StateT, ContextT] | None = None,
+    *,
+    state_schema: type[StateT] | None = None,
+    tools: list[BaseTool] | None = None,
+    can_jump_to: list[JumpTo] | None = None,
+    name: str | None = None,
+) -> (
+    Callable[[_CallableWithStateAndRuntime[StateT, ContextT]], AgentMiddleware[StateT, ContextT]]
+    | AgentMiddleware[StateT, ContextT]
+):
+    """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.
+        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"
+        name: Optional name for the generated middleware class. If not provided,
+            uses the decorated function's name.
 
-            wrapped = wrapped_with_runtime
-        else:
+    Returns:
+        Either an AgentMiddleware instance (if func is provided directly) or a decorator function
+        that can be applied to a function its wrapping.
 
-            def wrapped_without_runtime(
+    The decorated function should return:
+        - `dict[str, Any]` - State updates to merge into the agent state
+        - `Command` - A command to control flow (e.g., jump to different node)
+        - `None` - No state updates or flow control
+
+    Examples:
+        Basic usage:
+        ```python
+        @before_agent
+        def log_before_agent(state: AgentState, runtime: Runtime) -> None:
+            print(f"Starting agent with {len(state['messages'])} messages")
+        ```
+
+        With conditional jumping:
+        ```python
+        @before_agent(can_jump_to=["end"])
+        def conditional_before_agent(state: AgentState, runtime: Runtime) -> dict[str, Any] | None:
+            if some_condition(state):
+                return {"jump_to": "end"}
+            return None
+        ```
+
+        With custom state schema:
+        ```python
+        @before_agent(state_schema=MyCustomState)
+        def custom_before_agent(state: MyCustomState, runtime: Runtime) -> dict[str, Any]:
+            return {"custom_field": "initialized_value"}
+        ```
+    """
+
+    def decorator(
+        func: _CallableWithStateAndRuntime[StateT, ContextT],
+    ) -> AgentMiddleware[StateT, ContextT]:
+        is_async = iscoroutinefunction(func)
+
+        func_can_jump_to = (
+            can_jump_to if can_jump_to is not None else getattr(func, "__can_jump_to__", [])
+        )
+
+        if is_async:
+
+            async def async_wrapped(
                 self: AgentMiddleware[StateT, ContextT],  # noqa: ARG001
                 state: StateT,
+                runtime: Runtime[ContextT],
             ) -> dict[str, Any] | Command | None:
-                return func(state)  # type: ignore[call-arg]
+                return await func(state, runtime)  # type: ignore[misc]
+
+            # Preserve can_jump_to metadata on the wrapped function
+            if func_can_jump_to:
+                async_wrapped.__can_jump_to__ = func_can_jump_to  # type: ignore[attr-defined]
+
+            middleware_name = name or cast(
+                "str", getattr(func, "__name__", "BeforeAgentMiddleware")
+            )
+
+            return type(
+                middleware_name,
+                (AgentMiddleware,),
+                {
+                    "state_schema": state_schema or AgentState,
+                    "tools": tools or [],
+                    "abefore_agent": async_wrapped,
+                },
+            )()
+
+        def wrapped(
+            self: AgentMiddleware[StateT, ContextT],  # noqa: ARG001
+            state: StateT,
+            runtime: Runtime[ContextT],
+        ) -> dict[str, Any] | Command | None:
+            return func(state, runtime)  # type: ignore[return-value]
+
+        # Preserve can_jump_to metadata on the wrapped function
+        if func_can_jump_to:
+            wrapped.__can_jump_to__ = func_can_jump_to  # type: ignore[attr-defined]
+
+        # Use function name as default if no name provided
+        middleware_name = name or cast("str", getattr(func, "__name__", "BeforeAgentMiddleware"))
+
+        return type(
+            middleware_name,
+            (AgentMiddleware,),
+            {
+                "state_schema": state_schema or AgentState,
+                "tools": tools or [],
+                "before_agent": wrapped,
+            },
+        )()
+
+    if func is not None:
+        return decorator(func)
+    return decorator
+
+
+@overload
+def after_agent(
+    func: _CallableWithStateAndRuntime[StateT, ContextT],
+) -> AgentMiddleware[StateT, ContextT]: ...
+
+
+@overload
+def after_agent(
+    func: None = None,
+    *,
+    state_schema: type[StateT] | None = None,
+    tools: list[BaseTool] | None = None,
+    can_jump_to: list[JumpTo] | None = None,
+    name: str | None = None,
+) -> Callable[
+    [_CallableWithStateAndRuntime[StateT, ContextT]], AgentMiddleware[StateT, ContextT]
+]: ...
+
+
+def after_agent(
+    func: _CallableWithStateAndRuntime[StateT, ContextT] | None = None,
+    *,
+    state_schema: type[StateT] | None = None,
+    tools: list[BaseTool] | None = None,
+    can_jump_to: list[JumpTo] | None = None,
+    name: str | None = None,
+) -> (
+    Callable[[_CallableWithStateAndRuntime[StateT, ContextT]], AgentMiddleware[StateT, ContextT]]
+    | AgentMiddleware[StateT, ContextT]
+):
+    """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.
+        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"
+        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:
+        - `dict[str, Any]` - State updates to merge into the agent state
+        - `Command` - A command to control flow (e.g., jump to different node)
+        - `None` - No state updates or flow control
+
+    Examples:
+        Basic usage for logging agent completion:
+        ```python
+        @after_agent
+        def log_completion(state: AgentState, runtime: Runtime) -> None:
+            print(f"Agent completed with {len(state['messages'])} messages")
+        ```
+
+        With custom state schema:
+        ```python
+        @after_agent(state_schema=MyCustomState, name="MyAfterAgentMiddleware")
+        def custom_after_agent(state: MyCustomState, runtime: Runtime) -> dict[str, Any]:
+            return {"custom_field": "finalized_value"}
+        ```
+    """
+
+    def decorator(
+        func: _CallableWithStateAndRuntime[StateT, ContextT],
+    ) -> AgentMiddleware[StateT, ContextT]:
+        is_async = iscoroutinefunction(func)
+        # Extract can_jump_to from decorator parameter or from function metadata
+        func_can_jump_to = (
+            can_jump_to if can_jump_to is not None else getattr(func, "__can_jump_to__", [])
+        )
 
-            wrapped = wrapped_without_runtime  # type: ignore[assignment]
+        if is_async:
+
+            async def async_wrapped(
+                self: AgentMiddleware[StateT, ContextT],  # noqa: ARG001
+                state: StateT,
+                runtime: Runtime[ContextT],
+            ) -> dict[str, Any] | Command | None:
+                return await func(state, runtime)  # type: ignore[misc]
+
+            # Preserve can_jump_to metadata on the wrapped function
+            if func_can_jump_to:
+                async_wrapped.__can_jump_to__ = func_can_jump_to  # type: ignore[attr-defined]
+
+            middleware_name = name or cast("str", getattr(func, "__name__", "AfterAgentMiddleware"))
+
+            return type(
+                middleware_name,
+                (AgentMiddleware,),
+                {
+                    "state_schema": state_schema or AgentState,
+                    "tools": tools or [],
+                    "aafter_agent": async_wrapped,
+                },
+            )()
+
+        def wrapped(
+            self: AgentMiddleware[StateT, ContextT],  # noqa: ARG001
+            state: StateT,
+            runtime: Runtime[ContextT],
+        ) -> dict[str, Any] | Command | None:
+            return func(state, runtime)  # type: ignore[return-value]
+
+        # Preserve can_jump_to metadata on the wrapped function
+        if func_can_jump_to:
+            wrapped.__can_jump_to__ = func_can_jump_to  # type: ignore[attr-defined]
 
         # Use function name as default if no name provided
-        middleware_name = name or cast("str", getattr(func, "__name__", "AfterModelMiddleware"))
+        middleware_name = name or cast("str", getattr(func, "__name__", "AfterAgentMiddleware"))
 
         return type(
             middleware_name,
@@ -533,8 +992,130 @@ def after_model(
             {
                 "state_schema": state_schema or AgentState,
                 "tools": tools or [],
-                "after_model_jump_to": jump_to or [],
-                "after_model": wrapped,
+                "after_agent": wrapped,
+            },
+        )()
+
+    if func is not None:
+        return decorator(func)
+    return decorator
+
+
+@overload
+def dynamic_prompt(
+    func: _CallableReturningPromptString[StateT, ContextT],
+) -> AgentMiddleware[StateT, ContextT]: ...
+
+
+@overload
+def dynamic_prompt(
+    func: None = None,
+) -> Callable[
+    [_CallableReturningPromptString[StateT, ContextT]],
+    AgentMiddleware[StateT, ContextT],
+]: ...
+
+
+def dynamic_prompt(
+    func: _CallableReturningPromptString[StateT, ContextT] | None = None,
+) -> (
+    Callable[
+        [_CallableReturningPromptString[StateT, ContextT]],
+        AgentMiddleware[StateT, ContextT],
+    ]
+    | AgentMiddleware[StateT, ContextT]
+):
+    """Decorator used to dynamically generate system prompts for the model.
+
+    This is a convenience decorator that creates middleware using `modify_model_request`
+    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
+
+    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:
+        - `str` - The system prompt to use for the model request
+
+    Examples:
+        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")
+            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"])
+            if msg_count > 10:
+                return "You are in a long conversation. Be concise."
+            return "You are a helpful assistant."
+        ```
+
+        Using with agent:
+        ```python
+        agent = create_agent(model, middleware=[my_prompt])
+        ```
+    """
+
+    def decorator(
+        func: _CallableReturningPromptString[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:
+                prompt = await func(request, state, runtime)  # type: ignore[misc]
+                request.system_prompt = prompt
+                return request
+
+            middleware_name = cast("str", getattr(func, "__name__", "DynamicPromptMiddleware"))
+
+            return type(
+                middleware_name,
+                (AgentMiddleware,),
+                {
+                    "state_schema": AgentState,
+                    "tools": [],
+                    "amodify_model_request": 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))
+            request.system_prompt = prompt
+            return request
+
+        middleware_name = cast("str", getattr(func, "__name__", "DynamicPromptMiddleware"))
+
+        return type(
+            middleware_name,
+            (AgentMiddleware,),
+            {
+                "state_schema": AgentState,
+                "tools": [],
+                "modify_model_request": wrapped,
             },
         )()