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

revenium_middleware/litellm/client/validation.py

@@ -0,0 +1,207 @@
+"""
+Pydantic models for validating usage metadata in Revenium LiteLLM middleware.
+
+This module provides type-safe validation for metadata fields using Pydantic.
+It helps catch errors early in development and provides IDE autocomplete support.
+
+Example:
+    >>> from revenium_middleware.litellm.client import UsageMetadata
+    >>> 
+    >>> # Valid metadata
+    >>> metadata = UsageMetadata(
+    ...     organization_id="AcmeCorp",
+    ...     subscription_id="82764738",
+    ...     product_id="Platinum",
+    ...     trace_id="abc-123",
+    ...     agent="Lead Analyst",
+    ...     task_type="research"
+    ... )
+    >>> 
+    >>> # Type errors caught immediately
+    >>> try:
+    ...     bad_metadata = UsageMetadata(organization_id=12345)  # Wrong type
+    ... except ValidationError as e:
+    ...     print(e)
+"""
+
+try:
+    from pydantic import BaseModel, Field, field_validator, ConfigDict
+    from typing import Optional, Dict, Any
+    PYDANTIC_AVAILABLE = True
+except ImportError:
+    # Pydantic is optional - provide fallback
+    PYDANTIC_AVAILABLE = False
+    BaseModel = object  # type: ignore
+
+
+if PYDANTIC_AVAILABLE:
+    class SubscriberCredential(BaseModel):
+        """
+        Credential information for a subscriber.
+
+        Attributes:
+            name: An alias for an API key used by one or more users
+            value: The key value associated with the subscriber (e.g., an API key)
+        """
+        model_config = ConfigDict(extra="forbid")
+
+        name: Optional[str] = Field(
+            None,
+            description="An alias for an API key used by one or more users"
+        )
+        value: Optional[str] = Field(
+            None,
+            description="The key value associated with the subscriber (e.g., an API key)"
+        )
+
+
+    class Subscriber(BaseModel):
+        """
+        Subscriber information for tracking individual users.
+
+        Attributes:
+            id: The ID of the subscriber from non-Revenium systems
+            email: The email address of the subscriber
+            credential: Credential information (API key, etc.)
+        """
+        model_config = ConfigDict(extra="forbid")
+
+        id: Optional[str] = Field(
+            None,
+            description="The ID of the subscriber from non-Revenium systems"
+        )
+        email: Optional[str] = Field(
+            None,
+            description="The email address of the subscriber"
+        )
+        credential: Optional[SubscriberCredential] = Field(
+            None,
+            description="Credential information for the subscriber"
+        )
+
+
+    class UsageMetadata(BaseModel):
+        """
+        Complete usage metadata for Revenium metering.
+
+        All fields are optional. Adding them enables more detailed reporting
+        and analytics in Revenium.
+
+        Attributes:
+            trace_id: Unique identifier for a conversation or session
+            task_type: Classification of the AI operation by type of work
+            subscriber: Object containing subscriber information
+            organization_id: Customer or department ID from non-Revenium systems
+            subscription_id: Reference to a billing plan in non-Revenium systems
+            product_id: Your product or feature making the AI call
+            agent: Identifier for the specific AI agent
+            response_quality_score: The quality of the AI response (0.0 to 1.0)
+
+        Example:
+            >>> metadata = UsageMetadata(
+            ...     organization_id="AcmeCorp",
+            ...     subscription_id="82764738",
+            ...     product_id="Platinum",
+            ...     trace_id="abc-123",
+            ...     agent="Lead Analyst",
+            ...     task_type="market_research",
+            ...     response_quality_score=0.95
+            ... )
+            >>> metadata_dict = metadata.model_dump(exclude_none=True)
+        """
+        model_config = ConfigDict(extra="allow")  # Allow additional custom fields
+
+        trace_id: Optional[str] = Field(
+            None,
+            description="Unique identifier for a conversation or session"
+        )
+        task_type: Optional[str] = Field(
+            None,
+            description="Classification of the AI operation by type of work"
+        )
+        subscriber: Optional[Subscriber] = Field(
+            None,
+            description="Object containing subscriber information"
+        )
+        organization_id: Optional[str] = Field(
+            None,
+            description="Customer or department ID from non-Revenium systems"
+        )
+        subscription_id: Optional[str] = Field(
+            None,
+            description="Reference to a billing plan in non-Revenium systems"
+        )
+        product_id: Optional[str] = Field(
+            None,
+            description="Your product or feature making the AI call"
+        )
+        agent: Optional[str] = Field(
+            None,
+            description="Identifier for the specific AI agent"
+        )
+        response_quality_score: Optional[float] = Field(
+            None,
+            ge=0.0,
+            le=1.0,
+            description="The quality of the AI response (0.0 to 1.0)"
+        )
+
+        @field_validator('response_quality_score')
+        @classmethod
+        def validate_quality_score(cls, v: Optional[float]) -> Optional[float]:
+            """Validate that quality score is between 0 and 1."""
+            if v is not None and (v < 0.0 or v > 1.0):
+                raise ValueError('response_quality_score must be between 0.0 and 1.0')
+            return v
+
+        def to_dict(self) -> Dict[str, Any]:
+            """
+            Convert to dictionary, excluding None values.
+            
+            Returns:
+                Dict with only non-None values
+            """
+            return self.model_dump(exclude_none=True)
+
+else:
+    # Fallback when Pydantic is not installed
+    class UsageMetadata:  # type: ignore
+        """
+        Fallback UsageMetadata class when Pydantic is not installed.
+        
+        This provides basic functionality without validation. For full
+        type safety and validation, install pydantic:
+        
+            pip install "revenium-python-sdk[litellm]"
+        """
+        
+        def __init__(self, **kwargs):
+            """Initialize with any keyword arguments."""
+            self.__dict__.update(kwargs)
+        
+        def to_dict(self) -> Dict[str, Any]:
+            """Convert to dictionary, excluding None values."""
+            return {k: v for k, v in self.__dict__.items() if v is not None}
+    
+    class Subscriber:  # type: ignore
+        """Fallback Subscriber class when Pydantic is not installed."""
+        
+        def __init__(self, **kwargs):
+            """Initialize with any keyword arguments."""
+            self.__dict__.update(kwargs)
+    
+    class SubscriberCredential:  # type: ignore
+        """Fallback SubscriberCredential class when Pydantic is not installed."""
+        
+        def __init__(self, **kwargs):
+            """Initialize with any keyword arguments."""
+            self.__dict__.update(kwargs)
+
+
+__all__ = [
+    'UsageMetadata',
+    'Subscriber', 
+    'SubscriberCredential',
+    'PYDANTIC_AVAILABLE'
+]
+

