revenium-python-sdk 0.1.0__py3-none-any.whl

revenium_middleware/anthropic/trace_fields.py

@@ -0,0 +1,158 @@
+"""
+Trace visualization field capture and validation.
+
+This module provides functions to capture trace visualization fields from
+environment variables and validate them according to the specification.
+
+Shared functions are imported from _core.trace_fields. This module retains
+only Anthropic-specific functions: detect_vision_content and detect_operation_type.
+"""
+
+from typing import Optional, Dict, Any
+
+from revenium_middleware._core.trace_fields import (  # noqa: F401 — re-exported
+    TRACE_TYPE_MAX_LENGTH,
+    TRACE_NAME_MAX_LENGTH,
+    TRACE_TYPE_PATTERN,
+    get_environment,
+    get_region,
+    get_credential_alias,
+    get_trace_type,
+    get_trace_name,
+    get_parent_transaction_id,
+    get_transaction_name,
+    get_retry_number,
+    validate_trace_type,
+    validate_trace_name,
+)
+
+
+def detect_vision_content(messages: Optional[list] = None) -> bool:
+    """
+    Detect if messages contain vision/image content.
+
+    Anthropic vision content format:
+    {
+        "type": "image",
+        "source": {
+            "type": "base64" | "url",
+            "media_type": "image/jpeg",
+            "data": "..." # or "url": "..."
+        }
+    }
+
+    Args:
+        messages: List of message objects from Anthropic API request
+
+    Returns:
+        True if any message contains image content, False otherwise
+    """
+    if not messages:
+        return False
+
+    for message in messages:
+        if not isinstance(message, dict):
+            continue
+
+        content = message.get("content")
+        if content is None:
+            continue
+
+        # Content can be a string (no images) or a list of content blocks
+        if isinstance(content, list):
+            for block in content:
+                if isinstance(block, dict) and block.get("type") == "image":
+                    return True
+
+    return False
+
+
+def detect_operation_type(
+    provider: str,
+    endpoint: str,
+    request_body: Optional[Dict[str, Any]] = None
+) -> Dict[str, Optional[str]]:
+    """
+    Auto-detect operation type and subtype from provider, endpoint,
+    and request.
+
+    Args:
+        provider: Provider name (e.g., 'openai', 'azure_openai') or Provider enum
+        endpoint: API endpoint (e.g., '/chat/completions', '/embeddings')
+        request_body: Optional request body to check for tools/functions
+
+    Returns:
+        Dictionary with 'operationType' and 'operationSubtype' keys
+    """
+    # Handle Provider enum or string
+    if hasattr(provider, 'name'):
+        # It's a Provider enum, get the name (e.g., 'OPENAI', 'AZURE_OPENAI')
+        provider_str = provider.name
+    else:
+        provider_str = str(provider)
+
+    provider_lower = provider_str.lower()
+    request_body = request_body or {}
+
+    # OpenAI and Azure OpenAI
+    if provider_lower in ('openai', 'azure_openai', 'azure'):
+        # Chat completions
+        is_chat = (
+            'chat/completions' in endpoint or
+            endpoint.endswith('/chat/completions')
+        )
+        if is_chat:
+            # Check for tools or functions
+            has_tools = (
+                request_body.get('tools') or
+                request_body.get('functions')
+            )
+            if has_tools:
+                return {
+                    'operationType': 'TOOL_CALL',
+                    'operationSubtype': 'function_call'
+                }
+            return {
+                'operationType': 'CHAT',
+                'operationSubtype': None
+            }
+
+        # Embeddings
+        if 'embeddings' in endpoint or endpoint.endswith('/embeddings'):
+            return {
+                'operationType': 'EMBED',
+                'operationSubtype': None
+            }
+
+        # Moderations
+        if 'moderations' in endpoint or endpoint.endswith('/moderations'):
+            return {
+                'operationType': 'MODERATION',
+                'operationSubtype': None
+            }
+
+    # Anthropic
+    if provider_lower in ('anthropic', 'bedrock'):
+        # Messages endpoint (chat completions)
+        is_messages = (
+            '/messages' in endpoint or
+            endpoint.endswith('/messages')
+        )
+        if is_messages:
+            # Check for tools
+            has_tools = request_body.get('tools')
+            if has_tools:
+                return {
+                    'operationType': 'TOOL_CALL',
+                    'operationSubtype': None
+                }
+            return {
+                'operationType': 'CHAT',
+                'operationSubtype': None
+            }
+
+    # Default fallback
+    return {
+        'operationType': 'CHAT',
+        'operationSubtype': None
+    }

