aip-agents-binary 0.6.4__py3-none-any.whl

aip_agents/guardrails/manager.pyi

@@ -0,0 +1,42 @@
+from _typeshed import Incomplete
+from aip_agents.guardrails.engines.base import GuardrailEngine as GuardrailEngine
+from aip_agents.guardrails.schemas import GuardrailInput as GuardrailInput, GuardrailMode as GuardrailMode, GuardrailResult as GuardrailResult
+from typing import Any
+
+class GuardrailManager:
+    """Orchestrates multiple guardrail engines with fail-fast behavior.
+
+    The manager accepts one or more guardrail engines and runs them in sequence.
+    If any engine reports unsafe content, execution stops immediately (fail-fast)
+    and returns the violation result.
+
+    Attributes:
+        engines: List of guardrail engines to orchestrate
+    """
+    enabled: bool
+    engines: Incomplete
+    def __init__(self, engine: GuardrailEngine | list[GuardrailEngine] | None = None, engines: list[GuardrailEngine] | None = None) -> None:
+        """Initialize the GuardrailManager.
+
+        Args:
+            engine: Single guardrail engine to use
+            engines: List of guardrail engines to use
+
+        Raises:
+            ValueError: If both engine and engines are provided
+        """
+    def model_dump(self) -> dict[str, Any]:
+        """Serialize manager configuration into a JSON-compatible dictionary."""
+    async def check_content(self, content: str | GuardrailInput) -> GuardrailResult:
+        """Check content against all registered engines.
+
+        Executes engines in order with fail-fast behavior. If any engine
+        reports unsafe content, returns immediately with that result.
+
+        Args:
+            content: Content to check. Can be a string (treated as input-only)
+                    or GuardrailInput with input/output fields.
+
+        Returns:
+            GuardrailResult indicating if content passed all checks
+        """

aip_agents/guardrails/middleware.py

