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

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)

dao_ai/middleware/summarization.py

@@ -0,0 +1,197 @@
+"""
+Summarization middleware for DAO AI agents.
+
+This module provides a LoggingSummarizationMiddleware that extends LangChain's
+built-in SummarizationMiddleware with logging capabilities, and provides
+helper utilities for creating summarization middleware from DAO AI configuration.
+
+The middleware automatically:
+- Summarizes older messages using a separate LLM call when thresholds are exceeded
+- Replaces them with a summary message in State (permanently)
+- Keeps recent messages intact for context
+- Logs when summarization is triggered and completed
+
+Example:
+    from dao_ai.middleware import create_summarization_middleware
+    from dao_ai.config import ChatHistoryModel, LLMModel
+
+    chat_history = ChatHistoryModel(
+        model=LLMModel(name="gpt-4o-mini"),
+        max_tokens=256,
+        max_tokens_before_summary=4000,
+    )
+
+    middleware = create_summarization_middleware(chat_history)
+"""
+
+from typing import Any, Tuple
+
+from langchain.agents.middleware import SummarizationMiddleware
+from langchain_core.language_models import LanguageModelLike
+from langchain_core.messages import BaseMessage
+from langgraph.runtime import Runtime
+from loguru import logger
+
+from dao_ai.config import ChatHistoryModel
+
+__all__ = [
+    "SummarizationMiddleware",
+    "LoggingSummarizationMiddleware",
+    "create_summarization_middleware",
+]
+
+
+class LoggingSummarizationMiddleware(SummarizationMiddleware):
+    """
+    SummarizationMiddleware with logging for when summarization occurs.
+
+    This extends LangChain's SummarizationMiddleware to add logging at INFO level
+    when summarization is triggered and completed, providing visibility into
+    when conversation history is being summarized.
+
+    Logs include:
+    - Original message count and approximate token count (before summarization)
+    - New message count and approximate token count (after summarization)
+    - Number of messages that were summarized
+    """
+
+    def _log_summarization(
+        self,
+        original_message_count: int,
+        original_token_count: int,
+        result_messages: list[Any],
+    ) -> None:
+        """Log summarization details with before/after metrics."""
+        # Result messages: [RemoveMessage, summary_message, ...preserved_messages]
+        # New message count excludes RemoveMessage (index 0)
+        new_messages = [
+            msg for msg in result_messages if not self._is_remove_message(msg)
+        ]
+        new_message_count = len(new_messages)
+        new_token_count = self.token_counter(new_messages) if new_messages else 0
+
+        # Calculate how many messages were summarized
+        # preserved = new_messages - 1 (the summary message)
+        preserved_count = max(0, new_message_count - 1)
+        summarized_count = original_message_count - preserved_count
+
+        logger.info(
+            "Conversation summarized",
+            before_messages=original_message_count,
+            before_tokens=original_token_count,
+            after_messages=new_message_count,
+            after_tokens=new_token_count,
+            summarized_messages=summarized_count,
+        )
+        logger.debug(
+            "Summarization details",
+            trigger=self.trigger,
+            keep=self.keep,
+            preserved_messages=preserved_count,
+            token_reduction=original_token_count - new_token_count,
+        )
+
+    def _is_remove_message(self, msg: Any) -> bool:
+        """Check if a message is a RemoveMessage."""
+        return type(msg).__name__ == "RemoveMessage"
+
+    def before_model(
+        self, state: dict[str, Any], runtime: Runtime
+    ) -> dict[str, Any] | None:
+        """Process messages before model invocation, logging when summarization occurs."""
+        messages: list[BaseMessage] = state.get("messages", [])
+        original_message_count = len(messages)
+        original_token_count = self.token_counter(messages) if messages else 0
+
+        result = super().before_model(state, runtime)
+
+        if result is not None:
+            result_messages = result.get("messages", [])
+            self._log_summarization(
+                original_message_count,
+                original_token_count,
+                result_messages,
+            )
+
+        return result
+
+    async def abefore_model(
+        self, state: dict[str, Any], runtime: Runtime
+    ) -> dict[str, Any] | None:
+        """Process messages before model invocation (async), logging when summarization occurs."""
+        messages: list[BaseMessage] = state.get("messages", [])
+        original_message_count = len(messages)
+        original_token_count = self.token_counter(messages) if messages else 0
+
+        result = await super().abefore_model(state, runtime)
+
+        if result is not None:
+            result_messages = result.get("messages", [])
+            self._log_summarization(
+                original_message_count,
+                original_token_count,
+                result_messages,
+            )
+
+        return result
+
+
+def create_summarization_middleware(
+    chat_history: ChatHistoryModel,
+) -> LoggingSummarizationMiddleware:
+    """
+    Create a LoggingSummarizationMiddleware from DAO AI ChatHistoryModel configuration.
+
+    This factory function creates a LoggingSummarizationMiddleware instance
+    configured according to the DAO AI ChatHistoryModel settings. The middleware
+    includes logging at INFO level when summarization is triggered.
+
+    Args:
+        chat_history: ChatHistoryModel configuration for summarization
+
+    Returns:
+        List containing LoggingSummarizationMiddleware configured with the specified parameters
+
+    Example:
+        from dao_ai.config import ChatHistoryModel, LLMModel
+
+        chat_history = ChatHistoryModel(
+            model=LLMModel(name="gpt-4o-mini"),
+            max_tokens=256,
+            max_tokens_before_summary=4000,
+        )
+
+        middleware = create_summarization_middleware(chat_history)
+    """
+    logger.debug(
+        "Creating summarization middleware",
+        max_tokens=chat_history.max_tokens,
+        max_tokens_before_summary=chat_history.max_tokens_before_summary,
+        max_messages_before_summary=chat_history.max_messages_before_summary,
+    )
+
+    # Get the LLM model
+    model: LanguageModelLike = chat_history.model.as_chat_model()
+
+    # Determine trigger condition
+    # LangChain uses ("tokens", value) or ("messages", value) tuples
+    trigger: Tuple[str, int]
+    if chat_history.max_tokens_before_summary:
+        trigger = ("tokens", chat_history.max_tokens_before_summary)
+    elif chat_history.max_messages_before_summary:
+        trigger = ("messages", chat_history.max_messages_before_summary)
+    else:
+        # Default to a reasonable token threshold
+        trigger = ("tokens", chat_history.max_tokens * 10)
+
+    # Determine keep condition - how many recent messages/tokens to preserve
+    # Default to keeping enough for context
+    keep: Tuple[str, int] = ("tokens", chat_history.max_tokens)
+
+    logger.info("Summarization middleware configured", trigger=trigger, keep=keep)
+
+    return LoggingSummarizationMiddleware(
+        model=model,
+        trigger=trigger,
+        keep=keep,
+    )

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,
+    )