aip-agents-binary 0.0.0b2__py3-none-any.whl → 0.5.11__py3-none-any.whl

aip_agents/sentry/sentry.py

@@ -0,0 +1,151 @@
+"""This file contains the Sentry and OpenTelemetry configuration for BOSA SDK.
+
+Authors:
+    Saul Sayers (saul.sayers@gdplabs.id)
+"""
+
+import inspect
+import os
+
+from bosa_core.telemetry import (
+    FastAPIConfig,
+    OpenTelemetryConfig,
+    SentryConfig,
+    TelemetryConfig,
+    init_telemetry,
+)
+from bosa_core.telemetry.opentelemetry.instrument.functions import (
+    BOSAFunctionsInstrumentor,
+)
+from dotenv import load_dotenv
+from fastapi import FastAPI
+
+from aip_agents.agent import BaseAgent, GoogleADKAgent, LangChainAgent, LangGraphAgent
+from aip_agents.utils.logger import get_logger
+
+load_dotenv()
+
+
+logger = get_logger(__name__)
+
+SENTRY_DSN = os.getenv("SENTRY_DSN")
+SENTRY_ENVIRONMENT = os.getenv("SENTRY_ENVIRONMENT", "development")
+SENTRY_PROJECT = os.getenv("SENTRY_PROJECT")
+VERSION_NUMBER = os.getenv("VERSION_NUMBER", "0.0.0")
+BUILD_NUMBER = os.getenv("BUILD_NUMBER", "0")
+USE_OPENTELEMETRY = os.getenv("USE_OPENTELEMETRY", "true").lower() == "true"
+
+CLASSES_TO_INSTRUMENT = [
+    BaseAgent,
+    LangGraphAgent,
+    LangChainAgent,
+    GoogleADKAgent,
+]
+
+
+def get_all_methods(cls: type) -> list:
+    """Get all methods from a class.
+
+    Args:
+        cls: The class to get methods from.
+
+    Returns:
+        list: A list of methods.
+    """
+    methods = []
+    for name, member in inspect.getmembers(cls):
+        if name.startswith("_"):
+            continue  # skip dunder and private
+        if inspect.isfunction(member) or inspect.ismethod(member) or inspect.iscoroutinefunction(member):
+            methods.append(member)
+    return methods
+
+
+def instrument_bosa_functions() -> None:
+    """Instrument BOSA functions."""
+    if BOSAFunctionsInstrumentor is None:
+        return
+    agent_methods = []
+    for cls in CLASSES_TO_INSTRUMENT:
+        agent_methods.extend(get_all_methods(cls))
+    BOSAFunctionsInstrumentor().instrument(methods=agent_methods)
+
+
+def traces_sampler(*args) -> float:
+    """Determine appropriate sampling rate for Sentry transactions.
+
+    Args:
+        *args: Additional positional arguments
+
+    Returns:
+        float: Sampling rate between 0 and 1
+    """
+    # TODO: Dont sample healthcheck endpoints (return 0.0)
+    return 1.0
+
+
+def setup_sentry_with_open_telemetry(app: FastAPI) -> None:
+    """Configure telemetry with both Sentry and OpenTelemetry.
+
+    Args:
+        app: FastAPI application instance
+    """
+    try:
+        fastapi_config = FastAPIConfig(app)
+
+        # Create OpenTelemetry configuration
+        opentelemetry_init = OpenTelemetryConfig(
+            use_langchain=True,
+            fastapi_config=fastapi_config,
+        )
+
+        # Sentry configuration with OpenTelemetry
+        sentry_config = SentryConfig(
+            dsn=SENTRY_DSN,
+            traces_sampler=traces_sampler,
+            environment=SENTRY_ENVIRONMENT,
+            release=f"{SENTRY_PROJECT}@{VERSION_NUMBER}+{BUILD_NUMBER}",
+            send_default_pii=True,
+            open_telemetry_config=opentelemetry_init,
+        )
+
+        telemetry_config = TelemetryConfig(sentry_config=sentry_config)
+        init_telemetry(telemetry_config)
+        logger.info(f"Telemetry initialized with OpenTelemetry for environment: {SENTRY_ENVIRONMENT}")
+    except Exception as e:
+        logger.error(f"Failed to initialize telemetry with OpenTelemetry: {e}")
+
+
+def setup_sentry_only() -> None:
+    """Configure telemetry with Sentry only (no OpenTelemetry)."""
+    try:
+        # Sentry configuration without OpenTelemetry
+        sentry_config = SentryConfig(
+            dsn=SENTRY_DSN,
+            traces_sampler=traces_sampler,
+            environment=SENTRY_ENVIRONMENT,
+            release=f"{SENTRY_PROJECT}@{VERSION_NUMBER}+{BUILD_NUMBER}",
+            send_default_pii=True,
+        )
+
+        telemetry_config = TelemetryConfig(sentry_config=sentry_config)
+        init_telemetry(telemetry_config)
+        logger.info(f"Telemetry initialized with Sentry only for environment: {SENTRY_ENVIRONMENT}")
+    except Exception as e:
+        logger.error(f"Failed to initialize telemetry with Sentry only: {e}")
+
+
+def setup_telemetry(app: FastAPI) -> None:
+    """Configure and initialize telemetry based on configuration.
+
+    Args:
+        app: FastAPI application instance
+    """
+    if not SENTRY_DSN:
+        logger.warning("Warning: SENTRY_DSN not set. Sentry will not be enabled.")
+        return
+    if USE_OPENTELEMETRY:
+        setup_sentry_with_open_telemetry(app)
+    else:
+        setup_sentry_only()
+    instrument_bosa_functions()

