langchain 1.0.0a14__py3-none-any.whl → 1.0.0rc1__py3-none-any.whl

langchain/agents/middleware/model_call_limit.py

@@ -2,16 +2,33 @@
 
 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
+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 ModelCallLimitState(AgentState):
+    """State schema for ModelCallLimitMiddleware.
+
+    Extends AgentState with model call tracking fields.
+    """
+
+    thread_model_call_count: NotRequired[Annotated[int, PrivateStateAttr]]
+    run_model_call_count: NotRequired[Annotated[int, UntrackedValue, PrivateStateAttr]]
+
+
 def _build_limit_exceeded_message(
     thread_count: int,
     run_count: int,
@@ -69,7 +86,7 @@ class ModelCallLimitExceededError(Exception):
         super().__init__(msg)
 
 
-class ModelCallLimitMiddleware(AgentMiddleware):
+class ModelCallLimitMiddleware(AgentMiddleware[ModelCallLimitState, Any]):
     """Middleware that tracks model call counts and enforces limits.
 
     This middleware monitors the number of model calls made during agent execution
@@ -97,6 +114,8 @@ class ModelCallLimitMiddleware(AgentMiddleware):
         ```
     """
 
+    state_schema = ModelCallLimitState
+
     def __init__(
         self,
         *,
@@ -108,17 +127,16 @@ class ModelCallLimitMiddleware(AgentMiddleware):
 
         Args:
             thread_limit: Maximum number of model calls allowed per thread.
-                None means no limit. Defaults to `None`.
+                None means no limit.
             run_limit: Maximum number of model calls allowed per run.
-                None means no limit. Defaults to `None`.
+                None means no limit.
             exit_behavior: What to do when limits are exceeded.
                 - "end": Jump to the end of the agent execution and
                     inject an artificial AI message indicating that the limit was exceeded.
-                - "error": Raise a ModelCallLimitExceededError
-                Defaults to "end".
+                - "error": Raise a `ModelCallLimitExceededError`
 
         Raises:
-            ValueError: If both limits are None or if exit_behavior is invalid.
+            ValueError: If both limits are `None` or if `exit_behavior` is invalid.
         """
         super().__init__()
 
@@ -135,7 +153,7 @@ class ModelCallLimitMiddleware(AgentMiddleware):
         self.exit_behavior = exit_behavior
 
     @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: ModelCallLimitState, runtime: Runtime) -> dict[str, Any] | None:  # noqa: ARG002
         """Check model call limits before making a model call.
 
         Args:
@@ -175,3 +193,18 @@ class ModelCallLimitMiddleware(AgentMiddleware):
                 return {"jump_to": "end", "messages": [limit_ai_message]}
 
         return None
+
+    def after_model(self, state: ModelCallLimitState, runtime: Runtime) -> dict[str, Any] | None:  # noqa: ARG002
+        """Increment model call counts after a model call.
+
+        Args:
+            state: The current agent state.
+            runtime: The langgraph runtime.
+
+        Returns:
+            State updates with incremented call counts.
+        """
+        return {
+            "thread_model_call_count": state.get("thread_model_call_count", 0) + 1,
+            "run_model_call_count": state.get("run_model_call_count", 0) + 1,
+        }

langchain/agents/middleware/model_fallback.py

@@ -13,7 +13,7 @@ from langchain.agents.middleware.types import (
 from langchain.chat_models import init_chat_model
 
 if TYPE_CHECKING:
-    from collections.abc import Callable
+    from collections.abc import Awaitable, Callable
 
     from langchain_core.language_models.chat_models import BaseChatModel
 
@@ -75,8 +75,6 @@ class ModelFallbackMiddleware(AgentMiddleware):
 
         Args:
             request: Initial model request.
-            state: Current agent state.
-            runtime: LangGraph runtime.
             handler: Callback to execute the model.
 
         Returns:
@@ -102,3 +100,38 @@ class ModelFallbackMiddleware(AgentMiddleware):
                 continue
 
         raise last_exception
+
+    async def awrap_model_call(
+        self,
+        request: ModelRequest,
+        handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
+    ) -> ModelCallResult:
+        """Try fallback models in sequence on errors (async version).
+
+        Args:
+            request: Initial model request.
+            handler: Async callback to execute the model.
+
+        Returns:
+            AIMessage from successful model call.
+
+        Raises:
+            Exception: If all models fail, re-raises last exception.
+        """
+        # Try primary model first
+        last_exception: Exception
+        try:
+            return await handler(request)
+        except Exception as e:  # noqa: BLE001
+            last_exception = e
+
+        # Try fallback models
+        for fallback_model in self.models:
+            request.model = fallback_model
+            try:
+                return await handler(request)
+            except Exception as e:  # noqa: BLE001
+                last_exception = e
+                continue
+
+        raise last_exception

langchain/agents/middleware/pii.py

@@ -421,7 +421,7 @@ class PIIMiddleware(AgentMiddleware):
         - `credit_card`: Credit card numbers (validated with Luhn algorithm)
         - `ip`: IP addresses (validated with stdlib)
         - `mac_address`: MAC addresses
-        - `url`: URLs (both http/https and bare URLs)
+        - `url`: URLs (both `http`/`https` and bare URLs)
 
     Strategies:
         - `block`: Raise an exception when PII is detected
@@ -431,12 +431,12 @@ class PIIMiddleware(AgentMiddleware):
 
     Strategy Selection Guide:
 
-        | Strategy | Preserves Identity? | Best For                                |
-        | -------- | ------------------- | --------------------------------------- |
-        | `block`  | N/A                 | Avoid PII completely                    |
-        | `redact` | No                  | General compliance, log sanitization    |
-        | `mask`   | No                  | Human readability, customer service UIs |
-        | `hash`   | Yes (pseudonymous)  | Analytics, debugging                    |
+    | Strategy | Preserves Identity? | Best For                                |
+    | -------- | ------------------- | --------------------------------------- |
+    | `block`  | N/A                 | Avoid PII completely                    |
+    | `redact` | No                  | General compliance, log sanitization    |
+    | `mask`   | No                  | Human readability, customer service UIs |
+    | `hash`   | Yes (pseudonymous)  | Analytics, debugging                    |
 
     Example:
         ```python

langchain/agents/middleware/{planning.py → todo.py}

@@ -6,7 +6,7 @@ from __future__ import annotations
 from typing import TYPE_CHECKING, Annotated, Literal
 
 if TYPE_CHECKING:
-    from collections.abc import Callable
+    from collections.abc import Awaitable, Callable
 
 from langchain_core.messages import ToolMessage
 from langchain_core.tools import tool
@@ -126,7 +126,7 @@ def write_todos(todos: list[Todo], tool_call_id: Annotated[str, InjectedToolCall
     )
 
 
-class PlanningMiddleware(AgentMiddleware):
+class TodoListMiddleware(AgentMiddleware):
     """Middleware that provides todo list management capabilities to agents.
 
     This middleware adds a `write_todos` tool that allows agents to create and manage
@@ -139,10 +139,10 @@ class PlanningMiddleware(AgentMiddleware):
 
     Example:
         ```python
-        from langchain.agents.middleware.planning import PlanningMiddleware
+        from langchain.agents.middleware.todo import TodoListMiddleware
         from langchain.agents import create_agent
 
-        agent = create_agent("openai:gpt-4o", middleware=[PlanningMiddleware()])
+        agent = create_agent("openai:gpt-4o", middleware=[TodoListMiddleware()])
 
         # Agent now has access to write_todos tool and todo state tracking
         result = await agent.invoke({"messages": [HumanMessage("Help me refactor my codebase")]})
@@ -165,7 +165,7 @@ class PlanningMiddleware(AgentMiddleware):
         system_prompt: str = WRITE_TODOS_SYSTEM_PROMPT,
         tool_description: str = WRITE_TODOS_TOOL_DESCRIPTION,
     ) -> None:
-        """Initialize the PlanningMiddleware with optional custom prompts.
+        """Initialize the TodoListMiddleware with optional custom prompts.
 
         Args:
             system_prompt: Custom system prompt to guide the agent on using the todo tool.
@@ -204,3 +204,16 @@ class PlanningMiddleware(AgentMiddleware):
             else self.system_prompt
         )
         return handler(request)
+
+    async def awrap_model_call(
+        self,
+        request: ModelRequest,
+        handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
+    ) -> ModelCallResult:
+        """Update the system prompt to include the todo system prompt (async version)."""
+        request.system_prompt = (
+            request.system_prompt + "\n\n" + self.system_prompt
+            if request.system_prompt
+            else self.system_prompt
+        )
+        return await handler(request)

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,
         *,
@@ -181,7 +204,7 @@ class ToolCallLimitMiddleware(AgentMiddleware):
                 Defaults to "end".
 
         Raises:
-            ValueError: If both limits are None or if exit_behavior is invalid.
+            ValueError: If both limits are `None` or if `exit_behavior` is invalid.
         """
         super().__init__()
 
@@ -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/tool_emulator.py

@@ -123,7 +123,7 @@ class LLMToolEmulator(AgentMiddleware):
 
         # Extract tool information for emulation
         tool_args = request.tool_call["args"]
-        tool_description = request.tool.description
+        tool_description = request.tool.description if request.tool else "No description available"
 
         # Build prompt for emulator LLM
         prompt = (
@@ -175,7 +175,7 @@ class LLMToolEmulator(AgentMiddleware):
 
         # Extract tool information for emulation
         tool_args = request.tool_call["args"]
-        tool_description = request.tool.description
+        tool_description = request.tool.description if request.tool else "No description available"
 
         # Build prompt for emulator LLM
         prompt = (