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

revenium_middleware/__init__.py

@@ -0,0 +1,184 @@
+"""
+Revenium Middleware — the economic system of record for AI usage.
+
+This package provides core metering functionality and provider-specific
+middleware for deeply attributed AI usage metrics.
+"""
+import os
+import sys
+import logging
+import re
+
+
+class ReadableFormatter(logging.Formatter):
+    """
+    Custom formatter that improves readability of log messages.
+    - Adds visual separators for important events
+    - Uses colors if terminal supports them
+    - Keeps messages clean and scannable
+    """
+
+    # ANSI color codes
+    COLORS = {
+        'RESET': '\033[0m',
+        'BOLD': '\033[1m',
+        'DIM': '\033[2m',
+        'GREEN': '\033[92m',
+        'YELLOW': '\033[93m',
+        'BLUE': '\033[94m',
+        'CYAN': '\033[96m',
+        'RED': '\033[91m',
+    }
+
+    # Box drawing characters
+    SEPARATOR = "─" * 70
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.use_colors = self._supports_color()
+
+    def _supports_color(self):
+        """Check if terminal supports ANSI colors."""
+        return hasattr(sys.stdout, 'isatty') and sys.stdout.isatty()
+
+    def _colorize(self, text, color):
+        """Apply color to text if terminal supports it."""
+        if self.use_colors and color in self.COLORS:
+            return f"{self.COLORS[color]}{text}{self.COLORS['RESET']}"
+        return text
+
+    def _truncate_long_objects(self, message, max_length=200):
+        """Truncate very long object representations to keep logs readable."""
+        # Check if message contains object representations
+        if '(' in message and ')' in message and len(message) > max_length:
+            # Look for patterns like "ClassName(field=value, ...)"
+            # Match object representations
+            pattern = r'(\w+)\(([^)]{100,})\)'
+
+            def replacer(match):
+                class_name = match.group(1)
+                content = match.group(2)
+                # Truncate the content
+                if len(content) > 150:
+                    truncated = content[:150] + "..."
+                    return f"{class_name}({truncated})"
+                return match.group(0)
+
+            message = re.sub(pattern, replacer, message)
+
+        return message
+
+    def format(self, record):
+        # Format the base message
+        message = record.getMessage()
+
+        # Truncate very long object representations in DEBUG logs
+        if record.levelname == 'DEBUG':
+            message = self._truncate_long_objects(message)
+
+        # Add visual enhancements for key events
+        # Check log level first to ensure it takes precedence
+        if record.levelname == 'ERROR':
+            message = self._colorize(f"[ERROR] {message}", 'RED')
+        elif record.levelname == 'WARNING':
+            message = self._colorize(f"[WARNING] {message}", 'YELLOW')
+        elif "SUCCESS" in message or "successful" in message.lower():
+            message = self._colorize(f"[SUCCESS] {message}", 'GREEN')
+        elif "FAILURE" in message:
+            message = self._colorize(f"[ERROR] {message}", 'RED')
+        elif "Shutdown complete" in message:
+            # Make shutdown completion more visible
+            separator = self._colorize(self.SEPARATOR, 'GREEN')
+            success_msg = self._colorize(
+                "[COMPLETE] All operations finished successfully", 'GREEN'
+            )
+            message = f"\n{separator}\n{success_msg}\n{separator}"
+        elif "Shutdown initiated" in message:
+            separator = self._colorize(self.SEPARATOR, 'CYAN')
+            shutdown_msg = self._colorize('[SHUTDOWN] ' + message, 'CYAN')
+            message = f"\n{separator}\n{shutdown_msg}"
+
+        # Format timestamp and level
+        if record.levelname == 'DEBUG':
+            level = self._colorize('DEBUG', 'DIM')
+        elif record.levelname == 'INFO':
+            level = self._colorize('INFO', 'BLUE')
+        elif record.levelname == 'WARNING':
+            level = self._colorize('WARN', 'YELLOW')
+        elif record.levelname == 'ERROR':
+            level = self._colorize('ERROR', 'RED')
+        else:
+            level = record.levelname
+
+        # Build the final log message
+        timestamp = self.formatTime(record, '%H:%M:%S')
+        return f"{timestamp} [{level}] {message}"
+
+
+# Set up logger
+logger = logging.getLogger("revenium_middleware")
+log_level = os.environ.get("REVENIUM_LOG_LEVEL", "INFO").upper()
+try:
+    logger.setLevel(getattr(logging, log_level))
+except AttributeError:
+    logger.setLevel(logging.INFO)
+    logger.warning(f"Invalid log level: {log_level}, defaulting to INFO")
+
+# Configure a handler with the readable formatter if none exists
+if not logger.handlers and not logging.root.handlers:
+    handler = logging.StreamHandler()
+    formatter = ReadableFormatter()
+    handler.setFormatter(formatter)
+    logger.addHandler(handler)
+
+# Allow propagation to root logger for testing
+logger.propagate = True
+
+# Re-export everything from _core for backward compatibility.
+# Existing imports like `from revenium_middleware import client` keep working.
+from ._core import (  # noqa: E402
+    client,
+    run_async_in_thread,
+    shutdown_event,
+    revenium_meter,
+    revenium_metadata,
+    track_usage,
+    is_inside_decorated_function,
+    get_function_metadata,
+    set_decorated_context,
+    clear_decorated_context,
+    get_injected_metadata,
+    set_injected_metadata,
+    clear_injected_metadata,
+    merge_metadata,
+    is_selective_metering_enabled,
+)
+
+# Re-export tool metering utilities from revenium_metering (v6.8.2+)
+from revenium_metering import meter_tool, report_tool_call, configure  # noqa: E402
+
+__all__ = [
+    # Metering exports
+    "client",
+    "run_async_in_thread",
+    "shutdown_event",
+    # Decorator exports
+    "revenium_meter",
+    "revenium_metadata",
+    "track_usage",
+    # Context management exports
+    "is_inside_decorated_function",
+    "get_function_metadata",
+    "set_decorated_context",
+    "clear_decorated_context",
+    "get_injected_metadata",
+    "set_injected_metadata",
+    "clear_injected_metadata",
+    "merge_metadata",
+    # Config exports
+    "is_selective_metering_enabled",
+    # Tool metering exports (from revenium_metering)
+    "meter_tool",
+    "report_tool_call",
+    "configure",
+]

