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

langchain/agents/middleware/planning.py

@@ -0,0 +1,201 @@
+"""Planning and task management middleware for agents."""
+# ruff: noqa: E501
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Annotated, Literal
+
+from langchain_core.messages import ToolMessage
+from langchain_core.tools import tool
+from langgraph.types import Command
+from typing_extensions import NotRequired, TypedDict
+
+from langchain.agents.middleware.types import AgentMiddleware, AgentState, ModelRequest
+from langchain.tools import InjectedToolCallId
+
+if TYPE_CHECKING:
+    from langgraph.runtime import Runtime
+
+
+class Todo(TypedDict):
+    """A single todo item with content and status."""
+
+    content: str
+    """The content/description of the todo item."""
+
+    status: Literal["pending", "in_progress", "completed"]
+    """The current status of the todo item."""
+
+
+class PlanningState(AgentState):
+    """State schema for the todo middleware."""
+
+    todos: NotRequired[list[Todo]]
+    """List of todo items for tracking task progress."""
+
+
+WRITE_TODOS_TOOL_DESCRIPTION = """Use this tool to create and manage a structured task list for your current work session. This helps you track progress, organize complex tasks, and demonstrate thoroughness to the user.
+
+Only use this tool if you think it will be helpful in staying organized. If the user's request is trivial and takes less than 3 steps, it is better to NOT use this tool and just do the task directly.
+
+## When to Use This Tool
+Use this tool in these scenarios:
+
+1. Complex multi-step tasks - When a task requires 3 or more distinct steps or actions
+2. Non-trivial and complex tasks - Tasks that require careful planning or multiple operations
+3. User explicitly requests todo list - When the user directly asks you to use the todo list
+4. User provides multiple tasks - When users provide a list of things to be done (numbered or comma-separated)
+5. The plan may need future revisions or updates based on results from the first few steps
+
+## How to Use This Tool
+1. When you start working on a task - Mark it as in_progress BEFORE beginning work.
+2. After completing a task - Mark it as completed and add any new follow-up tasks discovered during implementation.
+3. You can also update future tasks, such as deleting them if they are no longer necessary, or adding new tasks that are necessary. Don't change previously completed tasks.
+4. You can make several updates to the todo list at once. For example, when you complete a task, you can mark the next task you need to start as in_progress.
+
+## When NOT to Use This Tool
+It is important to skip using this tool when:
+1. There is only a single, straightforward task
+2. The task is trivial and tracking it provides no benefit
+3. The task can be completed in less than 3 trivial steps
+4. The task is purely conversational or informational
+
+## Task States and Management
+
+1. **Task States**: Use these states to track progress:
+   - pending: Task not yet started
+   - in_progress: Currently working on (you can have multiple tasks in_progress at a time if they are not related to each other and can be run in parallel)
+   - completed: Task finished successfully
+
+2. **Task Management**:
+   - Update task status in real-time as you work
+   - Mark tasks complete IMMEDIATELY after finishing (don't batch completions)
+   - Complete current tasks before starting new ones
+   - Remove tasks that are no longer relevant from the list entirely
+   - IMPORTANT: When you write this todo list, you should mark your first task (or tasks) as in_progress immediately!.
+   - IMPORTANT: Unless all tasks are completed, you should always have at least one task in_progress to show the user that you are working on something.
+
+3. **Task Completion Requirements**:
+   - ONLY mark a task as completed when you have FULLY accomplished it
+   - If you encounter errors, blockers, or cannot finish, keep the task as in_progress
+   - When blocked, create a new task describing what needs to be resolved
+   - Never mark a task as completed if:
+     - There are unresolved issues or errors
+     - Work is partial or incomplete
+     - You encountered blockers that prevent completion
+     - You couldn't find necessary resources or dependencies
+     - Quality standards haven't been met
+
+4. **Task Breakdown**:
+   - Create specific, actionable items
+   - Break complex tasks into smaller, manageable steps
+   - 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."""
+
+WRITE_TODOS_SYSTEM_PROMPT = """## `write_todos`
+
+You have access to the `write_todos` tool to help you manage and plan complex objectives.
+Use this tool for complex objectives to ensure that you are tracking each necessary step and giving the user visibility into your progress.
+This tool is very helpful for planning complex objectives, and for breaking down these larger complex objectives into smaller steps.
+
+It is critical that you mark todos as completed as soon as you are done with a step. Do not batch up multiple steps before marking them as completed.
+For simple objectives that only require a few steps, it is better to just complete the objective directly and NOT use this tool.
+Writing todos takes time and tokens, use it when it is helpful for managing complex many-step problems! But not for simple few-step requests.
+
+## 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."""
+
+
+@tool(description=WRITE_TODOS_TOOL_DESCRIPTION)
+def write_todos(todos: list[Todo], tool_call_id: Annotated[str, InjectedToolCallId]) -> Command:
+    """Create and manage a structured task list for your current work session."""
+    return Command(
+        update={
+            "todos": todos,
+            "messages": [ToolMessage(f"Updated todo list to {todos}", tool_call_id=tool_call_id)],
+        }
+    )
+
+
+class PlanningMiddleware(AgentMiddleware):
+    """Middleware that provides todo list management capabilities to agents.
+
+    This middleware adds a `write_todos` tool that allows agents to create and manage
+    structured task lists for complex multi-step operations. It's designed to help
+    agents track progress, organize complex tasks, and provide users with visibility
+    into task completion status.
+
+    The middleware automatically injects system prompts that guide the agent on when
+    and how to use the todo functionality effectively.
+
+    Example:
+        ```python
+        from langchain.agents.middleware.planning import PlanningMiddleware
+        from langchain.agents import create_agent
+
+        agent = create_agent("openai:gpt-4o", middleware=[PlanningMiddleware()])
+
+        # Agent now has access to write_todos tool and todo state tracking
+        result = await agent.invoke({"messages": [HumanMessage("Help me refactor my codebase")]})
+
+        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
+
+    def __init__(
+        self,
+        *,
+        system_prompt: str = WRITE_TODOS_SYSTEM_PROMPT,
+        tool_description: str = WRITE_TODOS_TOOL_DESCRIPTION,
+    ) -> None:
+        """Initialize the PlanningMiddleware 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.
+        """
+        super().__init__()
+        self.system_prompt = system_prompt
+        self.tool_description = tool_description
+
+        # Dynamically create the write_todos tool with the custom description
+        @tool(description=self.tool_description)
+        def write_todos(
+            todos: list[Todo], tool_call_id: Annotated[str, InjectedToolCallId]
+        ) -> Command:
+            """Create and manage a structured task list for your current work session."""
+            return Command(
+                update={
+                    "todos": todos,
+                    "messages": [
+                        ToolMessage(f"Updated todo list to {todos}", tool_call_id=tool_call_id)
+                    ],
+                }
+            )
+
+        self.tools = [write_todos]
+
+    def modify_model_request(
+        self,
+        request: ModelRequest,
+        state: AgentState,  # noqa: ARG002
+        runtime: Runtime,  # noqa: ARG002
+    ) -> ModelRequest:
+        """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
+        )
+        return request

