codekingpro/portable-devtools

1"""Planning and task management middleware for agents."""2 3from collections.abc import Awaitable, Callable4from typing import Annotated, Any, Literal, cast5 6from langchain_core.messages import AIMessage, SystemMessage, ToolMessage7from langchain_core.tools import InjectedToolCallId, StructuredTool, tool8from langgraph.runtime import Runtime9from langgraph.types import Command10from pydantic import BaseModel11from typing_extensions import NotRequired, TypedDict, override12 13from langchain.agents.middleware.types import (14    AgentMiddleware,15    AgentState,16    ContextT,17    ModelRequest,18    ModelResponse,19    OmitFromInput,20    ResponseT,21)22from langchain.tools import ToolRuntime23 24 25class Todo(TypedDict):26    """A single todo item with content and status."""27 28    content: str29    """The content/description of the todo item."""30 31    status: Literal["pending", "in_progress", "completed"]32    """The current status of the todo item."""33 34 35class PlanningState(AgentState[ResponseT]):36    """State schema for the todo middleware.37 38    Type Parameters:39        ResponseT: The type of the structured response. Defaults to `Any`.40    """41 42    todos: Annotated[NotRequired[list[Todo]], OmitFromInput]43    """List of todo items for tracking task progress."""44 45 46class WriteTodosInput(BaseModel):47    """Input schema for the `write_todos` tool."""48 49    todos: list[Todo]50 51 52WRITE_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.53 54Only 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.55 56## When to Use This Tool57Use this tool in these scenarios:58 591. Complex multi-step tasks - When a task requires 3 or more distinct steps or actions602. Non-trivial and complex tasks - Tasks that require careful planning or multiple operations613. User explicitly requests todo list - When the user directly asks you to use the todo list624. User provides multiple tasks - When users provide a list of things to be done (numbered or comma-separated)635. The plan may need future revisions or updates based on results from the first few steps64 65## How to Use This Tool661. When you start working on a task - Mark it as in_progress BEFORE beginning work.672. After completing a task - Mark it as completed and add any new follow-up tasks discovered during implementation.683. 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.694. 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.70 71## When NOT to Use This Tool72It is important to skip using this tool when:731. There is only a single, straightforward task742. The task is trivial and tracking it provides no benefit753. The task can be completed in less than 3 trivial steps764. The task is purely conversational or informational77 78## Task States and Management79 801. **Task States**: Use these states to track progress:81   - pending: Task not yet started82   - 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)83   - completed: Task finished successfully84 852. **Task Management**:86   - Update task status in real-time as you work87   - Mark tasks complete IMMEDIATELY after finishing (don't batch completions)88   - Complete current tasks before starting new ones89   - Remove tasks that are no longer relevant from the list entirely90   - IMPORTANT: When you write this todo list, you should mark your first task (or tasks) as in_progress immediately!.91   - 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.92 933. **Task Completion Requirements**:94   - ONLY mark a task as completed when you have FULLY accomplished it95   - If you encounter errors, blockers, or cannot finish, keep the task as in_progress96   - When blocked, create a new task describing what needs to be resolved97   - Never mark a task as completed if:98     - There are unresolved issues or errors99     - Work is partial or incomplete100     - You encountered blockers that prevent completion101     - You couldn't find necessary resources or dependencies102     - Quality standards haven't been met103 1044. **Task Breakdown**:105   - Create specific, actionable items106   - Break complex tasks into smaller, manageable steps107   - Use clear, descriptive task names108 109Being proactive with task management demonstrates attentiveness and ensures you complete all requirements successfully110Remember: 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: E501111 112WRITE_TODOS_SYSTEM_PROMPT = """## `write_todos`113 114You have access to the `write_todos` tool to help you manage and plan complex objectives.115Use this tool for complex objectives to ensure that you are tracking each necessary step and giving the user visibility into your progress.116This tool is very helpful for planning complex objectives, and for breaking down these larger complex objectives into smaller steps.117 118It 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.119For simple objectives that only require a few steps, it is better to just complete the objective directly and NOT use this tool.120Writing todos takes time and tokens, use it when it is helpful for managing complex many-step problems! But not for simple few-step requests.121 122## Important To-Do List Usage Notes to Remember123- The `write_todos` tool should never be called multiple times in parallel.124- 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: E501125 126 127@tool(description=WRITE_TODOS_TOOL_DESCRIPTION)128def write_todos(129    todos: list[Todo], tool_call_id: Annotated[str, InjectedToolCallId]130) -> Command[Any]:131    """Create and manage a structured task list for your current work session."""132    return Command(133        update={134            "todos": todos,135            "messages": [ToolMessage(f"Updated todo list to {todos}", tool_call_id=tool_call_id)],136        }137    )138 139 140# Dynamically create the write_todos tool with the custom description141def _write_todos(142    runtime: ToolRuntime[ContextT, PlanningState[ResponseT]], todos: list[Todo]143) -> Command[Any]:144    """Create and manage a structured task list for your current work session."""145    return Command(146        update={147            "todos": todos,148            "messages": [149                ToolMessage(f"Updated todo list to {todos}", tool_call_id=runtime.tool_call_id)150            ],151        }152    )153 154 155async def _awrite_todos(156    runtime: ToolRuntime[ContextT, PlanningState[ResponseT]], todos: list[Todo]157) -> Command[Any]:158    """Create and manage a structured task list for your current work session."""159    return _write_todos(runtime, todos)160 161 162class TodoListMiddleware(AgentMiddleware[PlanningState[ResponseT], ContextT, ResponseT]):163    """Middleware that provides todo list management capabilities to agents.164 165    This middleware adds a `write_todos` tool that allows agents to create and manage166    structured task lists for complex multi-step operations. It's designed to help167    agents track progress, organize complex tasks, and provide users with visibility168    into task completion status.169 170    The middleware automatically injects system prompts that guide the agent on when171    and how to use the todo functionality effectively. It also enforces that the172    `write_todos` tool is called at most once per model turn, since the tool replaces173    the entire todo list and parallel calls would create ambiguity about precedence.174 175    Example:176        ```python177        from langchain.agents.middleware import TodoListMiddleware178        from langchain.agents import create_agent179 180        agent = create_agent("openai:gpt-4o", middleware=[TodoListMiddleware()])181 182        # Agent now has access to write_todos tool and todo state tracking183        result = await agent.invoke({"messages": [HumanMessage("Help me refactor my codebase")]})184 185        print(result["todos"])  # Array of todo items with status tracking186        ```187    """188 189    state_schema = PlanningState  # type: ignore[assignment]190 191    def __init__(192        self,193        *,194        system_prompt: str = WRITE_TODOS_SYSTEM_PROMPT,195        tool_description: str = WRITE_TODOS_TOOL_DESCRIPTION,196    ) -> None:197        """Initialize the `TodoListMiddleware` with optional custom prompts.198 199        Args:200            system_prompt: Custom system prompt to guide the agent on using the todo201                tool.202            tool_description: Custom description for the `write_todos` tool.203        """204        super().__init__()205        self.system_prompt = system_prompt206        self.tool_description = tool_description207 208        self.tools = [209            StructuredTool.from_function(210                name="write_todos",211                description=tool_description,212                func=_write_todos,213                coroutine=_awrite_todos,214                args_schema=WriteTodosInput,215                infer_schema=False,216            )217        ]218 219    def wrap_model_call(220        self,221        request: ModelRequest[ContextT],222        handler: Callable[[ModelRequest[ContextT]], ModelResponse[ResponseT]],223    ) -> ModelResponse[ResponseT] | AIMessage:224        """Update the system message to include the todo system prompt.225 226        Args:227            request: Model request to execute (includes state and runtime).228            handler: Async callback that executes the model request and returns229                `ModelResponse`.230 231        Returns:232            The model call result.233        """234        if request.system_message is not None:235            new_system_content = [236                *request.system_message.content_blocks,237                {"type": "text", "text": f"\n\n{self.system_prompt}"},238            ]239        else:240            new_system_content = [{"type": "text", "text": self.system_prompt}]241        new_system_message = SystemMessage(242            content=cast("list[str | dict[str, str]]", new_system_content)243        )244        return handler(request.override(system_message=new_system_message))245 246    async def awrap_model_call(247        self,248        request: ModelRequest[ContextT],249        handler: Callable[[ModelRequest[ContextT]], Awaitable[ModelResponse[ResponseT]]],250    ) -> ModelResponse[ResponseT] | AIMessage:251        """Update the system message to include the todo system prompt.252 253        Args:254            request: Model request to execute (includes state and runtime).255            handler: Async callback that executes the model request and returns256                `ModelResponse`.257 258        Returns:259            The model call result.260        """261        if request.system_message is not None:262            new_system_content = [263                *request.system_message.content_blocks,264                {"type": "text", "text": f"\n\n{self.system_prompt}"},265            ]266        else:267            new_system_content = [{"type": "text", "text": self.system_prompt}]268        new_system_message = SystemMessage(269            content=cast("list[str | dict[str, str]]", new_system_content)270        )271        return await handler(request.override(system_message=new_system_message))272 273    @override274    def after_model(275        self, state: PlanningState[ResponseT], runtime: Runtime[ContextT]276    ) -> dict[str, Any] | None:277        """Check for parallel write_todos tool calls and return errors if detected.278 279        The todo list is designed to be updated at most once per model turn. Since280        the `write_todos` tool replaces the entire todo list with each call, making281        multiple parallel calls would create ambiguity about which update should take282        precedence. This method prevents such conflicts by rejecting any response that283        contains multiple write_todos tool calls.284 285        Args:286            state: The current agent state containing messages.287            runtime: The LangGraph runtime instance.288 289        Returns:290            A dict containing error ToolMessages for each write_todos call if multiple291            parallel calls are detected, otherwise None to allow normal execution.292        """293        messages = state["messages"]294        if not messages:295            return None296 297        last_ai_msg = next((msg for msg in reversed(messages) if isinstance(msg, AIMessage)), None)298        if not last_ai_msg or not last_ai_msg.tool_calls:299            return None300 301        # Count write_todos tool calls302        write_todos_calls = [tc for tc in last_ai_msg.tool_calls if tc["name"] == "write_todos"]303 304        if len(write_todos_calls) > 1:305            # Create error tool messages for all write_todos calls306            error_messages = [307                ToolMessage(308                    content=(309                        "Error: The `write_todos` tool should never be called multiple times "310                        "in parallel. Please call it only once per model invocation to update "311                        "the todo list."312                    ),313                    tool_call_id=tc["id"],314                    status="error",315                )316                for tc in write_todos_calls317            ]318 319            # Keep the tool calls in the AI message but return error messages320            # This follows the same pattern as HumanInTheLoopMiddleware321            return {"messages": error_messages}322 323        return None324 325    @override326    async def aafter_model(327        self, state: PlanningState[ResponseT], runtime: Runtime[ContextT]328    ) -> dict[str, Any] | None:329        """Check for parallel write_todos tool calls and return errors if detected.330 331        Async version of `after_model`. The todo list is designed to be updated at332        most once per model turn. Since the `write_todos` tool replaces the entire333        todo list with each call, making multiple parallel calls would create ambiguity334        about which update should take precedence. This method prevents such conflicts335        by rejecting any response that contains multiple write_todos tool calls.336 337        Args:338            state: The current agent state containing messages.339            runtime: The LangGraph runtime instance.340 341        Returns:342            A dict containing error ToolMessages for each write_todos call if multiple343            parallel calls are detected, otherwise None to allow normal execution.344        """345        return self.after_model(state, runtime)346