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

revenium_middleware/google/common/protocols.py

@@ -0,0 +1,192 @@
+"""
+Protocol definitions for Google AI API responses.
+
+This module defines Protocol classes that describe the expected structure
+of API responses from Google AI and Vertex AI SDKs, improving type safety
+and enabling better static analysis.
+"""
+
+from typing import Protocol, Optional, List, Any, Union
+from typing_extensions import runtime_checkable
+
+
+@runtime_checkable
+class UsageMetadataProtocol(Protocol):
+    """Protocol for usage metadata in API responses."""
+
+    # Google AI SDK attributes
+    prompt_token_count: Optional[int]
+    candidates_token_count: Optional[int]
+    total_token_count: Optional[int]
+    cached_content_token_count: Optional[int]
+
+    # Alternative attribute names (for compatibility)
+    prompt_tokens: Optional[int]
+    completion_tokens: Optional[int]
+    total_tokens: Optional[int]
+    cached_tokens: Optional[int]
+    input_tokens: Optional[int]
+    output_tokens: Optional[int]
+
+
+@runtime_checkable
+class CandidateProtocol(Protocol):
+    """Protocol for candidate objects in API responses."""
+
+    finish_reason: Optional[str]
+    content: Optional[Any]
+    text: Optional[str]
+
+
+@runtime_checkable
+class EmbeddingProtocol(Protocol):
+    """Protocol for embedding objects."""
+
+    values: List[float]
+    statistics: Optional[Any]
+    _prediction_response: Optional[Any]
+
+
+@runtime_checkable
+class EmbeddingStatisticsProtocol(Protocol):
+    """Protocol for embedding statistics."""
+
+    token_count: Optional[int]
+    billableCharacterCount: Optional[int]
+
+
+@runtime_checkable
+class ChatResponseProtocol(Protocol):
+    """Protocol for chat/generation API responses."""
+
+    # Model information
+    model: Optional[str]
+    model_name: Optional[str]
+    model_version: Optional[str]
+    _model_name: Optional[str]
+
+    # Usage information
+    usage: Optional[UsageMetadataProtocol]
+    usage_metadata: Optional[UsageMetadataProtocol]
+
+    # Response content
+    text: Optional[str]
+    content: Optional[str]
+    candidates: Optional[List[CandidateProtocol]]
+
+    # Completion information
+    finish_reason: Optional[str]
+    stop_reason: Optional[str]
+
+
+@runtime_checkable
+class EmbeddingResponseProtocol(Protocol):
+    """Protocol for embedding API responses."""
+
+    # Model information
+    model: Optional[str]
+    model_name: Optional[str]
+
+    # Embeddings data
+    embeddings: Optional[List[EmbeddingProtocol]]
+    values: Optional[List[float]]  # For single embedding responses
+
+    # Usage information
+    usage: Optional[UsageMetadataProtocol]
+    usage_metadata: Optional[UsageMetadataProtocol]
+    statistics: Optional[EmbeddingStatisticsProtocol]
+
+
+@runtime_checkable
+class StreamChunkProtocol(Protocol):
+    """Protocol for streaming response chunks."""
+
+    # Model information
+    model: Optional[str]
+    model_version: Optional[str]
+
+    # Content
+    text: Optional[str]
+    content: Optional[str]
+    candidates: Optional[List[CandidateProtocol]]
+
+    # Usage information (typically in final chunk)
+    usage_metadata: Optional[UsageMetadataProtocol]
+
+    # Completion information
+    finish_reason: Optional[str]
+
+
+@runtime_checkable
+class ClientProtocol(Protocol):
+    """Protocol for Google AI client objects."""
+
+    _vertexai: Optional[bool]
+    _api_client: Optional[Any]
+
+
+# Type guards for runtime type checking
+def is_chat_response(response: Any) -> bool:
+    """Check if response matches ChatResponseProtocol."""
+    return isinstance(response, ChatResponseProtocol)
+
+
+def is_embedding_response(response: Any) -> bool:
+    """Check if response matches EmbeddingResponseProtocol."""
+    return isinstance(response, EmbeddingResponseProtocol)
+
+
+def is_stream_chunk(chunk: Any) -> bool:
+    """Check if chunk matches StreamChunkProtocol."""
+    return isinstance(chunk, StreamChunkProtocol)
+
+
+def has_usage_metadata(response: Any) -> bool:
+    """Check if response has usage metadata."""
+    return (
+        hasattr(response, "usage_metadata") and response.usage_metadata is not None
+    ) or (hasattr(response, "usage") and response.usage is not None)
+
+
+def has_token_counts(usage: Any) -> bool:
+    """Check if usage object has token count information."""
+    if not usage:
+        return False
+
+    # Check for any token count attributes
+    token_attrs = [
+        "total_token_count",
+        "total_tokens",
+        "prompt_token_count",
+        "prompt_tokens",
+        "input_tokens",
+        "candidates_token_count",
+        "completion_tokens",
+        "output_tokens",
+    ]
+
+    return any(
+        hasattr(usage, attr) and getattr(usage, attr, 0) > 0 for attr in token_attrs
+    )
+
+
+# Utility functions for safe attribute access
+def safe_getattr(obj: Any, attr: str, default: Any = None) -> Any:
+    """Safely get attribute with fallback to default."""
+    try:
+        return getattr(obj, attr, default)
+    except (AttributeError, TypeError):
+        return default
+
+
+def get_token_count(usage: Any, attr_names: List[str]) -> int:
+    """Get token count from usage object, trying multiple attribute names."""
+    if not usage:
+        return 0
+
+    for attr_name in attr_names:
+        value = safe_getattr(usage, attr_name, 0)
+        if isinstance(value, int) and value > 0:
+            return value
+
+    return 0

