langchain-dev-utils 1.3.7__py3-none-any.whl

langchain_dev_utils/agents/middleware/plan.py

@@ -0,0 +1,367 @@
+import json
+from typing import Awaitable, Callable, Literal, NotRequired, Optional, cast
+
+from langchain.agents.middleware import ModelRequest, ModelResponse
+from langchain.agents.middleware.types import (
+    AgentMiddleware,
+    AgentState,
+    ModelCallResult,
+)
+from langchain.tools import BaseTool, ToolRuntime, tool
+from langchain_core.messages import SystemMessage, ToolMessage
+from langgraph.types import Command
+from typing_extensions import TypedDict
+
+_DEFAULT_WRITE_PLAN_TOOL_DESCRIPTION = """Use this tool to create and manage a structured task list for complex or multi-step work. It helps you stay organized, track progress, and demonstrate to the user that you’re handling tasks systematically.
+
+## When to Use This Tool  
+Use this tool in the following scenarios:
+
+1. **Complex multi-step tasks** — when a task requires three or more distinct steps or actions.  
+2. **Non-trivial and complex tasks** — tasks that require careful planning or involve multiple operations.  
+3. **User explicitly requests a to-do list** — when the user directly asks you to use the to-do list feature.  
+4. **User provides multiple tasks** — when the user supplies a list of items to be done (e.g., numbered or comma-separated).  
+5. **The plan needs adjustment based on current execution** — when ongoing progress indicates the plan should be revised.
+
+## How to Use This Tool  
+1. **When starting a task** — before actually beginning work, invoke this tool with a task list (a list of strings). The first task will automatically be set to `in_progress`, and all others to `pending`.  
+2. **When updating the task list** — for example, after completing some tasks, if you find certain tasks are no longer needed, remove them; if new necessary tasks emerge, add them. However, **do not modify** tasks already marked as completed. In such cases, simply call this tool again with the updated task list.
+
+## When NOT to Use This Tool  
+Avoid using this tool in the following situations:  
+1. The task is a **single, straightforward action**.  
+2. The task is **too trivial**, and tracking it provides no benefit.  
+3. The task can be completed in **fewer than three simple steps**.  
+4. The current task list has been fully completed — in this case, use `finish_sub_plan()` to finalize.
+
+## How It Works  
+- **Input**: A parameter named `plan` containing a list of strings representing the tasks (e.g., `["Task 1", "Task 2", "Task 3"]`).  
+- **Automatic status assignment**:  
+  → First task: `in_progress`  
+  → Remaining tasks: `pending`  
+- When updating the plan, provide only the **next set of tasks to execute**. For example, if the next phase requires `["Task 4", "Task 5"]`, call this tool with `plan=["Task 4", "Task 5"]`.
+
+## Task States  
+- `pending`: Ready to start, awaiting execution  
+- `in_progress`: Currently being worked on  
+- `done`: Completed  
+
+## Best Practices  
+- Break large tasks into clear, actionable steps.  
+- Use specific and descriptive task names.  
+- Update the plan immediately if priorities shift or blockers arise.  
+- Never leave the plan empty — as long as unfinished tasks remain, at least one must be marked `in_progress`.  
+- Do not batch completions — mark each task as done immediately after finishing it.  
+- Remove irrelevant tasks entirely instead of leaving them in `pending` state.
+
+**Remember**: If a task is simple, just do it. This tool is meant to provide structure — not overhead.
+"""
+
+_DEFAULT_FINISH_SUB_PLAN_TOOL_DESCRIPTION = """This tool is used to mark the currently in-progress task in an existing task list as completed.
+
+## Functionality  
+- Marks the current task with status `in_progress` as `done`, and automatically sets the next task (previously `pending`) to `in_progress`.
+
+## When to Use  
+Use only when you have confirmed that the current task is truly finished.
+
+## Example  
+Before calling:
+```json
+[
+    {"content": "Task 1", "status": "done"},
+    {"content": "Task 2", "status": "in_progress"},
+    {"content": "Task 3", "status": "pending"}
+]
+```
+
+After calling `finish_sub_plan()`:
+```json
+[
+    {"content": "Task 1", "status": "done"},
+    {"content": "Task 2", "status": "done"},
+    {"content": "Task 3", "status": "in_progress"}
+]
+```
+
+**Note**:  
+- This tool is **only** for marking completion — do **not** use it to create or modify plans (use `write_plan` instead).  
+- Ensure the task is genuinely complete before invoking this function.  
+- No parameters are required — status updates are handled automatically.
+"""
+
+_DEFAULT_READ_PLAN_TOOL_DESCRIPTION = """
+Get all sub-plans with their current status.
+"""
+
+
+class Plan(TypedDict):
+    content: str
+    status: Literal["pending", "in_progress", "done"]
+
+
+class PlanState(AgentState):
+    plan: NotRequired[list[Plan]]
+
+
+class PlanToolDescription(TypedDict):
+    write_plan: NotRequired[str]
+    finish_sub_plan: NotRequired[str]
+    read_plan: NotRequired[str]
+
+
+def _create_write_plan_tool(
+    description: Optional[str] = None,
+) -> BaseTool:
+    """Create a tool for writing initial plan.
+
+    This function creates a tool that allows agents to write an initial plan
+    with a list of plans. The first plan in the plan will be marked as "in_progress"
+    and the rest as "pending".
+
+    Args:
+        description: The description of the tool. Uses default description if not provided.
+
+    Returns:
+        BaseTool: The tool for writing initial plan.
+    """
+
+    @tool(
+        description=description or _DEFAULT_WRITE_PLAN_TOOL_DESCRIPTION,
+    )
+    def write_plan(plan: list[str], runtime: ToolRuntime):
+        return Command(
+            update={
+                "plan": [
+                    {
+                        "content": content,
+                        "status": "pending" if index > 0 else "in_progress",
+                    }
+                    for index, content in enumerate(plan)
+                ],
+                "messages": [
+                    ToolMessage(
+                        content=f"Plan successfully written, please first execute the {plan[0]} sub-plan (no need to change the status to in_process)",
+                        tool_call_id=runtime.tool_call_id,
+                    )
+                ],
+            }
+        )
+
+    return write_plan
+
+
+def _create_finish_sub_plan_tool(
+    description: Optional[str] = None,
+) -> BaseTool:
+    """Create a tool for finishing sub-plan tasks.
+
+    This function creates a tool that allows agents to update the status of sub-plans in a plan.
+    Sub-plans can be marked as "done" to track progress.
+
+    Args:
+        description: The description of the tool. Uses default description if not provided.
+
+    Returns:
+        BaseTool: The tool for finishing sub-plan tasks.
+    """
+
+    @tool(
+        description=description or _DEFAULT_FINISH_SUB_PLAN_TOOL_DESCRIPTION,
+    )
+    def finish_sub_plan(
+        runtime: ToolRuntime,
+    ):
+        plan_list = runtime.state.get("plan", [])
+
+        sub_finish_plan = ""
+        sub_next_plan = ",all sub plan are done"
+        for plan in plan_list:
+            if plan["status"] == "in_progress":
+                plan["status"] = "done"
+                sub_finish_plan = f"finish sub plan:**{plan['content']}**"
+
+        for plan in plan_list:
+            if plan["status"] == "pending":
+                plan["status"] = "in_progress"
+                sub_next_plan = f",next plan:**{plan['content']}**"
+                break
+
+        return Command(
+            update={
+                "plan": plan_list,
+                "messages": [
+                    ToolMessage(
+                        content=sub_finish_plan + sub_next_plan,
+                        tool_call_id=runtime.tool_call_id,
+                    )
+                ],
+            }
+        )
+
+    return finish_sub_plan
+
+
+def _create_read_plan_tool(
+    description: Optional[str] = None,
+):
+    """Create a tool for reading all sub-plans.
+
+    This function creates a tool that allows agents to read all sub-plans
+    in the current plan with their status information.
+
+    Args:
+        description: The description of the tool. Uses default description if not provided.
+
+    Returns:
+        BaseTool: The tool for reading all sub-plans.
+    """
+
+    @tool(
+        description=description or _DEFAULT_READ_PLAN_TOOL_DESCRIPTION,
+    )
+    def read_plan(runtime: ToolRuntime):
+        plan_list = runtime.state.get("plan", [])
+        return json.dumps(plan_list)
+
+    return read_plan
+
+
+_PLAN_SYSTEM_PROMPT_NOT_READ_PLAN = """You can manage task plans using two simple tools:
+
+## write_plan
+- Use it to break complex tasks (3+ steps) into a clear, actionable list. Only include next steps to execute — the first becomes `"in_progress"`, the rest `"pending"`. Don’t use it for simple tasks (<3 steps).
+
+## finish_sub_plan
+- Call it **only when the current task is 100% done**. It automatically marks it `"done"` and promotes the next `"pending"` task to `"in_progress"`. No parameters needed. Never use it mid-task or if anything’s incomplete.
+Keep plans lean, update immediately, and never batch completions.
+
+**Note**: Make sure that all tasks end up with the status `"done"`.
+"""
+
+_PLAN_SYSTEM_PROMPT = """You can manage task plans using three simple tools:
+
+## write_plan
+- Use it to break complex tasks (3+ steps) into a clear, actionable list. Only include next steps to execute — the first becomes `"in_progress"`, the rest `"pending"`. Don’t use it for simple tasks (<3 steps).
+
+## finish_sub_plan
+- Call it **only when the current task is 100% done**. It automatically marks it `"done"` and promotes the next `"pending"` task to `"in_progress"`. No parameters needed. Never use it mid-task or if anything’s incomplete.
+
+## read_plan
+- Retrieve the full current plan list with statuses, especially when you forget which sub-plan you're supposed to execute next.
+- No parameters required—returns a current plan list with statuses.
+
+**Note**: Make sure that all tasks end up with the status `"done"`.
+"""
+
+
+class PlanMiddleware(AgentMiddleware):
+    """Middleware that provides plan management capabilities to agents.
+
+    This middleware adds a `write_plan` and `finish_sub_plan` (and `read_plan`
+    optional) tool that allows agents to create and manage structured plan 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
+    how to use the plan functionality effectively.
+
+    Args:
+        system_prompt: Custom system prompt to guide the agent on using the plan
+            tool. If not provided, uses the default `_PLAN_SYSTEM_PROMPT` or
+            `_PLAN_SYSTEM_PROMPT_NOT_READ_PLAN` based on the `use_read_plan_tool`
+            parameter.
+        custom_plan_tool_descriptions: Custom descriptions for the plan tools.
+            If not provided, uses the default descriptions.
+        use_read_plan_tool: Whether to use the `read_plan` tool.
+            If not provided, uses the default `True`.
+
+    Example:
+        ```python
+        from langchain_dev_utils.agents.middleware import PlanMiddleware
+        from langchain_dev_utils.agents import create_agent
+
+        agent = create_agent("vllm:qwen3-4b", middleware=[PlanMiddleware()])
+
+        # Agent now has access to write_plan tool and plan state tracking
+        result = await agent.invoke({"messages": [HumanMessage("Help me refactor my codebase")]})
+
+        print(result["plan"])  # Array of plan items with status tracking
+        ```
+    """
+
+    state_schema = PlanState
+
+    def __init__(
+        self,
+        *,
+        system_prompt: Optional[str] = None,
+        custom_plan_tool_descriptions: Optional[PlanToolDescription] = None,
+        use_read_plan_tool: bool = True,
+    ) -> None:
+        super().__init__()
+
+        if not custom_plan_tool_descriptions:
+            custom_plan_tool_descriptions = {}
+
+        write_plan_tool_description = custom_plan_tool_descriptions.get(
+            "write_plan",
+            _DEFAULT_WRITE_PLAN_TOOL_DESCRIPTION,
+        )
+        finish_sub_plan_tool_description = custom_plan_tool_descriptions.get(
+            "finish_sub_plan",
+            _DEFAULT_FINISH_SUB_PLAN_TOOL_DESCRIPTION,
+        )
+        read_plan_tool_description = custom_plan_tool_descriptions.get(
+            "read_plan",
+            _DEFAULT_READ_PLAN_TOOL_DESCRIPTION,
+        )
+
+        tools = [
+            _create_write_plan_tool(description=write_plan_tool_description),
+            _create_finish_sub_plan_tool(description=finish_sub_plan_tool_description),
+        ]
+
+        if use_read_plan_tool:
+            tools.append(_create_read_plan_tool(description=read_plan_tool_description))
+
+        if system_prompt is None:
+            if use_read_plan_tool:
+                system_prompt = _PLAN_SYSTEM_PROMPT
+            else:
+                system_prompt = _PLAN_SYSTEM_PROMPT_NOT_READ_PLAN
+
+        self.system_prompt = system_prompt
+        self.tools = tools
+
+    def _get_override_request(self, request: ModelRequest) -> ModelRequest:
+        """Add the plan system prompt to the system message."""
+        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 request.override(system_message=new_system_message)
+
+    def wrap_model_call(
+        self,
+        request: ModelRequest,
+        handler: Callable[[ModelRequest], ModelResponse],
+    ) -> ModelCallResult:
+        override_request = self._get_override_request(request)
+        return handler(override_request)
+
+    async def awrap_model_call(
+        self,
+        request: ModelRequest,
+        handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
+    ) -> ModelCallResult:
+        """Update the system message to include the plan system prompt."""
+        override_request = self._get_override_request(request)
+        return await handler(override_request)