revenium_middleware/litellm/proxy/__init__.py

@@ -0,0 +1,25 @@
+"""
+Revenium LiteLLM Proxy Middleware
+
+When you install and import this library, it will automatically hook
+LiteLLM proxy requests using a custom logger, and log token usage after
+each request. You can customize or extend this logging logic later
+to add user or organization metadata for metering purposes.
+"""
+
+import logging
+
+logger = logging.getLogger(__name__)
+
+try:
+    import litellm  # noqa: F401
+    from .middleware import MiddlewareHandler, proxy_handler_instance
+except ImportError:
+    logger.debug("LiteLLM SDK (litellm) not available, proxy middleware not loaded")
+    MiddlewareHandler = None  # type: ignore
+    proxy_handler_instance = None  # type: ignore
+
+__all__ = [
+    "MiddlewareHandler",
+    "proxy_handler_instance",
+]

revenium_middleware/litellm/proxy/middleware.py

@@ -0,0 +1,217 @@
+from litellm.integrations.custom_logger import CustomLogger
+from revenium_middleware import client, run_async_in_thread
+import logging
+
+logger = logging.getLogger("revenium_middleware.extension")
+
+
+def _extract_organization_name(headers, metadata):
+    """
+    Extract organization name with fallback priority.
+
+    Priority: x-revenium-organization-name > x-revenium-organization-id > user_api_key_team_alias
+
+    Args:
+        headers: Request headers dictionary
+        metadata: Request metadata dictionary
+
+    Returns:
+        Organization name string or None
+    """
+    organization_name = headers.get("x-revenium-organization-name")
+    if organization_name is None:
+        organization_name = headers.get("x-revenium-organization-id")
+    if organization_name is None:
+        organization_name = metadata.get('user_api_key_team_alias', '') or None
+    return organization_name
+
+
+def _extract_product_name(headers):
+    """
+    Extract product name with fallback priority.
+
+    Priority: x-revenium-product-name > x-revenium-product-id
+
+    Args:
+        headers: Request headers dictionary
+
+    Returns:
+        Product name string or None
+    """
+    product_name = headers.get("x-revenium-product-name")
+    if product_name is None:
+        product_name = headers.get("x-revenium-product-id")
+    return product_name
+
+
+class MiddlewareHandler(CustomLogger):
+    async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
+        # log: key, user, model, prompt, response, tokens, cost
+        # Access kwargs passed to litellm.completion()
+        # pprint.pprint(kwargs['litellm_params'])
+
+        model = kwargs.get("model", None)
+
+        # Access litellm_params passed to litellm.completion(), example access `metadata`
+        litellm_params = kwargs.get("litellm_params", {})
+        metadata = litellm_params.get("metadata", {})  # headers passed to LiteLLM proxy, can be found here
+        headers = metadata.get("headers", {})
+
+        response = response_obj
+        # tokens used in response
+        usage = response_obj["usage"]
+
+        # Create subscriber object from metadata and headers
+        subscriber = {}
+
+        # Extract subscriber information from metadata and headers
+        subscriber_id = metadata.get('x-revenium-subscriber-id', '') or headers.get("x-revenium-subscriber-id")
+        subscriber_email = metadata.get('user_api_key_user_email', '')
+        credential_name = metadata.get('user_api_key_alias', '')
+        credential_value = metadata.get('user_api_key_alias', '')
+
+        if subscriber_id:
+            subscriber["id"] = subscriber_id
+        if subscriber_email:
+            subscriber["email"] = subscriber_email
+        if credential_name or credential_value:
+            subscriber["credential"] = {
+                "name": credential_name,
+                "value": credential_value
+            }
+
+        # Extract organization and product names using helper functions
+        organization_name = _extract_organization_name(headers, metadata)
+        product_name = _extract_product_name(headers)
+
+        completion_args = {
+            "cache_creation_token_count": 0,
+            "cache_read_token_count": 0,
+            "input_token_cost": None,
+            "output_token_cost": None,
+            "total_cost": None,
+            "output_token_count": usage.completion_tokens,
+            "cost_type": "AI",
+            "model": model,
+            "input_token_count": usage.prompt_tokens,
+            "provider": "LITELLM",
+            "model_source": "LITELLM",
+            "reasoning_token_count": 0,
+            "request_time": start_time.strftime("%Y-%m-%dT%H:%M:%SZ"),
+            "response_time": end_time.strftime("%Y-%m-%dT%H:%M:%SZ"),
+            "completion_start_time": end_time.strftime("%Y-%m-%dT%H:%M:%SZ"),
+            "request_duration": (end_time - start_time).total_seconds() * 1000,  # Convert to milliseconds
+            "time_to_first_token": (end_time - start_time).total_seconds() * 1000,
+            "stop_reason": "END",
+            "total_token_count": usage.total_tokens,
+            "transaction_id": response.id,
+            "trace_id": headers.get("x-revenium-trace-id"),
+            "task_type": headers.get("x-revenium-task-type"),
+            "subscriber": subscriber if subscriber else None,
+            "organization_name": organization_name,
+            "subscription_id": headers.get("x-revenium-subscription-id"),
+            "product_name": product_name,
+            "agent": headers.get("x-revenium-agent"),
+            "response_quality_score": headers.get("x-revenium-response-quality-score"),
+            "is_streamed": metadata.get('hidden_params', {}).get('optional_params', {}).get('stream', False),
+            "operation_type": "CHAT",
+            "mediation_latency": metadata.get('hidden_params', {}).get('litellm_overhead_time_ms', 0),
+            "middleware_source": "PROXY",
+        }
+
+        logger.debug("Calling client.ai.create_completion with args: %s", completion_args)
+        result = client.ai.create_completion(**completion_args)
+        logger.debug("Result from create_completion: %s", result)
+
+        return
+
+    async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
+        # log: key, user, model, prompt, error, tokens, cost
+        # Access kwargs passed to litellm.completion()
+        # pprint.pprint(kwargs['litellm_params'])
+
+        model = kwargs.get("model", None)
+
+        # Access litellm_params passed to litellm.completion(), example access `metadata`
+        litellm_params = kwargs.get("litellm_params", {})
+        metadata = litellm_params.get("metadata", {})  # headers passed to LiteLLM proxy, can be found here
+        headers = metadata.get("headers", {})
+
+        # For failures, we may not have usage information
+        usage = getattr(response_obj, "usage", {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0})
+        if isinstance(usage, dict) is False:
+            usage = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
+
+        error_message = str(response_obj)
+        error_type = type(response_obj).__name__
+
+        # Create subscriber object from metadata and headers
+        subscriber = {}
+
+        # Extract subscriber information from metadata and headers
+        subscriber_id = metadata.get('x-revenium-subscriber-id', '') or headers.get("x-revenium-subscriber-id")
+        subscriber_email = metadata.get('user_api_key_user_email', '')
+        credential_name = metadata.get('user_api_key_alias', '')
+        credential_value = metadata.get('user_api_key_hash', '')
+
+        if subscriber_id:
+            subscriber["id"] = subscriber_id
+        if subscriber_email:
+            subscriber["email"] = subscriber_email
+        if credential_name or credential_value:
+            subscriber["credential"] = {
+                "name": credential_name,
+                "value": credential_value
+            }
+
+        # Extract organization and product names using helper functions
+        organization_name = _extract_organization_name(headers, metadata)
+        product_name = _extract_product_name(headers)
+
+        completion_args = {
+            "cache_creation_token_count": 0,
+            "cache_read_token_count": 0,
+            "input_token_cost": None,
+            "output_token_cost": None,
+            "total_cost": None,
+            "output_token_count": usage.get("completion_tokens", 0),
+            "cost_type": "AI",
+            "model": model,
+            "input_token_count": usage.get("prompt_tokens", 0),
+            "provider": "LITELLM",
+            "model_source": "LITELLM",
+            "reasoning_token_count": 0,
+            "request_time": start_time.strftime("%Y-%m-%dT%H:%M:%SZ"),
+            "response_time": end_time.strftime("%Y-%m-%dT%H:%M:%SZ"),
+            "completion_start_time": end_time.strftime("%Y-%m-%dT%H:%M:%SZ"),
+            "request_duration": (end_time - start_time).total_seconds() * 1000,  # Convert to milliseconds
+            "time_to_first_token": (end_time - start_time).total_seconds() * 1000,
+            # Time to first token in milliseconds
+            "stop_reason": "ERROR",
+            "total_token_count": usage.get("total_tokens", 0),
+            "transaction_id": getattr(response_obj, "id", "error-no-id"),
+            "trace_id": headers.get("x-revenium-trace-id"),
+            "task_type": headers.get("x-revenium-task-type"),
+            "subscriber": subscriber if subscriber else None,
+            "organization_name": organization_name,
+            "subscription_id": headers.get("x-revenium-subscription-id"),
+            "product_name": product_name,
+            "agent": headers.get("x-revenium-agent"),
+            "response_quality_score": headers.get("x-revenium-response-quality-score"),
+            "is_streamed": metadata.get('hidden_params', {}).get('optional_params', {}).get('stream', False),
+            "operation_type": "CHAT",
+            "middleware_source": "PROXY"
+        # "error_message": error_message,
+            # "error_type": error_type,
+            # "mediation_latency": metadata['hidden_params']['litellm_overhead_time_ms'],
+        }
+
+        logger.debug("Calling client.ai.create_completion with args (failure): %s", completion_args)
+        try:
+            result = client.ai.create_completion(**completion_args)
+            logger.debug("Result from create_completion (failure): %s", result)
+        except Exception as e:
+            logger.error("Error logging failure event: %s", e)
+
+
+proxy_handler_instance = MiddlewareHandler()

revenium_middleware/ollama/__init__.py

@@ -0,0 +1,28 @@
+"""
+Revenium Middleware for Ollama Python SDK.
+
+When you install and import this library, it will automatically hook
+ollama.chat, ollama.generate, and ollama.embed using wrapt, and log token usage after
+each request. You can customize or extend this logging logic later
+to add user or organization metadata for metering purposes.
+"""
+
+import logging
+
+logger = logging.getLogger(__name__)
+
+# Conditionally import middleware (requires ollama SDK)
+try:
+    import ollama as _ollama  # noqa: F401
+    from .middleware import chat_wrapper, generate_wrapper, embed_wrapper
+except ImportError:
+    logger.debug("Ollama SDK (ollama) not available, middleware not loaded")
+    chat_wrapper = None  # type: ignore
+    generate_wrapper = None  # type: ignore
+    embed_wrapper = None  # type: ignore
+
+__all__ = [
+    "chat_wrapper",
+    "generate_wrapper",
+    "embed_wrapper",
+]