revenium_middleware/_core/__init__.py

@@ -0,0 +1,65 @@
+"""
+Core functionality for Revenium middleware.
+
+This subpackage contains the foundational components shared across all
+provider-specific middleware implementations.
+"""
+
+from .metering import run_async_in_thread, shutdown_event, client
+from .context import (
+    is_inside_decorated_function,
+    get_function_metadata,
+    set_decorated_context,
+    clear_decorated_context,
+    get_injected_metadata,
+    set_injected_metadata,
+    clear_injected_metadata,
+    merge_metadata,
+)
+from .decorators import revenium_meter, revenium_metadata, track_usage
+from .config import is_selective_metering_enabled
+from .trace_fields import (
+    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,
+)
+
+__all__ = [
+    # Metering
+    "client",
+    "run_async_in_thread",
+    "shutdown_event",
+    # Decorators
+    "revenium_meter",
+    "revenium_metadata",
+    "track_usage",
+    # Context management
+    "is_inside_decorated_function",
+    "get_function_metadata",
+    "set_decorated_context",
+    "clear_decorated_context",
+    "get_injected_metadata",
+    "set_injected_metadata",
+    "clear_injected_metadata",
+    "merge_metadata",
+    # Config
+    "is_selective_metering_enabled",
+    # Trace fields
+    "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",
+]

revenium_middleware/_core/config.py