langchain_dev_utils/agents/middleware/summarization.py

@@ -0,0 +1,85 @@
+from typing import Any
+
+from langchain.agents.middleware.summarization import (
+    _DEFAULT_MESSAGES_TO_KEEP,
+    _DEFAULT_TRIM_TOKEN_LIMIT,
+    DEFAULT_SUMMARY_PROMPT,
+    ContextSize,
+    TokenCounter,
+)
+from langchain.agents.middleware.summarization import (
+    SummarizationMiddleware as _SummarizationMiddleware,
+)
+from langchain_core.messages.utils import count_tokens_approximately
+
+from langchain_dev_utils.chat_models.base import load_chat_model
+
+
+class SummarizationMiddleware(_SummarizationMiddleware):
+    """Initialize summarization middleware.
+
+    Args:
+        model: The language model to use for generating summaries.
+        trigger: One or more thresholds that trigger summarization.
+
+            Provide a single `ContextSize` tuple or a list of tuples, in which case
+            summarization runs when any threshold is breached.
+
+            Examples: `("messages", 50)`, `("tokens", 3000)`, `[("fraction", 0.8),
+                ("messages", 100)]`.
+        keep: Context retention policy applied after summarization.
+
+            Provide a `ContextSize` tuple to specify how much history to preserve.
+
+            Defaults to keeping the most recent 20 messages.
+
+            Examples: `("messages", 20)`, `("tokens", 3000)`, or
+                `("fraction", 0.3)`.
+        token_counter: Function to count tokens in messages.
+        summary_prompt: Prompt template for generating summaries.
+        trim_tokens_to_summarize: Maximum tokens to keep when preparing messages for
+            the summarization call.
+
+            Pass `None` to skip trimming entirely.
+
+    Examples:
+        ```python
+        from langchain_dev_utils.agents.middleware import SummarizationMiddleware
+
+        middleware = SummarizationMiddleware(
+            model="vllm:qwen3-4b",
+            trigger=("tokens", 100),
+            keep=("messages", 2),
+        )
+        ```
+    """
+
+    def __init__(
+        self,
+        model: str,
+        *,
+        trigger: ContextSize | list[ContextSize] | None = None,
+        keep: ContextSize = ("messages", _DEFAULT_MESSAGES_TO_KEEP),
+        token_counter: TokenCounter = count_tokens_approximately,
+        summary_prompt: str = DEFAULT_SUMMARY_PROMPT,
+        trim_tokens_to_summarize: int | None = _DEFAULT_TRIM_TOKEN_LIMIT,
+        **deprecated_kwargs: Any,
+    ) -> None:
+        chat_model = load_chat_model(model)
+
+        middleware_kwargs = {}
+        if trigger is not None:
+            middleware_kwargs["trigger"] = trigger
+        if keep is not None:
+            middleware_kwargs["keep"] = keep
+        if token_counter is not None:
+            middleware_kwargs["token_counter"] = token_counter
+        if summary_prompt is not None:
+            middleware_kwargs["summary_prompt"] = summary_prompt
+        if trim_tokens_to_summarize is not None:
+            middleware_kwargs["trim_tokens_to_summarize"] = trim_tokens_to_summarize
+
+        super().__init__(
+            model=chat_model,
+            **middleware_kwargs,
+        )