langchain/agents/middleware/prompt_caching.py

@@ -1,9 +1,11 @@
 """Anthropic prompt caching middleware."""
 
-from typing import Any, Literal
+from typing import Literal
 from warnings import warn
 
-from langchain.agents.middleware.types import AgentMiddleware, ModelRequest
+from langgraph.runtime import Runtime
+
+from langchain.agents.middleware.types import AgentMiddleware, AgentState, ModelRequest
 
 
 class AnthropicPromptCachingMiddleware(AgentMiddleware):
@@ -39,10 +41,11 @@ class AnthropicPromptCachingMiddleware(AgentMiddleware):
         self.min_messages_to_cache = min_messages_to_cache
         self.unsupported_model_behavior = unsupported_model_behavior
 
-    def modify_model_request(  # type: ignore[override]
+    def modify_model_request(
         self,
         request: ModelRequest,
-        state: dict[str, Any],  # noqa: ARG002
+        state: AgentState,  # noqa: ARG002
+        runtime: Runtime,  # noqa: ARG002
     ) -> ModelRequest:
         """Modify the model request to add cache control blocks."""
         try:

langchain/agents/middleware/summarization.py

@@ -16,6 +16,7 @@ from langchain_core.messages.utils import count_tokens_approximately, trim_messa
 from langgraph.graph.message import (
     REMOVE_ALL_MESSAGES,
 )
+from langgraph.runtime import Runtime
 
 from langchain.agents.middleware.types import AgentMiddleware, AgentState
 from langchain.chat_models import BaseChatModel, init_chat_model
@@ -98,7 +99,7 @@ class SummarizationMiddleware(AgentMiddleware):
         self.summary_prompt = summary_prompt
         self.summary_prefix = summary_prefix
 
-    def before_model(self, state: AgentState) -> dict[str, Any] | None:  # type: ignore[override]
+    def before_model(self, state: AgentState, runtime: Runtime) -> dict[str, Any] | None:  # noqa: ARG002
         """Process messages before model invocation, potentially triggering summarization."""
         messages = state["messages"]
         self._ensure_message_ids(messages)

langchain/agents/middleware/tool_call_limit.py

@@ -0,0 +1,260 @@
+"""Tool call limit middleware for agents."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Literal
+
+from langchain_core.messages import AIMessage, AnyMessage, HumanMessage
+
+from langchain.agents.middleware.types import AgentMiddleware, AgentState, hook_config
+
+if TYPE_CHECKING:
+    from langgraph.runtime import Runtime
+
+
+def _count_tool_calls_in_messages(messages: list[AnyMessage], tool_name: str | None = None) -> int:
+    """Count tool calls in a list of messages.
+
+    Args:
+        messages: List of messages to count tool calls in.
+        tool_name: If specified, only count calls to this specific tool.
+            If None, count all tool calls.
+
+    Returns:
+        The total number of tool calls (optionally filtered by tool_name).
+    """
+    count = 0
+    for message in messages:
+        if isinstance(message, AIMessage) and message.tool_calls:
+            if tool_name is None:
+                # Count all tool calls
+                count += len(message.tool_calls)
+            else:
+                # Count only calls to the specified tool
+                count += sum(1 for tc in message.tool_calls if tc["name"] == tool_name)
+    return count
+
+
+def _get_run_messages(messages: list[AnyMessage]) -> list[AnyMessage]:
+    """Get messages from the current run (after the last HumanMessage).
+
+    Args:
+        messages: Full list of messages.
+
+    Returns:
+        Messages from the current run (after last HumanMessage).
+    """
+    # Find the last HumanMessage
+    last_human_index = -1
+    for i in range(len(messages) - 1, -1, -1):
+        if isinstance(messages[i], HumanMessage):
+            last_human_index = i
+            break
+
+    # If no HumanMessage found, return all messages
+    if last_human_index == -1:
+        return messages
+
+    # Return messages after the last HumanMessage
+    return messages[last_human_index + 1 :]
+
+
+def _build_tool_limit_exceeded_message(
+    thread_count: int,
+    run_count: int,
+    thread_limit: int | None,
+    run_limit: int | None,
+    tool_name: str | None,
+) -> str:
+    """Build a message indicating which tool call limits were exceeded.
+
+    Args:
+        thread_count: Current thread tool call count.
+        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.
+
+    Returns:
+        A formatted message describing which limits were exceeded.
+    """
+    tool_desc = f"'{tool_name}' tool call" if tool_name else "Tool call"
+    exceeded_limits = []
+    if thread_limit is not None and thread_count >= thread_limit:
+        exceeded_limits.append(f"thread limit ({thread_count}/{thread_limit})")
+    if run_limit is not None and run_count >= run_limit:
+        exceeded_limits.append(f"run limit ({run_count}/{run_limit})")
+
+    return f"{tool_desc} limits exceeded: {', '.join(exceeded_limits)}"
+
+
+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.
+    """
+
+    def __init__(
+        self,
+        thread_count: int,
+        run_count: int,
+        thread_limit: int | None,
+        run_limit: int | None,
+        tool_name: str | None = None,
+    ) -> None:
+        """Initialize the exception with call count information.
+
+        Args:
+            thread_count: Current thread tool call count.
+            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.
+        """
+        self.thread_count = thread_count
+        self.run_count = run_count
+        self.thread_limit = thread_limit
+        self.run_limit = run_limit
+        self.tool_name = tool_name
+
+        msg = _build_tool_limit_exceeded_message(
+            thread_count, run_count, thread_limit, run_limit, tool_name
+        )
+        super().__init__(msg)
+
+
+class ToolCallLimitMiddleware(AgentMiddleware):
+    """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.
+
+    Run-level: The middleware counts tool calls made after the last HumanMessage,
+    representing the current run (invocation) of the agent.
+
+    Example:
+        ```python
+        from langchain.agents.middleware.tool_call_limit import ToolCallLimitMiddleware
+        from langchain.agents import create_agent
+
+        # Limit all tool calls globally
+        global_limiter = ToolCallLimitMiddleware(thread_limit=20, run_limit=10, exit_behavior="end")
+
+        # Limit a specific tool
+        search_limiter = ToolCallLimitMiddleware(
+            tool_name="search", thread_limit=5, run_limit=3, exit_behavior="end"
+        )
+
+        # Use both in the same agent
+        agent = create_agent("openai:gpt-4o", middleware=[global_limiter, search_limiter])
+
+        result = await agent.invoke({"messages": [HumanMessage("Help me with a task")]})
+        ```
+    """
+
+    def __init__(
+        self,
+        *,
+        tool_name: str | None = None,
+        thread_limit: int | None = None,
+        run_limit: int | None = None,
+        exit_behavior: Literal["end", "error"] = "end",
+    ) -> None:
+        """Initialize the tool call limit middleware.
+
+        Args:
+            tool_name: Name of the specific tool to limit. If None, limits apply
+                to all tools. Defaults to None.
+            thread_limit: Maximum number of tool calls allowed per thread.
+                None means no limit. Defaults to None.
+            run_limit: Maximum number of tool calls allowed per run.
+                None means no limit. Defaults to None.
+            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 ToolCallLimitExceededError
+                Defaults to "end".
+
+        Raises:
+            ValueError: If both limits are None or if exit_behavior is invalid.
+        """
+        super().__init__()
+
+        if thread_limit is None and run_limit is None:
+            msg = "At least one limit must be specified (thread_limit or run_limit)"
+            raise ValueError(msg)
+
+        if exit_behavior not in ("end", "error"):
+            msg = f"Invalid exit_behavior: {exit_behavior}. Must be 'end' or 'error'"
+            raise ValueError(msg)
+
+        self.tool_name = tool_name
+        self.thread_limit = thread_limit
+        self.run_limit = run_limit
+        self.exit_behavior = exit_behavior
+
+    @property
+    def name(self) -> str:
+        """The name of the middleware instance.
+
+        Includes the tool name if specified to allow multiple instances
+        of this middleware with different tool names.
+        """
+        base_name = self.__class__.__name__
+        if self.tool_name:
+            return f"{base_name}[{self.tool_name}]"
+        return base_name
+
+    @hook_config(can_jump_to=["end"])
+    def before_model(self, state: AgentState, 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.
+            runtime: The langgraph runtime.
+
+        Returns:
+            If limits are exceeded and exit_behavior is "end", returns
+            a Command to jump to the end with a limit exceeded message. Otherwise returns None.
+
+        Raises:
+            ToolCallLimitExceededError: If limits are exceeded and exit_behavior
+                is "error".
+        """
+        messages = state.get("messages", [])
+
+        # Count tool calls in entire thread
+        thread_count = _count_tool_calls_in_messages(messages, self.tool_name)
+
+        # 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)
+
+        # Check if any limits are exceeded
+        thread_limit_exceeded = self.thread_limit is not None and thread_count >= self.thread_limit
+        run_limit_exceeded = self.run_limit is not None and run_count >= self.run_limit
+
+        if thread_limit_exceeded or run_limit_exceeded:
+            if self.exit_behavior == "error":
+                raise ToolCallLimitExceededError(
+                    thread_count=thread_count,
+                    run_count=run_count,
+                    thread_limit=self.thread_limit,
+                    run_limit=self.run_limit,
+                    tool_name=self.tool_name,
+                )
+            if self.exit_behavior == "end":
+                # Create a message indicating the limit was exceeded
+                limit_message = _build_tool_limit_exceeded_message(
+                    thread_count, run_count, self.thread_limit, self.run_limit, self.tool_name
+                )
+                limit_ai_message = AIMessage(content=limit_message)
+
+                return {"jump_to": "end", "messages": [limit_ai_message]}
+
+        return None