langchain 1.0.0a13__py3-none-any.whl → 1.0.0a15__py3-none-any.whl

langchain/agents/middleware/tool_call_limit.py

@@ -2,16 +2,37 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Literal
+from typing import TYPE_CHECKING, Annotated, Any, Literal
 
 from langchain_core.messages import AIMessage, AnyMessage, HumanMessage
+from langgraph.channels.untracked_value import UntrackedValue
+from typing_extensions import NotRequired
 
-from langchain.agents.middleware.types import AgentMiddleware, AgentState, hook_config
+from langchain.agents.middleware.types import (
+    AgentMiddleware,
+    AgentState,
+    PrivateStateAttr,
+    hook_config,
+)
 
 if TYPE_CHECKING:
     from langgraph.runtime import Runtime
 
 
+class ToolCallLimitState(AgentState):
+    """State schema for ToolCallLimitMiddleware.
+
+    Extends AgentState with tool call tracking fields.
+
+    The count fields are dictionaries mapping tool names to execution counts.
+    This allows multiple middleware instances to track different tools independently.
+    The special key "__all__" is used for tracking all tool calls globally.
+    """
+
+    thread_tool_call_count: NotRequired[Annotated[dict[str, int], PrivateStateAttr]]
+    run_tool_call_count: NotRequired[Annotated[dict[str, int], UntrackedValue, PrivateStateAttr]]
+
+
 def _count_tool_calls_in_messages(messages: list[AnyMessage], tool_name: str | None = None) -> int:
     """Count tool calls in a list of messages.
 
@@ -124,18 +145,18 @@ class ToolCallLimitExceededError(Exception):
         super().__init__(msg)
 
 
-class ToolCallLimitMiddleware(AgentMiddleware):
+class ToolCallLimitMiddleware(AgentMiddleware[ToolCallLimitState, Any]):
     """Middleware that tracks tool call counts and enforces limits.
 
     This middleware monitors the number of tool calls made during agent execution
     and can terminate the agent when specified limits are reached. It supports
     both thread-level and run-level call counting with configurable exit behaviors.
 
-    Thread-level: The middleware counts all tool calls in the entire message history
-    and persists this count across multiple runs (invocations) of the agent.
+    Thread-level: The middleware tracks the total number of tool calls and persists
+    call count across multiple runs (invocations) of the agent.
 
-    Run-level: The middleware counts tool calls made after the last HumanMessage,
-    representing the current run (invocation) of the agent.
+    Run-level: The middleware tracks the number of tool calls made during a single
+    run (invocation) of the agent.
 
     Example:
         ```python
@@ -157,6 +178,8 @@ class ToolCallLimitMiddleware(AgentMiddleware):
         ```
     """
 
+    state_schema = ToolCallLimitState
+
     def __init__(
         self,
         *,
@@ -211,11 +234,11 @@ class ToolCallLimitMiddleware(AgentMiddleware):
         return base_name
 
     @hook_config(can_jump_to=["end"])
-    def before_model(self, state: AgentState, runtime: Runtime) -> dict[str, Any] | None:  # noqa: ARG002
+    def before_model(self, state: ToolCallLimitState, runtime: Runtime) -> dict[str, Any] | None:  # noqa: ARG002
         """Check tool call limits before making a model call.
 
         Args:
-            state: The current agent state containing messages.
+            state: The current agent state containing tool call counts.
             runtime: The langgraph runtime.
 
         Returns:
@@ -226,14 +249,14 @@ class ToolCallLimitMiddleware(AgentMiddleware):
             ToolCallLimitExceededError: If limits are exceeded and exit_behavior
                 is "error".
         """
-        messages = state.get("messages", [])
+        # Get the count key for this middleware instance
+        count_key = self.tool_name if self.tool_name else "__all__"
 
-        # Count tool calls in entire thread
-        thread_count = _count_tool_calls_in_messages(messages, self.tool_name)
+        thread_counts = state.get("thread_tool_call_count", {})
+        run_counts = state.get("run_tool_call_count", {})
 
-        # Count tool calls in current run (after last HumanMessage)
-        run_messages = _get_run_messages(messages)
-        run_count = _count_tool_calls_in_messages(run_messages, self.tool_name)
+        thread_count = thread_counts.get(count_key, 0)
+        run_count = run_counts.get(count_key, 0)
 
         # Check if any limits are exceeded
         thread_limit_exceeded = self.thread_limit is not None and thread_count >= self.thread_limit
@@ -258,3 +281,53 @@ class ToolCallLimitMiddleware(AgentMiddleware):
                 return {"jump_to": "end", "messages": [limit_ai_message]}
 
         return None
+
+    def after_model(self, state: ToolCallLimitState, runtime: Runtime) -> dict[str, Any] | None:  # noqa: ARG002
+        """Increment tool call counts after a model call (when tool calls are made).
+
+        Args:
+            state: The current agent state.
+            runtime: The langgraph runtime.
+
+        Returns:
+            State updates with incremented tool call counts if tool calls were made.
+        """
+        # Get the last AIMessage to check for tool calls
+        messages = state.get("messages", [])
+        if not messages:
+            return None
+
+        # Find the last AIMessage
+        last_ai_message = None
+        for message in reversed(messages):
+            if isinstance(message, AIMessage):
+                last_ai_message = message
+                break
+
+        if not last_ai_message or not last_ai_message.tool_calls:
+            return None
+
+        # Count relevant tool calls (filter by tool_name if specified)
+        tool_call_count = 0
+        for tool_call in last_ai_message.tool_calls:
+            if self.tool_name is None or tool_call["name"] == self.tool_name:
+                tool_call_count += 1
+
+        if tool_call_count == 0:
+            return None
+
+        # Get the count key for this middleware instance
+        count_key = self.tool_name if self.tool_name else "__all__"
+
+        # Get current counts
+        thread_counts = state.get("thread_tool_call_count", {}).copy()
+        run_counts = state.get("run_tool_call_count", {}).copy()
+
+        # Increment counts for this key
+        thread_counts[count_key] = thread_counts.get(count_key, 0) + tool_call_count
+        run_counts[count_key] = run_counts.get(count_key, 0) + tool_call_count
+
+        return {
+            "thread_tool_call_count": thread_counts,
+            "run_tool_call_count": run_counts,
+        }

langchain/agents/middleware/types.py

@@ -3,7 +3,7 @@
 from __future__ import annotations
 
 from collections.abc import Awaitable, Callable
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, replace
 from inspect import iscoroutinefunction
 from typing import (
     TYPE_CHECKING,
@@ -21,16 +21,15 @@ if TYPE_CHECKING:
 
     from langchain.tools.tool_node import ToolCallRequest
 
-# needed as top level import for pydantic schema generation on AgentState
+# Needed as top level import for Pydantic schema generation on AgentState
 from typing import TypeAlias
 
 from langchain_core.messages import AIMessage, AnyMessage, BaseMessage, ToolMessage  # 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.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
@@ -62,6 +61,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."""
@@ -76,6 +87,36 @@ class ModelRequest:
     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:
@@ -129,8 +170,6 @@ 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]):
@@ -263,18 +302,35 @@ class AgentMiddleware(Generic[StateT, ContextT]):
                 return AIMessage(content="Simplified response")
             ```
         """
-        raise NotImplementedError
+        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 awrap_model_call(
         self,
         request: ModelRequest,
         handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
     ) -> ModelCallResult:
-        """Async version of wrap_model_call.
+        """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:
             request: Model request to execute (includes state and runtime).
-            handler: Async callback that executes the model request.
+            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:
             ModelCallResult
@@ -291,7 +347,17 @@ class AgentMiddleware(Generic[StateT, ContextT]):
                             raise
             ```
         """
-        raise NotImplementedError
+        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."""
@@ -353,7 +419,77 @@ class AgentMiddleware(Generic[StateT, ContextT]):
                         continue
                     return result
         """
-        raise NotImplementedError
+        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
+            ```
+
+
+            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]):
@@ -1104,6 +1240,16 @@ def dynamic_prompt(
             request.system_prompt = prompt
             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"))
 
         return type(
@@ -1113,6 +1259,7 @@ def dynamic_prompt(
                 "state_schema": AgentState,
                 "tools": [],
                 "wrap_model_call": wrapped,
+                "awrap_model_call": async_wrapped_from_sync,
             },
         )()
 
@@ -1309,6 +1456,7 @@ def wrap_tool_call(
     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.
 
@@ -1316,13 +1464,6 @@ def wrap_tool_call(
         AgentMiddleware instance if func provided, otherwise a decorator.
 
     Examples:
-        Basic passthrough:
-        ```python
-        @wrap_tool_call
-        def passthrough(request, handler):
-            return handler(request)
-        ```
-
         Retry logic:
         ```python
         @wrap_tool_call
@@ -1336,6 +1477,18 @@ def wrap_tool_call(
                         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
@@ -1359,6 +1512,31 @@ def wrap_tool_call(
     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,

langchain/embeddings/__init__.py

@@ -3,10 +3,8 @@
 from langchain_core.embeddings import Embeddings
 
 from langchain.embeddings.base import init_embeddings
-from langchain.embeddings.cache import CacheBackedEmbeddings
 
 __all__ = [
-    "CacheBackedEmbeddings",
     "Embeddings",
     "init_embeddings",
 ]

langchain/messages/__init__.py

@@ -3,29 +3,61 @@
 from langchain_core.messages import (
     AIMessage,
     AIMessageChunk,
+    Annotation,
     AnyMessage,
+    AudioContentBlock,
+    Citation,
+    ContentBlock,
+    DataContentBlock,
+    FileContentBlock,
     HumanMessage,
+    ImageContentBlock,
     InvalidToolCall,
     MessageLikeRepresentation,
+    NonStandardAnnotation,
+    NonStandardContentBlock,
+    PlainTextContentBlock,
+    ReasoningContentBlock,
     RemoveMessage,
+    ServerToolCall,
+    ServerToolCallChunk,
+    ServerToolResult,
     SystemMessage,
+    TextContentBlock,
     ToolCall,
     ToolCallChunk,
     ToolMessage,
+    VideoContentBlock,
     trim_messages,
 )
 
 __all__ = [
     "AIMessage",
     "AIMessageChunk",
+    "Annotation",
     "AnyMessage",
+    "AudioContentBlock",
+    "Citation",
+    "ContentBlock",
+    "DataContentBlock",
+    "FileContentBlock",
     "HumanMessage",
+    "ImageContentBlock",
     "InvalidToolCall",
     "MessageLikeRepresentation",
+    "NonStandardAnnotation",
+    "NonStandardContentBlock",
+    "PlainTextContentBlock",
+    "ReasoningContentBlock",
     "RemoveMessage",
+    "ServerToolCall",
+    "ServerToolCallChunk",
+    "ServerToolResult",
     "SystemMessage",
+    "TextContentBlock",
     "ToolCall",
     "ToolCallChunk",
     "ToolMessage",
+    "VideoContentBlock",
     "trim_messages",
 ]

langchain/tools/__init__.py

@@ -8,11 +8,7 @@ from langchain_core.tools import (
     tool,
 )
 
-from langchain.tools.tool_node import (
-    InjectedState,
-    InjectedStore,
-    ToolNode,
-)
+from langchain.tools.tool_node import InjectedState, InjectedStore
 
 __all__ = [
     "BaseTool",
@@ -21,6 +17,5 @@ __all__ = [
     "InjectedToolArg",
     "InjectedToolCallId",
     "ToolException",
-    "ToolNode",
     "tool",
 ]