revenium_middleware/google/common/summary_printer.py

@@ -0,0 +1,271 @@
+"""
+Summary printer for terminal output of API usage metrics.
+
+This module provides functionality to display cost and metrics information
+after each API request in human-readable or JSON format.
+"""
+
+import json
+import logging
+import time
+import urllib.parse
+from dataclasses import dataclass
+from typing import Optional
+
+import requests
+
+from . import trace_fields
+
+logger = logging.getLogger("revenium_middleware.extension")
+
+
+@dataclass
+class CompletionMetrics:
+    """Data class to hold cost information from the Revenium API."""
+
+    total_cost: Optional[float] = None
+
+
+def fetch_completion_metrics(
+    transaction_id: str,
+    revenium_api_key: str,
+) -> Optional[CompletionMetrics]:
+    """
+    Fetch cost data from Revenium profitstream API.
+
+    Implements retry logic (3 attempts with 1-second delay).
+
+    Args:
+        transaction_id: The transaction ID to fetch metrics for
+        revenium_api_key: The Revenium API key for authentication
+
+    Returns:
+        CompletionMetrics with cost data, or None if fetch fails
+    """
+    team_id = trace_fields.get_team_id()
+    if not team_id:
+        logger.debug("No REVENIUM_TEAM_ID set, skipping cost fetch")
+        return CompletionMetrics(total_cost=None)
+
+    base_url = trace_fields.get_base_url()
+    # Use urlencode to properly encode query parameters
+    params = urllib.parse.urlencode({
+        'teamId': team_id,
+        'transactionId': transaction_id
+    })
+    url = (
+        f"{base_url}/profitstream/v2/api/sources/metrics/ai/completions"
+        f"?{params}"
+    )
+
+    headers = {
+        "Authorization": f"Bearer {revenium_api_key}",
+        "Content-Type": "application/json",
+    }
+
+    for attempt in range(trace_fields.SUMMARY_RETRY_ATTEMPTS):
+        try:
+            response = requests.get(
+                url,
+                headers=headers,
+                timeout=trace_fields.SUMMARY_API_TIMEOUT,
+            )
+
+            if response.status_code == 200:
+                data = response.json()
+                # Check both camelCase and snake_case
+                # Don't use 'or' to avoid treating 0.0 as falsy
+                total_cost = data.get("totalCost")
+                if total_cost is None:
+                    total_cost = data.get("total_cost")
+                return CompletionMetrics(total_cost=total_cost)
+            else:
+                logger.debug(
+                    f"Metrics API returned status {response.status_code} "
+                    f"on attempt {attempt + 1}"
+                )
+
+        except requests.exceptions.RequestException as e:
+            logger.debug(
+                f"Failed to fetch metrics on attempt {attempt + 1}: {e}"
+            )
+
+        # Wait before retry (except on last attempt)
+        if attempt < trace_fields.SUMMARY_RETRY_ATTEMPTS - 1:
+            time.sleep(trace_fields.SUMMARY_RETRY_DELAY)
+
+    logger.debug(
+        f"Failed to fetch metrics after {trace_fields.SUMMARY_RETRY_ATTEMPTS} attempts"
+    )
+    return None
+
+
+def format_and_print_json_summary(
+    model: str,
+    provider: str,
+    duration_seconds: float,
+    input_token_count: Optional[int],
+    output_token_count: Optional[int],
+    total_token_count: Optional[int],
+    cost: Optional[float],
+    trace_id: Optional[str],
+) -> None:
+    """
+    Print single-line JSON output for machine-readable summary.
+
+    Args:
+        model: Model name used
+        provider: Provider name (e.g., "GOOGLE")
+        duration_seconds: Request duration in seconds
+        input_token_count: Number of input tokens
+        output_token_count: Number of output tokens
+        total_token_count: Total token count
+        cost: Cost in dollars (or None if unavailable)
+        trace_id: Trace ID for the request
+    """
+    summary = {
+        "model": model,
+        "provider": provider,
+        "durationSeconds": round(duration_seconds, 3),
+        "inputTokenCount": input_token_count,
+        "outputTokenCount": output_token_count,
+        "totalTokenCount": total_token_count,
+    }
+
+    if cost is not None:
+        summary["cost"] = cost
+        summary["costStatus"] = "available"
+    else:
+        summary["costStatus"] = "unavailable"
+
+    if trace_id:
+        summary["traceId"] = trace_id
+
+    print(json.dumps(summary))
+
+
+def format_and_print_human_summary(
+    model: str,
+    provider: str,
+    duration_seconds: float,
+    input_token_count: Optional[int],
+    output_token_count: Optional[int],
+    total_token_count: Optional[int],
+    cost: Optional[float],
+    trace_id: Optional[str],
+    team_id: Optional[str],
+) -> None:
+    """
+    Print professional human-readable output.
+
+    NO EMOJIS - professional format only.
+
+    Args:
+        model: Model name used
+        provider: Provider name (e.g., "GOOGLE")
+        duration_seconds: Request duration in seconds
+        input_token_count: Number of input tokens
+        output_token_count: Number of output tokens
+        total_token_count: Total token count
+        cost: Cost in dollars (or None if unavailable)
+        trace_id: Trace ID for the request
+        team_id: Team ID (used for cost status message)
+    """
+    separator = "=" * 60
+
+    print(separator)
+    print("REVENIUM USAGE SUMMARY")
+    print(separator)
+    print(f"Model: {model}")
+    print(f"Provider: {provider}")
+    print(f"Duration: {duration_seconds:.2f}s")
+    print()
+    print("Token Usage:")
+    print(f"  Input Tokens:  {input_token_count or 0}")
+    print(f"  Output Tokens: {output_token_count or 0}")
+    print(f"  Total Tokens:  {total_token_count or 0}")
+    print()
+
+    if cost is not None:
+        print(f"Cost: ${cost:.6f}")
+    elif team_id:
+        print("Cost: Pending (aggregating... check Revenium dashboard)")
+    else:
+        print("Cost: Add REVENIUM_TEAM_ID to see pricing")
+
+    if trace_id:
+        print()
+        print(f"Trace ID: {trace_id}")
+
+    print(separator)
+
+
+def print_usage_summary(
+    model: str,
+    provider: str,
+    request_duration: int,
+    input_token_count: Optional[int],
+    output_token_count: Optional[int],
+    total_token_count: Optional[int],
+    transaction_id: str,
+    trace_id: Optional[str],
+    revenium_api_key: str,
+) -> None:
+    """
+    Main entry point for printing usage summary.
+
+    Fire-and-forget: Wrapped in try/except to never fail the main API call.
+
+    Args:
+        model: Model name used
+        provider: Provider name (e.g., "Google")
+        request_duration: Request duration in milliseconds
+        input_token_count: Number of input tokens
+        output_token_count: Number of output tokens
+        total_token_count: Total token count
+        transaction_id: Transaction ID for fetching metrics
+        trace_id: Trace ID for the request
+        revenium_api_key: Revenium API key for authentication
+    """
+    try:
+        # Check if summary output is enabled
+        summary_format = trace_fields.get_print_summary_config()
+        if summary_format is False:
+            return
+
+        # Convert duration from milliseconds to seconds
+        duration_seconds = request_duration / 1000.0
+
+        # Attempt to fetch cost metrics
+        metrics = fetch_completion_metrics(transaction_id, revenium_api_key)
+        cost = metrics.total_cost if metrics else None
+        team_id = trace_fields.get_team_id()
+
+        if summary_format == "json":
+            format_and_print_json_summary(
+                model=model,
+                provider=provider,
+                duration_seconds=duration_seconds,
+                input_token_count=input_token_count,
+                output_token_count=output_token_count,
+                total_token_count=total_token_count,
+                cost=cost,
+                trace_id=trace_id,
+            )
+        else:  # "human"
+            format_and_print_human_summary(
+                model=model,
+                provider=provider,
+                duration_seconds=duration_seconds,
+                input_token_count=input_token_count,
+                output_token_count=output_token_count,
+                total_token_count=total_token_count,
+                cost=cost,
+                trace_id=trace_id,
+                team_id=team_id,
+            )
+
+    except Exception as e:
+        # Fire-and-forget: Never fail the main operation
+        logger.debug(f"Failed to print summary: {e}")
+