revenium_middleware/google/__init__.py

@@ -0,0 +1,114 @@
+"""
+Revenium middleware for Google AI services.
+
+This package provides middleware to track and meter usage of Google AI services
+including Google AI SDK (google-genai) and Vertex AI SDK (vertexai).
+
+The middleware automatically wraps API calls to capture usage metrics and send
+them to Revenium for tracking and billing purposes.
+
+Usage:
+    Simply import this package before using Google AI services:
+
+    import revenium_middleware.google
+    from google import genai  # or import vertexai
+
+    # Your existing code works unchanged
+    client = genai.Client(api_key="your-key")
+    response = client.models.generate_content(...)
+"""
+
+import logging
+import os
+
+
+# Configure logging based on REVENIUM_LOG_LEVEL environment variable
+def _configure_logging():
+    """Configure logging for the Revenium middleware."""
+    log_level_str = os.getenv("REVENIUM_LOG_LEVEL", "INFO").upper()
+
+    # Map string levels to logging constants
+    level_mapping = {
+        "DEBUG": logging.DEBUG,
+        "INFO": logging.INFO,
+        "WARNING": logging.WARNING,
+        "ERROR": logging.ERROR,
+        "CRITICAL": logging.CRITICAL,
+    }
+
+    log_level = level_mapping.get(log_level_str, logging.INFO)
+
+    # Configure the revenium middleware logger
+    revenium_logger = logging.getLogger("revenium_middleware")
+    revenium_logger.setLevel(log_level)
+
+    # Only add handler if none exists to avoid duplicate logs
+    if not revenium_logger.handlers:
+        handler = logging.StreamHandler()
+        formatter = logging.Formatter("%(name)s - %(levelname)s - %(message)s")
+        handler.setFormatter(formatter)
+        revenium_logger.addHandler(handler)
+
+    return log_level
+
+
+# Configure logging
+_log_level = _configure_logging()
+
+# Set up logging
+logger = logging.getLogger(__name__)
+
+# Check if verbose startup logging is enabled
+_verbose_startup = os.getenv("REVENIUM_VERBOSE_STARTUP", "").lower() in (
+    "true",
+    "1",
+    "yes",
+)
+
+if _verbose_startup:
+    logger.info("Revenium middleware initialization starting")
+
+# Import common utilities (always available)
+from .common import utils
+
+# Import and activate middleware for available SDKs
+try:
+    # Try to import Google AI SDK and activate middleware
+    if _verbose_startup:
+        logger.debug("Attempting to import google.genai")
+    import google.genai
+
+    if _verbose_startup:
+        logger.debug("google.genai imported successfully, importing middleware")
+    from .google_ai import middleware as google_ai_middleware
+
+    logger.info("Google AI SDK middleware activated")
+except ImportError as e:
+    logger.debug("Google AI SDK (google-genai) not available: %s", e)
+
+try:
+    # Try to import Vertex AI SDK and activate middleware
+    if _verbose_startup:
+        logger.debug("Attempting to import vertexai")
+    import vertexai
+
+    if _verbose_startup:
+        logger.debug("vertexai imported successfully, importing middleware")
+    from .vertex_ai import middleware as vertex_ai_middleware
+
+    logger.info("Vertex AI SDK middleware activated")
+except ImportError as e:
+    logger.debug("Vertex AI SDK (vertexai) not available: %s", e)
+
+active_sdks = []
+if "google.genai" in globals():
+    active_sdks.append("google_ai")
+if "vertexai" in globals():
+    active_sdks.append("vertex_ai")
+
+logger.info(
+    "Revenium middleware activated for: %s",
+    ", ".join(active_sdks) if active_sdks else "none",
+)
+
+__all__ = ["utils"]

