langchain 0.3.27__py3-none-any.whl → 1.0.0__py3-none-any.whl

langchain/agents/middleware/summarization.py

@@ -0,0 +1,249 @@
+"""Summarization middleware."""
+
+import uuid
+from collections.abc import Callable, Iterable
+from typing import Any, cast
+
+from langchain_core.messages import (
+    AIMessage,
+    AnyMessage,
+    MessageLikeRepresentation,
+    RemoveMessage,
+    ToolMessage,
+)
+from langchain_core.messages.human import HumanMessage
+from langchain_core.messages.utils import count_tokens_approximately, trim_messages
+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
+
+TokenCounter = Callable[[Iterable[MessageLikeRepresentation]], int]
+
+DEFAULT_SUMMARY_PROMPT = """<role>
+Context Extraction Assistant
+</role>
+
+<primary_objective>
+Your sole objective in this task is to extract the highest quality/most relevant context from the conversation history below.
+</primary_objective>
+
+<objective_information>
+You're nearing the total number of input tokens you can accept, so you must extract the highest quality/most relevant pieces of information from your conversation history.
+This context will then overwrite the conversation history presented below. Because of this, ensure the context you extract is only the most important information to your overall goal.
+</objective_information>
+
+<instructions>
+The conversation history below will be replaced with the context you extract in this step. Because of this, you must do your very best to extract and record all of the most important context from the conversation history.
+You want to ensure that you don't repeat any actions you've already completed, so the context you extract from the conversation history should be focused on the most important information to your overall goal.
+</instructions>
+
+The user will message you with the full message history you'll be extracting context from, to then replace. Carefully read over it all, and think deeply about what information is most important to your overall goal that should be saved:
+
+With all of this in mind, please carefully read over the entire conversation history, and extract the most important and relevant context to replace it so that you can free up space in the conversation history.
+Respond ONLY with the extracted context. Do not include any additional information, or text before or after the extracted context.
+
+<messages>
+Messages to summarize:
+{messages}
+</messages>"""  # noqa: E501
+
+SUMMARY_PREFIX = "## Previous conversation summary:"
+
+_DEFAULT_MESSAGES_TO_KEEP = 20
+_DEFAULT_TRIM_TOKEN_LIMIT = 4000
+_DEFAULT_FALLBACK_MESSAGE_COUNT = 15
+_SEARCH_RANGE_FOR_TOOL_PAIRS = 5
+
+
+class SummarizationMiddleware(AgentMiddleware):
+    """Middleware that summarizes conversation history when token limits are approached.
+
+    This middleware monitors message token counts and automatically summarizes older
+    messages when a threshold is reached, preserving recent messages and maintaining
+    context continuity by ensuring AI/Tool message pairs remain together.
+    """
+
+    def __init__(
+        self,
+        model: str | BaseChatModel,
+        max_tokens_before_summary: int | None = None,
+        messages_to_keep: int = _DEFAULT_MESSAGES_TO_KEEP,
+        token_counter: TokenCounter = count_tokens_approximately,
+        summary_prompt: str = DEFAULT_SUMMARY_PROMPT,
+        summary_prefix: str = SUMMARY_PREFIX,
+    ) -> None:
+        """Initialize the summarization middleware.
+
+        Args:
+            model: The language model to use for generating summaries.
+            max_tokens_before_summary: Token threshold to trigger summarization.
+                If `None`, summarization is disabled.
+            messages_to_keep: Number of recent messages to preserve after summarization.
+            token_counter: Function to count tokens in messages.
+            summary_prompt: Prompt template for generating summaries.
+            summary_prefix: Prefix added to system message when including summary.
+        """
+        super().__init__()
+
+        if isinstance(model, str):
+            model = init_chat_model(model)
+
+        self.model = model
+        self.max_tokens_before_summary = max_tokens_before_summary
+        self.messages_to_keep = messages_to_keep
+        self.token_counter = token_counter
+        self.summary_prompt = summary_prompt
+        self.summary_prefix = summary_prefix
+
+    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)
+
+        total_tokens = self.token_counter(messages)
+        if (
+            self.max_tokens_before_summary is not None
+            and total_tokens < self.max_tokens_before_summary
+        ):
+            return None
+
+        cutoff_index = self._find_safe_cutoff(messages)
+
+        if cutoff_index <= 0:
+            return None
+
+        messages_to_summarize, preserved_messages = self._partition_messages(messages, cutoff_index)
+
+        summary = self._create_summary(messages_to_summarize)
+        new_messages = self._build_new_messages(summary)
+
+        return {
+            "messages": [
+                RemoveMessage(id=REMOVE_ALL_MESSAGES),
+                *new_messages,
+                *preserved_messages,
+            ]
+        }
+
+    def _build_new_messages(self, summary: str) -> list[HumanMessage]:
+        return [
+            HumanMessage(content=f"Here is a summary of the conversation to date:\n\n{summary}")
+        ]
+
+    def _ensure_message_ids(self, messages: list[AnyMessage]) -> None:
+        """Ensure all messages have unique IDs for the add_messages reducer."""
+        for msg in messages:
+            if msg.id is None:
+                msg.id = str(uuid.uuid4())
+
+    def _partition_messages(
+        self,
+        conversation_messages: list[AnyMessage],
+        cutoff_index: int,
+    ) -> tuple[list[AnyMessage], list[AnyMessage]]:
+        """Partition messages into those to summarize and those to preserve."""
+        messages_to_summarize = conversation_messages[:cutoff_index]
+        preserved_messages = conversation_messages[cutoff_index:]
+
+        return messages_to_summarize, preserved_messages
+
+    def _find_safe_cutoff(self, messages: list[AnyMessage]) -> int:
+        """Find safe cutoff point that preserves AI/Tool message pairs.
+
+        Returns the index where messages can be safely cut without separating
+        related AI and Tool messages. Returns 0 if no safe cutoff is found.
+        """
+        if len(messages) <= self.messages_to_keep:
+            return 0
+
+        target_cutoff = len(messages) - self.messages_to_keep
+
+        for i in range(target_cutoff, -1, -1):
+            if self._is_safe_cutoff_point(messages, i):
+                return i
+
+        return 0
+
+    def _is_safe_cutoff_point(self, messages: list[AnyMessage], cutoff_index: int) -> bool:
+        """Check if cutting at index would separate AI/Tool message pairs."""
+        if cutoff_index >= len(messages):
+            return True
+
+        search_start = max(0, cutoff_index - _SEARCH_RANGE_FOR_TOOL_PAIRS)
+        search_end = min(len(messages), cutoff_index + _SEARCH_RANGE_FOR_TOOL_PAIRS)
+
+        for i in range(search_start, search_end):
+            if not self._has_tool_calls(messages[i]):
+                continue
+
+            tool_call_ids = self._extract_tool_call_ids(cast("AIMessage", messages[i]))
+            if self._cutoff_separates_tool_pair(messages, i, cutoff_index, tool_call_ids):
+                return False
+
+        return True
+
+    def _has_tool_calls(self, message: AnyMessage) -> bool:
+        """Check if message is an AI message with tool calls."""
+        return (
+            isinstance(message, AIMessage) and hasattr(message, "tool_calls") and message.tool_calls  # type: ignore[return-value]
+        )
+
+    def _extract_tool_call_ids(self, ai_message: AIMessage) -> set[str]:
+        """Extract tool call IDs from an AI message."""
+        tool_call_ids = set()
+        for tc in ai_message.tool_calls:
+            call_id = tc.get("id") if isinstance(tc, dict) else getattr(tc, "id", None)
+            if call_id is not None:
+                tool_call_ids.add(call_id)
+        return tool_call_ids
+
+    def _cutoff_separates_tool_pair(
+        self,
+        messages: list[AnyMessage],
+        ai_message_index: int,
+        cutoff_index: int,
+        tool_call_ids: set[str],
+    ) -> bool:
+        """Check if cutoff separates an AI message from its corresponding tool messages."""
+        for j in range(ai_message_index + 1, len(messages)):
+            message = messages[j]
+            if isinstance(message, ToolMessage) and message.tool_call_id in tool_call_ids:
+                ai_before_cutoff = ai_message_index < cutoff_index
+                tool_before_cutoff = j < cutoff_index
+                if ai_before_cutoff != tool_before_cutoff:
+                    return True
+        return False
+
+    def _create_summary(self, messages_to_summarize: list[AnyMessage]) -> str:
+        """Generate summary for the given messages."""
+        if not messages_to_summarize:
+            return "No previous conversation history."
+
+        trimmed_messages = self._trim_messages_for_summary(messages_to_summarize)
+        if not trimmed_messages:
+            return "Previous conversation was too long to summarize."
+
+        try:
+            response = self.model.invoke(self.summary_prompt.format(messages=trimmed_messages))
+            return cast("str", response.content).strip()
+        except Exception as e:  # noqa: BLE001
+            return f"Error generating summary: {e!s}"
+
+    def _trim_messages_for_summary(self, messages: list[AnyMessage]) -> list[AnyMessage]:
+        """Trim messages to fit within summary generation limits."""
+        try:
+            return trim_messages(
+                messages,
+                max_tokens=_DEFAULT_TRIM_TOKEN_LIMIT,
+                token_counter=self.token_counter,
+                start_on="human",
+                strategy="last",
+                allow_partial=True,
+                include_system=True,
+            )
+        except Exception:  # noqa: BLE001
+            return messages[-_DEFAULT_FALLBACK_MESSAGE_COUNT:]

langchain/agents/middleware/todo.py

@@ -0,0 +1,219 @@
+"""Planning and task management middleware for agents."""
+# ruff: noqa: E501
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Annotated, Literal
+
+if TYPE_CHECKING:
+    from collections.abc import Awaitable, Callable
+
+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,
+    ModelCallResult,
+    ModelRequest,
+    ModelResponse,
+)
+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: 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 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
+        ```
+
+    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 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 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 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)