revenium_middleware/google/common/trace_fields.py

@@ -0,0 +1,205 @@
+"""
+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 Google-specific functions: detect_vision_content and detect_operation_type.
+"""
+
+import logging
+from typing import Optional, Dict, Any, Union, Literal
+
+from revenium_middleware._core.config import (
+    Config as _CoreConfig,
+    SummaryFormat,
+    parse_print_summary_value,
+    get_print_summary_config,
+    get_team_id,
+    get_base_url,
+)
+from revenium_middleware._core.trace_fields import (  # noqa: F401 — re-exported
+    TRACE_TYPE_MAX_LENGTH,
+    TRACE_NAME_MAX_LENGTH,
+    TRACE_TYPE_PATTERN,
+    ENV_REVENIUM_ENVIRONMENT,
+    ENV_ENVIRONMENT,
+    ENV_DEPLOYMENT_ENV,
+    ENV_REVENIUM_REGION,
+    ENV_AWS_REGION,
+    ENV_AWS_DEFAULT_REGION,
+    ENV_AZURE_REGION,
+    ENV_GCP_REGION,
+    ENV_GOOGLE_CLOUD_REGION,
+    ENV_REVENIUM_CREDENTIAL_ALIAS,
+    ENV_REVENIUM_TRACE_TYPE,
+    ENV_REVENIUM_TRACE_NAME,
+    ENV_REVENIUM_PARENT_TRANSACTION_ID,
+    ENV_REVENIUM_TRANSACTION_NAME,
+    ENV_REVENIUM_RETRY_NUMBER,
+    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,
+)
+
+logger = logging.getLogger(__name__)
+
+# Summary output configuration (re-exported from core)
+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
+
+SUMMARY_RETRY_ATTEMPTS = _CoreConfig.SUMMARY_RETRY_ATTEMPTS
+SUMMARY_RETRY_DELAY: float = 1.0  # Google uses 1.0s
+SUMMARY_API_TIMEOUT = _CoreConfig.SUMMARY_API_TIMEOUT
+DEFAULT_BASE_URL = _CoreConfig.DEFAULT_BASE_URL
+
+
+def detect_vision_content(contents: Any = None) -> bool:
+    """
+    Detect if contents contain vision/image content.
+
+    Google AI content formats that indicate vision input:
+    - Part with inline_data containing image/* or video/* MIME type
+    - Part with file_data containing image/* or video/* MIME type
+    - PIL Image objects passed directly
+    - google.genai types.Part with inline_data/file_data
+
+    Args:
+        contents: The contents parameter from generate_content() call.
+            Can be a string, list, Part, or Content object.
+
+    Returns:
+        True if any content contains image/video data, False otherwise.
+    """
+    if contents is None:
+        return False
+
+    # Normalize to a list for uniform processing
+    if not isinstance(contents, list):
+        contents = [contents]
+
+    for item in contents:
+        if _item_has_vision_content(item, _depth=0):
+            return True
+
+    return False
+
+
+_MAX_VISION_RECURSION_DEPTH = 10
+
+
+def _item_has_vision_content(item: Any, _depth: int = 0) -> bool:
+    """Check if a single content item contains vision data."""
+    if item is None or _depth > _MAX_VISION_RECURSION_DEPTH:
+        return False
+
+    # Check for PIL Image objects
+    try:
+        import PIL.Image
+        if isinstance(item, PIL.Image.Image):
+            return True
+    except ImportError:
+        pass
+
+    # Check for dict-based content (raw API format)
+    if isinstance(item, dict):
+        # Check inline_data
+        inline_data = item.get("inline_data")
+        if inline_data and isinstance(inline_data, dict):
+            mime_type = inline_data.get("mime_type", "")
+            if mime_type.startswith("image/") or mime_type.startswith("video/"):
+                return True
+
+        # Check file_data
+        file_data = item.get("file_data")
+        if file_data and isinstance(file_data, dict):
+            mime_type = file_data.get("mime_type", "")
+            if mime_type.startswith("image/") or mime_type.startswith("video/"):
+                return True
+
+        # Check parts list within a Content dict
+        parts = item.get("parts")
+        if parts and isinstance(parts, list):
+            for part in parts:
+                if _item_has_vision_content(part, _depth + 1):
+                    return True
+
+        # Check type field (e.g., {"type": "image", ...})
+        if item.get("type") in ("image", "video"):
+            return True
+
+        return False
+
+    # Check for google.genai types (Part objects with inline_data/file_data attributes)
+    if hasattr(item, "inline_data") and item.inline_data is not None:
+        mime_type = getattr(item.inline_data, "mime_type", "") or ""
+        if mime_type.startswith("image/") or mime_type.startswith("video/"):
+            return True
+
+    if hasattr(item, "file_data") and item.file_data is not None:
+        mime_type = getattr(item.file_data, "mime_type", "") or ""
+        if mime_type.startswith("image/") or mime_type.startswith("video/"):
+            return True
+
+    # Check Content objects with parts
+    if hasattr(item, "parts") and item.parts:
+        parts = item.parts if isinstance(item.parts, list) else [item.parts]
+        for part in parts:
+            if _item_has_vision_content(part, _depth + 1):
+                return True
+
+    # Check for bytes (raw image data passed directly)
+    if isinstance(item, bytes) and len(item) > 0:
+        # Check for common image magic bytes
+        if (item[:4] == b'\x89PNG' or  # PNG
+                item[:2] == b'\xff\xd8' or  # JPEG
+                item[:4] == b'GIF8' or  # GIF
+                (len(item) >= 12 and item[:4] == b'RIFF' and item[8:12] == b'WEBP')):  # WebP
+            return True
+
+    return False
+
+
+def detect_operation_type(
+    method_name: str,
+    request_body: Optional[Dict[str, Any]] = None
+) -> str:
+    """
+    Auto-detect operation type from method name and request.
+
+    Args:
+        method_name: API method name (e.g., 'generate_content', 'embed_content')
+        request_body: Optional request body to check for tools
+
+    Returns:
+        Operation type string ('CHAT', 'GENERATE', 'EMBED', 'TOOL_CALL')
+    """
+    request_body = request_body or {}
+
+    # Embeddings
+    if 'embed' in method_name.lower():
+        return 'EMBED'
+
+    # Content generation (chat/generate)
+    if 'generate' in method_name.lower() or 'chat' in method_name.lower():
+        # Check for tools in request
+        has_tools = request_body.get('tools')
+        if has_tools:
+            return 'TOOL_CALL'
+        return 'CHAT'
+
+    # Default fallback
+    return 'CHAT'
+
+
+# parse_print_summary_value, get_print_summary_config, get_team_id, get_base_url
+# are imported from revenium_middleware._core.config at the top of this file