225 lines
10 KiB
Python
225 lines
10 KiB
Python
"""Planning and task management middleware for agents."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import TYPE_CHECKING, Annotated, Literal, cast
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Awaitable, Callable
|
|
|
|
from langchain_core.messages import SystemMessage, 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,
|
|
ModelCallResult,
|
|
ModelRequest,
|
|
ModelResponse,
|
|
OmitFromInput,
|
|
)
|
|
from langchain.tools import InjectedToolCallId
|
|
|
|
|
|
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: Annotated[NotRequired[list[Todo]], OmitFromInput]
|
|
"""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.""" # noqa: E501
|
|
|
|
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.""" # noqa: E501
|
|
|
|
|
|
@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 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
|
|
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.todo import TodoListMiddleware
|
|
from langchain.agents import create_agent
|
|
|
|
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")]})
|
|
|
|
print(result["todos"]) # Array of todo items with status tracking
|
|
```
|
|
"""
|
|
|
|
state_schema = PlanningState
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
system_prompt: str = WRITE_TODOS_SYSTEM_PROMPT,
|
|
tool_description: str = WRITE_TODOS_TOOL_DESCRIPTION,
|
|
) -> None:
|
|
"""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.
|
|
"""
|
|
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 wrap_model_call(
|
|
self,
|
|
request: ModelRequest,
|
|
handler: Callable[[ModelRequest], ModelResponse],
|
|
) -> ModelCallResult:
|
|
"""Update the system message to include the todo system prompt."""
|
|
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.override(system_message=new_system_message))
|
|
|
|
async def awrap_model_call(
|
|
self,
|
|
request: ModelRequest,
|
|
handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
|
|
) -> ModelCallResult:
|
|
"""Update the system message to include the todo system prompt (async version)."""
|
|
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.override(system_message=new_system_message))
|