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

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

@@ -21,7 +21,6 @@ def create_factory_middleware(
     """
     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.
 
@@ -35,7 +34,7 @@ def create_factory_middleware(
         args: Arguments to pass to the factory function
 
     Returns:
-        An AgentMiddleware instance returned by the factory function
+        The AgentMiddleware instance returned by the factory function.
 
     Raises:
         ImportError: If the function cannot be loaded
@@ -59,9 +58,10 @@ def create_factory_middleware(
     factory: Callable[..., AgentMiddleware[AgentState, Context]] = load_function(
         function_name=function_name
     )
-    middleware: AgentMiddleware[AgentState, Context] = factory(**args)
+    middleware = factory(**args)
 
     logger.trace(
-        "Created middleware from factory", middleware_type=type(middleware).__name__
+        "Created middleware from factory",
+        middleware_type=type(middleware).__name__,
     )
     return middleware

dao_ai/middleware/guardrails.py

@@ -342,7 +342,7 @@ def create_guardrail_middleware(
         num_retries: Maximum number of retry attempts (default: 3)
 
     Returns:
-        GuardrailMiddleware configured with the specified parameters
+        List containing GuardrailMiddleware configured with the specified parameters
 
     Example:
         middleware = create_guardrail_middleware(
@@ -376,7 +376,7 @@ def create_content_filter_middleware(
         block_message: Message to return when content is blocked
 
     Returns:
-        ContentFilterMiddleware configured with the specified parameters
+        List containing ContentFilterMiddleware configured with the specified parameters
 
     Example:
         middleware = create_content_filter_middleware(
@@ -407,7 +407,7 @@ def create_safety_guardrail_middleware(
             defaults to gpt-4o-mini.
 
     Returns:
-        SafetyGuardrailMiddleware configured with the specified model
+        List containing SafetyGuardrailMiddleware configured with the specified model
 
     Example:
         from databricks_langchain import ChatDatabricks

dao_ai/middleware/human_in_the_loop.py

@@ -132,7 +132,7 @@ def create_human_in_the_loop_middleware(
         description_prefix: Message prefix shown when pausing for review
 
     Returns:
-        HumanInTheLoopMiddleware configured with the specified parameters
+        List containing HumanInTheLoopMiddleware configured with the specified parameters
 
     Example:
         from dao_ai.config import HumanInTheLoopModel
@@ -182,7 +182,8 @@ def create_hitl_middleware_from_tool_models(
         description_prefix: Message prefix shown when pausing for review
 
     Returns:
-        HumanInTheLoopMiddleware if any tools require approval, None otherwise
+        List containing HumanInTheLoopMiddleware if any tools require approval,
+        empty list otherwise
 
     Example:
         from dao_ai.config import ToolModel, PythonFunctionModel, HumanInTheLoopModel

dao_ai/middleware/message_validation.py

@@ -501,7 +501,7 @@ def create_user_id_validation_middleware() -> UserIdValidationMiddleware:
     and format of user_id in the runtime context.
 
     Returns:
-        UserIdValidationMiddleware instance
+        List containing UserIdValidationMiddleware instance
 
     Example:
         middleware = create_user_id_validation_middleware()
@@ -518,7 +518,7 @@ def create_thread_id_validation_middleware() -> ThreadIdValidationMiddleware:
     of thread_id in the runtime context.
 
     Returns:
-        ThreadIdValidationMiddleware instance
+        List containing ThreadIdValidationMiddleware instance
 
     Example:
         middleware = create_thread_id_validation_middleware()
@@ -550,7 +550,7 @@ def create_custom_field_validation_middleware(
             optionally 'description', 'required', and 'example_value' keys.
 
     Returns:
-        CustomFieldValidationMiddleware configured with the specified fields
+        List containing CustomFieldValidationMiddleware configured with the specified fields
 
     Example:
         middleware = create_custom_field_validation_middleware(
@@ -577,7 +577,7 @@ def create_filter_last_human_message_middleware() -> FilterLastHumanMessageMiddl
     process only the latest user input without conversation history.
 
     Returns:
-        FilterLastHumanMessageMiddleware instance
+        List containing FilterLastHumanMessageMiddleware instance
 
     Example:
         middleware = create_filter_last_human_message_middleware()

dao_ai/middleware/model_call_limit.py

@@ -0,0 +1,77 @@
+"""
+Model call limit middleware for DAO AI agents.
+
+Limits the number of model (LLM) calls to prevent infinite loops or excessive costs.
+
+Example:
+    from dao_ai.middleware import create_model_call_limit_middleware
+
+    # Limit model calls per run and thread
+    middleware = create_model_call_limit_middleware(
+        thread_limit=10,
+        run_limit=5,
+    )
+"""
+
+from __future__ import annotations
+
+from typing import Literal
+
+from langchain.agents.middleware import ModelCallLimitMiddleware
+from loguru import logger
+
+__all__ = [
+    "ModelCallLimitMiddleware",
+    "create_model_call_limit_middleware",
+]
+
+
+def create_model_call_limit_middleware(
+    thread_limit: int | None = None,
+    run_limit: int | None = None,
+    exit_behavior: Literal["error", "end"] = "end",
+) -> ModelCallLimitMiddleware:
+    """
+    Create a ModelCallLimitMiddleware to limit LLM API calls.
+
+    Prevents runaway agents from making too many API calls and helps
+    enforce cost controls on production deployments.
+
+    Args:
+        thread_limit: Max model calls per thread (conversation).
+            Requires checkpointer. None = no limit.
+        run_limit: Max model calls per run (single invocation).
+            None = no limit.
+        exit_behavior: What to do when limit hit:
+            - "end": Stop execution gracefully (default)
+            - "error": Raise ModelCallLimitExceededError immediately
+
+    Returns:
+        List containing ModelCallLimitMiddleware instance
+
+    Raises:
+        ValueError: If no limits specified
+
+    Example:
+        # Limit to 5 model calls per run, 10 per thread
+        limiter = create_model_call_limit_middleware(
+            run_limit=5,
+            thread_limit=10,
+            exit_behavior="end",
+        )
+    """
+    if thread_limit is None and run_limit is None:
+        raise ValueError("At least one of thread_limit or run_limit must be specified.")
+
+    logger.debug(
+        "Creating model call limit middleware",
+        thread_limit=thread_limit,
+        run_limit=run_limit,
+        exit_behavior=exit_behavior,
+    )
+
+    return ModelCallLimitMiddleware(
+        thread_limit=thread_limit,
+        run_limit=run_limit,
+        exit_behavior=exit_behavior,
+    )

dao_ai/middleware/model_retry.py

@@ -0,0 +1,121 @@
+"""
+Model retry middleware for DAO AI agents.
+
+Automatically retries failed model (LLM) calls with configurable exponential backoff.
+
+Example:
+    from dao_ai.middleware import create_model_retry_middleware
+
+    # Retry failed model calls with exponential backoff
+    middleware = create_model_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 ModelRetryMiddleware
+from loguru import logger
+
+__all__ = [
+    "ModelRetryMiddleware",
+    "create_model_retry_middleware",
+]
+
+
+def create_model_retry_middleware(
+    max_retries: int = 3,
+    backoff_factor: float = 2.0,
+    initial_delay: float = 1.0,
+    max_delay: float | None = None,
+    jitter: bool = False,
+    retry_on: tuple[type[Exception], ...] | Callable[[Exception], bool] | None = None,
+    on_failure: Literal["continue", "error"] | Callable[[Exception], str] = "continue",
+) -> ModelRetryMiddleware:
+    """
+    Create a ModelRetryMiddleware for automatic model call retries.
+
+    Handles transient failures in model API calls with exponential backoff.
+    Useful for handling rate limits, network issues, and temporary outages.
+
+    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.
+        retry_on: When to retry:
+            - None: Retry on all errors (default)
+            - tuple of Exception types: Retry only on these
+            - callable: Function(exception) -> bool for custom logic
+        on_failure: Behavior when all retries exhausted:
+            - "continue": Return AIMessage with error, let agent continue (default)
+            - "error": Re-raise exception, stop execution
+            - callable: Function(exception) -> str for custom error message
+
+    Returns:
+        List containing ModelRetryMiddleware instance
+
+    Example:
+        # Basic retry with defaults
+        retry = create_model_retry_middleware()
+
+        # Custom backoff for rate limits
+        retry = create_model_retry_middleware(
+            max_retries=5,
+            backoff_factor=2.0,
+            initial_delay=1.0,
+            max_delay=60.0,
+            jitter=True,
+        )
+
+        # Retry only on specific exceptions, fail hard
+        retry = create_model_retry_middleware(
+            max_retries=3,
+            retry_on=(RateLimitError, TimeoutError),
+            on_failure="error",
+        )
+
+        # Custom retry logic
+        def should_retry(error: Exception) -> bool:
+            return "rate_limit" in str(error).lower()
+
+        retry = create_model_retry_middleware(
+            max_retries=5,
+            retry_on=should_retry,
+        )
+    """
+    logger.debug(
+        "Creating model retry middleware",
+        max_retries=max_retries,
+        backoff_factor=backoff_factor,
+        initial_delay=initial_delay,
+        max_delay=max_delay,
+        jitter=jitter,
+        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 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 ModelRetryMiddleware(**kwargs)

dao_ai/middleware/pii.py

@@ -0,0 +1,157 @@
+"""
+PII detection middleware for DAO AI agents.
+
+Detects and handles Personally Identifiable Information (PII) in conversations
+using configurable strategies (redact, mask, hash, block).
+
+Example:
+    from dao_ai.middleware import create_pii_middleware
+
+    # Redact emails in user input
+    middleware = create_pii_middleware(
+        pii_type="email",
+        strategy="redact",
+        apply_to_input=True,
+    )
+"""
+
+from __future__ import annotations
+
+from typing import Any, Callable, Literal, Pattern
+
+from langchain.agents.middleware import PIIMiddleware
+from loguru import logger
+
+__all__ = [
+    "PIIMiddleware",
+    "create_pii_middleware",
+]
+
+# Type alias for PII detector
+PIIDetector = str | Pattern[str] | Callable[[str], list[dict[str, str | int]]]
+
+# Built-in PII types
+BUILTIN_PII_TYPES = frozenset({"email", "credit_card", "ip", "mac_address", "url"})
+
+
+def create_pii_middleware(
+    pii_type: str,
+    strategy: Literal["redact", "mask", "hash", "block"] = "redact",
+    detector: PIIDetector | None = None,
+    apply_to_input: bool = True,
+    apply_to_output: bool = False,
+    apply_to_tool_results: bool = False,
+) -> PIIMiddleware:
+    """
+    Create a PIIMiddleware for detecting and handling PII.
+
+    Detects Personally Identifiable Information in conversations and handles
+    it according to the specified strategy. Useful for compliance, privacy,
+    and sanitizing logs.
+
+    Built-in PII types:
+    - email: Email addresses
+    - credit_card: Credit card numbers (Luhn validated)
+    - ip: IP addresses
+    - mac_address: MAC addresses
+    - url: URLs
+
+    Args:
+        pii_type: Type of PII to detect. Use built-in types (email, credit_card,
+            ip, mac_address, url) or custom type names with a detector.
+        strategy: How to handle detected PII:
+            - "redact": Replace with [REDACTED_{TYPE}] (default)
+            - "mask": Partially obscure (e.g., ****-****-****-1234)
+            - "hash": Replace with deterministic hash
+            - "block": Raise exception when detected
+        detector: Custom detector for non-built-in types. Can be:
+            - str: Regex pattern string
+            - re.Pattern: Compiled regex pattern
+            - Callable: Function(content: str) -> list[dict] with keys:
+                - text: The matched text
+                - start: Start index
+                - end: End index
+            Default None (uses built-in detector for built-in types).
+        apply_to_input: Check user messages before model call. Default True.
+        apply_to_output: Check AI messages after model call. Default False.
+        apply_to_tool_results: Check tool results after execution. Default False.
+
+    Returns:
+        List containing PIIMiddleware instance
+
+    Raises:
+        ValueError: If custom pii_type without detector, or invalid strategy
+
+    Example:
+        # Redact emails in input
+        email_redactor = create_pii_middleware(
+            pii_type="email",
+            strategy="redact",
+            apply_to_input=True,
+        )
+
+        # Mask credit cards
+        card_masker = create_pii_middleware(
+            pii_type="credit_card",
+            strategy="mask",
+            apply_to_input=True,
+            apply_to_output=True,
+        )
+
+        # Block API keys with custom regex
+        api_key_blocker = create_pii_middleware(
+            pii_type="api_key",
+            detector=r"sk-[a-zA-Z0-9]{32}",
+            strategy="block",
+        )
+
+        # Custom SSN detector with validation
+        def detect_ssn(content: str) -> list[dict]:
+            matches = []
+            pattern = r"\\d{3}-\\d{2}-\\d{4}"
+            for match in re.finditer(pattern, content):
+                ssn = match.group(0)
+                first_three = int(ssn[:3])
+                if first_three not in [0, 666] and not (900 <= first_three <= 999):
+                    matches.append({
+                        "text": ssn,
+                        "start": match.start(),
+                        "end": match.end(),
+                    })
+            return matches
+
+        ssn_hasher = create_pii_middleware(
+            pii_type="ssn",
+            detector=detect_ssn,
+            strategy="hash",
+        )
+    """
+    # Validate: custom types require detector
+    if pii_type not in BUILTIN_PII_TYPES and detector is None:
+        raise ValueError(
+            f"Custom PII type '{pii_type}' requires a detector. "
+            f"Built-in types are: {', '.join(sorted(BUILTIN_PII_TYPES))}"
+        )
+
+    logger.debug(
+        "Creating PII middleware",
+        pii_type=pii_type,
+        strategy=strategy,
+        has_custom_detector=detector is not None,
+        apply_to_input=apply_to_input,
+        apply_to_output=apply_to_output,
+        apply_to_tool_results=apply_to_tool_results,
+    )
+
+    # Build kwargs
+    kwargs: dict[str, Any] = {
+        "strategy": strategy,
+        "apply_to_input": apply_to_input,
+        "apply_to_output": apply_to_output,
+        "apply_to_tool_results": apply_to_tool_results,
+    }
+
+    if detector is not None:
+        kwargs["detector"] = detector
+
+    return PIIMiddleware(pii_type, **kwargs)