@@ -0,0 +1,199 @@
+"""GuardrailMiddleware for integrating guardrails into agent execution.
+
+This module provides GuardrailMiddleware that hooks into the agent execution
+flow to automatically check content before and after model invocations.
+
+Authors:
+    Reinhart Linanda (reinhart.linanda@gdplabs.id)
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from langchain_core.messages import AIMessage, HumanMessage
+
+from aip_agents.guardrails.exceptions import GuardrailViolationError
+from aip_agents.guardrails.schemas import GuardrailInput
+from aip_agents.middleware.base import AgentMiddleware, ModelRequest
+
+if TYPE_CHECKING:
+    from aip_agents.guardrails.manager import GuardrailManager
+
+
+class GuardrailMiddleware(AgentMiddleware):
+    """Middleware that integrates guardrails into agent execution.
+
+    This middleware wraps a GuardrailManager and automatically checks content
+    at appropriate points during agent execution:
+
+    - Before model invocation: checks user input from messages
+    - After model invocation: checks AI output from messages
+
+    If unsafe content is detected, raises GuardrailViolationError to stop execution.
+
+    Attributes:
+        guardrail_manager: The GuardrailManager to use for content checking
+    """
+
+    def __init__(self, guardrail_manager: GuardrailManager) -> None:
+        """Initialize the GuardrailMiddleware.
+
+        Args:
+            guardrail_manager: GuardrailManager instance to use for checking
+        """
+        self.guardrail_manager = guardrail_manager
+
+    @property
+    def tools(self) -> list:
+        """Guardrails are passive filters and don't contribute tools."""
+        return []
+
+    @property
+    def system_prompt_additions(self) -> str | None:
+        """Guardrails are passive filters and don't modify system prompts."""
+        return None
+
+    async def abefore_model(self, state: dict[str, Any]) -> dict[str, Any]:
+        """Asynchronously check user input before model invocation.
+
+        Extracts the last user message from state and checks it with guardrails.
+        If unsafe, raises GuardrailViolationError to stop execution.
+
+        Args:
+            state: Current agent state containing messages and context
+
+        Returns:
+            Empty dict (no state modifications needed)
+
+        Raises:
+            GuardrailViolationError: If user input violates safety policies
+        """
+        # Extract last user message from state
+        messages = state.get("messages", [])
+        user_input = self._extract_last_user_message(messages)
+
+        if user_input is not None:
+            # Check input content
+            result = await self.guardrail_manager.check_content(user_input)
+
+            if not result.is_safe:
+                raise GuardrailViolationError(result)
+
+        return {}
+
+    def before_model(self, state: dict[str, Any]) -> dict[str, Any]:
+        """Check user input before model invocation (synchronous wrapper).
+
+        Note:
+            This is a synchronous wrapper for the async `abefore_model()` method.
+            LangGraph agents primarily use `abefore_model()` in async contexts.
+            This method should rarely be called directly. If called from an async
+            context with a running event loop, it will attempt to handle it,
+            but `abefore_model()` should be preferred.
+
+        Args:
+            state: Current agent state containing messages and context
+
+        Returns:
+            Empty dict (no state modifications needed)
+
+        Raises:
+            GuardrailViolationError: If user input violates safety policies
+        """
+        import asyncio
+
+        user_input = self._extract_last_user_message(state.get("messages", []))
+        if user_input is None:
+            return {}
+
+        # Check if we're in an async context with a running loop
+        try:
+            loop = asyncio.get_running_loop()
+            if loop.is_running():
+                # We're in an async context with a running loop
+                # Use nest_asyncio to allow nested event loops
+                # This enables calling asyncio.run() from within a running loop
+                import nest_asyncio
+
+                nest_asyncio.apply()
+                result = asyncio.run(self.guardrail_manager.check_content(user_input))
+            else:
+                # Loop exists but not running - safe to use asyncio.run()
+                result = asyncio.run(self.guardrail_manager.check_content(user_input))
+        except RuntimeError:
+            # No running loop - safe to use asyncio.run()
+            result = asyncio.run(self.guardrail_manager.check_content(user_input))
+
+        if not result.is_safe:
+            raise GuardrailViolationError(result)
+
+        return {}
+
+    def modify_model_request(self, request: ModelRequest, state: dict[str, Any]) -> ModelRequest:
+        """Guardrails don't modify model requests."""
+        return request
+
+    async def aafter_model(self, state: dict[str, Any]) -> dict[str, Any]:
+        """Asynchronously check AI output after model invocation.
+
+        Extracts the last AI message from state and checks it with guardrails.
+        If unsafe, raises GuardrailViolationError to stop execution.
+
+        Args:
+            state: Current agent state after model invocation
+
+        Returns:
+            Empty dict (no state modifications needed)
+
+        Raises:
+            GuardrailViolationError: If AI output violates safety policies
+        """
+        # Extract last AI message from state
+        messages = state.get("messages", [])
+        ai_output = self._extract_last_ai_message(messages)
+
+        if ai_output is not None:
+            # Check output content
+            result = await self.guardrail_manager.check_content(GuardrailInput(input=None, output=ai_output))
+
+            if not result.is_safe:
+                raise GuardrailViolationError(result)
+
+        return {}
+
+    def after_model(self, state: dict[str, Any]) -> dict[str, Any]:
+        """Check AI output after model invocation (synchronous wrapper)."""
+        return {}
+
+    def _extract_last_user_message(self, messages: list) -> str | None:
+        """Extract the last user message from a list of messages.
+
+        Searches backwards through messages to find the most recent HumanMessage.
+
+        Args:
+            messages: List of message objects
+
+        Returns:
+            Content of the last user message, or None if not found
+        """
+        for message in reversed(messages):
+            if isinstance(message, HumanMessage) and message.content:
+                return str(message.content)
+        return None
+
+    def _extract_last_ai_message(self, messages: list) -> str | None:
+        """Extract the last AI message from a list of messages.
+
+        Searches backwards through messages to find the most recent AIMessage.
+
+        Args:
+            messages: List of message objects
+
+        Returns:
+            Content of the last AI message, or None if not found
+        """
+        for message in reversed(messages):
+            if isinstance(message, AIMessage) and message.content:
+                return str(message.content)
+        return None

aip_agents/guardrails/middleware.pyi