@@ -0,0 +1,165 @@
+"""
+Shared configuration constants and settings for Revenium middleware.
+
+This module is the single source of truth for configuration values shared
+across all provider middlewares. Provider-specific config files extend
+this via class inheritance and re-export symbols for backward compatibility.
+"""
+
+import logging
+import os
+from typing import Set, Literal, Union, Optional
+
+logger = logging.getLogger(__name__)
+
+# Type alias for summary format
+SummaryFormat = Literal["human", "json"]
+
+
+class Config:
+    """Shared configuration constants for all Revenium middleware providers."""
+
+    # Threading and async timeouts
+    THREAD_JOIN_TIMEOUT: float = 5.0
+    API_REQUEST_TIMEOUT: float = 30.0
+    BACKGROUND_THREAD_TIMEOUT: float = 5.0
+
+    # Cache settings
+    PROVIDER_CACHE_TTL: int = 3600  # 1 hour
+    MODEL_CACHE_TTL: int = 3600     # 1 hour
+
+    # Logging and security
+    MAX_LOG_STRING_LENGTH: int = 100
+    MAX_SANITIZATION_DEPTH: int = 3
+
+    # Stream processing
+    STREAM_CHUNK_BUFFER_SIZE: int = 1000
+    STREAM_TIMEOUT: float = 30.0
+
+    # Retry and resilience
+    MAX_RETRY_ATTEMPTS: int = 3
+    RETRY_BACKOFF_FACTOR: float = 1.5
+    INITIAL_RETRY_DELAY: float = 0.1
+
+    # Core environment variable names
+    ENV_REVENIUM_API_KEY: str = "REVENIUM_METERING_API_KEY"
+    ENV_LOG_LEVEL: str = "REVENIUM_LOG_LEVEL"
+
+    # Trace visualization environment variables
+    ENV_REVENIUM_ENVIRONMENT: str = "REVENIUM_ENVIRONMENT"
+    ENV_ENVIRONMENT: str = "ENVIRONMENT"
+    ENV_DEPLOYMENT_ENV: str = "DEPLOYMENT_ENV"
+    ENV_REVENIUM_REGION: str = "REVENIUM_REGION"
+    ENV_REVENIUM_CREDENTIAL_ALIAS: str = "REVENIUM_CREDENTIAL_ALIAS"
+    ENV_REVENIUM_TRACE_TYPE: str = "REVENIUM_TRACE_TYPE"
+    ENV_REVENIUM_TRACE_NAME: str = "REVENIUM_TRACE_NAME"
+    ENV_REVENIUM_PARENT_TRANSACTION_ID: str = "REVENIUM_PARENT_TRANSACTION_ID"
+    ENV_REVENIUM_TRANSACTION_NAME: str = "REVENIUM_TRANSACTION_NAME"
+    ENV_REVENIUM_RETRY_NUMBER: str = "REVENIUM_RETRY_NUMBER"
+
+    # Prompt capture settings
+    ENV_REVENIUM_CAPTURE_PROMPTS: str = "REVENIUM_CAPTURE_PROMPTS"
+    CAPTURE_PROMPTS: bool = os.getenv("REVENIUM_CAPTURE_PROMPTS", "false").lower() in ("true", "1", "yes", "on")
+    MAX_PROMPT_LENGTH: int = 50_000  # Maximum characters per prompt field
+
+    # Terminal summary settings
+    ENV_REVENIUM_PRINT_SUMMARY: str = "REVENIUM_PRINT_SUMMARY"
+    ENV_REVENIUM_TEAM_ID: str = "REVENIUM_TEAM_ID"
+    ENV_REVENIUM_BASE_URL: str = "REVENIUM_METERING_BASE_URL"
+    SUMMARY_RETRY_ATTEMPTS: int = 3
+    SUMMARY_RETRY_DELAY: float = 2.0
+    SUMMARY_API_TIMEOUT: float = 5.0
+    DEFAULT_BASE_URL: str = "https://api.revenium.ai"
+
+
+class SecurityConfig:
+    """Security-related configuration shared across all providers."""
+
+    SENSITIVE_FIELDS: Set[str] = {
+        'api_key', 'subscriber_credential', 'subscriber_email', 'messages',
+        'input', 'content', 'prompt', 'text', 'data', 'authorization',
+        'x-api-key', 'bearer', 'token', 'password', 'secret', 'key',
+        'credential', 'auth', 'private',
+        'system_prompt', 'input_messages', 'output_response'
+    }
+
+    SENSITIVE_PATTERNS: Set[str] = {
+        'sk-', 'pk-', 'Bearer ', 'Basic ', 'Token '
+    }
+
+
+def is_selective_metering_enabled() -> bool:
+    """
+    Check if selective metering is enabled.
+
+    When enabled, only functions decorated with @revenium_meter will be metered.
+    When disabled (default), all API calls are metered automatically.
+
+    The setting is controlled by the REVENIUM_SELECTIVE_METERING environment variable.
+    Accepted values for enabled: "true", "1", "yes", "on" (case-insensitive)
+    """
+    env_value = os.environ.get("REVENIUM_SELECTIVE_METERING", "false").lower()
+    return env_value in ("true", "1", "yes", "on")
+
+
+def get_config_value(key: str, default: any = None) -> any:
+    """Get configuration value from environment or use default."""
+    return os.getenv(key, default)
+
+
+def is_debug_enabled() -> bool:
+    """Check if debug logging is enabled."""
+    log_level = get_config_value(Config.ENV_LOG_LEVEL, "INFO").upper()
+    return log_level == "DEBUG"
+
+
+def get_timeout_config() -> dict:
+    """Get base timeout-related configuration."""
+    return {
+        'thread_join': Config.THREAD_JOIN_TIMEOUT,
+        'api_request': Config.API_REQUEST_TIMEOUT,
+        'background_thread': Config.BACKGROUND_THREAD_TIMEOUT,
+        'stream': Config.STREAM_TIMEOUT,
+    }
+
+
+def parse_print_summary_value(value: Optional[str]) -> Union[bool, SummaryFormat]:
+    """
+    Parse REVENIUM_PRINT_SUMMARY environment variable value.
+
+    Returns:
+        False if disabled, 'human' or 'json' if enabled
+    """
+    if value is None:
+        return False
+
+    value_lower = value.lower().strip()
+
+    if value_lower in ('false', '0', 'no', 'off', 'disabled', ''):
+        return False
+    elif value_lower in ('true', '1', 'yes', 'on', 'enabled', 'human'):
+        return 'human'
+    elif value_lower == 'json':
+        return 'json'
+    else:
+        logger.warning(
+            f"Invalid REVENIUM_PRINT_SUMMARY value '{value}'. "
+            f"Expected 'true', 'human', 'json', or 'false'. Defaulting to disabled."
+        )
+        return False
+
+
+def get_print_summary_config() -> Union[bool, SummaryFormat]:
+    """Get print summary configuration from environment."""
+    value = get_config_value(Config.ENV_REVENIUM_PRINT_SUMMARY)
+    return parse_print_summary_value(value)
+
+
+def get_team_id() -> Optional[str]:
+    """Get Revenium team ID from environment."""
+    return get_config_value(Config.ENV_REVENIUM_TEAM_ID)
+
+
+def get_base_url() -> str:
+    """Get Revenium base URL from environment (defaults to https://api.revenium.ai)."""
+    return get_config_value(Config.ENV_REVENIUM_BASE_URL, Config.DEFAULT_BASE_URL)

