dao-ai 0.1.2__py3-none-any.whl → 0.1.20__py3-none-any.whl

dao_ai/middleware/summarization.py

@@ -150,7 +150,7 @@ def create_summarization_middleware(
         chat_history: ChatHistoryModel configuration for summarization
 
     Returns:
-        LoggingSummarizationMiddleware configured with the specified parameters
+        List containing LoggingSummarizationMiddleware configured with the specified parameters
 
     Example:
         from dao_ai.config import ChatHistoryModel, LLMModel

dao_ai/middleware/tool_call_limit.py

@@ -0,0 +1,210 @@
+"""
+Tool call limit middleware for DAO AI agents.
+
+This module provides a factory for creating LangChain's ToolCallLimitMiddleware
+from DAO AI configuration.
+
+Example:
+    from dao_ai.middleware import create_tool_call_limit_middleware
+
+    # Global limit across all tools
+    middleware = create_tool_call_limit_middleware(
+        thread_limit=20,
+        run_limit=10,
+    )
+
+    # Limit specific tool by name
+    search_limiter = create_tool_call_limit_middleware(
+        tool="search_web",
+        run_limit=3,
+        exit_behavior="continue",
+    )
+"""
+
+from __future__ import annotations
+
+from typing import Any, Literal
+
+from langchain.agents.middleware import ToolCallLimitMiddleware
+from langchain_core.tools import BaseTool
+from loguru import logger
+
+from dao_ai.config import BaseFunctionModel, ToolModel
+
+__all__ = [
+    "ToolCallLimitMiddleware",
+    "create_tool_call_limit_middleware",
+]
+
+
+def _resolve_tool(tool: str | ToolModel | dict[str, Any]) -> list[str]:
+    """
+    Resolve tool argument to a list of actual tool names.
+
+    Args:
+        tool: String name, ToolModel, or dict to resolve
+
+    Returns:
+        List of tool name strings
+
+    Raises:
+        ValueError: If dict cannot be converted to ToolModel
+        TypeError: If tool is not a supported type
+    """
+    # String: return as single-item list
+    if isinstance(tool, str):
+        return [tool]
+
+    # Dict: convert to ToolModel first
+    if isinstance(tool, dict):
+        try:
+            tool_model = ToolModel(**tool)
+        except Exception as e:
+            raise ValueError(
+                f"Failed to construct ToolModel from dict: {e}\n"
+                f"Dict must have 'name' and 'function' keys."
+            ) from e
+    elif isinstance(tool, ToolModel):
+        tool_model = tool
+    else:
+        raise TypeError(
+            f"tool must be str, ToolModel, or dict, got {type(tool).__name__}"
+        )
+
+    # Extract tool names from ToolModel
+    return _extract_tool_names(tool_model)
+
+
+def _extract_tool_names(tool_model: ToolModel) -> list[str]:
+    """
+    Extract actual tool names from a ToolModel.
+
+    A single ToolModel can produce multiple tools (e.g., UC functions).
+    Falls back to ToolModel.name if extraction fails.
+    """
+    function = tool_model.function
+
+    # String function references can't be introspected
+    if not isinstance(function, BaseFunctionModel):
+        logger.debug(
+            "Cannot extract names from string function, using ToolModel.name",
+            tool_model_name=tool_model.name,
+        )
+        return [tool_model.name]
+
+    # Try to extract names from created tools
+    try:
+        tool_names = [
+            tool.name
+            for tool in function.as_tools()
+            if isinstance(tool, BaseTool) and tool.name
+        ]
+        if tool_names:
+            logger.trace(
+                "Extracted tool names",
+                tool_model_name=tool_model.name,
+                tool_names=tool_names,
+            )
+            return tool_names
+    except Exception as e:
+        logger.warning(
+            "Error extracting tool names from ToolModel",
+            tool_model_name=tool_model.name,
+            error=str(e),
+        )
+
+    # Fallback to ToolModel.name
+    logger.debug(
+        "Falling back to ToolModel.name",
+        tool_model_name=tool_model.name,
+    )
+    return [tool_model.name]
+
+
+def create_tool_call_limit_middleware(
+    tool: str | ToolModel | dict[str, Any] | None = None,
+    thread_limit: int | None = None,
+    run_limit: int | None = None,
+    exit_behavior: Literal["continue", "error", "end"] = "continue",
+) -> ToolCallLimitMiddleware:
+    """
+    Create a ToolCallLimitMiddleware with graceful termination support.
+
+    Factory for LangChain's ToolCallLimitMiddleware that supports DAO AI
+    configuration types.
+
+    Args:
+        tool: Tool to limit. Can be:
+            - None: Global limit on all tools
+            - str: Limit specific tool by name
+            - ToolModel: Limit tool(s) from DAO AI config
+            - dict: Tool config dict (converted to ToolModel)
+        thread_limit: Max calls per thread (conversation). Requires checkpointer.
+        run_limit: Max calls per run (single invocation).
+        exit_behavior: What to do when limit hit:
+            - "continue": Block tool with error message, let agent continue
+            - "error": Raise ToolCallLimitExceededError immediately
+            - "end": Stop execution gracefully (single-tool only)
+
+    Returns:
+        A ToolCallLimitMiddleware instance. If ToolModel produces multiple tools,
+        only the first tool is used (with a warning logged).
+
+    Raises:
+        ValueError: If no limits specified, or invalid dict
+        TypeError: If tool is unsupported type
+
+    Example:
+        # Global limit
+        limiter = create_tool_call_limit_middleware(run_limit=10)
+
+        # Tool-specific limit
+        limiter = create_tool_call_limit_middleware(
+            tool="search_web",
+            run_limit=3,
+            exit_behavior="continue",
+        )
+    """
+    if thread_limit is None and run_limit is None:
+        raise ValueError("At least one of thread_limit or run_limit must be specified.")
+
+    # Global limit: no tool parameter
+    if tool is None:
+        logger.debug(
+            "Creating global tool call limit",
+            thread_limit=thread_limit,
+            run_limit=run_limit,
+            exit_behavior=exit_behavior,
+        )
+        return ToolCallLimitMiddleware(
+            thread_limit=thread_limit,
+            run_limit=run_limit,
+            exit_behavior=exit_behavior,
+        )
+
+    # Resolve to list of tool names
+    names = _resolve_tool(tool)
+
+    # Use first tool name (warn if multiple)
+    tool_name = names[0]
+    if len(names) > 1:
+        logger.warning(
+            "ToolModel resolved to multiple tool names, using first only",
+            tool_names=names,
+            using=tool_name,
+        )
+
+    logger.debug(
+        "Creating tool call limit middleware",
+        tool_name=tool_name,
+        thread_limit=thread_limit,
+        run_limit=run_limit,
+        exit_behavior=exit_behavior,
+    )
+
+    return ToolCallLimitMiddleware(
+        tool_name=tool_name,
+        thread_limit=thread_limit,
+        run_limit=run_limit,
+        exit_behavior=exit_behavior,
+    )

dao_ai/middleware/tool_retry.py

@@ -0,0 +1,174 @@
+"""
+Tool retry middleware for DAO AI agents.
+
+Automatically retries failed tool calls with configurable exponential backoff.
+
+Example:
+    from dao_ai.middleware import create_tool_retry_middleware
+
+    # Retry failed tool calls with exponential backoff
+    middleware = create_tool_retry_middleware(
+        max_retries=3,
+        backoff_factor=2.0,
+        initial_delay=1.0,
+    )
+"""
+
+from __future__ import annotations
+
+from typing import Any, Callable, Literal
+
+from langchain.agents.middleware import ToolRetryMiddleware
+from langchain_core.tools import BaseTool
+from loguru import logger
+
+from dao_ai.config import BaseFunctionModel, ToolModel
+
+__all__ = [
+    "ToolRetryMiddleware",
+    "create_tool_retry_middleware",
+]
+
+
+def _resolve_tools(
+    tools: list[str | ToolModel | dict[str, Any]] | None,
+) -> list[str] | None:
+    """
+    Resolve tool specs to a list of tool name strings.
+
+    Returns None if tools is None (apply to all tools).
+    """
+    if tools is None:
+        return None
+
+    result: list[str] = []
+    for tool in tools:
+        if isinstance(tool, str):
+            result.append(tool)
+        elif isinstance(tool, dict):
+            try:
+                tool_model = ToolModel(**tool)
+                result.extend(_extract_tool_names(tool_model))
+            except Exception as e:
+                raise ValueError(f"Failed to construct ToolModel from dict: {e}") from e
+        elif isinstance(tool, ToolModel):
+            result.extend(_extract_tool_names(tool))
+        else:
+            raise TypeError(
+                f"Tool must be str, ToolModel, or dict, got {type(tool).__name__}"
+            )
+
+    return result if result else None
+
+
+def _extract_tool_names(tool_model: ToolModel) -> list[str]:
+    """Extract tool names from ToolModel, falling back to ToolModel.name."""
+    function = tool_model.function
+
+    if not isinstance(function, BaseFunctionModel):
+        return [tool_model.name]
+
+    try:
+        tool_names = [
+            tool.name
+            for tool in function.as_tools()
+            if isinstance(tool, BaseTool) and tool.name
+        ]
+        return tool_names if tool_names else [tool_model.name]
+    except Exception:
+        return [tool_model.name]
+
+
+def create_tool_retry_middleware(
+    max_retries: int = 3,
+    backoff_factor: float = 2.0,
+    initial_delay: float = 1.0,
+    max_delay: float | None = None,
+    jitter: bool = False,
+    tools: list[str | ToolModel | dict[str, Any]] | None = None,
+    retry_on: tuple[type[Exception], ...] | Callable[[Exception], bool] | None = None,
+    on_failure: Literal["continue", "error"] | Callable[[Exception], str] = "continue",
+) -> ToolRetryMiddleware:
+    """
+    Create a ToolRetryMiddleware for automatic tool call retries.
+
+    Handles transient failures in external API calls with exponential backoff.
+
+    Args:
+        max_retries: Max retry attempts after initial call. Default 3.
+        backoff_factor: Multiplier for exponential backoff. Default 2.0.
+            Delay = initial_delay * (backoff_factor ** retry_number)
+            Set to 0.0 for constant delay.
+        initial_delay: Initial delay in seconds before first retry. Default 1.0.
+        max_delay: Max delay in seconds (caps exponential growth). None = no cap.
+        jitter: Add ±25% random jitter to avoid thundering herd. Default False.
+        tools: List of tools to apply retry to. Can be:
+            - None: Apply to all tools (default)
+            - list of str: Tool names
+            - list of ToolModel: DAO AI tool models
+            - list of dict: Tool config dicts
+        retry_on: When to retry:
+            - None: Retry on all errors (default)
+            - tuple of Exception types: Retry only on these
+            - callable: Function(exception) -> bool
+        on_failure: Behavior when all retries exhausted:
+            - "continue": Return error message, let agent continue (default)
+            - "error": Re-raise exception, stop execution
+            - callable: Function(exception) -> str for custom message
+
+    Returns:
+        List containing ToolRetryMiddleware instance
+
+    Example:
+        # Basic retry with defaults
+        retry = create_tool_retry_middleware()
+
+        # Retry specific tools with custom backoff
+        retry = create_tool_retry_middleware(
+            max_retries=5,
+            backoff_factor=1.5,
+            initial_delay=0.5,
+            tools=["search_web", "query_database"],
+        )
+
+        # Retry only on specific exceptions
+        retry = create_tool_retry_middleware(
+            max_retries=3,
+            retry_on=(TimeoutError, ConnectionError),
+            on_failure="error",
+        )
+    """
+    tool_names = _resolve_tools(tools)
+
+    logger.debug(
+        "Creating tool retry middleware",
+        max_retries=max_retries,
+        backoff_factor=backoff_factor,
+        initial_delay=initial_delay,
+        max_delay=max_delay,
+        jitter=jitter,
+        tools=tool_names or "all",
+        on_failure=on_failure if isinstance(on_failure, str) else "custom",
+    )
+
+    # Build kwargs
+    kwargs: dict[str, Any] = {
+        "max_retries": max_retries,
+        "backoff_factor": backoff_factor,
+        "initial_delay": initial_delay,
+        "on_failure": on_failure,
+    }
+
+    if tool_names is not None:
+        kwargs["tools"] = tool_names
+
+    if max_delay is not None:
+        kwargs["max_delay"] = max_delay
+
+    if jitter:
+        kwargs["jitter"] = jitter
+
+    if retry_on is not None:
+        kwargs["retry_on"] = retry_on
+
+    return ToolRetryMiddleware(**kwargs)

dao_ai/middleware/tool_selector.py

@@ -0,0 +1,129 @@
+"""
+Tool selector middleware for intelligently filtering tools before LLM calls.
+
+This middleware uses an LLM to select relevant tools from a large set, improving
+performance and accuracy by reducing context size and improving focus.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+from langchain.agents.middleware import LLMToolSelectorMiddleware
+from langchain_core.language_models import LanguageModelLike
+from loguru import logger
+
+from dao_ai.config import ToolModel
+
+
+def create_llm_tool_selector_middleware(
+    model: LanguageModelLike,
+    max_tools: int = 3,
+    always_include: list[str | ToolModel | dict[str, Any]] | None = None,
+) -> LLMToolSelectorMiddleware:
+    """
+    Create an LLMToolSelectorMiddleware for intelligent tool selection.
+
+    Uses an LLM to analyze the current query and select the most relevant tools
+    before calling the main model. This is particularly useful for agents with
+    many tools (10+) where most aren't relevant for any given query.
+
+    Benefits:
+    - Reduces token usage by filtering irrelevant tools
+    - Improves model focus and accuracy
+    - Optimizes cost for agents with large tool sets
+    - Maintains context window efficiency
+
+    Args:
+        model: The LLM to use for tool selection. Typically a smaller, faster
+            model like "gpt-4o-mini" or similar.
+        max_tools: Maximum number of tools to select for each query.
+            Default 3. Adjust based on your use case - higher values
+            increase context but improve tool coverage.
+        always_include: List of tools that should always be included regardless
+            of the LLM's selection. Can be:
+            - str: Tool name
+            - ToolModel: Full tool configuration
+            - dict: Tool configuration dictionary
+            Use this for critical tools that should always be available.
+
+    Returns:
+        LLMToolSelectorMiddleware configured with the specified parameters
+
+    Example:
+        from dao_ai.middleware import create_llm_tool_selector_middleware
+        from dao_ai.llms import create_llm
+
+        # Use a fast, cheap model for tool selection
+        selector_llm = create_llm("databricks-gpt-4o-mini")
+
+        middleware = create_llm_tool_selector_middleware(
+            model=selector_llm,
+            max_tools=3,
+            always_include=["search_web"],  # Always include search
+        )
+
+    Use Cases:
+        - Large tool sets (10+ tools) where most are specialized
+        - Cost optimization by reducing tokens in main model calls
+        - Improved accuracy by reducing tool confusion
+        - Dynamic tool filtering based on query relevance
+
+    Note:
+        The selector model makes an additional LLM call for each agent turn.
+        Choose a fast, inexpensive model to minimize latency and cost overhead.
+    """
+    # Extract tool names from always_include
+    always_include_names: list[str] = []
+    if always_include:
+        always_include_names = _resolve_tool_names(always_include)
+
+    logger.debug(
+        "Creating LLM tool selector middleware",
+        max_tools=max_tools,
+        always_include_count=len(always_include_names),
+        always_include=always_include_names,
+    )
+
+    return LLMToolSelectorMiddleware(
+        model=model,
+        max_tools=max_tools,
+        always_include=always_include_names if always_include_names else None,
+    )
+
+
+def _resolve_tool_names(tools: list[str | ToolModel | dict[str, Any]]) -> list[str]:
+    """
+    Extract tool names from a list of tool specifications.
+
+    Args:
+        tools: List of tool specifications (strings, ToolModels, or dicts)
+
+    Returns:
+        List of tool names as strings
+    """
+    names: list[str] = []
+
+    for tool_spec in tools:
+        if isinstance(tool_spec, str):
+            # Simple string tool name
+            names.append(tool_spec)
+        elif isinstance(tool_spec, ToolModel):
+            # ToolModel - use its name
+            names.append(tool_spec.name)
+        elif isinstance(tool_spec, dict):
+            # Dictionary - try to extract name
+            if "name" in tool_spec:
+                names.append(tool_spec["name"])
+            else:
+                logger.warning(
+                    "Tool dict missing 'name' field, skipping",
+                    tool_spec=tool_spec,
+                )
+        else:
+            logger.warning(
+                "Unknown tool specification type, skipping",
+                tool_spec_type=type(tool_spec).__name__,
+            )
+
+    return names