aip_agents/sentry/sentry.pyi

@@ -1,7 +1,7 @@
 from _typeshed import Incomplete
 from aip_agents.agent import BaseAgent as BaseAgent, GoogleADKAgent as GoogleADKAgent, LangChainAgent as LangChainAgent, LangGraphAgent as LangGraphAgent
-from aip_agents.utils.logger_manager import LoggerManager as LoggerManager
-from fastapi import FastAPI as FastAPI
+from aip_agents.utils.logger import get_logger as get_logger
+from fastapi import FastAPI
 
 logger: Incomplete
 SENTRY_DSN: Incomplete

aip_agents/storage/__init__.py

@@ -0,0 +1,41 @@
+"""Storage module for aip_agents.
+
+This module provides comprehensive storage functionality including object storage
+clients, storage providers for tool outputs, and configuration management.
+
+Authors:
+    Fachriza Adhiatma (fachriza.d.adhiatma@gdplabs.id)
+    Christian Trisno Sen Long Chen (christian.t.s.l.chen@gdplabs.id)
+"""
+
+# Object storage clients
+from aip_agents.storage.base import BaseObjectStorageClient
+from aip_agents.storage.clients.minio_client import MinioConfig, MinioObjectStorage
+
+# Configuration and factory
+from aip_agents.storage.config import (
+    StorageConfig,
+    StorageProviderFactory,
+    StorageType,
+)
+
+# Storage providers
+from aip_agents.storage.providers.base import BaseStorageProvider, StorageError
+from aip_agents.storage.providers.memory import InMemoryStorageProvider
+from aip_agents.storage.providers.object_storage import ObjectStorageProvider
+
+__all__ = [
+    # Object storage
+    "BaseObjectStorageClient",
+    "MinioConfig",
+    "MinioObjectStorage",
+    # Storage providers
+    "BaseStorageProvider",
+    "StorageError",
+    "InMemoryStorageProvider",
+    "ObjectStorageProvider",
+    # Configuration
+    "StorageConfig",
+    "StorageType",
+    "StorageProviderFactory",
+]

aip_agents/storage/base.py