langchain_dev_utils/agents/middleware/tool_call_repair.py

@@ -0,0 +1,96 @@
+from typing import Any, Awaitable, Callable, cast
+
+from langchain.agents.middleware import AgentMiddleware, ModelRequest, ModelResponse
+from langchain.agents.middleware.types import ModelCallResult
+from langchain_core.messages import AIMessage, BaseMessage
+
+from langchain_dev_utils._utils import _check_pkg_install
+
+
+class ToolCallRepairMiddleware(AgentMiddleware):
+    """Middleware to repair invalid tool calls in AIMessages.
+
+    This middleware attempts to repair JSON-formatted tool arguments in
+    AIMessages that have invalid tool calls. It uses the `json_repair`
+    package to fix common JSON errors.
+
+    Example:
+        ```python
+        from langchain_dev_utils.agents.middleware import ToolCallRepairMiddleware
+
+        middleware = ToolCallRepairMiddleware()
+        ```
+    """
+
+    def _repair_msgs(self, messages: list[BaseMessage]) -> list[BaseMessage]:
+        _check_pkg_install("json_repair")
+        from json import JSONDecodeError
+
+        from json_repair import loads
+
+        results = []
+        for msg in messages:
+            if (
+                isinstance(msg, AIMessage)
+                and hasattr(msg, "invalid_tool_calls")
+                and len(msg.invalid_tool_calls) > 0
+            ):
+                new_invalid_toolcalls = []
+                new_tool_calls = [*msg.tool_calls]
+
+                for invalid_tool_call in msg.invalid_tool_calls:
+                    args = invalid_tool_call.get("args")
+                    if args:
+                        try:
+                            args = cast(dict[str, Any], loads(args))
+                            new_tool_calls.append(
+                                {
+                                    "name": invalid_tool_call.get(
+                                        "name",
+                                    )
+                                    or "",
+                                    "id": invalid_tool_call.get("id", ""),
+                                    "type": "tool_call",
+                                    "args": args,
+                                }
+                            )
+                        except JSONDecodeError:
+                            new_invalid_toolcalls.append(invalid_tool_call)
+                    else:
+                        new_invalid_toolcalls.append(invalid_tool_call)
+
+                new_msg = msg.model_copy(
+                    update={
+                        "tool_calls": new_tool_calls,
+                        "invalid_tool_calls": new_invalid_toolcalls,
+                    }
+                )
+                results.append(new_msg)
+            else:
+                results.append(msg)
+
+        return results
+
+    def wrap_model_call(
+        self, request: ModelRequest, handler: Callable[[ModelRequest], ModelResponse]
+    ) -> ModelCallResult:
+        response = handler(request)
+        results = self._repair_msgs(response.result)
+
+        return ModelResponse(
+            result=results,
+            structured_response=response.structured_response,
+        )
+
+    async def awrap_model_call(
+        self,
+        request: ModelRequest,
+        handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
+    ) -> ModelCallResult:
+        response = await handler(request)
+        results = self._repair_msgs(response.result)
+
+        return ModelResponse(
+            result=results,
+            structured_response=response.structured_response,
+        )