@@ -0,0 +1,87 @@
+from _typeshed import Incomplete
+from aip_agents.guardrails.exceptions import GuardrailViolationError as GuardrailViolationError
+from aip_agents.guardrails.manager import GuardrailManager as GuardrailManager
+from aip_agents.guardrails.schemas import GuardrailInput as GuardrailInput
+from aip_agents.middleware.base import AgentMiddleware as AgentMiddleware, ModelRequest as ModelRequest
+from typing import Any
+
+class GuardrailMiddleware(AgentMiddleware):
+    """Middleware that integrates guardrails into agent execution.
+
+    This middleware wraps a GuardrailManager and automatically checks content
+    at appropriate points during agent execution:
+
+    - Before model invocation: checks user input from messages
+    - After model invocation: checks AI output from messages
+
+    If unsafe content is detected, raises GuardrailViolationError to stop execution.
+
+    Attributes:
+        guardrail_manager: The GuardrailManager to use for content checking
+    """
+    guardrail_manager: Incomplete
+    def __init__(self, guardrail_manager: GuardrailManager) -> None:
+        """Initialize the GuardrailMiddleware.
+
+        Args:
+            guardrail_manager: GuardrailManager instance to use for checking
+        """
+    @property
+    def tools(self) -> list:
+        """Guardrails are passive filters and don't contribute tools."""
+    @property
+    def system_prompt_additions(self) -> str | None:
+        """Guardrails are passive filters and don't modify system prompts."""
+    async def abefore_model(self, state: dict[str, Any]) -> dict[str, Any]:
+        """Asynchronously check user input before model invocation.
+
+        Extracts the last user message from state and checks it with guardrails.
+        If unsafe, raises GuardrailViolationError to stop execution.
+
+        Args:
+            state: Current agent state containing messages and context
+
+        Returns:
+            Empty dict (no state modifications needed)
+
+        Raises:
+            GuardrailViolationError: If user input violates safety policies
+        """
+    def before_model(self, state: dict[str, Any]) -> dict[str, Any]:
+        """Check user input before model invocation (synchronous wrapper).
+
+        Note:
+            This is a synchronous wrapper for the async `abefore_model()` method.
+            LangGraph agents primarily use `abefore_model()` in async contexts.
+            This method should rarely be called directly. If called from an async
+            context with a running event loop, it will attempt to handle it,
+            but `abefore_model()` should be preferred.
+
+        Args:
+            state: Current agent state containing messages and context
+
+        Returns:
+            Empty dict (no state modifications needed)
+
+        Raises:
+            GuardrailViolationError: If user input violates safety policies
+        """
+    def modify_model_request(self, request: ModelRequest, state: dict[str, Any]) -> ModelRequest:
+        """Guardrails don't modify model requests."""
+    async def aafter_model(self, state: dict[str, Any]) -> dict[str, Any]:
+        """Asynchronously check AI output after model invocation.
+
+        Extracts the last AI message from state and checks it with guardrails.
+        If unsafe, raises GuardrailViolationError to stop execution.
+
+        Args:
+            state: Current agent state after model invocation
+
+        Returns:
+            Empty dict (no state modifications needed)
+
+        Raises:
+            GuardrailViolationError: If AI output violates safety policies
+        """
+    def after_model(self, state: dict[str, Any]) -> dict[str, Any]:
+        """Check AI output after model invocation (synchronous wrapper)."""

aip_agents/guardrails/schemas.py

@@ -0,0 +1,63 @@
+"""Schemas for guardrail input, output, and configuration.
+
+This module defines the data structures used throughout the guardrails system,
+including input/output schemas and configuration objects.
+
+Authors:
+    Reinhart Linanda (reinhart.linanda@gdplabs.id)
+"""
+
+from enum import StrEnum
+
+from pydantic import BaseModel, ConfigDict
+
+
+class GuardrailMode(StrEnum):
+    """Modes determining what content an engine checks."""
+
+    INPUT_ONLY = "input_only"
+    OUTPUT_ONLY = "output_only"
+    INPUT_OUTPUT = "input_output"
+    DISABLED = "disabled"
+
+
+class GuardrailInput(BaseModel):
+    """Input schema for guardrail checks.
+
+    Attributes:
+        input: User input content to check (queries, prompts, context)
+        output: AI output content to check (responses, generated text)
+    """
+
+    model_config = ConfigDict(extra="forbid")
+
+    input: str | None = None
+    output: str | None = None
+
+
+class GuardrailResult(BaseModel):
+    """Result schema returned by guardrail engines and managers.
+
+    Attributes:
+        is_safe: Whether the content passed all checks
+        reason: Explanation when content is blocked (None if safe)
+        filtered_content: Cleaned/sanitized content if engine provides it
+    """
+
+    model_config = ConfigDict(extra="forbid")
+
+    is_safe: bool
+    reason: str | None = None
+    filtered_content: str | None = None
+
+
+class BaseGuardrailEngineConfig(BaseModel):
+    """Base configuration for guardrail engines.
+
+    Attributes:
+        guardrail_mode: What content this engine should check
+    """
+
+    model_config = ConfigDict(extra="forbid")
+
+    guardrail_mode: GuardrailMode = GuardrailMode.INPUT_OUTPUT

aip_agents/guardrails/schemas.pyi

@@ -0,0 +1,43 @@
+from _typeshed import Incomplete
+from enum import StrEnum
+from pydantic import BaseModel
+
+class GuardrailMode(StrEnum):
+    """Modes determining what content an engine checks."""
+    INPUT_ONLY: str
+    OUTPUT_ONLY: str
+    INPUT_OUTPUT: str
+    DISABLED: str
+
+class GuardrailInput(BaseModel):
+    """Input schema for guardrail checks.
+
+    Attributes:
+        input: User input content to check (queries, prompts, context)
+        output: AI output content to check (responses, generated text)
+    """
+    model_config: Incomplete
+    input: str | None
+    output: str | None
+
+class GuardrailResult(BaseModel):
+    """Result schema returned by guardrail engines and managers.
+
+    Attributes:
+        is_safe: Whether the content passed all checks
+        reason: Explanation when content is blocked (None if safe)
+        filtered_content: Cleaned/sanitized content if engine provides it
+    """
+    model_config: Incomplete
+    is_safe: bool
+    reason: str | None
+    filtered_content: str | None
+
+class BaseGuardrailEngineConfig(BaseModel):
+    """Base configuration for guardrail engines.
+
+    Attributes:
+        guardrail_mode: What content this engine should check
+    """
+    model_config: Incomplete
+    guardrail_mode: GuardrailMode

