aip-agents-binary 0.5.24__py3-none-any.whl → 0.5.25b1__py3-none-any.whl

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/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']

aip_agents/mcp/client/base_mcp_client.pyi

@@ -0,0 +1,148 @@
+from _typeshed import Incomplete
+from abc import ABC, abstractmethod
+from aip_agents.mcp.client.persistent_session import PersistentMCPSession as PersistentMCPSession
+from aip_agents.mcp.client.session_pool import MCPSessionPool as MCPSessionPool
+from aip_agents.utils.logger import get_logger as get_logger
+from gllm_tools.mcp.client.client import MCPClient
+from gllm_tools.mcp.client.config import MCPConfiguration as MCPConfiguration
+from mcp.types import CallToolResult, Tool as Tool
+from typing import Any
+
+logger: Incomplete
+
+class BaseMCPClient(MCPClient, ABC):
+    """Base MCP Client with persistent session management for aip-agents.
+
+    This class provides:
+    - Persistent session management across tool calls
+    - One-time tool registration and caching
+    - Automatic connection health monitoring and reconnection
+    - Centralized cleanup of all MCP resources
+    - Generic tool discovery from all configured MCP servers
+
+    Subclasses should implement SDK-specific tool conversion in _process_tool() method.
+    """
+    session_pool: Incomplete
+    def __init__(self, servers: dict[str, MCPConfiguration]) -> None:
+        """Initialize the base MCP client with session pool.
+
+        Args:
+            servers (dict[str, MCPConfiguration]): Dictionary of MCP server configurations by server name
+        """
+    async def initialize(self) -> None:
+        """Initialize all MCP sessions and cache tools once.
+
+        This method is idempotent and only performs initialization if not already done.
+        It establishes persistent connections to all configured MCP servers and caches
+        available tools for efficient access.
+
+        Raises:
+            Exception: If any session initialization fails
+        """
+    @abstractmethod
+    async def get_tools(self, server: str | None = None) -> list[Any]:
+        """Get framework-specific tools from MCP servers.
+
+        This method must be implemented by subclasses to provide framework-specific
+        tool conversion (e.g., StructuredTool for LangChain, FunctionTool for Google ADK).
+
+        Args:
+            server (str | None): Optional server name to filter tools from a specific server.
+                    If None, returns tools from all configured servers.
+
+        Returns:
+            list[Any]: List of framework-specific tool objects.
+        """
+    async def get_raw_mcp_tools(self, server: str | None = None) -> list[Tool]:
+        """Get raw MCP tools - for subclasses to perform framework-specific conversions.
+
+        This method provides access to the cached raw MCP Tool objects.
+        Subclasses use this to convert to framework-specific tools.
+
+        Args:
+            server (str | None): Optional server name to filter tools from a specific server.
+                    If None, returns tools from all configured servers.
+
+        Returns:
+            list[Tool]: List of raw MCP Tool objects. Empty list if not initialized or no tools available.
+        """
+    def get_tools_count(self, server: str | None = None) -> int:
+        """Get count of raw MCP tools without expensive copying.
+
+        This is an efficient way to get tool counts for logging/metrics
+        without the overhead of copying tool lists.
+
+        Args:
+            server (str | None): Optional server name to filter tools from a specific server.
+                    If None, returns count from all configured servers.
+
+        Returns:
+            int: Count of raw MCP tools available.
+        """
+    async def get_session(self, server_name: str) -> PersistentMCPSession:
+        """Get a persistent session for a specific server.
+
+        Args:
+            server_name (str): The name of the MCP server
+
+        Returns:
+            PersistentMCPSession: Persistent MCP session for the specified server
+
+        Raises:
+            KeyError: If the server is not configured or no active session exists
+        """
+    async def call_tool(self, server_name: str, tool_name: str, arguments: dict[str, Any]) -> CallToolResult:
+        """Execute a tool on a specific MCP server using persistent session.
+
+        Args:
+            server_name (str): The MCP server to execute the tool on
+            tool_name (str): The name of the tool to execute
+            arguments (dict[str, Any]): Arguments for the tool execution
+
+        Returns:
+            CallToolResult: The result of the tool execution
+
+        Raises:
+            KeyError: If the server doesn't exist
+            Exception: If tool execution fails
+        """
+    async def read_resource(self, server_name: str, resource_uri: str):
+        """Read an MCP resource from a specific server using persistent session.
+
+        Args:
+            server_name (str): The MCP server to read the resource from
+            resource_uri (str): The URI of the resource to read
+
+        Returns:
+            Any: The resource content
+
+        Raises:
+            KeyError: If the server doesn't exist
+            Exception: If resource reading fails
+        """
+    async def cleanup(self) -> None:
+        """Clean up all MCP resources and close sessions.
+
+        This method properly closes all persistent sessions and cleans up resources.
+        It can be called multiple times safely.
+        """
+    @property
+    def is_initialized(self) -> bool:
+        """Check if the client is fully initialized.
+
+        Returns:
+            bool: True if sessions are initialized and tools are cached, False otherwise
+        """
+    @property
+    def active_sessions(self) -> dict[str, PersistentMCPSession]:
+        """Get all active persistent sessions.
+
+        Returns:
+            dict[str, PersistentMCPSession]: Dictionary of active sessions by server name
+        """
+    def get_session_count(self) -> int:
+        """Get the number of active MCP sessions.
+
+        Returns:
+            Number of active persistent sessions
+        """

aip_agents/mcp/client/connection_manager.pyi

@@ -0,0 +1,48 @@
+from _typeshed import Incomplete
+from aip_agents.mcp.client.transports import TransportType as TransportType, create_transport as create_transport
+from aip_agents.utils.logger import get_logger as get_logger
+from gllm_tools.mcp.client.config import MCPConfiguration
+from typing import Any
+
+logger: Incomplete
+
+class MCPConnectionManager:
+    """Manages MCP connection lifecycle following mcp-use patterns.
+
+    This connection manager handles the transport connection lifecycle in a background
+    task to avoid cancel scope issues during cleanup. It supports automatic transport
+    negotiation (HTTP -> SSE fallback) and graceful shutdown. Invalid explicit transports
+    are normalized via aliases (e.g., 'streamable_http' -> 'http') or fall back to
+    auto-detection with a warning.
+    """
+    TRANSPORT_ALIASES: Incomplete
+    server_name: Incomplete
+    config: Incomplete
+    transport_type: Incomplete
+    max_retries: Incomplete
+    initial_retry_delay: Incomplete
+    def __init__(self, server_name: str, config: MCPConfiguration) -> None:
+        """Initialize connection manager.
+
+        Args:
+            server_name (str): Name of the MCP server
+            config (MCPConfiguration): MCP server configuration
+        """
+    async def start(self) -> tuple[Any, Any]:
+        """Start connection in background task.
+
+        Returns:
+            tuple[Any, Any]: Tuple of (read_stream, write_stream) for ClientSession
+
+        Raises:
+            Exception: If connection establishment fails
+        """
+    async def stop(self) -> None:
+        """Stop connection gracefully."""
+    @property
+    def is_connected(self) -> bool:
+        """Check if connection is active.
+
+        Returns:
+            bool: True if connected, False otherwise
+        """

aip_agents/mcp/client/google_adk/__init__.pyi

@@ -0,0 +1,3 @@
+from aip_agents.mcp.client.google_adk.client import GoogleADKMCPClient as GoogleADKMCPClient
+
+__all__ = ['GoogleADKMCPClient']