langchain-dev-utils 1.2.6__py3-none-any.whl → 1.2.8__py3-none-any.whl

langchain_dev_utils/__init__.py

@@ -1 +1 @@
-__version__ = "1.2.6"
+__version__ = "1.2.8"

langchain_dev_utils/_utils.py

@@ -1,13 +1,17 @@
 from importlib import util
+from typing import Literal
 
 from pydantic import BaseModel
 
 
-def _check_langchain_openai_install() -> None:
-    if not util.find_spec("langchain_openai"):
-        msg = (
-            "Please install langchain_dev_utils[standard],when use 'openai-compatible'"
-        )
+def _check_pkg_install(
+    pkg: Literal["langchain_openai", "json_repair"],
+) -> None:
+    if not util.find_spec(pkg):
+        if pkg == "langchain_openai":
+            msg = "Please install langchain_dev_utils[standard],when use 'openai-compatible'"
+        else:
+            msg = "Please install langchain_dev_utils[standard] to use ToolCallRepairMiddleware."
         raise ImportError(msg)
 
 

langchain_dev_utils/agents/__init__.py

@@ -1,5 +1,4 @@
 from .factory import create_agent
 from .wrap import wrap_agent_as_tool
 
-
 __all__ = ["create_agent", "wrap_agent_as_tool"]

langchain_dev_utils/agents/factory.py

