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

revenium_middleware/google/vertex_ai/provider.py

@@ -0,0 +1,99 @@
+"""
+Provider detection and configuration for Vertex AI SDK support.
+
+This module handles provider detection for the native Vertex AI SDK (vertexai package).
+All Vertex AI SDK usage reports as "Google" provider for unified analytics.
+"""
+
+import os
+import logging
+from ..common import ProviderMetadata
+
+logger = logging.getLogger("revenium_middleware.extension")
+
+
+def detect_provider() -> str:
+    """
+    Detect Vertex AI provider configuration.
+
+    For the native Vertex AI SDK, we always return "vertex_ai" since
+    this module only handles the vertexai package.
+
+    Returns:
+        Always returns "vertex_ai" for this module
+    """
+    logger.debug("Vertex AI SDK provider detected")
+    return "vertex_ai"
+
+
+def get_provider_metadata() -> ProviderMetadata:
+    """
+    Get provider metadata for Vertex AI SDK usage records.
+
+    Returns standardized "Google" provider name for unified analytics.
+
+    Returns:
+        ProviderMetadata with "Google" provider name
+    """
+    return ProviderMetadata.for_vertex_ai_sdk()
+
+
+def validate_vertex_ai_configuration() -> bool:
+    """
+    Validate that Vertex AI is properly configured.
+
+    Checks for required environment variables and configuration.
+
+    Returns:
+        True if Vertex AI is properly configured, False otherwise
+    """
+    # Check for required Vertex AI configuration
+    project_id = os.getenv("GOOGLE_CLOUD_PROJECT")
+    location = os.getenv("GOOGLE_CLOUD_LOCATION")
+
+    if not project_id:
+        logger.warning(
+            "GOOGLE_CLOUD_PROJECT environment variable not set for Vertex AI"
+        )
+        return False
+
+    if not location:
+        logger.warning(
+            "GOOGLE_CLOUD_LOCATION environment variable not set for Vertex AI"
+        )
+        return False
+
+    logger.debug(
+        f"Vertex AI configuration validated: project={project_id}, location={location}"
+    )
+    return True
+
+
+def get_vertex_ai_config() -> dict:
+    """
+    Get Vertex AI configuration from environment variables.
+
+    Returns:
+        Dictionary with Vertex AI configuration
+    """
+    return {
+        "project_id": os.getenv("GOOGLE_CLOUD_PROJECT"),
+        "location": os.getenv("GOOGLE_CLOUD_LOCATION", "us-central1"),
+        "credentials_path": os.getenv("GOOGLE_APPLICATION_CREDENTIALS"),
+    }
+
+
+def is_vertex_ai_available() -> bool:
+    """
+    Check if Vertex AI SDK is available and properly configured.
+
+    Returns:
+        True if Vertex AI SDK is available and configured, False otherwise
+    """
+    try:
+        import vertexai
+
+        return validate_vertex_ai_configuration()
+    except ImportError:
+        logger.debug("Vertex AI SDK not available")
+        return False

revenium_middleware/litellm/__init__.py

@@ -0,0 +1,25 @@
+"""
+Revenium Middleware for LiteLLM.
+
+Provides both client-side (wrapt-based) and proxy-side (custom logger) middleware
+for metering LiteLLM usage.
+"""
+
+import logging
+
+logger = logging.getLogger(__name__)
+
+# Conditionally import subpackages
+try:
+    import litellm as _litellm  # noqa: F401
+    from . import client
+    from . import proxy
+except ImportError:
+    logger.debug("LiteLLM SDK (litellm) not available, middleware not loaded")
+    client = None  # type: ignore
+    proxy = None  # type: ignore
+
+__all__ = [
+    "client",
+    "proxy",
+]

revenium_middleware/litellm/client/__init__.py

