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

revenium_middleware/openai/prompt_extractor.py

@@ -0,0 +1,173 @@
+"""
+Prompt extraction utilities for capturing AI prompts and responses.
+
+This module provides functions to extract and truncate prompts from OpenAI API
+requests and responses for optional storage in Revenium analytics.
+"""
+
+import json
+import logging
+from typing import Dict, Any, Optional, List
+
+from .config import Config
+from revenium_middleware._core.prompt_extraction import (  # noqa: F401 — re-exported
+    extract_streaming_response_content,
+)
+
+logger = logging.getLogger("revenium_middleware.extension")
+
+
+def extract_prompts_from_request(kwargs: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Extract system prompt and input messages from OpenAI API request.
+    
+    Args:
+        kwargs: The kwargs dict passed to the OpenAI API call
+        
+    Returns:
+        Dict containing:
+            - systemPrompt: String or None (system message content)
+            - inputMessages: JSON string or None (non-system messages)
+            - promptsTruncated: Boolean (True if any field was truncated)
+    """
+    messages = kwargs.get('messages', [])
+    
+    if not messages:
+        return {
+            'systemPrompt': None,
+            'inputMessages': None,
+            'promptsTruncated': False
+        }
+    
+    system_prompt = None
+    user_messages = []
+    prompts_truncated = False
+    
+    # Separate system message from other messages
+    for msg in messages:
+        if not isinstance(msg, dict):
+            continue
+            
+        role = msg.get('role')
+        if role == 'system':
+            # Extract system prompt (take first system message if multiple)
+            if system_prompt is None:
+                content = msg.get('content', '')
+                if isinstance(content, str):
+                    system_prompt = content
+                elif isinstance(content, list):
+                    # Handle content as array (multimodal)
+                    system_prompt = json.dumps(content)
+        else:
+            # Collect all non-system messages
+            user_messages.append(msg)
+    
+    # Apply truncation to system prompt
+    # Keep total length at MAX_PROMPT_LENGTH by subtracting marker length
+    if system_prompt and len(system_prompt) > Config.MAX_PROMPT_LENGTH:
+        marker = "...[TRUNCATED]"
+        marker_len = len(marker)
+        truncate_at = Config.MAX_PROMPT_LENGTH - marker_len
+        system_prompt = system_prompt[:truncate_at] + marker
+        prompts_truncated = True
+        logger.debug(f"System prompt truncated to {Config.MAX_PROMPT_LENGTH} characters")
+    
+    # Convert user messages to JSON string
+    input_messages = None
+    if user_messages:
+        try:
+            # First, truncate individual message contents to avoid invalid JSON
+            # when the total exceeds the limit
+            marker = "...[TRUNCATED]"
+            marker_len = len(marker)
+            truncated_messages = []
+            for msg in user_messages:
+                truncated_msg = msg.copy()
+                content = msg.get('content', '')
+                if isinstance(content, str) and len(content) > Config.MAX_PROMPT_LENGTH // 2:
+                    # Truncate individual messages to half the limit to be safe
+                    truncate_at = (Config.MAX_PROMPT_LENGTH // 2) - marker_len
+                    truncated_msg['content'] = content[:truncate_at] + marker
+                    prompts_truncated = True
+                truncated_messages.append(truncated_msg)
+
+            input_messages = json.dumps(truncated_messages, ensure_ascii=False)
+
+            # Apply final truncation if still too long (keeps valid JSON structure)
+            if len(input_messages) > Config.MAX_PROMPT_LENGTH:
+                truncate_at = Config.MAX_PROMPT_LENGTH - marker_len
+                input_messages = input_messages[:truncate_at] + marker
+                prompts_truncated = True
+                logger.debug(
+                    f"Input messages truncated to "
+                    f"{Config.MAX_PROMPT_LENGTH} characters"
+                )
+        except (TypeError, ValueError) as e:
+            logger.warning(f"Failed to serialize input messages to JSON: {e}")
+            input_messages = None
+    
+    return {
+        'systemPrompt': system_prompt,
+        'inputMessages': input_messages,
+        'promptsTruncated': prompts_truncated
+    }
+
+
+def extract_response_content(response: Any, prompts_truncated: bool = False) -> Dict[str, Any]:
+    """
+    Extract output response content from OpenAI API response.
+    
+    Args:
+        response: OpenAI API response object (ChatCompletion or similar)
+        prompts_truncated: Whether prompts were already truncated (from request)
+        
+    Returns:
+        Dict containing:
+            - outputResponse: String or None (assistant response content)
+            - promptsTruncated: Boolean (True if any field was truncated)
+    """
+    output_response = None
+    was_truncated = prompts_truncated
+    
+    try:
+        # Extract content from response.choices[0].message.content
+        if hasattr(response, 'choices') and response.choices and len(response.choices) > 0:
+            first_choice = response.choices[0]
+            
+            if hasattr(first_choice, 'message') and hasattr(first_choice.message, 'content'):
+                content = first_choice.message.content
+                
+                if content:
+                    if isinstance(content, str):
+                        output_response = content
+                    else:
+                        # Handle non-string content (e.g., structured output)
+                        output_response = json.dumps(content, ensure_ascii=False)
+                    
+                    # Apply truncation - keep total at MAX_PROMPT_LENGTH
+                    if len(output_response) > Config.MAX_PROMPT_LENGTH:
+                        marker = "...[TRUNCATED]"
+                        marker_len = len(marker)
+                        truncate_at = Config.MAX_PROMPT_LENGTH - marker_len
+                        output_response = (
+                            output_response[:truncate_at]
+                            + marker
+                        )
+                        was_truncated = True
+                        logger.debug(
+                            f"Output response truncated to "
+                            f"{Config.MAX_PROMPT_LENGTH} characters"
+                        )
+    except Exception as e:
+        logger.warning(f"Failed to extract response content: {e}")
+        output_response = None
+    
+    return {
+        'outputResponse': output_response,
+        'promptsTruncated': was_truncated
+    }
+
+
+
+# extract_streaming_response_content is imported from _core.prompt_extraction
+

revenium_middleware/openai/provider.py

@@ -0,0 +1,170 @@
+"""
+Provider detection and configuration for Azure OpenAI support.
+
+This module handles detection of Azure OpenAI vs standard OpenAI based on:
+1. Environment variables (AZURE_OPENAI_ENDPOINT)
+2. Base URL substring matching ("azure")
+3. Client instance type detection
+
+The detection is lazy-loaded and only impacts Azure users.
+"""
+
+import os
+import logging
+import threading
+from enum import Enum, auto
+from typing import Optional, Any
+
+logger = logging.getLogger("revenium_middleware.extension")
+
+
+class Provider(Enum):
+    """Supported AI providers."""
+    OPENAI = auto()
+    AZURE_OPENAI = auto()
+    OLLAMA = auto()  # Existing provider support
+
+
+def detect_provider(client: Optional[Any] = None, base_url: Optional[str] = None) -> Provider:
+    """
+    Detect which AI provider is being used based on available information.
+
+    Detection priority:
+    1. Client instance type (AzureOpenAI) - most reliable
+    2. Base URL substring matching ("azure")
+    3. Environment variables (AZURE_OPENAI_ENDPOINT) - only if client suggests Azure
+    4. Default to OpenAI
+
+    Args:
+        client: OpenAI client instance (may be AzureOpenAI)
+        base_url: Base URL for API calls
+
+    Returns:
+        Provider enum indicating detected provider
+    """
+    logger.debug("Detecting AI provider...")
+
+    # 1. Check client instance type first (most reliable)
+    if client and hasattr(client, '__class__'):
+        client_class_name = client.__class__.__name__
+        if "Azure" in client_class_name:
+            logger.debug(f"Azure provider detected via client type: {client_class_name}")
+            return Provider.AZURE_OPENAI
+    
+    # 2. Check base URL for Azure substring (broader than just azure.com)
+    if base_url and "azure" in str(base_url).lower():
+        logger.debug(f"Azure provider detected via base_url substring: {base_url}")
+        return Provider.AZURE_OPENAI
+
+    # 3. Check for client base_url if not provided directly
+    if client and hasattr(client, 'base_url') and client.base_url:
+        if "azure" in str(client.base_url).lower():
+            logger.debug(f"Azure provider detected via client.base_url: {client.base_url}")
+            return Provider.AZURE_OPENAI
+
+    # 4. Check for OLLAMA via base URL patterns
+    if base_url and ("localhost:11434" in str(base_url) or "ollama" in str(base_url).lower()):
+        logger.debug(f"OLLAMA provider detected via base_url: {base_url}")
+        return Provider.OLLAMA
+
+    if client and hasattr(client, 'base_url') and client.base_url:
+        client_url = str(client.base_url).lower()
+        if "localhost:11434" in client_url or "ollama" in client_url:
+            logger.debug(f"OLLAMA provider detected via client.base_url: {client.base_url}")
+            return Provider.OLLAMA
+
+    # 5. Check environment variables only if we have some indication this might be Azure
+    # (This prevents false positives when both Azure and OpenAI configs are present)
+    if client and hasattr(client, '__class__') and "Azure" in str(type(client)):
+        azure_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT")
+        if azure_endpoint:
+            logger.debug(f"Azure provider detected via AZURE_OPENAI_ENDPOINT: {azure_endpoint}")
+            return Provider.AZURE_OPENAI
+
+    # 6. Default to OpenAI
+    logger.debug("Defaulting to OpenAI provider")
+    return Provider.OPENAI
+
+
+def get_provider_metadata(provider: Provider) -> dict:
+    """
+    Get provider-specific metadata for usage records.
+    
+    Args:
+        provider: Detected provider
+        
+    Returns:
+        Dictionary with provider and model_source fields
+    """
+    if provider == Provider.AZURE_OPENAI:
+        return {
+            "provider": "Azure",
+            "model_source": "OPENAI"
+        }
+    elif provider == Provider.OLLAMA:
+        return {
+            "provider": "OLLAMA", 
+            "model_source": "OLLAMA"
+        }
+    else:  # OPENAI
+        return {
+            "provider": "OPENAI",
+            "model_source": "OPENAI"
+        }
+
+
+def is_azure_provider(provider: Provider) -> bool:
+    """
+    Check if the provider is Azure OpenAI.
+    
+    Args:
+        provider: Provider to check
+        
+    Returns:
+        True if Azure OpenAI, False otherwise
+    """
+    return provider == Provider.AZURE_OPENAI
+
+
+# Global provider cache to avoid repeated detection
+_detected_provider: Optional[Provider] = None
+_provider_detection_attempted: bool = False
+_provider_lock = threading.Lock()  # Thread safety for global state
+
+
+def get_or_detect_provider(client: Optional[Any] = None, base_url: Optional[str] = None,
+                          force_redetect: bool = False) -> Provider:
+    """
+    Get cached provider or detect if not already done.
+
+    This provides lazy loading - detection only happens when needed and is cached.
+    Thread-safe implementation prevents race conditions in multi-threaded environments.
+
+    Args:
+        client: OpenAI client instance
+        base_url: Base URL for API calls
+        force_redetect: Force re-detection even if cached
+
+    Returns:
+        Detected provider
+    """
+    global _detected_provider, _provider_detection_attempted
+
+    # Thread-safe provider detection with double-checked locking pattern
+    if force_redetect or not _provider_detection_attempted:
+        with _provider_lock:
+            # Double-check inside the lock to prevent race conditions
+            if force_redetect or not _provider_detection_attempted:
+                _detected_provider = detect_provider(client, base_url)
+                _provider_detection_attempted = True
+                logger.debug(f"Provider detection completed: {_detected_provider}")
+
+    return _detected_provider
+
+
+def reset_provider_cache():
+    """Reset provider detection cache. Useful for testing. Thread-safe implementation."""
+    global _detected_provider, _provider_detection_attempted
+    with _provider_lock:
+        _detected_provider = None
+        _provider_detection_attempted = False

revenium_middleware/openai/summary_printer.py

@@ -0,0 +1,292 @@
+"""
+Terminal summary printer for Revenium OpenAI middleware.
+
+This module provides functionality to print cost/metrics summaries to the terminal
+after each API request. Supports both human-readable and JSON output formats.
+
+Fetches cost data from Revenium's traces API and formats for console display.
+"""
+
+import json
+import logging
+import time
+from typing import Optional, Dict, Any
+from urllib.parse import urlencode
+from urllib.request import Request, urlopen
+from urllib.error import URLError, HTTPError
+
+from .config import (
+    Config,
+    SummaryFormat,
+    get_print_summary_config,
+    get_team_id,
+    get_base_url,
+)
+
+logger = logging.getLogger("revenium_middleware.summary_printer")
+
+
+class CompletionMetrics:
+    """Completion metrics from Revenium API."""
+
+    def __init__(self, total_cost: Optional[float] = None):
+        self.total_cost = total_cost
+
+
+def fetch_completion_metrics(
+    transaction_id: str,
+    revenium_api_key: str,
+) -> Optional[CompletionMetrics]:
+    """
+    Fetch metrics from Revenium completions API.
+
+    Args:
+        transaction_id: The transaction ID to fetch metrics for
+        revenium_api_key: Revenium API key for authentication
+
+    Returns:
+        CompletionMetrics if successful, None otherwise
+    """
+    team_id = get_team_id()
+    if not team_id:
+        logger.debug(
+            "Team ID not configured, skipping cost retrieval for summary"
+        )
+        return None
+
+    base_url = get_base_url().rstrip("/")
+    # Note: profitstream API uses a different path structure than the metering API
+    url = f"{base_url}/profitstream/v2/api/sources/metrics/ai/completions"
+    params = {
+        "teamId": team_id,
+        "transactionId": transaction_id,
+    }
+    url_with_params = f"{url}?{urlencode(params)}"
+
+    logger.debug(f"Fetching completion metrics from {url_with_params}")
+
+    max_retries = Config.SUMMARY_RETRY_ATTEMPTS
+    retry_delay = Config.SUMMARY_RETRY_DELAY
+
+    for attempt in range(max_retries):
+        try:
+            request = Request(url_with_params)
+            request.add_header("Authorization", revenium_api_key)
+            request.add_header("Content-Type", "application/json")
+
+            with urlopen(request, timeout=Config.SUMMARY_API_TIMEOUT) as response:
+                if response.status == 200:
+                    data = json.loads(response.read().decode())
+                    embedded = data.get("_embedded", {})
+                    metrics_list = embedded.get("aICompletionMetricResourceList", [])
+
+                    if metrics_list:
+                        first_metric = metrics_list[0]
+                        total_cost = first_metric.get("totalCost")
+                        logger.debug(f"Retrieved cost: {total_cost}")
+                        return CompletionMetrics(total_cost=total_cost)
+
+                    logger.debug(
+                        f"No metrics found yet (attempt {attempt + 1}/{max_retries})"
+                    )
+
+            if attempt < max_retries - 1:
+                logger.debug(
+                    f"Waiting for metrics to aggregate "
+                    f"(attempt {attempt + 1}/{max_retries})..."
+                )
+                time.sleep(retry_delay)
+
+        except (HTTPError, URLError) as e:
+            logger.debug(
+                f"Failed to fetch trace metrics: {e} "
+                f"(attempt {attempt + 1}/{max_retries})"
+            )
+            if attempt < max_retries - 1:
+                time.sleep(retry_delay)
+        except Exception as e:
+            logger.debug(
+                f"Unexpected error fetching metrics: {e} "
+                f"(attempt {attempt + 1}/{max_retries})"
+            )
+            if attempt < max_retries - 1:
+                time.sleep(retry_delay)
+
+    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],
+    trace_id: Optional[str],
+    metrics: Optional[CompletionMetrics],
+) -> None:
+    """
+    Format and print summary in JSON format.
+
+    Args:
+        model: Model name
+        provider: Provider name
+        duration_seconds: Request duration in seconds
+        input_token_count: Input token count
+        output_token_count: Output token count
+        total_token_count: Total token count
+        trace_id: Trace ID
+        metrics: Optional completion metrics from API
+    """
+    team_id = get_team_id()
+
+    summary: Dict[str, Any] = {
+        "model": model,
+        "provider": provider,
+        "durationSeconds": round(duration_seconds, 2),
+        "inputTokenCount": input_token_count,
+        "outputTokenCount": output_token_count,
+        "totalTokenCount": total_token_count,
+        "cost": metrics.total_cost if metrics and metrics.total_cost is not None else None,
+    }
+
+    # Add cost status if cost is null
+    if summary["cost"] is None:
+        summary["costStatus"] = "pending" if team_id else "unavailable"
+
+    # Add trace ID if present
+    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],
+    trace_id: Optional[str],
+    metrics: Optional[CompletionMetrics],
+) -> None:
+    """
+    Format and print summary in human-readable format.
+
+    Args:
+        model: Model name
+        provider: Provider name
+        duration_seconds: Request duration in seconds
+        input_token_count: Input token count
+        output_token_count: Output token count
+        total_token_count: Total token count
+        trace_id: Trace ID
+        metrics: Optional completion metrics from API
+    """
+    team_id = get_team_id()
+
+    print("=" * 60)
+    print("REVENIUM USAGE SUMMARY")
+    print("=" * 60)
+    print(f"Model: {model}")
+    print(f"Provider: {provider}")
+    print(f"Duration: {duration_seconds:.2f}s")
+
+    print("\nToken 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):,}")
+
+    if metrics and metrics.total_cost is not None:
+        print(f"\nCost: ${metrics.total_cost:.6f}")
+    else:
+        if team_id:
+            print(
+                "\nCost: Pending (aggregating... check Revenium dashboard)"
+            )
+        else:
+            print(
+                "\nCost: Add REVENIUM_TEAM_ID to see pricing "
+                "(find your team ID in the Revenium web app)"
+            )
+
+    if trace_id:
+        print(f"\nTrace ID: {trace_id}")
+
+    print("=" * 60 + "\n")
+
+
+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: Optional[str],
+    trace_id: Optional[str],
+    revenium_api_key: str,
+) -> None:
+    """
+    Print usage summary to console (fire-and-forget).
+
+    This function is called after tracking is complete to optionally display
+    a formatted cost/metrics summary in the terminal.
+
+    Args:
+        model: Model name
+        provider: Provider name
+        request_duration: Request duration in milliseconds
+        input_token_count: Input token count
+        output_token_count: Output token count
+        total_token_count: Total token count
+        transaction_id: Transaction ID for fetching metrics
+        trace_id: Trace ID for display
+        revenium_api_key: Revenium API key for authentication
+    """
+    print_summary = get_print_summary_config()
+    if not print_summary:
+        return
+
+    # Determine format
+    format_type: SummaryFormat = "human" if print_summary is True else print_summary
+
+    duration_seconds = request_duration / 1000.0
+
+    # Fetch metrics if team_id and transaction_id are available
+    metrics: Optional[CompletionMetrics] = None
+    team_id = get_team_id()
+    if team_id and transaction_id:
+        try:
+            metrics = fetch_completion_metrics(transaction_id, revenium_api_key)
+        except Exception as e:
+            logger.debug(f"Failed to fetch metrics: {e}")
+
+    # Print summary in the appropriate format
+    try:
+        if format_type == "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,
+                trace_id=trace_id,
+                metrics=metrics,
+            )
+        else:
+            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,
+                trace_id=trace_id,
+                metrics=metrics,
+            )
+    except Exception as e:
+        logger.debug(f"Failed to format and print summary: {e}")
+