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

dao_ai/middleware/guardrails.py

@@ -0,0 +1,420 @@
+"""
+Guardrail middleware for DAO AI agents.
+
+This module provides middleware implementations for applying guardrails
+to agent responses, including LLM-based judging and content validation.
+
+Factory functions are provided for consistent configuration via the
+DAO AI middleware factory pattern.
+"""
+
+from typing import Any, Optional
+
+from langchain.agents.middleware import hook_config
+from langchain_core.language_models import LanguageModelLike
+from langchain_core.messages import AIMessage, BaseMessage, HumanMessage
+from langgraph.runtime import Runtime
+from loguru import logger
+from openevals.llm import create_llm_as_judge
+
+from dao_ai.messages import last_ai_message, last_human_message
+from dao_ai.middleware.base import AgentMiddleware
+from dao_ai.state import AgentState, Context
+
+
+def _extract_text_content(message: BaseMessage) -> str:
+    """
+    Extract text content from a message, handling both string and list formats.
+
+    Args:
+        message: The message to extract text from
+
+    Returns:
+        The extracted text content as a string
+    """
+    content = message.content
+
+    if isinstance(content, str):
+        return content
+    elif isinstance(content, list):
+        # Extract text from content blocks (e.g., Claude's structured content)
+        text_parts = []
+        for block in content:
+            if isinstance(block, dict) and block.get("type") == "text":
+                text_parts.append(block.get("text", ""))
+            elif isinstance(block, str):
+                text_parts.append(block)
+        return " ".join(text_parts)
+    else:
+        return str(content)
+
+
+__all__ = [
+    "GuardrailMiddleware",
+    "ContentFilterMiddleware",
+    "SafetyGuardrailMiddleware",
+    "create_guardrail_middleware",
+    "create_content_filter_middleware",
+    "create_safety_guardrail_middleware",
+]
+
+
+class GuardrailMiddleware(AgentMiddleware[AgentState, Context]):
+    """
+    Middleware that applies LLM-based guardrails to agent responses.
+
+    Uses an LLM judge to evaluate responses against a prompt/criteria and
+    can request improvements if the response doesn't meet the criteria.
+
+    This is equivalent to the previous reflection_guardrail pattern but
+    implemented as middleware for better composability.
+
+    Args:
+        guardrail_name: Name identifying this guardrail
+        model: The LLM to use for evaluation
+        prompt: The evaluation prompt/criteria
+        num_retries: Maximum number of retry attempts (default: 3)
+    """
+
+    def __init__(
+        self,
+        name: str,
+        model: LanguageModelLike,
+        prompt: str,
+        num_retries: int = 3,
+    ):
+        super().__init__()
+        self.guardrail_name = name
+        self.model = model
+        self.prompt = prompt
+        self.num_retries = num_retries
+        self._retry_count = 0
+
+    @property
+    def name(self) -> str:
+        """Return the guardrail name for middleware identification."""
+        return self.guardrail_name
+
+    def after_model(
+        self, state: AgentState, runtime: Runtime[Context]
+    ) -> dict[str, Any] | None:
+        """
+        Evaluate the model's response using an LLM judge.
+
+        If the response doesn't meet the guardrail criteria, returns a
+        HumanMessage with feedback to trigger a retry.
+        """
+        messages: list[BaseMessage] = state.get("messages", [])
+
+        if not messages:
+            return None
+
+        ai_message: AIMessage | None = last_ai_message(messages)
+        human_message: HumanMessage | None = last_human_message(messages)
+
+        if not ai_message or not human_message:
+            return None
+
+        # Skip evaluation if the AI message has tool calls (not the final response yet)
+        if ai_message.tool_calls:
+            logger.trace(
+                "Guardrail skipping evaluation - AI message contains tool calls",
+                guardrail_name=self.guardrail_name,
+            )
+            return None
+
+        # Skip evaluation if the AI message has no content to evaluate
+        if not ai_message.content:
+            logger.trace(
+                "Guardrail skipping evaluation - AI message has no content",
+                guardrail_name=self.guardrail_name,
+            )
+            return None
+
+        # Extract text content from messages (handles both string and structured content)
+        human_content = _extract_text_content(human_message)
+        ai_content = _extract_text_content(ai_message)
+
+        logger.debug(
+            "Evaluating response with guardrail",
+            guardrail_name=self.guardrail_name,
+            input_length=len(human_content),
+            output_length=len(ai_content),
+        )
+
+        evaluator = create_llm_as_judge(
+            prompt=self.prompt,
+            judge=self.model,
+        )
+
+        eval_result = evaluator(inputs=human_content, outputs=ai_content)
+
+        if eval_result["score"]:
+            logger.debug(
+                "Response approved by guardrail",
+                guardrail_name=self.guardrail_name,
+                comment=eval_result["comment"],
+            )
+            self._retry_count = 0
+            return None
+        else:
+            self._retry_count += 1
+            comment: str = eval_result["comment"]
+
+            if self._retry_count >= self.num_retries:
+                logger.warning(
+                    "Guardrail failed - max retries reached",
+                    guardrail_name=self.guardrail_name,
+                    retry_count=self._retry_count,
+                    max_retries=self.num_retries,
+                    critique=comment,
+                )
+                self._retry_count = 0
+
+                # Add system message to inform user of guardrail failure
+                failure_message = (
+                    f"⚠️ **Quality Check Failed**\n\n"
+                    f"The response did not meet the '{self.guardrail_name}' quality standards "
+                    f"after {self.num_retries} attempts.\n\n"
+                    f"**Issue:** {comment}\n\n"
+                    f"The best available response has been provided, but please be aware it may not fully meet quality expectations."
+                )
+                return {"messages": [AIMessage(content=failure_message)]}
+
+            logger.warning(
+                "Guardrail requested improvements",
+                guardrail_name=self.guardrail_name,
+                retry=self._retry_count,
+                max_retries=self.num_retries,
+                critique=comment,
+            )
+
+            content: str = "\n".join([str(human_message.content), comment])
+            return {"messages": [HumanMessage(content=content)]}
+
+
+class ContentFilterMiddleware(AgentMiddleware[AgentState, Context]):
+    """
+    Middleware that filters responses containing banned keywords.
+
+    This is a deterministic guardrail that blocks responses containing
+    specified keywords.
+
+    Args:
+        banned_keywords: List of keywords to block
+        block_message: Message to return when content is blocked
+    """
+
+    def __init__(
+        self,
+        banned_keywords: list[str],
+        block_message: str = "I cannot provide that response. Please rephrase your request.",
+    ):
+        super().__init__()
+        self.banned_keywords = [kw.lower() for kw in banned_keywords]
+        self.block_message = block_message
+
+    @hook_config(can_jump_to=["end"])
+    def before_agent(
+        self, state: AgentState, runtime: Runtime[Context]
+    ) -> dict[str, Any] | None:
+        """Block requests containing banned keywords."""
+        messages: list[BaseMessage] = state.get("messages", [])
+
+        if not messages:
+            return None
+
+        first_message = messages[0]
+        if not isinstance(first_message, HumanMessage):
+            return None
+
+        content = str(first_message.content).lower()
+
+        for keyword in self.banned_keywords:
+            if keyword in content:
+                logger.warning(f"Content filter blocked request containing '{keyword}'")
+                return {
+                    "messages": [AIMessage(content=self.block_message)],
+                    "jump_to": "end",
+                }
+
+        return None
+
+    def after_model(
+        self, state: AgentState, runtime: Runtime[Context]
+    ) -> dict[str, Any] | None:
+        """Block responses containing banned keywords."""
+        messages: list[BaseMessage] = state.get("messages", [])
+
+        if not messages:
+            return None
+
+        last_message: AIMessage | None = last_ai_message(messages)
+        if not last_message:
+            return None
+
+        content = str(last_message.content).lower()
+
+        for keyword in self.banned_keywords:
+            if keyword in content:
+                logger.warning("Content filter blocked response", keyword=keyword)
+                # Modify the last message content
+                last_message.content = self.block_message
+                return None
+
+        return None
+
+
+class SafetyGuardrailMiddleware(AgentMiddleware[AgentState, Context]):
+    """
+    Model-based safety guardrail middleware.
+
+    Uses an LLM to evaluate whether responses are safe and appropriate.
+    This is applied after the agent completes to validate the final output.
+
+    Args:
+        safety_model: The LLM to use for safety evaluation
+    """
+
+    def __init__(self, safety_model: Optional[LanguageModelLike] = None):
+        super().__init__()
+        self._safety_model = safety_model
+
+    @property
+    def safety_model(self) -> LanguageModelLike:
+        """Lazily initialize the safety model."""
+        if self._safety_model is None:
+            from langchain.chat_models import init_chat_model
+
+            self._safety_model = init_chat_model("gpt-4o-mini")
+        return self._safety_model
+
+    def after_agent(
+        self, state: AgentState, runtime: Runtime[Context]
+    ) -> dict[str, Any] | None:
+        """Evaluate response safety before returning to user."""
+        messages: list[BaseMessage] = state.get("messages", [])
+
+        if not messages:
+            return None
+
+        last_message: AIMessage | None = last_ai_message(messages)
+        if not last_message:
+            return None
+
+        safety_prompt = f"""Evaluate if this response is safe and appropriate.
+        Respond with only 'SAFE' or 'UNSAFE'.
+        
+        Response: {last_message.content}"""
+
+        result = self.safety_model.invoke([{"role": "user", "content": safety_prompt}])
+
+        if "UNSAFE" in str(result.content):
+            logger.warning("Safety guardrail blocked unsafe response")
+            last_message.content = (
+                "I cannot provide that response. Please rephrase your request."
+            )
+
+        return None
+
+
+# =============================================================================
+# Factory Functions
+# =============================================================================
+
+
+def create_guardrail_middleware(
+    name: str,
+    model: LanguageModelLike,
+    prompt: str,
+    num_retries: int = 3,
+) -> GuardrailMiddleware:
+    """
+    Create a GuardrailMiddleware instance.
+
+    Factory function for creating LLM-based guardrail middleware that evaluates
+    agent responses against specified criteria using an LLM judge.
+
+    Args:
+        name: Name identifying this guardrail
+        model: The LLM to use for evaluation
+        prompt: The evaluation prompt/criteria
+        num_retries: Maximum number of retry attempts (default: 3)
+
+    Returns:
+        List containing GuardrailMiddleware configured with the specified parameters
+
+    Example:
+        middleware = create_guardrail_middleware(
+            name="tone_check",
+            model=ChatDatabricks(endpoint="databricks-meta-llama-3-3-70b-instruct"),
+            prompt="Evaluate if the response is professional and helpful.",
+            num_retries=2,
+        )
+    """
+    logger.trace("Creating guardrail middleware", guardrail_name=name)
+    return GuardrailMiddleware(
+        name=name,
+        model=model,
+        prompt=prompt,
+        num_retries=num_retries,
+    )
+
+
+def create_content_filter_middleware(
+    banned_keywords: list[str],
+    block_message: str = "I cannot provide that response. Please rephrase your request.",
+) -> ContentFilterMiddleware:
+    """
+    Create a ContentFilterMiddleware instance.
+
+    Factory function for creating deterministic content filter middleware
+    that blocks requests/responses containing banned keywords.
+
+    Args:
+        banned_keywords: List of keywords to block
+        block_message: Message to return when content is blocked
+
+    Returns:
+        List containing ContentFilterMiddleware configured with the specified parameters
+
+    Example:
+        middleware = create_content_filter_middleware(
+            banned_keywords=["password", "secret", "api_key"],
+            block_message="I cannot discuss sensitive credentials.",
+        )
+    """
+    logger.trace(
+        "Creating content filter middleware", keywords_count=len(banned_keywords)
+    )
+    return ContentFilterMiddleware(
+        banned_keywords=banned_keywords,
+        block_message=block_message,
+    )
+
+
+def create_safety_guardrail_middleware(
+    safety_model: Optional[LanguageModelLike] = None,
+) -> SafetyGuardrailMiddleware:
+    """
+    Create a SafetyGuardrailMiddleware instance.
+
+    Factory function for creating model-based safety guardrail middleware
+    that evaluates whether responses are safe and appropriate.
+
+    Args:
+        safety_model: The LLM to use for safety evaluation. If not provided,
+            defaults to gpt-4o-mini.
+
+    Returns:
+        List containing SafetyGuardrailMiddleware configured with the specified model
+
+    Example:
+        from databricks_langchain import ChatDatabricks
+
+        middleware = create_safety_guardrail_middleware(
+            safety_model=ChatDatabricks(endpoint="databricks-meta-llama-3-3-70b-instruct"),
+        )
+    """
+    logger.trace("Creating safety guardrail middleware")
+    return SafetyGuardrailMiddleware(safety_model=safety_model)