@@ -0,0 +1,81 @@
+"""
+Revenium LiteLLM Client Middleware
+
+When you install and import this library, it will automatically hook
+litellm.completion using wrapt, and log token usage after each request.
+
+New in v0.2.0:
+- Context-based metadata injection via metadata_context
+- Type-safe validation with UsageMetadata (requires pydantic)
+- Decorator-based metadata injection (8 decorators available)
+- Framework integrations (CrewAI integration available)
+
+Basic Usage:
+    >>> import revenium_middleware.litellm.client.middleware
+    >>> import litellm
+    >>> response = litellm.completion(
+    ...     model="gpt-4",
+    ...     messages=[{"role": "user", "content": "Hello"}],
+    ...     usage_metadata={"agent": "my-agent"}
+    ... )
+
+Context-Based Usage:
+    >>> from revenium_middleware.litellm.client import metadata_context
+    >>> with metadata_context.set(agent="Lead Analyst", task_type="research"):
+    ...     response = litellm.completion(...)  # Metadata auto-injected
+"""
+import logging
+
+logger = logging.getLogger(__name__)
+
+# Import components that don't require litellm SDK
+from .context import metadata_context, MetadataContext
+from .validation import UsageMetadata, Subscriber, SubscriberCredential, PYDANTIC_AVAILABLE
+from .decorators import (
+    track_agent,
+    track_task,
+    track_trace,
+    track_organization,
+    track_subscription,
+    track_product,
+    track_subscriber,
+    track_quality
+)
+from .hooks import (
+    register_metadata_hook,
+    unregister_metadata_hook,
+    clear_metadata_hooks,
+    get_registered_hooks
+)
+
+# Conditionally import middleware (requires litellm SDK)
+try:
+    import litellm as _litellm  # noqa: F401
+    from .middleware import completion_wrapper
+except ImportError:
+    logger.debug("LiteLLM SDK (litellm) not available, middleware not loaded")
+    completion_wrapper = None  # type: ignore
+
+__version__ = "0.2.0"
+
+__all__ = [
+    'completion_wrapper',
+    'metadata_context',
+    'MetadataContext',
+    'UsageMetadata',
+    'Subscriber',
+    'SubscriberCredential',
+    'PYDANTIC_AVAILABLE',
+    'track_agent',
+    'track_task',
+    'track_trace',
+    'track_organization',
+    'track_subscription',
+    'track_product',
+    'track_subscriber',
+    'track_quality',
+    'register_metadata_hook',
+    'unregister_metadata_hook',
+    'clear_metadata_hooks',
+    'get_registered_hooks',
+]

revenium_middleware/litellm/client/config.py

@@ -0,0 +1,53 @@
+"""
+Configuration for Revenium LiteLLM middleware.
+
+Imports shared config from core and adds LiteLLM-specific settings.
+"""
+
+import os
+from typing import Optional
+
+from revenium_middleware._core.config import (  # noqa: F401
+    Config as _CoreConfig,
+    SummaryFormat,
+    parse_print_summary_value,
+    get_print_summary_config,
+    get_team_id,
+    get_base_url,
+)
+
+# Re-export environment variable names for backward compatibility
+ENV_REVENIUM_PRINT_SUMMARY = _CoreConfig.ENV_REVENIUM_PRINT_SUMMARY
+ENV_REVENIUM_TEAM_ID = _CoreConfig.ENV_REVENIUM_TEAM_ID
+ENV_REVENIUM_METERING_BASE_URL = _CoreConfig.ENV_REVENIUM_BASE_URL
+ENV_REVENIUM_METERING_API_KEY = _CoreConfig.ENV_REVENIUM_API_KEY
+
+# Summary settings (LiteLLM uses a shorter retry delay)
+SUMMARY_RETRY_ATTEMPTS = _CoreConfig.SUMMARY_RETRY_ATTEMPTS
+SUMMARY_RETRY_DELAY: float = 1.0  # LiteLLM uses 1.0s vs core's 2.0s
+SUMMARY_API_TIMEOUT = _CoreConfig.SUMMARY_API_TIMEOUT
+
+DEFAULT_BASE_URL = _CoreConfig.DEFAULT_BASE_URL
+
+
+def get_api_key() -> Optional[str]:
+    """Get the Revenium API key from environment."""
+    return os.getenv(ENV_REVENIUM_METERING_API_KEY)
+
+
+__all__ = [
+    'ENV_REVENIUM_PRINT_SUMMARY',
+    'ENV_REVENIUM_TEAM_ID',
+    'ENV_REVENIUM_METERING_BASE_URL',
+    'ENV_REVENIUM_METERING_API_KEY',
+    'SUMMARY_RETRY_ATTEMPTS',
+    'SUMMARY_RETRY_DELAY',
+    'SUMMARY_API_TIMEOUT',
+    'DEFAULT_BASE_URL',
+    'SummaryFormat',
+    'parse_print_summary_value',
+    'get_print_summary_config',
+    'get_team_id',
+    'get_base_url',
+    'get_api_key',
+]

revenium_middleware/litellm/client/context.py

