langchain 1.0.5__py3-none-any.whl → 1.2.4__py3-none-any.whl

langchain/agents/middleware/todo.py

@@ -1,17 +1,18 @@
 """Planning and task management middleware for agents."""
-# ruff: noqa: E501
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Annotated, Literal
+from typing import TYPE_CHECKING, Annotated, Any, Literal, cast
 
 if TYPE_CHECKING:
     from collections.abc import Awaitable, Callable
 
-from langchain_core.messages import ToolMessage
+    from langgraph.runtime import Runtime
+
+from langchain_core.messages import AIMessage, SystemMessage, ToolMessage
 from langchain_core.tools import tool
 from langgraph.types import Command
-from typing_extensions import NotRequired, TypedDict
+from typing_extensions import NotRequired, TypedDict, override
 
 from langchain.agents.middleware.types import (
     AgentMiddleware,
@@ -34,7 +35,7 @@ class Todo(TypedDict):
     """The current status of the todo item."""
 
 
-class PlanningState(AgentState):
+class PlanningState(AgentState[Any]):
     """State schema for the todo middleware."""
 
     todos: Annotated[NotRequired[list[Todo]], OmitFromInput]
@@ -99,7 +100,7 @@ It is important to skip using this tool when:
    - Use clear, descriptive task names
 
 Being proactive with task management demonstrates attentiveness and ensures you complete all requirements successfully
-Remember: If you only need to make a few tool calls to complete a task, and it is clear what you need to do, it is better to just do the task directly and NOT call this tool at all."""
+Remember: If you only need to make a few tool calls to complete a task, and it is clear what you need to do, it is better to just do the task directly and NOT call this tool at all."""  # noqa: E501
 
 WRITE_TODOS_SYSTEM_PROMPT = """## `write_todos`
 
@@ -113,11 +114,13 @@ Writing todos takes time and tokens, use it when it is helpful for managing comp
 
 ## Important To-Do List Usage Notes to Remember
 - The `write_todos` tool should never be called multiple times in parallel.
-- Don't be afraid to revise the To-Do list as you go. New information may reveal new tasks that need to be done, or old tasks that are irrelevant."""
+- Don't be afraid to revise the To-Do list as you go. New information may reveal new tasks that need to be done, or old tasks that are irrelevant."""  # noqa: E501
 
 
 @tool(description=WRITE_TODOS_TOOL_DESCRIPTION)
-def write_todos(todos: list[Todo], tool_call_id: Annotated[str, InjectedToolCallId]) -> Command:
+def write_todos(
+    todos: list[Todo], tool_call_id: Annotated[str, InjectedToolCallId]
+) -> Command[Any]:
     """Create and manage a structured task list for your current work session."""
     return Command(
         update={
@@ -136,7 +139,9 @@ class TodoListMiddleware(AgentMiddleware):
     into task completion status.
 
     The middleware automatically injects system prompts that guide the agent on when
-    and how to use the todo functionality effectively.
+    and how to use the todo functionality effectively. It also enforces that the
+    `write_todos` tool is called at most once per model turn, since the tool replaces
+    the entire todo list and parallel calls would create ambiguity about precedence.
 
     Example:
         ```python
@@ -150,12 +155,6 @@ class TodoListMiddleware(AgentMiddleware):
 
         print(result["todos"])  # Array of todo items with status tracking
         ```
-
-    Args:
-        system_prompt: Custom system prompt to guide the agent on using the todo tool.
-            If not provided, uses the default `WRITE_TODOS_SYSTEM_PROMPT`.
-        tool_description: Custom description for the write_todos tool.
-            If not provided, uses the default `WRITE_TODOS_TOOL_DESCRIPTION`.
     """
 
     state_schema = PlanningState
@@ -166,11 +165,12 @@ class TodoListMiddleware(AgentMiddleware):
         system_prompt: str = WRITE_TODOS_SYSTEM_PROMPT,
         tool_description: str = WRITE_TODOS_TOOL_DESCRIPTION,
     ) -> None:
-        """Initialize the TodoListMiddleware 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.
-            tool_description: Custom description for the write_todos tool.
+            system_prompt: Custom system prompt to guide the agent on using the todo
+                tool.
+            tool_description: Custom description for the `write_todos` tool.
         """
         super().__init__()
         self.system_prompt = system_prompt
@@ -180,7 +180,7 @@ class TodoListMiddleware(AgentMiddleware):
         @tool(description=self.tool_description)
         def write_todos(
             todos: list[Todo], tool_call_id: Annotated[str, InjectedToolCallId]
-        ) -> Command:
+        ) -> Command[Any]:
             """Create and manage a structured task list for your current work session."""
             return Command(
                 update={
@@ -198,23 +198,121 @@ class TodoListMiddleware(AgentMiddleware):
         request: ModelRequest,
         handler: Callable[[ModelRequest], ModelResponse],
     ) -> ModelCallResult:
-        """Update the system prompt to include the todo system prompt."""
-        request.system_prompt = (
-            request.system_prompt + "\n\n" + self.system_prompt
-            if request.system_prompt
-            else self.system_prompt
+        """Update the system message to include the todo system prompt.
+
+        Args:
+            request: Model request to execute (includes state and runtime).
+            handler: Async callback that executes the model request and returns
+                `ModelResponse`.
+
+        Returns:
+            The model call result.
+        """
+        if request.system_message is not None:
+            new_system_content = [
+                *request.system_message.content_blocks,
+                {"type": "text", "text": f"\n\n{self.system_prompt}"},
+            ]
+        else:
+            new_system_content = [{"type": "text", "text": self.system_prompt}]
+        new_system_message = SystemMessage(
+            content=cast("list[str | dict[str, str]]", new_system_content)
         )
-        return handler(request)
+        return handler(request.override(system_message=new_system_message))
 
     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
+        """Update the system message to include the todo system prompt.
+
+        Args:
+            request: Model request to execute (includes state and runtime).
+            handler: Async callback that executes the model request and returns
+                `ModelResponse`.
+
+        Returns:
+            The model call result.
+        """
+        if request.system_message is not None:
+            new_system_content = [
+                *request.system_message.content_blocks,
+                {"type": "text", "text": f"\n\n{self.system_prompt}"},
+            ]
+        else:
+            new_system_content = [{"type": "text", "text": self.system_prompt}]
+        new_system_message = SystemMessage(
+            content=cast("list[str | dict[str, str]]", new_system_content)
         )
-        return await handler(request)
+        return await handler(request.override(system_message=new_system_message))
+
+    @override
+    def after_model(self, state: AgentState[Any], runtime: Runtime) -> dict[str, Any] | None:
+        """Check for parallel write_todos tool calls and return errors if detected.
+
+        The todo list is designed to be updated at most once per model turn. Since
+        the `write_todos` tool replaces the entire todo list with each call, making
+        multiple parallel calls would create ambiguity about which update should take
+        precedence. This method prevents such conflicts by rejecting any response that
+        contains multiple write_todos tool calls.
+
+        Args:
+            state: The current agent state containing messages.
+            runtime: The LangGraph runtime instance.
+
+        Returns:
+            A dict containing error ToolMessages for each write_todos call if multiple
+            parallel calls are detected, otherwise None to allow normal execution.
+        """
+        messages = state["messages"]
+        if not messages:
+            return None
+
+        last_ai_msg = next((msg for msg in reversed(messages) if isinstance(msg, AIMessage)), None)
+        if not last_ai_msg or not last_ai_msg.tool_calls:
+            return None
+
+        # Count write_todos tool calls
+        write_todos_calls = [tc for tc in last_ai_msg.tool_calls if tc["name"] == "write_todos"]
+
+        if len(write_todos_calls) > 1:
+            # Create error tool messages for all write_todos calls
+            error_messages = [
+                ToolMessage(
+                    content=(
+                        "Error: The `write_todos` tool should never be called multiple times "
+                        "in parallel. Please call it only once per model invocation to update "
+                        "the todo list."
+                    ),
+                    tool_call_id=tc["id"],
+                    status="error",
+                )
+                for tc in write_todos_calls
+            ]
+
+            # Keep the tool calls in the AI message but return error messages
+            # This follows the same pattern as HumanInTheLoopMiddleware
+            return {"messages": error_messages}
+
+        return None
+
+    @override
+    async def aafter_model(self, state: AgentState[Any], runtime: Runtime) -> dict[str, Any] | None:
+        """Check for parallel write_todos tool calls and return errors if detected.
+
+        Async version of `after_model`. The todo list is designed to be updated at
+        most once per model turn. Since the `write_todos` tool replaces the entire
+        todo list with each call, making multiple parallel calls would create ambiguity
+        about which update should take precedence. This method prevents such conflicts
+        by rejecting any response that contains multiple write_todos tool calls.
+
+        Args:
+            state: The current agent state containing messages.
+            runtime: The LangGraph runtime instance.
+
+        Returns:
+            A dict containing error ToolMessages for each write_todos call if multiple
+            parallel calls are detected, otherwise None to allow normal execution.
+        """
+        return self.after_model(state, runtime)

langchain/agents/middleware/tool_call_limit.py

@@ -7,7 +7,7 @@ from typing import TYPE_CHECKING, Annotated, Any, Generic, Literal
 from langchain_core.messages import AIMessage, ToolCall, ToolMessage
 from langgraph.channels.untracked_value import UntrackedValue
 from langgraph.typing import ContextT
-from typing_extensions import NotRequired
+from typing_extensions import NotRequired, override
 
 from langchain.agents.middleware.types import (
     AgentMiddleware,
@@ -23,22 +23,23 @@ if TYPE_CHECKING:
 ExitBehavior = Literal["continue", "error", "end"]
 """How to handle execution when tool call limits are exceeded.
 
-- `"continue"`: Block exceeded tools with error messages, let other tools continue (default)
-- `"error"`: Raise a `ToolCallLimitExceededError` exception
-- `"end"`: Stop execution immediately, injecting a ToolMessage and an AI message
-    for the single tool call that exceeded the limit. Raises `NotImplementedError`
-    if there are other pending tool calls (due to parallel tool calling).
+- `'continue'`: Block exceeded tools with error messages, let other tools continue
+    (default)
+- `'error'`: Raise a `ToolCallLimitExceededError` exception
+- `'end'`: Stop execution immediately, injecting a `ToolMessage` and an `AIMessage` for
+    the single tool call that exceeded the limit. Raises `NotImplementedError` if there
+    are other pending tool calls (due to parallel tool calling).
 """
 
 
 class ToolCallLimitState(AgentState[ResponseT], Generic[ResponseT]):
-    """State schema for ToolCallLimitMiddleware.
+    """State schema for `ToolCallLimitMiddleware`.
 
-    Extends AgentState with tool call tracking fields.
+    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.
+    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]]
@@ -46,13 +47,13 @@ class ToolCallLimitState(AgentState[ResponseT], Generic[ResponseT]):
 
 
 def _build_tool_message_content(tool_name: str | None) -> str:
-    """Build the error message content for ToolMessage when limit is exceeded.
+    """Build the error message content for `ToolMessage` when limit is exceeded.
 
     This message is sent to the model, so it should not reference thread/run concepts
     that the model has no notion of.
 
     Args:
-        tool_name: Tool name being limited (if specific tool), or None for all tools.
+        tool_name: Tool name being limited (if specific tool), or `None` for all tools.
 
     Returns:
         A concise message instructing the model not to call the tool again.
@@ -70,7 +71,7 @@ def _build_final_ai_message_content(
     run_limit: int | None,
     tool_name: str | None,
 ) -> str:
-    """Build the final AI message content for 'end' behavior.
+    """Build the final AI message content for `'end'` behavior.
 
     This message is displayed to the user, so it should include detailed information
     about which limits were exceeded.
@@ -80,7 +81,7 @@ def _build_final_ai_message_content(
         run_count: Current run tool call count.
         thread_limit: Thread tool call limit (if set).
         run_limit: Run tool call limit (if set).
-        tool_name: Tool name being limited (if specific tool), or None for all tools.
+        tool_name: Tool name being limited (if specific tool), or `None` for all tools.
 
     Returns:
         A formatted message describing which limits were exceeded.
@@ -100,8 +101,8 @@ def _build_final_ai_message_content(
 class ToolCallLimitExceededError(Exception):
     """Exception raised when tool call limits are exceeded.
 
-    This exception is raised when the configured exit behavior is 'error'
-    and either the thread or run tool call limit has been exceeded.
+    This exception is raised when the configured exit behavior is `'error'` and either
+    the thread or run tool call limit has been exceeded.
     """
 
     def __init__(
@@ -145,48 +146,53 @@ class ToolCallLimitMiddleware(
 
     Configuration:
         - `exit_behavior`: How to handle when limits are exceeded
-          - `"continue"`: Block exceeded tools, let execution continue (default)
-          - `"error"`: Raise an exception
-          - `"end"`: Stop immediately with a ToolMessage + AI message for the single
-            tool call that exceeded the limit (raises `NotImplementedError` if there
-            are other pending tool calls (due to parallel tool calling).
+            - `'continue'`: Block exceeded tools, let execution continue (default)
+            - `'error'`: Raise an exception
+            - `'end'`: Stop immediately with a `ToolMessage` + AI message for the single
+                tool call that exceeded the limit (raises `NotImplementedError` if there
+                are other pending tool calls (due to parallel tool calling).
 
     Examples:
-        Continue execution with blocked tools (default):
-        ```python
-        from langchain.agents.middleware.tool_call_limit import ToolCallLimitMiddleware
-        from langchain.agents import create_agent
-
-        # Block exceeded tools but let other tools and model continue
-        limiter = ToolCallLimitMiddleware(
-            thread_limit=20,
-            run_limit=10,
-            exit_behavior="continue",  # default
-        )
+        !!! example "Continue execution with blocked tools (default)"
+
+            ```python
+            from langchain.agents.middleware.tool_call_limit import ToolCallLimitMiddleware
+            from langchain.agents import create_agent
+
+            # Block exceeded tools but let other tools and model continue
+            limiter = ToolCallLimitMiddleware(
+                thread_limit=20,
+                run_limit=10,
+                exit_behavior="continue",  # default
+            )
+
+            agent = create_agent("openai:gpt-4o", middleware=[limiter])
+            ```
+
+        !!! example "Stop immediately when limit exceeded"
 
-        agent = create_agent("openai:gpt-4o", middleware=[limiter])
-        ```
+            ```python
+            # End execution immediately with an AI message
+            limiter = ToolCallLimitMiddleware(run_limit=5, exit_behavior="end")
 
-        Stop immediately when limit exceeded:
-        ```python
-        # End execution immediately with an AI message
-        limiter = ToolCallLimitMiddleware(run_limit=5, exit_behavior="end")
+            agent = create_agent("openai:gpt-4o", middleware=[limiter])
+            ```
 
-        agent = create_agent("openai:gpt-4o", middleware=[limiter])
-        ```
+        !!! example "Raise exception on limit"
 
-        Raise exception on limit:
-        ```python
-        # Strict limit with exception handling
-        limiter = ToolCallLimitMiddleware(tool_name="search", thread_limit=5, exit_behavior="error")
+            ```python
+            # Strict limit with exception handling
+            limiter = ToolCallLimitMiddleware(
+                tool_name="search", thread_limit=5, exit_behavior="error"
+            )
 
-        agent = create_agent("openai:gpt-4o", middleware=[limiter])
+            agent = create_agent("openai:gpt-4o", middleware=[limiter])
 
-        try:
-            result = await agent.invoke({"messages": [HumanMessage("Task")]})
-        except ToolCallLimitExceededError as e:
-            print(f"Search limit exceeded: {e}")
-        ```
+            try:
+                result = await agent.invoke({"messages": [HumanMessage("Task")]})
+            except ToolCallLimitExceededError as e:
+                print(f"Search limit exceeded: {e}")
+            ```
 
     """
 
@@ -204,23 +210,24 @@ class ToolCallLimitMiddleware(
 
         Args:
             tool_name: Name of the specific tool to limit. If `None`, limits apply
-                to all tools. Defaults to `None`.
+                to all tools.
             thread_limit: Maximum number of tool calls allowed per thread.
-                `None` means no limit. Defaults to `None`.
+                `None` means no limit.
             run_limit: Maximum number of tool calls allowed per run.
-                `None` means no limit. Defaults to `None`.
+                `None` means no limit.
             exit_behavior: How to handle when limits are exceeded.
-                - `"continue"`: Block exceeded tools with error messages, let other
-                  tools continue. Model decides when to end. (default)
-                - `"error"`: Raise a `ToolCallLimitExceededError` exception
-                - `"end"`: Stop execution immediately with a ToolMessage + AI message
-                  for the single tool call that exceeded the limit. Raises
-                  `NotImplementedError` if there are multiple parallel tool
-                  calls to other tools or multiple pending tool calls.
+
+                - `'continue'`: Block exceeded tools with error messages, let other
+                    tools continue. Model decides when to end.
+                - `'error'`: Raise a `ToolCallLimitExceededError` exception
+                - `'end'`: Stop execution immediately with a `ToolMessage` + AI message
+                    for the single tool call that exceeded the limit. Raises
+                    `NotImplementedError` if there are multiple parallel tool
+                    calls to other tools or multiple pending tool calls.
 
         Raises:
-            ValueError: If both limits are `None`, if exit_behavior is invalid,
-                or if run_limit exceeds thread_limit.
+            ValueError: If both limits are `None`, if `exit_behavior` is invalid,
+                or if `run_limit` exceeds `thread_limit`.
         """
         super().__init__()
 
@@ -293,7 +300,8 @@ class ToolCallLimitMiddleware(
             run_count: Current run call count.
 
         Returns:
-            Tuple of (allowed_calls, blocked_calls, final_thread_count, final_run_count).
+            Tuple of `(allowed_calls, blocked_calls, final_thread_count,
+                final_run_count)`.
         """
         allowed_calls: list[ToolCall] = []
         blocked_calls: list[ToolCall] = []
@@ -314,10 +322,11 @@ class ToolCallLimitMiddleware(
         return allowed_calls, blocked_calls, temp_thread_count, temp_run_count
 
     @hook_config(can_jump_to=["end"])
+    @override
     def after_model(
         self,
         state: ToolCallLimitState[ResponseT],
-        runtime: Runtime[ContextT],  # noqa: ARG002
+        runtime: Runtime[ContextT],
     ) -> dict[str, Any] | None:
         """Increment tool call counts after a model call and check limits.
 
@@ -327,13 +336,13 @@ class ToolCallLimitMiddleware(
 
         Returns:
             State updates with incremented tool call counts. If limits are exceeded
-            and exit_behavior is "end", also includes a jump to end with a ToolMessage
-            and AI message for the single exceeded tool call.
+                and exit_behavior is `'end'`, also includes a jump to end with a
+                `ToolMessage` and AI message for the single exceeded tool call.
 
         Raises:
-            ToolCallLimitExceededError: If limits are exceeded and exit_behavior
-                is "error".
-            NotImplementedError: If limits are exceeded, exit_behavior is "end",
+            ToolCallLimitExceededError: If limits are exceeded and `exit_behavior`
+                is `'error'`.
+            NotImplementedError: If limits are exceeded, `exit_behavior` is `'end'`,
                 and there are multiple tool calls.
         """
         # Get the last AIMessage to check for tool calls
@@ -352,7 +361,7 @@ class ToolCallLimitMiddleware(
             return None
 
         # Get the count key for this middleware instance
-        count_key = self.tool_name if self.tool_name else "__all__"
+        count_key = self.tool_name or "__all__"
 
         # Get current counts
         thread_counts = state.get("thread_tool_call_count", {}).copy()
@@ -452,3 +461,28 @@ class ToolCallLimitMiddleware(
             "run_tool_call_count": run_counts,
             "messages": artificial_messages,
         }
+
+    @hook_config(can_jump_to=["end"])
+    async def aafter_model(
+        self,
+        state: ToolCallLimitState[ResponseT],
+        runtime: Runtime[ContextT],
+    ) -> dict[str, Any] | None:
+        """Async increment tool call counts after a model call and check limits.
+
+        Args:
+            state: The current agent state.
+            runtime: The langgraph runtime.
+
+        Returns:
+            State updates with incremented tool call counts. If limits are exceeded
+                and exit_behavior is `'end'`, also includes a jump to end with a
+                `ToolMessage` and AI message for the single exceeded tool call.
+
+        Raises:
+            ToolCallLimitExceededError: If limits are exceeded and `exit_behavior`
+                is `'error'`.
+            NotImplementedError: If limits are exceeded, `exit_behavior` is `'end'`,
+                and there are multiple tool calls.
+        """
+        return self.after_model(state, runtime)