langchain_dev_utils/agents/middleware/tool_emulator.py

@@ -0,0 +1,60 @@
+from langchain.agents.middleware.tool_emulator import (
+    LLMToolEmulator as _LLMToolEmulator,
+)
+from langchain_core.tools import BaseTool
+
+from langchain_dev_utils.chat_models.base import load_chat_model
+
+
+class LLMToolEmulator(_LLMToolEmulator):
+    """Middleware that emulates specified tools using an LLM instead of executing them.
+
+    This middleware allows selective emulation of tools for testing purposes.
+    By default (when tools=None), all tools are emulated. You can specify which
+    tools to emulate by passing a list of tool names or BaseTool instances.
+
+    Args:
+        tools: List of tool names (str) or BaseTool instances to emulate.
+            If None (default), ALL tools will be emulated.
+            If empty list, no tools will be emulated.
+        model: Model to use for emulation. Must be a string identifier.
+
+    Examples:
+        # Emulate all tools (default behavior):
+        ```python
+        from langchain_dev_utils.agents import create_agent
+        from langchain_dev_utils.agents.middleware import LLMToolEmulator
+
+        middleware = LLMToolEmulator(
+            model="vllm:qwen3-4b"
+        )
+
+        agent = create_agent(
+            model="vllm:qwen3-4b",
+            tools=[get_weather, get_user_location, calculator],
+            middleware=[middleware],
+        )
+        ```
+
+        # Emulate specific tools by name:
+        ```python
+        middleware = LLMToolEmulator(model="vllm:qwen3-4b", tools=["get_weather", "get_user_location"])
+        ```
+
+        # Emulate specific tools by passing tool instances:
+        ```python
+        middleware = LLMToolEmulator(model="vllm:qwen3-4b", tools=[get_weather, get_user_location])
+        ```
+    """
+
+    def __init__(
+        self,
+        *,
+        model: str,
+        tools: list[str | BaseTool] | None = None,
+    ) -> None:
+        chat_model = load_chat_model(model)
+        super().__init__(
+            model=chat_model,
+            tools=tools,
+        )