dao_ai/middleware/human_in_the_loop.py

@@ -0,0 +1,233 @@
+"""
+Human-in-the-loop middleware for DAO AI agents.
+
+This module provides utilities for creating HITL middleware from DAO AI configuration.
+It re-exports LangChain's built-in HumanInTheLoopMiddleware.
+
+LangChain's HumanInTheLoopMiddleware automatically:
+- Pauses agent execution for human approval of tool calls
+- Allows humans to approve, edit, or reject tool calls
+- Uses LangGraph's interrupt mechanism for persistence across pauses
+
+Example:
+    from dao_ai.middleware import create_human_in_the_loop_middleware
+
+    middleware = create_human_in_the_loop_middleware(
+        interrupt_on={"send_email": True, "delete_record": True},
+    )
+"""
+
+from typing import Any, Sequence
+
+from langchain.agents.middleware import HumanInTheLoopMiddleware
+from langchain.agents.middleware.human_in_the_loop import (
+    Action,
+    ActionRequest,
+    ApproveDecision,
+    Decision,
+    DecisionType,
+    EditDecision,
+    HITLRequest,
+    HITLResponse,
+    InterruptOnConfig,
+    RejectDecision,
+    ReviewConfig,
+)
+from loguru import logger
+
+from dao_ai.config import HumanInTheLoopModel, ToolModel
+
+__all__ = [
+    # LangChain middleware
+    "HumanInTheLoopMiddleware",
+    # LangChain HITL types
+    "Action",
+    "ActionRequest",
+    "ApproveDecision",
+    "Decision",
+    "DecisionType",
+    "EditDecision",
+    "HITLRequest",
+    "HITLResponse",
+    "InterruptOnConfig",
+    "RejectDecision",
+    "ReviewConfig",
+    # DAO AI helper functions and models
+    "create_human_in_the_loop_middleware",
+    "create_hitl_middleware_from_tool_models",
+]
+
+
+def _hitl_config_to_allowed_decisions(
+    hitl_config: HumanInTheLoopModel,
+) -> list[DecisionType]:
+    """
+    Extract allowed decisions from HumanInTheLoopModel.
+
+    LangChain's HumanInTheLoopMiddleware supports 3 decision types:
+    - "approve": Execute tool with original arguments
+    - "edit": Modify arguments before execution
+    - "reject": Skip execution with optional feedback message
+
+    Args:
+        hitl_config: HumanInTheLoopModel with allowed_decisions
+
+    Returns:
+        List of allowed decision types (e.g., ["approve", "edit", "reject"])
+    """
+    return hitl_config.allowed_decisions  # type: ignore
+
+
+def _config_to_interrupt_on_entry(
+    config: HumanInTheLoopModel | bool,
+) -> dict[str, Any] | bool:
+    """
+    Convert a HITL config value to interrupt_on entry format.
+
+    Args:
+        config: HumanInTheLoopModel, True, or False
+
+    Returns:
+        dict with allowed_decisions and optional description, True, or False
+    """
+    if config is False:
+        return False
+    if config is True:
+        return {"allowed_decisions": ["approve", "edit", "reject"]}
+    if isinstance(config, HumanInTheLoopModel):
+        interrupt_entry: dict[str, Any] = {
+            "allowed_decisions": _hitl_config_to_allowed_decisions(config)
+        }
+        # If review_prompt is provided, use it as the description
+        if config.review_prompt is not None:
+            interrupt_entry["description"] = config.review_prompt
+        return interrupt_entry
+
+    logger.warning(
+        "Unknown HITL config type, defaulting to True",
+        config_type=type(config).__name__,
+    )
+    return True
+
+
+def create_human_in_the_loop_middleware(
+    interrupt_on: dict[str, HumanInTheLoopModel | bool | dict[str, Any]],
+    description_prefix: str = "Tool execution pending approval",
+) -> HumanInTheLoopMiddleware:
+    """
+    Create a HumanInTheLoopMiddleware instance.
+
+    Factory function for creating LangChain's built-in HumanInTheLoopMiddleware.
+    Accepts HumanInTheLoopModel, bool, or raw dict configurations per tool.
+
+    Note: This middleware requires a checkpointer to be configured on the agent.
+
+    Args:
+        interrupt_on: Dictionary mapping tool names to HITL configuration.
+            Each tool can be configured with:
+            - HumanInTheLoopModel: Full configuration with custom settings
+            - True: Enable HITL with default settings (approve, edit, reject)
+            - False: Disable HITL for this tool
+            - dict: Raw interrupt_on config (e.g., {"allowed_decisions": [...]})
+        description_prefix: Message prefix shown when pausing for review
+
+    Returns:
+        List containing HumanInTheLoopMiddleware configured with the specified parameters
+
+    Example:
+        from dao_ai.config import HumanInTheLoopModel
+
+        middleware = create_human_in_the_loop_middleware(
+            interrupt_on={
+                "send_email": HumanInTheLoopModel(review_prompt="Review email"),
+                "delete_record": True,
+                "search": False,
+            },
+        )
+    """
+    # Convert HumanInTheLoopModel entries to dict format
+    normalized_interrupt_on: dict[str, Any] = {}
+    for tool_name, config in interrupt_on.items():
+        if isinstance(config, (HumanInTheLoopModel, bool)):
+            normalized_interrupt_on[tool_name] = _config_to_interrupt_on_entry(config)
+        else:
+            # Already in dict format
+            normalized_interrupt_on[tool_name] = config
+
+    logger.debug(
+        "Creating HITL middleware",
+        tools_count=len(normalized_interrupt_on),
+        tools=list(normalized_interrupt_on.keys()),
+    )
+
+    return HumanInTheLoopMiddleware(
+        interrupt_on=normalized_interrupt_on,
+        description_prefix=description_prefix,
+    )
+
+
+def create_hitl_middleware_from_tool_models(
+    tool_models: Sequence[ToolModel],
+    description_prefix: str = "Tool execution pending approval",
+) -> HumanInTheLoopMiddleware | None:
+    """
+    Create HumanInTheLoopMiddleware from ToolModel configurations.
+
+    Scans tool_models for those with human_in_the_loop configured and
+    creates the appropriate middleware. This is the primary entry point
+    used by the agent node creation.
+
+    Args:
+        tool_models: List of ToolModel configurations from agent config
+        description_prefix: Message prefix shown when pausing for review
+
+    Returns:
+        List containing HumanInTheLoopMiddleware if any tools require approval,
+        empty list otherwise
+
+    Example:
+        from dao_ai.config import ToolModel, PythonFunctionModel, HumanInTheLoopModel
+
+        tool_models = [
+            ToolModel(
+                name="email_tool",
+                function=PythonFunctionModel(
+                    name="send_email",
+                    human_in_the_loop=HumanInTheLoopModel(
+                        review_prompt="Review this email",
+                    ),
+                ),
+            ),
+        ]
+
+        middleware = create_hitl_middleware_from_tool_models(tool_models)
+    """
+    from dao_ai.config import BaseFunctionModel
+
+    interrupt_on: dict[str, HumanInTheLoopModel] = {}
+
+    for tool_model in tool_models:
+        function = tool_model.function
+
+        if not isinstance(function, BaseFunctionModel):
+            continue
+
+        hitl_config: HumanInTheLoopModel | None = function.human_in_the_loop
+        if not hitl_config:
+            continue
+
+        # Get tool names created by this function
+        for func_tool in function.as_tools():
+            tool_name: str | None = getattr(func_tool, "name", None)
+            if tool_name:
+                interrupt_on[tool_name] = hitl_config
+                logger.trace("Tool configured for HITL", tool_name=tool_name)
+
+    if not interrupt_on:
+        logger.trace("No tools require HITL - returning None")
+        return None
+
+    return create_human_in_the_loop_middleware(
+        interrupt_on=interrupt_on,
+        description_prefix=description_prefix,
+    )