@@ -0,0 +1,85 @@
+"""Base class for object storage clients.
+
+Authors:
+    Fachriza Adhiatma (fachriza.d.adhiatma@gdplabs.id)
+    Christian Trisno Sen Long Chen (christian.t.s.l.chen@gdplabs.id)
+"""
+
+from abc import ABC, abstractmethod
+from typing import BinaryIO
+
+
+class BaseObjectStorageClient(ABC):
+    """Abstract base class for object storage clients."""
+
+    @abstractmethod
+    def object_exists(self, object_key: str) -> bool:
+        """Check if an object exists in the storage.
+
+        Args:
+            object_key: The key of the object to check.
+
+        Returns:
+            True if the object exists, False otherwise.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def upload(
+        self,
+        object_key: str,
+        file_stream: bytes | BinaryIO,
+        filename: str | None = None,
+        content_type: str | None = None,
+        metadata: dict[str, str] | None = None,
+    ) -> str:
+        """Upload data to object storage.
+
+        Args:
+            object_key: The key of the object in the storage
+            file_stream: The binary data to upload
+            filename: The name of the file
+            content_type: The content type of the file
+            metadata: Additional metadata to store with the object
+
+        Returns:
+            The key of the uploaded object
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def get(self, object_key: str) -> bytes:
+        """Get data from object storage.
+
+        Args:
+            object_key: The key of the object in the storage
+
+        Returns:
+            The binary data of the object
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def delete(self, object_key: str) -> None:
+        """Delete data from object storage.
+
+        Args:
+            object_key: The key of the object in the storage
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def generate_presigned_url(
+        self, object_key: str, expires: int = 24, response_headers: dict[str, str] | None = None
+    ) -> str:
+        """Generate a presigned URL for accessing the object.
+
+        Args:
+            object_key: The key of the object in the storage
+            expires: The number of hours the URL is valid for
+            response_headers: Additional headers to include in the response
+
+        Returns:
+            The presigned URL
+        """
+        raise NotImplementedError

aip_agents/storage/base.pyi

@@ -1,8 +1,7 @@
-import abc
 from abc import ABC, abstractmethod
 from typing import BinaryIO
 
-class BaseObjectStorageClient(ABC, metaclass=abc.ABCMeta):
+class BaseObjectStorageClient(ABC):
     """Abstract base class for object storage clients."""
     @abstractmethod
     def object_exists(self, object_key: str) -> bool:

aip_agents/storage/clients/__init__.py

@@ -0,0 +1,12 @@
+"""Storage client implementations.
+
+This module contains concrete implementations of object storage clients
+for various backends like MinIO, AWS S3, etc.
+
+Authors:
+    Fachriza Adhiatma (fachriza.d.adhiatma@gdplabs.id)
+"""
+
+from aip_agents.storage.clients.minio_client import MinioConfig, MinioObjectStorage
+
+__all__ = ["MinioConfig", "MinioObjectStorage"]

aip_agents/storage/clients/minio_client.py

