aip-agents-binary 0.5.21__py3-none-macosx_13_0_arm64.whl → 0.6.8__py3-none-macosx_13_0_arm64.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__.py

@@ -7,8 +7,44 @@ Authors:
     Putu Ravindra Wiguna (putu.r.wiguna@gdplabs.id)
 """
 
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
 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
+
+if TYPE_CHECKING:
+    from aip_agents.mcp.client.google_adk.client import GoogleADKMCPClient
+    from aip_agents.mcp.client.langchain.client import LangchainMCPClient
 
 __all__ = ["GoogleADKMCPClient", "LangchainMCPClient", "BaseMCPClient"]
+
+
+def __getattr__(name: str) -> Any:
+    """Lazy import of MCP client implementations.
+
+    This avoids importing heavy dependencies (Google ADK, Vertex AI, etc.)
+    when they are not needed.
+
+    Args:
+        name: Attribute name to import.
+
+    Returns:
+        The requested class.
+
+    Raises:
+        AttributeError: If attribute is not found.
+    """
+    if name == "GoogleADKMCPClient":
+        from aip_agents.mcp.client.google_adk.client import (
+            GoogleADKMCPClient as _GoogleADKMCPClient,
+        )
+
+        return _GoogleADKMCPClient
+    elif name == "LangchainMCPClient":
+        from aip_agents.mcp.client.langchain.client import (
+            LangchainMCPClient as _LangchainMCPClient,
+        )
+
+        return _LangchainMCPClient
+    raise AttributeError(f"module '{__name__}' has no attribute '{name}'")

aip_agents/mcp/client/connection_manager.py

@@ -60,6 +60,9 @@ class MCPConnectionManager:
     async def start(self) -> tuple[Any, Any]:
         """Start connection in background task.
 
+        For HTTP/SSE transports, establishes connection directly to avoid anyio context issues.
+        For stdio transport, uses background task to manage subprocess lifecycle.
+
         Returns:
             tuple[Any, Any]: Tuple of (read_stream, write_stream) for ClientSession
 
@@ -67,6 +70,17 @@ class MCPConnectionManager:
             Exception: If connection establishment fails
         """
         logger.debug(f"Starting connection manager for {self.server_name}")
+
+        # Determine transport type first
+        self.transport_type = self._get_transport_type()
+
+        # For HTTP/SSE: connect directly (no background task needed)
+        # This avoids anyio.BrokenResourceError when streams cross task boundaries
+        if self.transport_type in (TransportType.HTTP, TransportType.SSE):
+            await self._establish_connection()
+            return self._connection
+
+        # For stdio: use background task to manage subprocess
         self._task = asyncio.create_task(self._connection_task())
         await self._ready_event.wait()
 
@@ -78,6 +92,20 @@ class MCPConnectionManager:
     async def stop(self) -> None:
         """Stop connection gracefully."""
         logger.debug(f"Stopping connection manager for {self.server_name}")
+
+        # For HTTP/SSE (no background task), just close transport
+        if self.transport_type in (TransportType.HTTP, TransportType.SSE):
+            if self._transport:
+                try:
+                    close_result = self._transport.close()
+                    if inspect.isawaitable(close_result):
+                        await close_result
+                except Exception as exc:
+                    logger.warning(f"Failed to close transport cleanly for {self.server_name}: {exc}")
+            self._connection = None
+            return
+
+        # For stdio (with background task), wait for task to finish
         if self._task and not self._task.done():
             self._stop_event.set()
             try:
@@ -94,6 +122,11 @@ class MCPConnectionManager:
         Returns:
             bool: True if connected, False otherwise
         """
+        # For HTTP/SSE (no background task), just check if connection exists
+        if self.transport_type in (TransportType.HTTP, TransportType.SSE):
+            return self._connection is not None
+
+        # For stdio (with background task), check task status too
         return (
             self._connection is not None
             and self._task is not None
@@ -144,7 +177,9 @@ class MCPConnectionManager:
         Raises:
             ConnectionError: If all connection attempts fail
         """
-        self.transport_type = self._get_transport_type()
+        # transport_type may already be set by start() for HTTP/SSE
+        if not self.transport_type:
+            self.transport_type = self._get_transport_type()
         details = f"URL: {self.config.get('url', 'N/A')}, Command: {self.config.get('command', 'N/A')}"
         logger.info(f"Establishing connection to {self.server_name} via {self.transport_type} ({details})")
 

aip_agents/mcp/client/connection_manager.pyi

@@ -31,6 +31,9 @@ class MCPConnectionManager:
     async def start(self) -> tuple[Any, Any]:
         """Start connection in background task.
 
+        For HTTP/SSE transports, establishes connection directly to avoid anyio context issues.
+        For stdio transport, uses background task to manage subprocess lifecycle.
+
         Returns:
             tuple[Any, Any]: Tuple of (read_stream, write_stream) for ClientSession