revenium_middleware/google/common/__init__.py

@@ -0,0 +1,127 @@
+"""
+Common utilities and shared types for Google AI middleware.
+
+This module contains shared functionality used by both Google AI SDK
+and Vertex AI SDK middleware implementations.
+"""
+
+from .types import (
+    OperationType,
+    Provider,
+    ProviderMetadata,
+    UsageData,
+    TokenCounts,
+    normalize_stop_reason,
+    GOOGLE_AI_STOP_REASONS,
+    VERTEX_AI_STOP_REASONS,
+)
+
+from .utils import (
+    generate_transaction_id,
+    format_timestamp,
+    calculate_duration_ms,
+    log_token_usage,
+    log_image_usage,
+    log_video_usage,
+    create_metering_call,
+    create_image_metering_call,
+    create_video_metering_call,
+    extract_model_name,
+    extract_token_counts,
+    create_usage_data,
+    is_debug_logging_enabled,
+)
+
+from .exceptions import (
+    ReveniumMiddlewareError,
+    MeteringError,
+    TokenExtractionError,
+    ProviderDetectionError,
+    ConfigurationError,
+    StreamingError,
+    APIResponseError,
+    handle_metering_error,
+    safe_extract,
+)
+
+from .protocols import (
+    UsageMetadataProtocol,
+    CandidateProtocol,
+    EmbeddingProtocol,
+    ChatResponseProtocol,
+    EmbeddingResponseProtocol,
+    StreamChunkProtocol,
+    ClientProtocol,
+    is_chat_response,
+    is_embedding_response,
+    is_stream_chunk,
+    has_usage_metadata,
+    has_token_counts,
+    safe_getattr,
+    get_token_count,
+)
+
+from .summary_printer import (
+    CompletionMetrics,
+    fetch_completion_metrics,
+    format_and_print_json_summary,
+    format_and_print_human_summary,
+    print_usage_summary,
+)
+
+__all__ = [
+    # Types
+    "OperationType",
+    "Provider",
+    "ProviderMetadata",
+    "UsageData",
+    "TokenCounts",
+    "normalize_stop_reason",
+    "GOOGLE_AI_STOP_REASONS",
+    "VERTEX_AI_STOP_REASONS",
+    # Utils
+    "generate_transaction_id",
+    "format_timestamp",
+    "calculate_duration_ms",
+    "log_token_usage",
+    "log_image_usage",
+    "log_video_usage",
+    "create_metering_call",
+    "create_image_metering_call",
+    "create_video_metering_call",
+    "extract_model_name",
+    "extract_token_counts",
+    "create_usage_data",
+    "is_debug_logging_enabled",
+    # Exceptions
+    "ReveniumMiddlewareError",
+    "MeteringError",
+    "TokenExtractionError",
+    "ProviderDetectionError",
+    "ConfigurationError",
+    "StreamingError",
+    "APIResponseError",
+    "handle_metering_error",
+    "safe_extract",
+    # Protocols
+    "UsageMetadataProtocol",
+    "CandidateProtocol",
+    "EmbeddingProtocol",
+    "ChatResponseProtocol",
+    "EmbeddingResponseProtocol",
+    "StreamChunkProtocol",
+    "ClientProtocol",
+    "is_chat_response",
+    "is_embedding_response",
+    "is_stream_chunk",
+    "has_usage_metadata",
+    "has_token_counts",
+    "safe_getattr",
+    "get_token_count",
+    # Summary Printer
+    "CompletionMetrics",
+    "fetch_completion_metrics",
+    "format_and_print_json_summary",
+    "format_and_print_human_summary",
+    "print_usage_summary",
+]