@@ -1,5 +1,4 @@
-from typing import Any, Sequence
-from typing import Callable
+from typing import Any, Callable, Sequence
 
 from langchain.agents import create_agent as _create_agent
 from langchain.agents.middleware.types import (
@@ -63,21 +62,14 @@ def create_agent(  # noqa: PLR0915
     Example:
         >>> from langchain_dev_utils.chat_models import register_model_provider
         >>> from langchain_dev_utils.agents import create_agent
-        >>> from langchain_core.tools import tool
-        >>> import datetime
         >>>
-        >>> # Register a model provider
+        # Register a model provider, must be done before creating the agent
         >>> register_model_provider(
         ...     provider_name="vllm",
         ...     chat_model="openai-compatible",
         ...     base_url="http://localhost:8000/v1",
         ... )
         >>>
-        >>> @tool
-        ... def get_current_time() -> str:
-        ...     \"\"\"Get current time.\"\"\"
-        ...     return str(datetime.datetime.now().timestamp())
-        >>>
         >>> agent = create_agent(
         ...     "vllm:qwen3-4b",
         ...     tools=[get_current_time],

langchain_dev_utils/agents/file_system.py

@@ -1,5 +1,5 @@
-from typing import Annotated, Literal, Optional
 import warnings
+from typing import Annotated, Literal, Optional
 
 from langchain.tools import BaseTool, ToolRuntime, tool
 from langchain_core.messages import ToolMessage

langchain_dev_utils/agents/middleware/__init__.py

@@ -7,6 +7,7 @@ from .plan import (
     create_write_plan_tool,
 )
 from .summarization import SummarizationMiddleware
+from .tool_call_repair import ToolCallRepairMiddleware
 from .tool_emulator import LLMToolEmulator
 from .tool_selection import LLMToolSelectorMiddleware
 
@@ -20,4 +21,5 @@ __all__ = [
     "ModelFallbackMiddleware",
     "LLMToolEmulator",
     "ModelRouterMiddleware",
+    "ToolCallRepairMiddleware",
 ]

langchain_dev_utils/agents/middleware/model_fallback.py

@@ -26,7 +26,7 @@ class ModelFallbackMiddleware(_ModelFallbackMiddleware):
         )
 
         agent = create_agent(
-            model="vllm:qwen3-4b",#Primary model
+            model="vllm:qwen3-4b", #Primary model
             middleware=[fallback],
         )
 

langchain_dev_utils/agents/middleware/model_router.py

@@ -68,11 +68,15 @@ class ModelRouterState(AgentState):
 
 
 class ModelRouterMiddleware(AgentMiddleware):
-    """Model routing middleware that automatically selects the most suitable model based on input content.
+    """Model routing middleware that automatically selects the most suitable model
+    based on input content.
 
     Args:
-        router_model: Model identifier used for routing selection, it can be a model name or a BaseChatModel instance
-        model_list: List of available routing models, each containing model_name, model_description, tools(Optional), model_kwargs(Optional), model_system_prompt(Optional)
+        router_model: Model identifier used for routing selection, it can be a
+            model name or a BaseChatModel instance
+        model_list: List of available routing models, each containing model_name,
+            model_description, tools(Optional), model_kwargs(Optional),
+            model_system_prompt(Optional)
         router_prompt: Routing prompt template, uses default template if None
 
     Examples:
@@ -145,9 +149,7 @@ class ModelRouterMiddleware(AgentMiddleware):
         model_name = await self._aselect_model(state["messages"])
         return {"router_model_selection": model_name}
 
-    def wrap_model_call(
-        self, request: ModelRequest, handler: Callable[[ModelRequest], ModelResponse]
-    ) -> ModelCallResult:
+    def _get_override_kwargs(self, request: ModelRequest) -> dict[str, Any]:
         model_dict = {
             item["model_name"]: {
                 "tools": item.get("tools", None),
@@ -159,49 +161,38 @@ class ModelRouterMiddleware(AgentMiddleware):
         select_model_name = request.state.get("router_model_selection", "default-model")
 
         override_kwargs = {}
-        if select_model_name != "default-model":
-            if select_model_name in model_dict:
-                model_values = model_dict.get(select_model_name, {})
-                if model_values["kwargs"] is not None:
-                    model = load_chat_model(select_model_name, **model_values["kwargs"])
-                else:
-                    model = load_chat_model(select_model_name)
-                override_kwargs["model"] = model
-                if model_values["tools"] is not None:
-                    override_kwargs["tools"] = model_values["tools"]
-                if model_values["system_prompt"] is not None:
-                    override_kwargs["system_message"] = SystemMessage(
-                        content=model_values["system_prompt"]
-                    )
-        return handler(request.override(**override_kwargs))
+        if select_model_name != "default-model" and select_model_name in model_dict:
+            model_values = model_dict.get(select_model_name, {})
+            if model_values["kwargs"] is not None:
+                model = load_chat_model(select_model_name, **model_values["kwargs"])
+            else:
+                model = load_chat_model(select_model_name)
+            override_kwargs["model"] = model
+            if model_values["tools"] is not None:
+                override_kwargs["tools"] = model_values["tools"]
+            if model_values["system_prompt"] is not None:
+                override_kwargs["system_message"] = SystemMessage(
+                    content=model_values["system_prompt"]
+                )
+
+        return override_kwargs
+
+    def wrap_model_call(
+        self, request: ModelRequest, handler: Callable[[ModelRequest], ModelResponse]
+    ) -> ModelCallResult:
+        override_kwargs = self._get_override_kwargs(request)
+        if override_kwargs:
+            return handler(request.override(**override_kwargs))
+        else:
+            return handler(request)
 
     async def awrap_model_call(
         self,
         request: ModelRequest,
         handler: Callable[[ModelRequest], Awaitable[ModelResponse]],
     ) -> ModelCallResult:
-        model_dict = {
-            item["model_name"]: {
-                "tools": item.get("tools", None),
-                "kwargs": item.get("model_kwargs", None),
-                "system_prompt": item.get("model_system_prompt", None),
-            }
-            for item in self.model_list
-        }
-        select_model_name = request.state.get("router_model_selection", "default-model")
-        override_kwargs = {}
-        if select_model_name != "default-model":
-            if select_model_name in model_dict:
-                model_values = model_dict.get(select_model_name, {})
-                if model_values["kwargs"] is not None:
-                    model = load_chat_model(select_model_name, **model_values["kwargs"])
-                else:
-                    model = load_chat_model(select_model_name)
-                override_kwargs["model"] = model
-                if model_values["tools"] is not None:
-                    override_kwargs["tools"] = model_values["tools"]
-                if model_values["system_prompt"] is not None:
-                    override_kwargs["system_message"] = SystemMessage(
-                        content=model_values["system_prompt"]
-                    )
-        return await handler(request.override(**override_kwargs))
+        override_kwargs = self._get_override_kwargs(request)
+        if override_kwargs:
+            return await handler(request.override(**override_kwargs))
+        else:
+            return await handler(request)

langchain_dev_utils/agents/middleware/plan.py

@@ -1,6 +1,5 @@
 import json
-from typing import Awaitable, Callable, Literal, Optional, cast
-from typing import NotRequired
+from typing import Awaitable, Callable, Literal, NotRequired, Optional, cast
 
 from langchain.agents.middleware import ModelRequest, ModelResponse
 from langchain.agents.middleware.types import (
@@ -160,7 +159,8 @@ def create_finish_sub_plan_tool(
 ) -> 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.
+    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.
@@ -273,16 +273,20 @@ _PLAN_SYSTEM_PROMPT = """You can manage task plans using three simple tools:
 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.
+    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.
+    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.
+        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.
         write_plan_tool_description: Description of the `write_plan` tool.
             If not provided, uses the default `_DEFAULT_WRITE_PLAN_TOOL_DESCRIPTION`.
         finish_sub_plan_tool_description: Description of the `finish_sub_plan` tool.
@@ -291,7 +295,7 @@ class PlanMiddleware(AgentMiddleware):
             If not provided, uses the default `_DEFAULT_READ_PLAN_TOOL_DESCRIPTION`.
         use_read_plan_tool: Whether to use the `read_plan` tool.
             If not provided, uses the default `True`.
-        message_key: The key of the message to be updated. Defaults to "messages".
+
     Example:
         ```python
         from langchain_dev_utils.agents.middleware import PlanMiddleware
@@ -316,7 +320,6 @@ class PlanMiddleware(AgentMiddleware):
         finish_sub_plan_tool_description: Optional[str] = None,
         read_plan_tool_description: Optional[str] = None,
         use_read_plan_tool: bool = True,
-        message_key: Optional[str] = None,
     ) -> None:
         super().__init__()
 
@@ -332,12 +335,8 @@ class PlanMiddleware(AgentMiddleware):
         )
 
         tools = [
-            create_write_plan_tool(
-                description=write_plan_tool_description, message_key=message_key
-            ),
-            create_finish_sub_plan_tool(
-                description=finish_sub_plan_tool_description, message_key=message_key
-            ),
+            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:

langchain_dev_utils/agents/middleware/summarization.py

@@ -1,12 +1,14 @@
 from typing import Any
 
 from langchain.agents.middleware.summarization import (
-    ContextSize,
-    DEFAULT_SUMMARY_PROMPT,
-    SummarizationMiddleware as _SummarizationMiddleware,
-    TokenCounter,
     _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
 

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

@@ -20,7 +20,7 @@ class LLMToolEmulator(_LLMToolEmulator):
         model: Model to use for emulation. Must be a string identifier.
 
     Examples:
-        Emulate all tools (default behavior):
+        # Emulate all tools (default behavior):
         ```python
         from langchain_dev_utils.agents import create_agent
         from langchain_dev_utils.agents.middleware import LLMToolEmulator
@@ -36,12 +36,12 @@ class LLMToolEmulator(_LLMToolEmulator):
         )
         ```
 
-        Emulate specific tools by name:
+        # 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:
+        # Emulate specific tools by passing tool instances:
         ```python
         middleware = LLMToolEmulator(model="vllm:qwen3-4b", tools=[get_weather, get_user_location])
         ```

langchain_dev_utils/agents/middleware/tool_selection.py

@@ -30,7 +30,7 @@ class LLMToolSelectorMiddleware(_LLMToolSelectorMiddleware):
             against the max_tools limit.
 
     Examples:
-        Basic usage with tool limit:
+        # Basic usage with tool limit:
         ```python
         from langchain_dev_utils.agents.middleware import LLMToolSelectorMiddleware
 
@@ -40,7 +40,7 @@ class LLMToolSelectorMiddleware(_LLMToolSelectorMiddleware):
         )
         ```
 
-        With always-included tools:
+        # With always-included tools:
         ```python
         middleware = LLMToolSelectorMiddleware(
             model="vllm:qwen3-4b",
@@ -49,7 +49,7 @@ class LLMToolSelectorMiddleware(_LLMToolSelectorMiddleware):
         )
         ```
 
-        With custom system prompt:
+        # With custom system prompt:
         ```python
         custom_prompt = "Select tools that can help answer user questions about data."
         middleware = LLMToolSelectorMiddleware(

langchain_dev_utils/agents/plan.py

@@ -1,5 +1,5 @@
-from typing import Literal, Optional
 import warnings
+from typing import Literal, Optional
 
 from langchain.tools import BaseTool, ToolRuntime, tool
 from langchain_core.messages import ToolMessage

langchain_dev_utils/agents/wrap.py

@@ -54,30 +54,18 @@ def wrap_agent_as_tool(
         BaseTool: The wrapped agent as a tool
 
     Example:
-        >>> import datetime
-        >>> from langchain_core.messages import HumanMessage
-        >>> from langchain_core.tools import tool
         >>> from langchain_dev_utils.agents import wrap_agent_as_tool, create_agent
-
-        >>> @tool
-        ... def get_current_time() -> str:
-        ...     "\"\"Get the current timestamp.\"\"\"
-        ...     return str(datetime.datetime.now().timestamp())
-
-        >>> # Define an agent for querying the time
-        >>> time_agent = create_agent(
-        ...     "vllm:qwen3-4b", tools=[get_current_time], name="time_agent"
+        >>>
+        >>> call_time_agent_tool = wrap_agent_as_tool(
+        ...     time_agent,
+        ...     tool_name="call_time_agent",
+        ...     tool_description="Used to invoke the time sub-agent to perform time-related tasks"
         ... )
-        >>> tool = wrap_agent_as_tool(
-        ...     time_agent, tool_name="call_time_agent", tool_description="Invoke the time agent"
-        ... )
-        >>> print(tool)
-
-        >>> # Use it as a tool
-        >>> agent = create_agent("vllm:qwen3-4b", tools=[tool], name="agent")
+        >>>
+        >>> agent = create_agent("vllm:qwen3-4b", tools=[call_time_agent_tool], name="agent")
 
         >>> response = agent.invoke({"messages": [HumanMessage(content="What time is it now?")]})
-        >>> print(response)
+        >>> response
     """
     if agent.name is None:
         raise ValueError("Agent name must not be None")