dao-ai 0.0.28__py3-none-any.whl → 0.1.5__py3-none-any.whl

dao_ai/middleware/base.py

@@ -0,0 +1,50 @@
+"""
+Base classes and types for DAO AI middleware.
+
+This module re-exports LangChain's middleware types for convenience.
+Use LangChainAgentMiddleware directly with DAO AI's state and context types.
+
+Example:
+    from langchain.agents.middleware import AgentMiddleware
+    from dao_ai.state import AgentState, Context
+    from langgraph.runtime import Runtime
+
+    class MyMiddleware(AgentMiddleware[AgentState, Context]):
+        def before_model(
+            self,
+            state: AgentState,
+            runtime: Runtime[Context]
+        ) -> dict[str, Any] | None:
+            print(f"About to call model with {len(state['messages'])} messages")
+            return None
+"""
+
+from langchain.agents.middleware import (
+    AgentMiddleware,
+    ModelRequest,
+    after_agent,
+    after_model,
+    before_agent,
+    before_model,
+    dynamic_prompt,
+    wrap_model_call,
+    wrap_tool_call,
+)
+from langchain.agents.middleware.types import ModelResponse
+
+# Re-export LangChain types for convenience
+__all__ = [
+    # Base middleware class
+    "AgentMiddleware",
+    # Types
+    "ModelRequest",
+    "ModelResponse",
+    # Decorators
+    "before_agent",
+    "before_model",
+    "after_agent",
+    "after_model",
+    "wrap_model_call",
+    "wrap_tool_call",
+    "dynamic_prompt",
+]

dao_ai/middleware/context_editing.py

@@ -0,0 +1,230 @@
+"""
+Context editing middleware for DAO AI agents.
+
+Manages conversation context by clearing older tool call outputs when token limits
+are reached, while preserving recent results.
+
+Example:
+    from dao_ai.middleware import create_context_editing_middleware
+
+    # Clear old tool outputs when context exceeds 100k tokens
+    middleware = create_context_editing_middleware(
+        trigger=100000,
+        keep=3,
+    )
+"""
+
+from __future__ import annotations
+
+from typing import Any, Literal
+
+from langchain.agents.middleware import ClearToolUsesEdit, ContextEditingMiddleware
+from langchain_core.tools import BaseTool
+from loguru import logger
+
+from dao_ai.config import BaseFunctionModel, ToolModel
+
+__all__ = [
+    "ContextEditingMiddleware",
+    "ClearToolUsesEdit",
+    "create_context_editing_middleware",
+    "create_clear_tool_uses_edit",
+]
+
+
+def _resolve_tool_names(
+    tools: list[str | ToolModel | dict[str, Any]] | None,
+) -> list[str]:
+    """Resolve tool specs to a list of tool name strings."""
+    if tools is None:
+        return []
+
+    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
+
+
+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_clear_tool_uses_edit(
+    trigger: int = 100000,
+    keep: int = 3,
+    clear_at_least: int = 0,
+    clear_tool_inputs: bool = False,
+    exclude_tools: list[str | ToolModel | dict[str, Any]] | None = None,
+    placeholder: str = "[cleared]",
+) -> ClearToolUsesEdit:
+    """
+    Create a ClearToolUsesEdit for use with ContextEditingMiddleware.
+
+    This edit strategy clears older tool results when the conversation exceeds
+    a token threshold, while preserving recent results.
+
+    Args:
+        trigger: Token count that triggers the edit. When conversation exceeds
+            this, older tool outputs are cleared. Default 100000.
+        keep: Number of most recent tool results to preserve. These are never
+            cleared. Default 3.
+        clear_at_least: Minimum tokens to reclaim when edit runs.
+            0 means clear as much as needed. Default 0.
+        clear_tool_inputs: Whether to clear tool call arguments on AI messages.
+            When True, tool call arguments are replaced with empty objects.
+            Default False.
+        exclude_tools: Tools to never clear. Can be:
+            - list of str: Tool names
+            - list of ToolModel: DAO AI tool models
+            - list of dict: Tool config dicts
+            Default None (no exclusions).
+        placeholder: Text inserted for cleared tool outputs.
+            Default "[cleared]".
+
+    Returns:
+        ClearToolUsesEdit instance
+
+    Example:
+        edit = create_clear_tool_uses_edit(
+            trigger=50000,
+            keep=5,
+            clear_tool_inputs=True,
+            exclude_tools=["important_tool"],
+        )
+    """
+    excluded = _resolve_tool_names(exclude_tools) if exclude_tools else []
+
+    logger.debug(
+        "Creating ClearToolUsesEdit",
+        trigger=trigger,
+        keep=keep,
+        clear_at_least=clear_at_least,
+        clear_tool_inputs=clear_tool_inputs,
+        exclude_tools=excluded or "none",
+        placeholder=placeholder,
+    )
+
+    return ClearToolUsesEdit(
+        trigger=trigger,
+        keep=keep,
+        clear_at_least=clear_at_least,
+        clear_tool_inputs=clear_tool_inputs,
+        exclude_tools=excluded,
+        placeholder=placeholder,
+    )
+
+
+def create_context_editing_middleware(
+    trigger: int = 100000,
+    keep: int = 3,
+    clear_at_least: int = 0,
+    clear_tool_inputs: bool = False,
+    exclude_tools: list[str | ToolModel | dict[str, Any]] | None = None,
+    placeholder: str = "[cleared]",
+    token_count_method: Literal["approximate", "model"] = "approximate",
+) -> ContextEditingMiddleware:
+    """
+    Create a ContextEditingMiddleware with ClearToolUsesEdit.
+
+    Manages conversation context by clearing older tool call outputs when token
+    limits are reached. Useful for long conversations with many tool calls that
+    exceed context window limits.
+
+    Use cases:
+    - Long conversations with many tool calls exceeding token limits
+    - Reducing token costs by removing older irrelevant tool outputs
+    - Maintaining only the most recent N tool results in context
+
+    Args:
+        trigger: Token count that triggers clearing. When conversation exceeds
+            this threshold, older tool outputs are cleared. Default 100000.
+        keep: Number of most recent tool results to always preserve.
+            These are never cleared. Default 3.
+        clear_at_least: Minimum tokens to reclaim when edit runs.
+            0 means clear as much as needed. Default 0.
+        clear_tool_inputs: Whether to also clear tool call arguments on AI
+            messages. When True, replaces arguments with empty objects.
+            Default False (preserves tool call context).
+        exclude_tools: Tools to never clear outputs from. Can be:
+            - list of str: Tool names
+            - list of ToolModel: DAO AI tool models
+            - list of dict: Tool config dicts
+            Default None (no exclusions).
+        placeholder: Text inserted for cleared tool outputs.
+            Default "[cleared]".
+        token_count_method: How to count tokens:
+            - "approximate": Fast estimation (default)
+            - "model": Accurate count using model tokenizer
+
+    Returns:
+        List containing ContextEditingMiddleware instance
+
+    Example:
+        # Basic usage - clear old tool outputs after 100k tokens
+        middleware = create_context_editing_middleware(
+            trigger=100000,
+            keep=3,
+        )
+
+        # Aggressive clearing with exclusions
+        middleware = create_context_editing_middleware(
+            trigger=50000,
+            keep=5,
+            clear_tool_inputs=True,
+            exclude_tools=["important_tool", "critical_search"],
+            placeholder="[output cleared to save context]",
+        )
+
+        # Accurate token counting
+        middleware = create_context_editing_middleware(
+            trigger=100000,
+            keep=3,
+            token_count_method="model",
+        )
+    """
+    edit = create_clear_tool_uses_edit(
+        trigger=trigger,
+        keep=keep,
+        clear_at_least=clear_at_least,
+        clear_tool_inputs=clear_tool_inputs,
+        exclude_tools=exclude_tools,
+        placeholder=placeholder,
+    )
+
+    logger.debug(
+        "Creating ContextEditingMiddleware",
+        token_count_method=token_count_method,
+    )
+
+    return ContextEditingMiddleware(
+        edits=[edit],
+        token_count_method=token_count_method,
+    )