revenium_middleware/google/common/exceptions.py

@@ -0,0 +1,137 @@
+"""
+Custom exception classes for Revenium Google AI middleware.
+
+This module defines a hierarchy of exceptions for different error scenarios
+that can occur during middleware operation.
+"""
+
+from typing import Optional, Any
+
+
+class ReveniumMiddlewareError(Exception):
+    """Base exception for all Revenium middleware errors."""
+
+    def __init__(self, message: str, details: Optional[dict] = None):
+        super().__init__(message)
+        self.message = message
+        self.details = details or {}
+
+    def __str__(self) -> str:
+        if self.details:
+            return f"{self.message} (Details: {self.details})"
+        return self.message
+
+
+class MeteringError(ReveniumMiddlewareError):
+    """Raised when metering operations fail."""
+
+    def __init__(
+        self,
+        message: str,
+        transaction_id: Optional[str] = None,
+        api_response: Optional[Any] = None,
+        **kwargs,
+    ):
+        super().__init__(message, kwargs)
+        self.transaction_id = transaction_id
+        self.api_response = api_response
+
+
+class TokenExtractionError(ReveniumMiddlewareError):
+    """Raised when token count extraction fails."""
+
+    def __init__(
+        self,
+        message: str,
+        response_type: Optional[str] = None,
+        operation_type: Optional[str] = None,
+        **kwargs,
+    ):
+        super().__init__(message, kwargs)
+        self.response_type = response_type
+        self.operation_type = operation_type
+
+
+class ProviderDetectionError(ReveniumMiddlewareError):
+    """Raised when provider detection fails."""
+
+    def __init__(self, message: str, available_sdks: Optional[list] = None, **kwargs):
+        super().__init__(message, kwargs)
+        self.available_sdks = available_sdks or []
+
+
+class ConfigurationError(ReveniumMiddlewareError):
+    """Raised when configuration is invalid or missing."""
+
+    def __init__(self, message: str, missing_config: Optional[list] = None, **kwargs):
+        super().__init__(message, kwargs)
+        self.missing_config = missing_config or []
+
+
+class StreamingError(ReveniumMiddlewareError):
+    """Raised when streaming operations fail."""
+
+    def __init__(
+        self,
+        message: str,
+        chunk_count: Optional[int] = None,
+        stream_state: Optional[str] = None,
+        **kwargs,
+    ):
+        super().__init__(message, kwargs)
+        self.chunk_count = chunk_count
+        self.stream_state = stream_state
+
+
+class APIResponseError(ReveniumMiddlewareError):
+    """Raised when API response is malformed or unexpected."""
+
+    def __init__(
+        self,
+        message: str,
+        response_data: Optional[Any] = None,
+        expected_format: Optional[str] = None,
+        **kwargs,
+    ):
+        super().__init__(message, kwargs)
+        self.response_data = response_data
+        self.expected_format = expected_format
+
+
+# Utility functions for error handling
+def handle_metering_error(func):
+    """Decorator to handle metering errors gracefully."""
+
+    def wrapper(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except MeteringError:
+            # Re-raise metering errors as-is
+            raise
+        except Exception as e:
+            # Convert other exceptions to MeteringError
+            raise MeteringError(
+                f"Unexpected error in {func.__name__}: {str(e)}",
+                details={"original_error": type(e).__name__},
+            ) from e
+
+    return wrapper
+
+
+def safe_extract(func):
+    """Decorator to handle extraction errors gracefully."""
+
+    def wrapper(*args, **kwargs):
+        try:
+            return func(*args, **kwargs)
+        except (TokenExtractionError, APIResponseError):
+            # Re-raise extraction errors as-is
+            raise
+        except Exception as e:
+            # Convert other exceptions to TokenExtractionError
+            raise TokenExtractionError(
+                f"Unexpected error in {func.__name__}: {str(e)}",
+                details={"original_error": type(e).__name__},
+            ) from e
+
+    return wrapper