@@ -0,0 +1,318 @@
+"""Minio object storage implementation.
+
+Authors:
+    Fachriza Adhiatma (fachriza.d.adhiatma@gdplabs.id)
+"""
+
+import datetime
+import io
+import os
+from dataclasses import dataclass
+from typing import BinaryIO
+
+from minio import Minio, S3Error
+from urllib3.exceptions import HTTPError
+
+from aip_agents.storage.base import BaseObjectStorageClient
+
+S3_ERR_NO_SUCH_KEY = "NoSuchKey"
+S3_ERR_NO_SUCH_BUCKET = "NoSuchBucket"
+S3_ERR_ACCESS_DENIED = "AccessDenied"
+
+
+@dataclass
+class MinioConfig:
+    """Configuration for MinIO object storage client.
+
+    Attributes:
+        endpoint: MinIO server endpoint URL
+        access_key: Access key for authentication
+        secret_key: Secret key for authentication
+        bucket: Bucket name to use for storage
+        secure: Whether to use HTTPS (defaults to True)
+    """
+
+    endpoint: str
+    access_key: str
+    secret_key: str
+    bucket: str
+    secure: bool = True
+
+    @classmethod
+    def from_env(cls) -> "MinioConfig":
+        """Create MinioConfig from environment variables.
+
+        Expected environment variables:
+        - OBJECT_STORAGE_URL
+        - OBJECT_STORAGE_USER
+        - OBJECT_STORAGE_PASSWORD
+        - OBJECT_STORAGE_BUCKET
+        - OBJECT_STORAGE_SECURE (optional, defaults to True)
+
+        Returns:
+            MinioConfig instance
+
+        Raises:
+            ValueError: If required environment variables are not set
+        """
+        endpoint = os.getenv("OBJECT_STORAGE_URL")
+        access_key = os.getenv("OBJECT_STORAGE_USER")
+        secret_key = os.getenv("OBJECT_STORAGE_PASSWORD")
+        bucket = os.getenv("OBJECT_STORAGE_BUCKET")
+        secure_str = os.getenv("OBJECT_STORAGE_SECURE")
+
+        if not all([endpoint, access_key, secret_key, bucket]):
+            raise ValueError(
+                "MinIO configuration incomplete. Required environment variables: "
+                "OBJECT_STORAGE_URL, OBJECT_STORAGE_USER, "
+                "OBJECT_STORAGE_PASSWORD, OBJECT_STORAGE_BUCKET"
+            )
+
+        secure = secure_str.lower() in ("true", "1", "yes", "on")
+
+        return cls(endpoint=endpoint, access_key=access_key, secret_key=secret_key, bucket=bucket, secure=secure)
+
+
+class MinioObjectStorage(BaseObjectStorageClient):
+    """Implementation of ObjectStorageInterface using Minio."""
+
+    def __init__(self, config: MinioConfig | None = None, ensure_bucket: bool = True):
+        """Initialize MinioObjectStorage with configuration.
+
+        Args:
+            config: MinioConfig instance. If None, will attempt to load from environment variables.
+            ensure_bucket: Whether to ensure bucket exists during initialization (optional). Defaults to True.
+
+        Raises:
+            ValueError: If configuration is invalid or incomplete
+        """
+        if config is None:
+            config = MinioConfig.from_env()
+
+        self.config = config
+        self.client = Minio(
+            endpoint=config.endpoint,
+            access_key=config.access_key,
+            secret_key=config.secret_key,
+            secure=config.secure,
+        )
+        self.bucket = config.bucket
+        self._bucket_exists = None
+
+        if ensure_bucket:
+            self._ensure_bucket_exists()
+
+    def _ensure_bucket_exists(self) -> None:
+        """Ensure the bucket exists, creating it if necessary.
+
+        Uses caching to avoid repeated bucket existence checks.
+
+        Raises:
+            PermissionError: If user lacks permission to create the bucket
+            ConnectionError: If there's a network/connection issue creating the bucket
+        """
+        if self._bucket_exists is None:
+            try:
+                self._bucket_exists = self.client.bucket_exists(self.bucket)
+                if not self._bucket_exists:
+                    self.client.make_bucket(self.bucket)
+                    self._bucket_exists = True
+            except S3Error as e:
+                if e.code == S3_ERR_ACCESS_DENIED:  # pragma: no cover - difficult to test with mocks
+                    raise PermissionError(
+                        f"Access denied: Insufficient permissions to create bucket '{self.bucket}'. "
+                        f"Please check your AWS/MinIO user permissions."
+                    ) from e
+                else:
+                    raise ConnectionError(f"Failed to create bucket '{self.bucket}': {str(e)}") from e
+
+    def _prepare_file_stream(self, file_stream: bytes | BinaryIO) -> tuple[BinaryIO, int]:
+        """Prepare file stream for upload and validate its content.
+
+        Args:
+            file_stream: The file stream to prepare (bytes or file-like object)
+
+        Returns:
+            Tuple containing:
+                - Prepared file stream
+                - Length of the file content
+
+        Raises:
+            ValueError: If the file stream is empty or invalid
+        """
+        if isinstance(file_stream, bytes):
+            file_stream = io.BytesIO(file_stream)
+            file_length = len(file_stream.getvalue())
+        else:
+            # For file-like objects, seek to end to get size, then reset
+            try:
+                current_pos = file_stream.tell()
+                file_stream.seek(0, 2)  # Seek to end
+                file_length = file_stream.tell()
+                file_stream.seek(current_pos)
+            except (OSError, AttributeError) as e:
+                raise ValueError(f"Invalid file stream: {str(e)}") from e
+
+        if file_length == 0:
+            raise ValueError("File stream is empty.")
+
+        return file_stream, file_length
+
+    def upload(
+        self,
+        object_key: str,
+        file_stream: bytes | BinaryIO,
+        filename: str | None = None,
+        content_type: str | None = None,
+        metadata: dict[str, str] | None = None,
+    ) -> str:
+        """Upload data to Minio object storage.
+
+        Args:
+            object_key: The key to store the object under
+            file_stream: The data to upload (bytes or file-like object)
+            filename: The filename of the data (optional)
+            content_type: The content type of the data (optional)
+            metadata: Additional metadata to store with the object (optional)
+
+        Returns:
+            The object key of the uploaded data
+
+        Raises:
+            ValueError: If the file stream is empty or invalid
+            ConnectionError: If there's an issue connecting to Minio
+            Exception: For other unexpected errors during upload
+        """
+        if self._bucket_exists is None:
+            self._ensure_bucket_exists()
+        prepared_stream, file_length = self._prepare_file_stream(file_stream)
+
+        try:
+            self.client.put_object(
+                bucket_name=self.bucket,
+                object_name=object_key,
+                data=prepared_stream,
+                length=file_length,
+                content_type=content_type,
+                metadata=metadata,
+            )
+        except S3Error as e:
+            raise ConnectionError(f"Failed to upload object '{object_key}' to Minio: {str(e)}") from e
+        except HTTPError as e:
+            raise ConnectionError(f"A network error occurred during object upload: {str(e)}") from e
+
+        return object_key
+
+    def get(self, object_key: str) -> bytes:
+        """Get data from Minio object storage.
+
+        Args:
+            object_key: The key of the object to retrieve
+
+        Returns:
+            The object data as bytes
+
+        Raises:
+            KeyError: If the object is not found
+            ConnectionError: If there's a network or connection issue
+        """
+        try:
+            with self.client.get_object(bucket_name=self.bucket, object_name=object_key) as response:
+                return response.read()
+        except S3Error as e:
+            if e.code == S3_ERR_NO_SUCH_KEY:
+                raise KeyError(f"Object with key '{object_key}' not found") from e
+            raise ConnectionError(f"Failed to get object '{object_key}': {str(e)}") from e
+        except HTTPError as e:
+            raise ConnectionError(f"A network error occurred during object retrieval: {str(e)}") from e
+
+    def delete(self, object_key: str) -> None:
+        """Delete data from Minio object storage.
+
+        Args:
+            object_key: The key of the object to delete
+
+        Raises:
+            ConnectionError: If there's a network or connection issue
+        """
+        try:
+            self.client.remove_object(bucket_name=self.bucket, object_name=object_key)
+        except S3Error as e:
+            if e.code != S3_ERR_NO_SUCH_KEY:  # Don't raise error if object doesn't exist
+                raise ConnectionError(f"Failed to delete object '{object_key}': {str(e)}") from e
+        except HTTPError as e:
+            raise ConnectionError(f"A network error occurred during file deletion: {str(e)}") from e
+
+    def object_exists(self, object_key: str) -> bool:
+        """Check if an object exists in the MinIO bucket.
+
+        Args:
+            object_key: The key of the object to check.
+
+        Returns:
+            True if the object exists, False otherwise.
+
+        Raises:
+            ConnectionError: If there's a network or connection issue (excluding not found errors)
+        """
+        try:
+            self.client.stat_object(self.bucket, object_key)
+            return True
+        except S3Error as e:
+            if e.code in (S3_ERR_NO_SUCH_KEY, S3_ERR_NO_SUCH_BUCKET):
+                return False
+            raise ConnectionError(f"Failed to check object existence '{object_key}': {str(e)}") from e
+        except HTTPError as e:
+            raise ConnectionError(f"A network error occurred while checking object existence: {str(e)}") from e
+
+    def list_objects(self, prefix: str = "") -> list[str]:
+        """List objects in the bucket with optional prefix filter.
+
+        Args:
+            prefix: Optional prefix to filter objects
+
+        Returns:
+            List of object keys
+
+        Raises:
+            ConnectionError: If there's an issue listing objects
+        """
+        try:
+            objects = self.client.list_objects(self.bucket, prefix=prefix)
+            return [obj.object_name for obj in objects]
+        except S3Error as e:
+            raise ConnectionError(f"Failed to list objects in bucket '{self.bucket}': {str(e)}") from e
+        except HTTPError as e:
+            raise ConnectionError(f"A network error occurred while listing objects: {str(e)}") from e
+
+    def generate_presigned_url(
+        self, object_key: str, expires: int = 24, response_headers: dict[str, str] | None = None
+    ) -> str:
+        """Generate a presigned URL for accessing the object.
+
+        Args:
+            object_key: The key of the object
+            expires: Expiration time in hours (defaults to 24)
+            response_headers: Additional response headers (optional)
+
+        Returns:
+            A presigned URL
+
+        Raises:
+            ValueError: If expiration time is not positive
+            ConnectionError: If there's an issue generating the URL
+        """
+        if expires <= 0:
+            raise ValueError("Expiration time must be positive")
+
+        try:
+            return self.client.presigned_get_object(
+                bucket_name=self.bucket,
+                object_name=object_key,
+                expires=datetime.timedelta(hours=expires),
+                response_headers=response_headers,
+            )
+        except S3Error as e:
+            raise ConnectionError(f"Failed to generate presigned URL for '{object_key}': {str(e)}") from e
+        except HTTPError as e:
+            raise ConnectionError(f"A network error occurred while generating presigned URL: {str(e)}") from e

