langchain-dev-utils 1.3.1__tar.gz → 1.3.3__tar.gz

{langchain_dev_utils-1.3.1 → langchain_dev_utils-1.3.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: langchain-dev-utils
-Version: 1.3.1
+Version: 1.3.3
 Summary: A practical utility library for LangChain and LangGraph development
 Project-URL: Source Code, https://github.com/TBice123123/langchain-dev-utils
 Project-URL: repository, https://github.com/TBice123123/langchain-dev-utils

{langchain_dev_utils-1.3.1 → langchain_dev_utils-1.3.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "langchain-dev-utils"
-version = "1.3.1"
+version = "1.3.3"
 description = "A practical utility library for LangChain and LangGraph development"
 readme = "README.md"
 authors = [{ name = "tiebingice", email = "tiebingice123@outlook.com" }]

langchain_dev_utils-1.3.3/src/langchain_dev_utils/__init__.py

@@ -0,0 +1 @@
+__version__ = "1.3.3"

langchain_dev_utils-1.3.3/src/langchain_dev_utils/agents/__init__.py

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

{langchain_dev_utils-1.3.1 → langchain_dev_utils-1.3.3}/src/langchain_dev_utils/agents/middleware/__init__.py

@@ -1,13 +1,8 @@
 from .format_prompt import format_prompt
-from .handoffs import HandoffsAgentMiddleware, create_handoffs_tool
+from .handoffs import HandoffAgentMiddleware
 from .model_fallback import ModelFallbackMiddleware
 from .model_router import ModelRouterMiddleware
-from .plan import (
-    PlanMiddleware,
-    create_finish_sub_plan_tool,
-    create_read_plan_tool,
-    create_write_plan_tool,
-)
+from .plan import PlanMiddleware
 from .summarization import SummarizationMiddleware
 from .tool_call_repair import ToolCallRepairMiddleware
 from .tool_emulator import LLMToolEmulator
@@ -17,14 +12,10 @@ __all__ = [
     "SummarizationMiddleware",
     "LLMToolSelectorMiddleware",
     "PlanMiddleware",
-    "create_finish_sub_plan_tool",
-    "create_read_plan_tool",
-    "create_write_plan_tool",
     "ModelFallbackMiddleware",
     "LLMToolEmulator",
     "ModelRouterMiddleware",
     "ToolCallRepairMiddleware",
     "format_prompt",
-    "create_handoffs_tool",
-    "HandoffsAgentMiddleware",
+    "HandoffAgentMiddleware",
 ]

{langchain_dev_utils-1.3.1 → langchain_dev_utils-1.3.3}/src/langchain_dev_utils/agents/middleware/handoffs.py

@@ -1,4 +1,4 @@
-from typing import Any, Awaitable, Callable
+from typing import Any, Awaitable, Callable, Literal
 
 from langchain.agents import AgentState
 from langchain.agents.middleware import AgentMiddleware, ModelRequest, ModelResponse
@@ -19,35 +19,24 @@ class MultiAgentState(AgentState):
 class AgentConfig(TypedDict):
     model: NotRequired[str | BaseChatModel]
     prompt: str | SystemMessage
-    tools: list[BaseTool | dict[str, Any]]
+    tools: NotRequired[list[BaseTool | dict[str, Any]]]
     default: NotRequired[bool]
+    handoffs: list[str] | Literal["all"]
 
 
-def create_handoffs_tool(
-    agent_name: str,
-    tool_name: Optional[str] = None,
-    tool_description: Optional[str] = None,
-):
+def _create_handoffs_tool(agent_name: str, tool_description: Optional[str] = None):
     """Create a tool for handoffs to a specified agent.
 
     Args:
         agent_name (str): The name of the agent to transfer to.
-        tool_name (Optional[str], optional): The name of the tool. Defaults to None.
-        tool_description (Optional[str], optional): The description of the tool. Defaults to None.
 
     Returns:
         BaseTool: A tool instance for handoffs to the specified agent.
-
-    Example:
-        Basic usage
-        >>> from langchain_dev_utils.agents.middleware import create_handoffs_tool
-        >>> handoffs_tool = create_handoffs_tool("time_agent")
     """
-    if tool_name is None:
-        tool_name = f"transfer_to_{agent_name}"
-        if not tool_name.endswith("_agent"):
-            tool_name += "_agent"
 
+    tool_name = f"transfer_to_{agent_name}"
+    if not tool_name.endswith("_agent"):
+        tool_name += "_agent"
     if tool_description is None:
         tool_description = f"Transfer to the {agent_name}"
 
@@ -75,30 +64,107 @@ def _get_default_active_agent(state: dict[str, AgentConfig]) -> Optional[str]:
     return None
 
 
-class HandoffsAgentMiddleware(AgentMiddleware):
+def _transform_agent_config(
+    config: dict[str, AgentConfig],
+    handoffs_tools: list[BaseTool],
+) -> dict[str, AgentConfig]:
+    """Transform the agent config to add handoffs tools.
+
+    Args:
+        config (dict[str, AgentConfig]): The agent config.
+        handoffs_tools (list[BaseTool]): The list of handoffs tools.
+
+    Returns:
+        dict[str, AgentConfig]: The transformed agent config.
+    """
+
+    new_config = {}
+    for agent_name, _cfg in config.items():
+        new_config[agent_name] = {}
+
+        if "model" in _cfg:
+            new_config[agent_name]["model"] = _cfg["model"]
+        if "prompt" in _cfg:
+            new_config[agent_name]["prompt"] = _cfg["prompt"]
+        if "default" in _cfg:
+            new_config[agent_name]["default"] = _cfg["default"]
+        if "tools" in _cfg:
+            new_config[agent_name]["tools"] = _cfg["tools"]
+
+        handoffs = _cfg.get("handoffs", [])
+        if handoffs == "all":
+            handoff_tools = [
+                handoff_tool
+                for handoff_tool in handoffs_tools
+                if handoff_tool.name != f"transfer_to_{agent_name}"
+            ]
+        else:
+            if not isinstance(handoffs, list):
+                raise ValueError(
+                    f"handoffs for agent {agent_name} must be a list of agent names or 'all'"
+                )
+
+            handoff_tools = [
+                handoff_tool
+                for handoff_tool in handoffs_tools
+                if handoff_tool.name
+                in [
+                    f"transfer_to_{_handoff_agent_name}"
+                    for _handoff_agent_name in handoffs
+                ]
+            ]
+
+        new_config[agent_name]["tools"] = [
+            *new_config[agent_name].get("tools", []),
+            *handoff_tools,
+        ]
+    return new_config
+
+
+class HandoffAgentMiddleware(AgentMiddleware):
     """Agent middleware for switching between multiple agents.
     This middleware dynamically replaces model call parameters based on the currently active agent configuration, enabling seamless switching between different agents.
 
     Args:
         agents_config (dict[str, AgentConfig]): A dictionary of agent configurations.
+        custom_handoffs_tool_descriptions (Optional[dict[str, str]]): A dictionary of custom tool descriptions for handoffs tools. Defaults to None.
 
     Examples:
         ```python
-        from langchain_dev_utils.agents.middleware import HandoffsAgentMiddleware
-        middleware = HandoffsAgentMiddleware(agents_config)
+        from langchain_dev_utils.agents.middleware import HandoffAgentMiddleware
+        middleware = HandoffAgentMiddleware(agents_config)
         ```
     """
 
     state_schema = MultiAgentState
 
-    def __init__(self, agents_config: dict[str, AgentConfig]):
+    def __init__(
+        self,
+        agents_config: dict[str, AgentConfig],
+        custom_handoffs_tool_descriptions: Optional[dict[str, str]] = None,
+    ) -> None:
         default_agent_name = _get_default_active_agent(agents_config)
         if default_agent_name is None:
             raise ValueError(
                 "No default agent found, you must set one by set default=True"
             )
+
+        if custom_handoffs_tool_descriptions is None:
+            custom_handoffs_tool_descriptions = {}
+
+        handoffs_tools = [
+            _create_handoffs_tool(
+                agent_name,
+                custom_handoffs_tool_descriptions.get(agent_name),
+            )
+            for agent_name in agents_config.keys()
+        ]
         self.default_agent_name = default_agent_name
-        self.agents_config = agents_config
+        self.agents_config = _transform_agent_config(
+            agents_config,
+            handoffs_tools,
+        )
+        self.tools = handoffs_tools
 
     def _get_active_agent_config(self, request: ModelRequest) -> dict[str, Any]:
         active_agent_name = request.state.get("active_agent", self.default_agent_name)

{langchain_dev_utils-1.3.1 → langchain_dev_utils-1.3.3}/src/langchain_dev_utils/agents/middleware/plan.py

@@ -104,9 +104,14 @@ class PlanState(AgentState):
     plan: NotRequired[list[Plan]]
 
 
-def create_write_plan_tool(
+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,
-    message_key: Optional[str] = None,
 ) -> BaseTool:
     """Create a tool for writing initial plan.
 
@@ -116,22 +121,15 @@ def create_write_plan_tool(
 
     Args:
         description: The description of the tool. Uses default description if not provided.
-        message_key: The key of the message to be updated. Defaults to "messages".
 
     Returns:
         BaseTool: The tool for writing initial plan.
-
-    Example:
-        Basic usage:
-        >>> from langchain_dev_utils.agents.middleware import create_write_plan_tool
-        >>> write_plan_tool = create_write_plan_tool()
     """
 
     @tool(
         description=description or _DEFAULT_WRITE_PLAN_TOOL_DESCRIPTION,
     )
     def write_plan(plan: list[str], runtime: ToolRuntime):
-        msg_key = message_key or "messages"
         return Command(
             update={
                 "plan": [
@@ -141,7 +139,7 @@ def create_write_plan_tool(
                     }
                     for index, content in enumerate(plan)
                 ],
-                msg_key: [
+                "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,
@@ -153,9 +151,8 @@ def create_write_plan_tool(
     return write_plan
 
 
-def create_finish_sub_plan_tool(
+def _create_finish_sub_plan_tool(
     description: Optional[str] = None,
-    message_key: Optional[str] = None,
 ) -> BaseTool:
     """Create a tool for finishing sub-plan tasks.
 
@@ -164,15 +161,9 @@ def create_finish_sub_plan_tool(
 
     Args:
         description: The description of the tool. Uses default description if not provided.
-        message_key: The key of the message to be updated. Defaults to "messages".
 
     Returns:
         BaseTool: The tool for finishing sub-plan tasks.
-
-    Example:
-        Basic usage:
-        >>> from langchain_dev_utils.agents.middleware import create_finish_sub_plan_tool
-        >>> finish_sub_plan_tool = create_finish_sub_plan_tool()
     """
 
     @tool(
@@ -181,7 +172,6 @@ def create_finish_sub_plan_tool(
     def finish_sub_plan(
         runtime: ToolRuntime,
     ):
-        msg_key = message_key or "messages"
         plan_list = runtime.state.get("plan", [])
 
         sub_finish_plan = ""
@@ -200,7 +190,7 @@ def create_finish_sub_plan_tool(
         return Command(
             update={
                 "plan": plan_list,
-                msg_key: [
+                "messages": [
                     ToolMessage(
                         content=sub_finish_plan + sub_next_plan,
                         tool_call_id=runtime.tool_call_id,
@@ -212,7 +202,7 @@ def create_finish_sub_plan_tool(
     return finish_sub_plan
 
 
-def create_read_plan_tool(
+def _create_read_plan_tool(
     description: Optional[str] = None,
 ):
     """Create a tool for reading all sub-plans.
@@ -225,11 +215,6 @@ def create_read_plan_tool(
 
     Returns:
         BaseTool: The tool for reading all sub-plans.
-
-    Example:
-        Basic usage:
-        >>> from langchain_dev_utils.agents.middleware import create_read_plan_tool
-        >>> read_plan_tool = create_read_plan_tool()
     """
 
     @tool(
@@ -287,12 +272,8 @@ class PlanMiddleware(AgentMiddleware):
             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.
-            If not provided, uses the default `_DEFAULT_FINISH_SUB_PLAN_TOOL_DESCRIPTION`.
-        read_plan_tool_description: Description of the `read_plan` tool.
-            If not provided, uses the default `_DEFAULT_READ_PLAN_TOOL_DESCRIPTION`.
+        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`.
 
@@ -316,31 +297,34 @@ class PlanMiddleware(AgentMiddleware):
         self,
         *,
         system_prompt: Optional[str] = None,
-        write_plan_tool_description: Optional[str] = None,
-        finish_sub_plan_tool_description: Optional[str] = None,
-        read_plan_tool_description: Optional[str] = None,
+        custom_plan_tool_descriptions: Optional[PlanToolDescription] = None,
         use_read_plan_tool: bool = True,
     ) -> None:
         super().__init__()
 
-        write_plan_tool_description = (
-            write_plan_tool_description or _DEFAULT_WRITE_PLAN_TOOL_DESCRIPTION
+        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 = (
-            finish_sub_plan_tool_description
-            or _DEFAULT_FINISH_SUB_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 = (
-            read_plan_tool_description or _DEFAULT_READ_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),
+            _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))
+            tools.append(_create_read_plan_tool(description=read_plan_tool_description))
 
         if system_prompt is None:
             if use_read_plan_tool:

langchain_dev_utils-1.3.3/src/langchain_dev_utils/agents/wrap.py

@@ -0,0 +1,274 @@
+import asyncio
+from typing import Any, Awaitable, Callable, Optional, cast
+
+from langchain.tools import ToolRuntime
+from langchain_core.messages import AnyMessage, HumanMessage
+from langchain_core.tools import BaseTool, StructuredTool
+from langgraph.graph.state import CompiledStateGraph
+
+from langchain_dev_utils.message_convert import format_sequence
+
+
+def _process_input(request: str, runtime: ToolRuntime) -> str:
+    return request
+
+
+def _process_output(
+    request: str, response: list[AnyMessage], runtime: ToolRuntime
+) -> Any:
+    return response[-1].content
+
+
+def wrap_agent_as_tool(
+    agent: CompiledStateGraph,
+    tool_name: Optional[str] = None,
+    tool_description: Optional[str] = None,
+    pre_input_hooks: Optional[
+        tuple[
+            Callable[[str, ToolRuntime], str],
+            Callable[[str, ToolRuntime], Awaitable[str]],
+        ]
+        | Callable[[str, ToolRuntime], str]
+    ] = None,
+    post_output_hooks: Optional[
+        tuple[
+            Callable[[str, list[AnyMessage], ToolRuntime], Any],
+            Callable[[str, list[AnyMessage], ToolRuntime], Awaitable[Any]],
+        ]
+        | Callable[[str, list[AnyMessage], ToolRuntime], Any]
+    ] = None,
+) -> BaseTool:
+    """Wraps an agent as a tool
+
+    Args:
+        agent: The agent to wrap
+        tool_name: The name of the tool
+        tool_description: The description of the tool
+        pre_input_hooks: Hooks to run before the input is processed
+        post_output_hooks: Hooks to run after the output is processed
+
+    Returns:
+        BaseTool: The wrapped agent as a tool
+
+    Example:
+        >>> from langchain_dev_utils.agents import wrap_agent_as_tool, create_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"
+        ... )
+        >>>
+        >>> agent = create_agent("vllm:qwen3-4b", tools=[call_time_agent_tool], name="agent")
+
+        >>> response = agent.invoke({"messages": [HumanMessage(content="What time is it now?")]})
+        >>> response
+    """
+    if agent.name is None:
+        raise ValueError("Agent name must not be None")
+
+    process_input = _process_input
+    process_input_async = _process_input
+    process_output = _process_output
+    process_output_async = _process_output
+
+    if pre_input_hooks:
+        if isinstance(pre_input_hooks, tuple):
+            process_input = pre_input_hooks[0]
+            process_input_async = pre_input_hooks[1]
+        else:
+            process_input = pre_input_hooks
+            process_input_async = pre_input_hooks
+
+    if post_output_hooks:
+        if isinstance(post_output_hooks, tuple):
+            process_output = post_output_hooks[0]
+            process_output_async = post_output_hooks[1]
+        else:
+            process_output = post_output_hooks
+            process_output_async = post_output_hooks
+
+    def call_agent(
+        request: str,
+        runtime: ToolRuntime,
+    ) -> str:
+        request = process_input(request, runtime) if process_input else request
+
+        messages = [HumanMessage(content=request)]
+        response = agent.invoke({"messages": messages})
+
+        response = process_output(request, response["messages"], runtime)
+        return response
+
+    async def acall_agent(
+        request: str,
+        runtime: ToolRuntime,
+    ) -> str:
+        if asyncio.iscoroutinefunction(process_input_async):
+            request = await process_input_async(request, runtime)
+        else:
+            request = cast(str, process_input_async(request, runtime))
+
+        messages = [HumanMessage(content=request)]
+        response = await agent.ainvoke({"messages": messages})
+
+        if asyncio.iscoroutinefunction(process_output_async):
+            response = await process_output_async(
+                request, response["messages"], runtime
+            )
+        else:
+            response = process_output(request, response["messages"], runtime)
+
+        return response
+
+    if tool_name is None:
+        tool_name = f"transfor_to_{agent.name}"
+        if not tool_name.endswith("_agent"):
+            tool_name += "_agent"
+
+    if tool_description is None:
+        tool_description = f"This tool transforms input to {agent.name}"
+
+    return StructuredTool.from_function(
+        func=call_agent,
+        coroutine=acall_agent,
+        name=tool_name,
+        description=tool_description,
+    )
+
+
+def wrap_all_agents_as_tool(
+    agents: list[CompiledStateGraph],
+    tool_name: Optional[str] = None,
+    tool_description: Optional[str] = None,
+    pre_input_hooks: Optional[
+        tuple[
+            Callable[[str, ToolRuntime], str],
+            Callable[[str, ToolRuntime], Awaitable[str]],
+        ]
+        | Callable[[str, ToolRuntime], str]
+    ] = None,
+    post_output_hooks: Optional[
+        tuple[
+            Callable[[str, list[AnyMessage], ToolRuntime], Any],
+            Callable[[str, list[AnyMessage], ToolRuntime], Awaitable[Any]],
+        ]
+        | Callable[[str, list[AnyMessage], ToolRuntime], Any]
+    ] = None,
+) -> BaseTool:
+    """Wraps all agents as single tool
+
+    Args:
+        agents: The agents to wrap
+        tool_name: The name of the tool, default to "task"
+        tool_description: The description of the tool
+        pre_input_hooks: Hooks to run before the input is processed
+        post_output_hooks: Hooks to run after the output is processed
+
+    Returns:
+        BaseTool: The wrapped agents as single tool
+
+    Example:
+        >>> from langchain_dev_utils.agents import wrap_all_agents_as_tool, create_agent
+        >>>
+        >>> call_time_agent_tool = wrap_all_agents_as_tool(
+        ...     [time_agent,weather_agent],
+        ...     tool_name="call_sub_agents",
+        ...     tool_description="Used to invoke the sub-agents to perform tasks"
+        ... )
+        >>>
+        >>> agent = create_agent("vllm:qwen3-4b", tools=[call_sub_agents_tool], name="agent")
+
+        >>> response = agent.invoke({"messages": [HumanMessage(content="What time is it now?")]})
+        >>> response
+    """
+    if len(agents) <= 1:
+        raise ValueError("At least more than one agent must be provided")
+
+    agents_map = {}
+
+    for agent in agents:
+        if agent.name is None:
+            raise ValueError("Agent name must not be provided")
+        if agent.name in agents_map:
+            raise ValueError("Agent name must be unique")
+        agents_map[agent.name] = agent
+
+    process_input = _process_input
+    process_input_async = _process_input
+    process_output = _process_output
+    process_output_async = _process_output
+
+    if pre_input_hooks:
+        if isinstance(pre_input_hooks, tuple):
+            process_input = pre_input_hooks[0]
+            process_input_async = pre_input_hooks[1]
+        else:
+            process_input = pre_input_hooks
+            process_input_async = pre_input_hooks
+
+    if post_output_hooks:
+        if isinstance(post_output_hooks, tuple):
+            process_output = post_output_hooks[0]
+            process_output_async = post_output_hooks[1]
+        else:
+            process_output = post_output_hooks
+            process_output_async = post_output_hooks
+
+    def call_agent(
+        agent_name: str,
+        description: str,
+        runtime: ToolRuntime,
+    ) -> str:
+        task_description = (
+            process_input(description, runtime) if process_input else description
+        )
+
+        if agent_name not in agents_map:
+            raise ValueError(f"Agent {agent_name} not found")
+
+        messages = [HumanMessage(content=task_description)]
+        response = agents_map[agent_name].invoke({"messages": messages})
+
+        response = process_output(task_description, response["messages"], runtime)
+        return response
+
+    async def acall_agent(
+        agent_name: str,
+        description: str,
+        runtime: ToolRuntime,
+    ) -> str:
+        if asyncio.iscoroutinefunction(process_input_async):
+            task_description = await process_input_async(description, runtime)
+        else:
+            task_description = cast(str, process_input_async(description, runtime))
+
+        if agent_name not in agents_map:
+            raise ValueError(f"Agent {agent_name} not found")
+
+        messages = [HumanMessage(content=task_description)]
+        response = await agents_map[agent_name].ainvoke({"messages": messages})
+
+        if asyncio.iscoroutinefunction(process_output_async):
+            response = await process_output_async(
+                task_description, response["messages"], runtime
+            )
+        else:
+            response = process_output(task_description, response["messages"], runtime)
+
+        return response
+
+    if tool_name is None:
+        tool_name = "task"
+
+    if tool_description is None:
+        tool_description = (
+            "Launch an ephemeral subagent for a task.\nAvailable agents:\n "
+            + format_sequence(list(agents_map.keys()), with_num=True)
+        )
+    return StructuredTool.from_function(
+        func=call_agent,
+        coroutine=acall_agent,
+        name=tool_name,
+        description=tool_description,
+    )