aip_agents/guardrails/utils.py

@@ -0,0 +1,45 @@
+"""Utility functions for guardrail mode conversion.
+
+This module provides utilities for converting between aip-agents GuardrailMode
+and gllm-guardrail GuardrailMode enums.
+
+Authors:
+    Reinhart Linanda (reinhart.linanda@gdplabs.id)
+"""
+
+from typing import Any
+
+from aip_agents.guardrails.schemas import GuardrailMode
+
+
+def convert_guardrail_mode_to_gl_sdk(mode: GuardrailMode) -> Any:
+    """Convert aip-agents GuardrailMode to gllm-guardrail GuardrailMode.
+
+    This function performs lazy import of gllm-guardrail to support optional
+    dependencies. The conversion is necessary because we maintain our own
+    GuardrailMode enum for API consistency while wrapping the external library.
+
+    Args:
+        mode: The aip-agents GuardrailMode to convert
+
+    Returns:
+        The corresponding gllm-guardrail GuardrailMode enum value
+
+    Raises:
+        ImportError: If gllm-guardrail is not installed
+    """
+    try:
+        from gllm_guardrail.constants import GuardrailMode as GLGuardrailMode  # pragma: no cover
+    except ImportError as e:  # pragma: no cover
+        raise ImportError(
+            "gllm-guardrail is required for guardrails. Install with: pip install 'aip-agents[guardrails]'"
+        ) from e  # pragma: no cover
+
+    mode_mapping = {
+        GuardrailMode.INPUT_ONLY: GLGuardrailMode.INPUT_ONLY,
+        GuardrailMode.OUTPUT_ONLY: GLGuardrailMode.OUTPUT_ONLY,
+        GuardrailMode.INPUT_OUTPUT: GLGuardrailMode.BOTH,
+        GuardrailMode.DISABLED: GLGuardrailMode.DISABLED,
+    }
+
+    return mode_mapping[mode]

aip_agents/guardrails/utils.pyi

@@ -0,0 +1,19 @@
+from aip_agents.guardrails.schemas import GuardrailMode as GuardrailMode
+from typing import Any
+
+def convert_guardrail_mode_to_gl_sdk(mode: GuardrailMode) -> Any:
+    """Convert aip-agents GuardrailMode to gllm-guardrail GuardrailMode.
+
+    This function performs lazy import of gllm-guardrail to support optional
+    dependencies. The conversion is necessary because we maintain our own
+    GuardrailMode enum for API consistency while wrapping the external library.
+
+    Args:
+        mode: The aip-agents GuardrailMode to convert
+
+    Returns:
+        The corresponding gllm-guardrail GuardrailMode enum value
+
+    Raises:
+        ImportError: If gllm-guardrail is not installed
+    """

aip_agents/mcp/__init__.py

@@ -0,0 +1 @@
+"""Model Context Protocol related classes and functions."""

aip_agents/mcp/client/__init__.py

@@ -0,0 +1,14 @@
+"""MCP Client.
+
+This module provides a adapter client for interacting with MCP servers.
+
+Authors:
+    Fachriza Dian Adhiatma (fachriza.d.adhiatma@gdplabs.id)
+    Putu Ravindra Wiguna (putu.r.wiguna@gdplabs.id)
+"""
+
+from aip_agents.mcp.client.base_mcp_client import BaseMCPClient
+from aip_agents.mcp.client.google_adk.client import GoogleADKMCPClient
+from aip_agents.mcp.client.langchain.client import LangchainMCPClient
+
+__all__ = ["GoogleADKMCPClient", "LangchainMCPClient", "BaseMCPClient"]

aip_agents/mcp/client/__init__.pyi

@@ -0,0 +1,5 @@
+from aip_agents.mcp.client.base_mcp_client import BaseMCPClient as BaseMCPClient
+from aip_agents.mcp.client.google_adk.client import GoogleADKMCPClient as GoogleADKMCPClient
+from aip_agents.mcp.client.langchain.client import LangchainMCPClient as LangchainMCPClient
+
+__all__ = ['GoogleADKMCPClient', 'LangchainMCPClient', 'BaseMCPClient']