dao_ai/middleware/core.py

@@ -0,0 +1,67 @@
+"""
+Core middleware utilities for DAO AI.
+
+This module provides the factory function for creating middleware instances
+from fully qualified function names.
+"""
+
+from typing import Any, Callable
+
+from langchain.agents.middleware import AgentMiddleware
+from loguru import logger
+
+from dao_ai.state import AgentState, Context
+from dao_ai.utils import load_function
+
+
+def create_factory_middleware(
+    function_name: str,
+    args: dict[str, Any] | None = None,
+) -> AgentMiddleware[AgentState, Context]:
+    """
+    Create middleware from a factory function.
+
+    This factory function dynamically loads a Python function and calls it
+    with the provided arguments to create a middleware instance.
+
+    The factory function should return a middleware object compatible with
+    LangChain's create_agent middleware parameter (AgentMiddleware or any
+    callable/object that implements the middleware interface).
+
+    Args:
+        function_name: Fully qualified name of the factory function
+                       (e.g., 'my_module.create_custom_middleware')
+        args: Arguments to pass to the factory function
+
+    Returns:
+        The AgentMiddleware instance returned by the factory function.
+
+    Raises:
+        ImportError: If the function cannot be loaded
+
+    Example:
+        # Factory function in my_module.py:
+        def create_custom_middleware(threshold: float = 0.5) -> AgentMiddleware[AgentState, Context]:
+            return MyCustomMiddleware(threshold=threshold)
+
+        # Usage:
+        middleware = create_factory_middleware(
+            function_name="my_module.create_custom_middleware",
+            args={"threshold": 0.8}
+        )
+    """
+    if args is None:
+        args = {}
+
+    logger.trace("Creating factory middleware", function_name=function_name, args=args)
+
+    factory: Callable[..., AgentMiddleware[AgentState, Context]] = load_function(
+        function_name=function_name
+    )
+    middleware = factory(**args)
+
+    logger.trace(
+        "Created middleware from factory",
+        middleware_type=type(middleware).__name__,
+    )
+    return middleware