@@ -0,0 +1,198 @@
+"""
+Context-based metadata injection for Revenium LiteLLM middleware.
+
+This module provides thread-safe context management for injecting metadata into
+LiteLLM completion calls without requiring explicit kwargs. It uses Python's
+contextvars module to maintain isolated context per execution thread/task.
+
+Example:
+    >>> from revenium_middleware.litellm.client import metadata_context
+    >>> 
+    >>> # Using context manager
+    >>> with metadata_context.set(agent="Lead Analyst", task_type="research"):
+    ...     response = litellm.completion(model="gpt-4", messages=[...])
+    ...     # Metadata automatically injected
+    >>> 
+    >>> # Using direct API
+    >>> metadata_context.update(trace_id="abc-123")
+    >>> response = litellm.completion(model="gpt-4", messages=[...])
+    >>> metadata_context.clear()
+"""
+
+import contextvars
+from typing import Dict, Any, Optional
+
+
+# Thread-safe context variable for storing metadata
+_metadata_context: contextvars.ContextVar[Dict[str, Any]] = contextvars.ContextVar(
+    'revenium_metadata', 
+    default={}
+)
+
+
+class MetadataContext:
+    """
+    Thread-safe metadata context manager for Revenium middleware.
+    
+    This class provides methods to set, get, update, and clear metadata that
+    will be automatically injected into LiteLLM completion calls. It uses
+    Python's contextvars module to ensure thread-safety and proper isolation
+    in async contexts.
+    
+    All metadata set through this API will be merged with explicit usage_metadata
+    kwargs, with explicit kwargs taking precedence.
+    """
+    
+    @staticmethod
+    def set(**kwargs) -> '_MetadataContextManager':
+        """
+        Set metadata for the current context using a context manager.
+        
+        This method returns a context manager that sets the metadata on entry
+        and restores the previous metadata on exit. This is the recommended
+        way to set metadata for a specific scope.
+        
+        Args:
+            **kwargs: Metadata fields to set (e.g., agent="Lead Analyst",
+                     task_type="research", trace_id="abc-123")
+        
+        Returns:
+            _MetadataContextManager: Context manager for scoped metadata
+        
+        Example:
+            >>> with metadata_context.set(agent="Analyst", task_type="research"):
+            ...     # All LiteLLM calls here get this metadata
+            ...     response = litellm.completion(...)
+        """
+        return _MetadataContextManager(kwargs)
+    
+    @staticmethod
+    def get() -> Dict[str, Any]:
+        """
+        Get the current context metadata.
+        
+        Returns a copy of the current metadata dictionary. Modifying the
+        returned dictionary will not affect the context.
+        
+        Returns:
+            Dict[str, Any]: Copy of current metadata
+        
+        Example:
+            >>> metadata_context.update(agent="Analyst")
+            >>> current = metadata_context.get()
+            >>> print(current)
+            {'agent': 'Analyst'}
+        """
+        return _metadata_context.get().copy()
+    
+    @staticmethod
+    def update(**kwargs) -> None:
+        """
+        Update the current context metadata.
+        
+        This merges the provided metadata with existing metadata in the
+        current context. Unlike set(), this does not use a context manager
+        and the metadata persists until explicitly cleared.
+        
+        Args:
+            **kwargs: Metadata fields to update
+        
+        Example:
+            >>> metadata_context.update(agent="Analyst")
+            >>> metadata_context.update(task_type="research")
+            >>> # Both fields are now set
+            >>> metadata_context.clear()  # Clean up when done
+        """
+        current = _metadata_context.get().copy()
+        current.update(kwargs)
+        _metadata_context.set(current)
+    
+    @staticmethod
+    def clear(*keys: str) -> None:
+        """
+        Clear specific metadata fields or all metadata.
+        
+        If no keys are provided, clears all metadata. If keys are provided,
+        only those specific fields are removed.
+        
+        Args:
+            *keys: Optional field names to clear. If not provided, clears all.
+        
+        Example:
+            >>> metadata_context.update(agent="A", task_type="T", trace_id="123")
+            >>> metadata_context.clear("agent")  # Remove only agent
+            >>> metadata_context.clear()  # Remove all remaining
+        """
+        if not keys:
+            # Clear all metadata
+            _metadata_context.set({})
+        else:
+            # Clear specific keys
+            current = _metadata_context.get().copy()
+            for key in keys:
+                current.pop(key, None)
+            _metadata_context.set(current)
+
+
+class _MetadataContextManager:
+    """
+    Context manager for scoped metadata injection.
+    
+    This class implements the context manager protocol to provide scoped
+    metadata that is automatically cleaned up on exit. It saves the previous
+    metadata state on entry and restores it on exit.
+    
+    This class should not be instantiated directly. Use MetadataContext.set()
+    instead.
+    """
+    
+    def __init__(self, metadata: Dict[str, Any]):
+        """
+        Initialize the context manager.
+        
+        Args:
+            metadata: Metadata dictionary to set in this context
+        """
+        self.metadata = metadata
+        self.token: Optional[contextvars.Token] = None
+    
+    def __enter__(self) -> '_MetadataContextManager':
+        """
+        Enter the context and set metadata.
+        
+        Saves the current metadata state and sets the new metadata.
+        
+        Returns:
+            Self for use in with statements
+        """
+        # Merge new metadata with existing context
+        current = _metadata_context.get().copy()
+        current.update(self.metadata)
+        self.token = _metadata_context.set(current)
+        return self
+    
+    def __exit__(self, exc_type, exc_val, exc_tb) -> bool:
+        """
+        Exit the context and restore previous metadata.
+        
+        Restores the metadata state that existed before entering this context.
+        
+        Args:
+            exc_type: Exception type if an exception occurred
+            exc_val: Exception value if an exception occurred
+            exc_tb: Exception traceback if an exception occurred
+        
+        Returns:
+            False to propagate any exception that occurred
+        """
+        if self.token is not None:
+            _metadata_context.reset(self.token)
+        return False
+
+
+# Global instance for convenient access
+metadata_context = MetadataContext()
+
+
+__all__ = ['MetadataContext', 'metadata_context']
+