langchain_dev_utils/agents/middleware/tool_selection.py

@@ -0,0 +1,82 @@
+from typing import Optional
+
+from langchain.agents.middleware.tool_selection import (
+    LLMToolSelectorMiddleware as _LLMToolSelectorMiddleware,
+)
+
+from langchain_dev_utils.chat_models.base import load_chat_model
+
+
+class LLMToolSelectorMiddleware(_LLMToolSelectorMiddleware):
+    """Intelligent tool selection middleware using LLM to filter relevant tools.
+
+    This middleware leverages a language model to analyze user queries and select
+    only the most relevant tools from a potentially large toolset. This optimization
+    reduces token consumption and improves model performance by focusing on appropriate tools.
+
+    The selection process analyzes the user's request and matches it against available
+    tool descriptions to determine relevance.
+
+    Args:
+        model: String identifier for the model to use for tool selection.
+            Must be a valid model identifier that can be loaded by load_chat_model().
+        system_prompt: Custom instructions for the selection model. If not provided,
+            uses the default selection prompt from the parent class.
+        max_tools: Maximum number of tools to select and pass to the main model.
+            If the LLM selects more tools than this limit, only the first max_tools
+            tools will be used. If None, no limit is applied.
+        always_include: List of tool names that must always be included in the
+            selection regardless of the LLM's decision. These tools do not count
+            against the max_tools limit.
+
+    Examples:
+        # Basic usage with tool limit:
+        ```python
+        from langchain_dev_utils.agents.middleware import LLMToolSelectorMiddleware
+
+        middleware = LLMToolSelectorMiddleware(
+            model="vllm:qwen3-4b",
+            max_tools=3
+        )
+        ```
+
+        # With always-included tools:
+        ```python
+        middleware = LLMToolSelectorMiddleware(
+            model="vllm:qwen3-4b",
+            max_tools=5,
+            always_include=["search", "calculator"]
+        )
+        ```
+
+        # With custom system prompt:
+        ```python
+        custom_prompt = "Select tools that can help answer user questions about data."
+        middleware = LLMToolSelectorMiddleware(
+            model="vllm:qwen3-4b",
+            system_prompt=custom_prompt
+        )
+        ```
+    """
+
+    def __init__(
+        self,
+        *,
+        model: str,
+        system_prompt: Optional[str] = None,
+        max_tools: Optional[int] = None,
+        always_include: Optional[list[str]] = None,
+    ) -> None:
+        chat_model = load_chat_model(model)
+
+        tool_selector_kwargs = {}
+        if system_prompt is not None:
+            tool_selector_kwargs["system_prompt"] = system_prompt
+        if max_tools is not None:
+            tool_selector_kwargs["max_tools"] = max_tools
+        if always_include is not None:
+            tool_selector_kwargs["always_include"] = always_include
+        super().__init__(
+            model=chat_model,
+            **tool_selector_kwargs,
+        )