aip_agents/storage/config.py

@@ -0,0 +1,62 @@
+"""Storage configuration system.
+
+This module keeps factory logic while delegating data schemas to ``aip_agents.schema.storage``.
+"""
+
+from aip_agents.schema.storage import OBJECT_STORAGE_PREFIX, StorageConfig, StorageType
+from aip_agents.storage.base import BaseObjectStorageClient
+from aip_agents.storage.clients.minio_client import MinioConfig, MinioObjectStorage
+from aip_agents.storage.providers.base import BaseStorageProvider
+from aip_agents.storage.providers.memory import InMemoryStorageProvider
+from aip_agents.storage.providers.object_storage import ObjectStorageProvider
+
+__all__ = ["OBJECT_STORAGE_PREFIX", "StorageConfig", "StorageType", "StorageProviderFactory"]
+
+
+class StorageProviderFactory:
+    """Factory for creating storage providers based on configuration."""
+
+    @staticmethod
+    def create(
+        config: StorageConfig, object_storage_client: BaseObjectStorageClient | None = None
+    ) -> BaseStorageProvider:
+        """Create storage provider based on configuration.
+
+        Args:
+            config (StorageConfig): Storage configuration object.
+            object_storage_client (BaseObjectStorageClient | None, optional): Optional object storage client for object storage type.
+
+        Returns:
+            BaseStorageProvider: The created storage provider instance.
+        """
+        if config.storage_type == StorageType.MEMORY:
+            return InMemoryStorageProvider()
+
+        if config.storage_type == StorageType.OBJECT_STORAGE:
+            if object_storage_client is None:
+                try:
+                    minio_config = MinioConfig.from_env()
+                    object_storage_client = MinioObjectStorage(minio_config)
+                except ValueError as exc:
+                    raise ValueError(
+                        f"Object storage client required but not provided and cannot create from env: {exc}"
+                    ) from exc
+
+            return ObjectStorageProvider(
+                client=object_storage_client, prefix=config.object_prefix, use_json=config.object_use_json
+            )
+
+        raise ValueError(f"Unknown storage type: {config.storage_type}")
+
+    @staticmethod
+    def create_from_env(object_storage_client: BaseObjectStorageClient | None = None) -> BaseStorageProvider:
+        """Create storage provider from environment variables.
+
+        Args:
+            object_storage_client (BaseObjectStorageClient | None, optional): Optional object storage client for object storage type.
+
+        Returns:
+            BaseStorageProvider: The created storage provider instance.
+        """
+        config = StorageConfig.from_env()
+        return StorageProviderFactory.create(config, object_storage_client)