revenium_middleware/_core/context.py

@@ -0,0 +1,109 @@
+"""
+Context tracking for selective metering and metadata injection with decorators.
+
+This module provides thread-safe and async-safe context tracking to determine
+whether code is currently executing inside a decorated function that should be metered,
+and to store metadata that should be injected into API calls.
+"""
+
+import contextvars
+from typing import Optional, Dict, Any
+
+# Context variable to track if we're inside a decorated function
+_decorated_function_context: contextvars.ContextVar[bool] = contextvars.ContextVar(
+    'revenium_decorated_function', default=False
+)
+
+# Context variable to store metadata from the current decorated function
+_function_metadata_context: contextvars.ContextVar[Optional[Dict[str, Any]]] = contextvars.ContextVar(
+    'revenium_function_metadata', default=None
+)
+
+# Context variable to store injected metadata from @revenium_metadata decorator
+_injected_metadata_context: contextvars.ContextVar[Optional[Dict[str, Any]]] = contextvars.ContextVar(
+    'revenium_injected_metadata', default=None
+)
+
+
+def is_inside_decorated_function() -> bool:
+    """
+    Check if code is currently executing inside a decorated function.
+    
+    Returns:
+        True if inside a decorated function, False otherwise
+    """
+    return _decorated_function_context.get()
+
+
+def get_function_metadata() -> Optional[Dict[str, Any]]:
+    """
+    Get metadata from the current decorated function context.
+    
+    Returns:
+        Dictionary of metadata or None if not in decorated function
+    """
+    return _function_metadata_context.get()
+
+
+def set_decorated_context(is_decorated: bool, metadata: Optional[Dict[str, Any]] = None) -> None:
+    """
+    Set the decorated function context.
+    
+    Args:
+        is_decorated: Whether we're inside a decorated function
+        metadata: Optional metadata from the decorator
+    """
+    _decorated_function_context.set(is_decorated)
+    _function_metadata_context.set(metadata)
+
+
+def clear_decorated_context() -> None:
+    """Clear the decorated function context."""
+    _decorated_function_context.set(False)
+    _function_metadata_context.set(None)
+
+
+def get_injected_metadata() -> Optional[Dict[str, Any]]:
+    """
+    Get metadata from the current @revenium_metadata decorator context.
+
+    Returns:
+        Dictionary of injected metadata or None if not in decorated function
+    """
+    return _injected_metadata_context.get()
+
+
+def set_injected_metadata(metadata: Optional[Dict[str, Any]]) -> None:
+    """
+    Set the injected metadata context.
+
+    Args:
+        metadata: Dictionary of metadata to inject into API calls
+    """
+    _injected_metadata_context.set(metadata)
+
+
+def clear_injected_metadata() -> None:
+    """Clear the injected metadata context."""
+    _injected_metadata_context.set(None)
+
+
+def merge_metadata(api_metadata: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
+    """
+    Merge injected metadata with API-level metadata.
+
+    API-level metadata takes precedence over injected metadata.
+
+    Args:
+        api_metadata: Metadata passed directly to the API call
+
+    Returns:
+        Merged metadata dictionary with API-level metadata taking precedence
+    """
+    injected = get_injected_metadata() or {}
+    api = api_metadata or {}
+
+    # Start with injected metadata, then override with API-level metadata
+    merged = {**injected, **api}
+    return merged
+