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

langchain_dev_utils/agents/middleware/handoffs.py

@@ -0,0 +1,214 @@
+from typing import Any, Awaitable, Callable, Literal, cast
+
+from langchain.agents import AgentState
+from langchain.agents.middleware import AgentMiddleware, ModelRequest, ModelResponse
+from langchain.agents.middleware.types import ModelCallResult
+from langchain.tools import BaseTool, ToolRuntime, tool
+from langchain_core.language_models.chat_models import BaseChatModel
+from langchain_core.messages import SystemMessage, ToolMessage
+from langgraph.types import Command
+from typing_extensions import NotRequired, Optional, TypedDict
+
+from langchain_dev_utils.chat_models import load_chat_model
+
+
+class MultiAgentState(AgentState):
+    active_agent: NotRequired[str]
+
+
+class AgentConfig(TypedDict):
+    model: NotRequired[str | BaseChatModel]
+    prompt: str | SystemMessage
+    tools: NotRequired[list[BaseTool | dict[str, Any]]]
+    default: NotRequired[bool]
+    handoffs: list[str] | Literal["all"]
+
+
+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.
+
+    Returns:
+        BaseTool: A tool instance for handoffs to the specified 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}"
+
+    @tool(name_or_callable=tool_name, description=tool_description)
+    def handoffs_tool(runtime: ToolRuntime) -> Command:
+        return Command(
+            update={
+                "messages": [
+                    ToolMessage(
+                        content=f"Transferred to {agent_name}",
+                        tool_call_id=runtime.tool_call_id,
+                    )
+                ],
+                "active_agent": agent_name,
+            }
+        )
+
+    return handoffs_tool
+
+
+def _get_default_active_agent(state: dict[str, AgentConfig]) -> Optional[str]:
+    for agent_name, config in state.items():
+        if config.get("default", False):
+            return agent_name
+    return None
+
+
+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.
+        handoffs_tool_overrides (Optional[dict[str, BaseTool]]): A dictionary of handoffs tools to override. Defaults to None.
+
+    Examples:
+        ```python
+        from langchain_dev_utils.agents.middleware import HandoffAgentMiddleware
+        middleware = HandoffAgentMiddleware(agents_config)
+        ```
+    """
+
+    state_schema = MultiAgentState
+
+    def __init__(
+        self,
+        agents_config: dict[str, AgentConfig],
+        custom_handoffs_tool_descriptions: Optional[dict[str, str]] = None,
+        handoffs_tool_overrides: Optional[dict[str, BaseTool]] = 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 = {}
+
+        if handoffs_tool_overrides is None:
+            handoffs_tool_overrides = {}
+
+        handoffs_tools = []
+        for agent_name in agents_config.keys():
+            if not handoffs_tool_overrides.get(agent_name):
+                handoffs_tools.append(
+                    _create_handoffs_tool(
+                        agent_name,
+                        custom_handoffs_tool_descriptions.get(agent_name),
+                    )
+                )
+            else:
+                handoffs_tools.append(
+                    cast(BaseTool, handoffs_tool_overrides.get(agent_name))
+                )
+
+        self.default_agent_name = default_agent_name
+        self.agents_config = _transform_agent_config(
+            agents_config,
+            handoffs_tools,
+        )
+        self.tools = handoffs_tools
+
+    def _get_override_request(self, request: ModelRequest) -> ModelRequest:
+        active_agent_name = request.state.get("active_agent", self.default_agent_name)
+
+        _config = self.agents_config[active_agent_name]
+
+        params = {}
+        if _config.get("model"):
+            model = _config.get("model")
+            if isinstance(model, str):
+                model = load_chat_model(model)
+            params["model"] = model
+        if _config.get("prompt"):
+            params["system_prompt"] = _config.get("prompt")
+        if _config.get("tools"):
+            params["tools"] = _config.get("tools")
+
+        if params:
+            return request.override(**params)
+        else:
+            return request
+
+    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:
+        override_request = self._get_override_request(request)
+        return await handler(override_request)

langchain_dev_utils/agents/middleware/model_fallback.py

@@ -0,0 +1,49 @@
+from langchain.agents.middleware.model_fallback import (
+    ModelFallbackMiddleware as _ModelFallbackMiddleware,
+)
+
+from langchain_dev_utils.chat_models.base import load_chat_model
+
+
+class ModelFallbackMiddleware(_ModelFallbackMiddleware):
+    """Automatic fallback to alternative models on errors.
+
+    Retries failed model calls with alternative models in sequence until
+    success or all models exhausted. Primary model specified in create_agent().
+
+    Args:
+        first_model: The first model to try on error. Must be a string identifier.
+        additional_models: Additional models to try in sequence on error.
+
+    Example:
+        ```python
+        from langchain_dev_utils.agents.middleware import ModelFallbackMiddleware
+        from langchain_dev_utils.agents import create_agent
+
+        fallback = ModelFallbackMiddleware(
+            "vllm:qwen3-8b", ## Try first on error
+            "vllm:gpt-oss-20b", #Then this
+        )
+
+        agent = create_agent(
+            model="vllm:qwen3-4b", #Primary model
+            middleware=[fallback],
+        )
+
+        # If primary fails: tries qwen3-8b, then gpt-oss-20b
+        result = await agent.invoke({"messages": [HumanMessage("Hello")]})
+        ```
+    """
+
+    def __init__(
+        self,
+        first_model: str,
+        *additional_models: str,
+    ) -> None:
+        first_chat_model = load_chat_model(first_model)
+
+        additional_chat_models = [load_chat_model(model) for model in additional_models]
+        super().__init__(
+            first_chat_model,
+            *additional_chat_models,
+        )

langchain_dev_utils/agents/middleware/model_router.py

@@ -0,0 +1,200 @@
+from typing import Annotated, Any, Awaitable, Callable, NotRequired, Optional, cast
+
+from langchain.agents import AgentState
+from langchain.agents.middleware import AgentMiddleware, ModelRequest, ModelResponse
+from langchain.agents.middleware.types import ModelCallResult, OmitFromInput
+from langchain_core.language_models import BaseChatModel
+from langchain_core.messages import AnyMessage, SystemMessage
+from langchain_core.tools import BaseTool
+from langgraph.runtime import Runtime
+from pydantic import BaseModel, Field
+from typing_extensions import TypedDict
+
+from langchain_dev_utils.chat_models import load_chat_model
+from langchain_dev_utils.message_convert import format_sequence
+
+
+class ModelDict(TypedDict):
+    model_name: str
+    model_description: str
+    tools: NotRequired[list[BaseTool | dict[str, Any]]]
+    model_kwargs: NotRequired[dict[str, Any]]
+    model_instance: NotRequired[BaseChatModel]
+    model_system_prompt: NotRequired[str]
+
+
+class SelectModel(BaseModel):
+    """Tool for model selection - Must call this tool to return the finally selected model"""
+
+    model_name: str = Field(
+        ...,
+        description="Selected model name (must be the full model name, for example, openai:gpt-4o)",
+    )
+
+
+_ROUTER_MODEL_PROMPT = """
+# Role Description
+You are an intelligent routing model, specializing in analyzing task requirements and matching appropriate AI models.
+
+# Available Model List
+{model_card}
+
+# Core Responsibilities
+1. **Task Analysis**: Deeply understand the type, complexity, and special needs of the user's task
+2. **Model Matching**: Select the most appropriate model based on the task characteristics
+3. **Tool Call**: **Must call SelectModel tool** to return the final selection
+
+# ⚠️ Important Instructions
+**After completing the analysis, you must immediately call the SelectModel tool to return the model selection result.**
+**This is the only way to output, and you are forbidden to return the result in any other form.**
+
+# Selection Standards
+- Consider the type of the task (dialogue, inference, creation, analysis, etc.)
+- Evaluate the complexity of the task
+- Match the professional ability and applicable scenario of the model
+- Ensure that the model's ability matches the task requirements to a high degree
+
+# Execution Process
+1. Analyze user task requirements
+2. Compare the capabilities of available models
+3. Call **SelectModel tool** to submit selection
+4. Task completion
+
+Strictly adhere to tool call requirements!
+"""
+
+
+class ModelRouterState(AgentState):
+    router_model_selection: Annotated[str, OmitFromInput]
+
+
+class ModelRouterMiddleware(AgentMiddleware):
+    """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_instance(Optional), model_system_prompt(Optional)
+        router_prompt: Routing prompt template, uses default template if None
+
+    Examples:
+        ```python
+        from langchain_dev_utils.agents.middleware import ModelRouterMiddleware
+        middleware = ModelRouterMiddleware(
+            router_model="vllm:qwen3-4b",
+            model_list=model_list
+        )
+        ```
+    """
+
+    state_schema = ModelRouterState
+
+    def __init__(
+        self,
+        router_model: str | BaseChatModel,
+        model_list: list[ModelDict],
+        router_prompt: Optional[str] = None,
+    ) -> None:
+        super().__init__()
+        if isinstance(router_model, BaseChatModel):
+            self.router_model = router_model.with_structured_output(SelectModel)
+        else:
+            self.router_model = load_chat_model(router_model).with_structured_output(
+                SelectModel
+            )
+        self.model_list = model_list
+
+        if router_prompt is None:
+            router_prompt = _ROUTER_MODEL_PROMPT.format(
+                model_card=format_sequence(
+                    [
+                        f"model_name:\n {model['model_name']}\n model_description:\n {model['model_description']}"
+                        for model in model_list
+                    ],
+                    with_num=True,
+                )
+            )
+
+        self.router_prompt = router_prompt
+
+    def _select_model(self, messages: list[AnyMessage]):
+        response = cast(
+            SelectModel,
+            self.router_model.invoke(
+                [SystemMessage(content=self.router_prompt), *messages]
+            ),
+        )
+        return response.model_name if response is not None else "default-model"
+
+    async def _aselect_model(self, messages: list[AnyMessage]):
+        response = cast(
+            SelectModel,
+            await self.router_model.ainvoke(
+                [SystemMessage(content=self.router_prompt), *messages]
+            ),
+        )
+        return response.model_name if response is not None else "default-model"
+
+    def before_agent(
+        self, state: ModelRouterState, runtime: Runtime
+    ) -> dict[str, Any] | None:
+        model_name = self._select_model(state["messages"])
+        return {"router_model_selection": model_name}
+
+    async def abefore_agent(
+        self, state: ModelRouterState, runtime: Runtime
+    ) -> dict[str, Any] | None:
+        model_name = await self._aselect_model(state["messages"])
+        return {"router_model_selection": model_name}
+
+    def _get_override_request(self, request: ModelRequest) -> ModelRequest:
+        model_dict = {
+            item["model_name"]: {
+                "tools": item.get("tools", None),
+                "kwargs": item.get("model_kwargs", None),
+                "system_prompt": item.get("model_system_prompt", None),
+                "model_instance": item.get("model_instance", 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" and select_model_name in model_dict:
+            model_values = model_dict.get(select_model_name, {})
+            if model_values["model_instance"] is not None:
+                model = model_values["model_instance"]
+            else:
+                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"]
+                )
+
+        if override_kwargs:
+            return request.override(**override_kwargs)
+        else:
+            return request
+
+    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:
+        override_request = self._get_override_request(request)
+        return await handler(override_request)