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

revenium_middleware/anthropic/prompt_extractor.py

@@ -0,0 +1,178 @@
+"""
+Prompt extraction utilities for capturing AI prompts and responses.
+
+This module provides functions to extract and truncate prompts from Anthropic 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 Anthropic API request.
+    
+    Args:
+        kwargs: The kwargs dict passed to the Anthropic API call
+        
+    Returns:
+        Dict containing:
+            - systemPrompt: String or None (system message content)
+            - inputMessages: JSON string or None (user/assistant messages)
+            - promptsTruncated: Boolean (True if any field was truncated)
+    """
+    # Extract system prompt (Anthropic has separate 'system' parameter)
+    system_prompt = kwargs.get('system')
+    messages = kwargs.get('messages', [])
+    prompts_truncated = False
+    
+    # Apply truncation to system prompt
+    if system_prompt:
+        if isinstance(system_prompt, list):
+            # System can be an array of text blocks
+            try:
+                system_prompt = json.dumps(system_prompt, ensure_ascii=False)
+            except (TypeError, ValueError) as e:
+                logger.warning(f"Failed to serialize system prompt to JSON: {e}")
+                system_prompt = str(system_prompt)
+        
+        if 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 messages to JSON string
+    input_messages = None
+    if messages:
+        try:
+            # First, truncate individual message contents to avoid invalid JSON
+            marker = "...[TRUNCATED]"
+            marker_len = len(marker)
+            truncated_messages = []
+            
+            for msg in messages:
+                if not isinstance(msg, dict):
+                    continue
+                    
+                truncated_msg = msg.copy()
+                content = msg.get('content', '')
+                
+                # Handle content as string
+                if isinstance(content, str) and len(content) > Config.MAX_PROMPT_LENGTH // 2:
+                    truncate_at = (Config.MAX_PROMPT_LENGTH // 2) - marker_len
+                    truncated_msg['content'] = content[:truncate_at] + marker
+                    prompts_truncated = True
+                # Handle content as array (multimodal - text, images, etc.)
+                elif isinstance(content, list):
+                    truncated_content = []
+                    for block in content:
+                        if isinstance(block, dict):
+                            truncated_block = block.copy()
+                            # Truncate text blocks
+                            if block.get('type') == 'text':
+                                text = block.get('text', '')
+                                if len(text) > Config.MAX_PROMPT_LENGTH // 2:
+                                    truncate_at = (Config.MAX_PROMPT_LENGTH // 2) - marker_len
+                                    truncated_block['text'] = text[:truncate_at] + marker
+                                    prompts_truncated = True
+                            truncated_content.append(truncated_block)
+                        else:
+                            truncated_content.append(block)
+                    truncated_msg['content'] = truncated_content
+                
+                truncated_messages.append(truncated_msg)
+
+            input_messages = json.dumps(truncated_messages, ensure_ascii=False)
+
+            # Apply final truncation if still too long
+            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 Anthropic API response.
+    
+    Args:
+        response: Anthropic API response object (Message)
+        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.content (array of content blocks)
+        if hasattr(response, 'content') and response.content:
+            content_blocks = []
+
+            for block in response.content:
+                if hasattr(block, 'type'):
+                    if block.type == 'text' and hasattr(block, 'text'):
+                        content_blocks.append(block.text)
+                    else:
+                        # Handle other content types (tool_use, etc.)
+                        try:
+                            content_blocks.append(json.dumps(block.__dict__, ensure_ascii=False))
+                        except:
+                            content_blocks.append(str(block))
+
+            # Join all content blocks
+            if content_blocks:
+                output_response = '\n'.join(content_blocks)
+
+                # Apply truncation
+                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/anthropic/provider.py

@@ -0,0 +1,141 @@
+"""
+Provider detection and configuration for AWS Bedrock support.
+
+This module handles detection of AWS Bedrock vs standard Anthropic based on:
+1. Client instance type (boto3 bedrock-runtime client)
+2. Base URL substring matching ("amazonaws.com")
+3. Default to Anthropic
+
+The detection is simple and focused on the MVP requirements.
+"""
+
+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."""
+    ANTHROPIC = auto()
+    BEDROCK = auto()
+
+
+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 (boto3 bedrock-runtime) - most reliable
+    2. Base URL substring matching ("amazonaws.com")
+    3. Default to Anthropic
+
+    Args:
+        client: Client instance (may be boto3 bedrock-runtime client)
+        base_url: Base URL for API calls
+
+    Returns:
+        Provider enum indicating detected provider
+    """
+    logger.debug("Detecting AI provider...")
+
+    # 1. Check if client is boto3 bedrock-runtime client (most reliable)
+    if client and hasattr(client, "meta"):
+        try:
+            if hasattr(client.meta, "service_model") and \
+               client.meta.service_model.service_name == "bedrock-runtime":
+                logger.debug("Bedrock provider detected via boto3 client service_name")
+                return Provider.BEDROCK
+        except AttributeError:
+            # If meta doesn't have service_model, continue to next check
+            pass
+    
+    # 2. Check base URL for AWS substring
+    if base_url and "amazonaws.com" in str(base_url).lower():
+        logger.debug(f"Bedrock provider detected via base_url: {base_url}")
+        return Provider.BEDROCK
+
+    # 3. Check for client base_url if not provided directly
+    if client and hasattr(client, 'base_url') and client.base_url:
+        if "amazonaws.com" in str(client.base_url).lower():
+            logger.debug(f"Bedrock provider detected via client.base_url: {client.base_url}")
+            return Provider.BEDROCK
+
+    # 4. Default to Anthropic
+    logger.debug("Defaulting to Anthropic provider")
+    return Provider.ANTHROPIC
+
+
+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.BEDROCK:
+        return {
+            "provider": "AWS",
+            "model_source": "ANTHROPIC"
+        }
+    else:  # ANTHROPIC
+        return {
+            "provider": "ANTHROPIC",
+            "model_source": "ANTHROPIC"
+        }
+
+
+def is_bedrock_provider(provider: Provider) -> bool:
+    """
+    Check if the provider is AWS Bedrock.
+    
+    Args:
+        provider: Provider to check
+        
+    Returns:
+        True if AWS Bedrock, False otherwise
+    """
+    return provider == Provider.BEDROCK
+
+
+# Thread-local storage for provider cache to ensure thread safety
+_thread_local = threading.local()
+
+
+def _get_thread_cache():
+    """Get thread-local cache, initializing if necessary."""
+    if not hasattr(_thread_local, 'detected_provider'):
+        _thread_local.detected_provider = None
+        _thread_local.detection_attempted = False
+    return _thread_local
+
+
+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
+    per thread for thread safety.
+
+    Args:
+        client: Client instance
+        base_url: Base URL for API calls
+        force_redetect: Force re-detection even if cached
+
+    Returns:
+        Detected provider
+    """
+    cache = _get_thread_cache()
+
+    if force_redetect or not cache.detection_attempted:
+        cache.detected_provider = detect_provider(client, base_url)
+        cache.detection_attempted = True
+        logger.debug(f"Provider detection completed: {cache.detected_provider}")
+
+    return cache.detected_provider

revenium_middleware/anthropic/summary_printer.py

@@ -0,0 +1,286 @@
+"""
+Summary printer module for Revenium Anthropic middleware.
+
+This module provides terminal summary output functionality that displays
+cost and metrics information after each API request in either JSON or
+human-readable format.
+"""
+
+import json
+import logging
+import time
+from dataclasses import dataclass
+from typing import Optional
+import urllib.request
+import urllib.error
+
+from .config import (
+    Config,
+    get_base_url,
+    get_print_summary_config,
+    get_team_id,
+)
+
+logger = logging.getLogger("revenium_middleware.summary_printer")
+
+
+@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 2-second delay).
+
+    Args:
+        transaction_id: The transaction ID to fetch metrics for
+        revenium_api_key: The Revenium API key for authentication
+
+    Returns:
+        CompletionMetrics if successful, None if fetch fails or cost not available
+    """
+    team_id = get_team_id()
+    if not team_id:
+        logger.debug("No team ID configured, skipping cost fetch")
+        return None
+
+    base_url = get_base_url()
+    url = (
+        f"{base_url}/profitstream/v2/api/sources/metrics/ai/completions"
+        f"?teamId={team_id}&transactionId={transaction_id}"
+    )
+
+    headers = {
+        "Authorization": f"Bearer {revenium_api_key}",
+        "Content-Type": "application/json",
+    }
+
+    for attempt in range(Config.SUMMARY_RETRY_ATTEMPTS):
+        try:
+            request = urllib.request.Request(url, headers=headers, method="GET")
+            with urllib.request.urlopen(
+                request, timeout=Config.SUMMARY_API_TIMEOUT
+            ) as response:
+                data = json.loads(response.read().decode("utf-8"))
+
+                # Extract cost from response
+                total_cost = data.get("totalCost")
+                if total_cost is not None:
+                    return CompletionMetrics(total_cost=float(total_cost))
+
+                # Cost might be in a nested structure
+                if isinstance(data, dict):
+                    # Try common field names
+                    for field in ["cost", "total_cost", "totalCost"]:
+                        if field in data and data[field] is not None:
+                            return CompletionMetrics(total_cost=float(data[field]))
+
+                return CompletionMetrics(total_cost=None)
+
+        except urllib.error.HTTPError as e:
+            logger.debug(f"HTTP error fetching metrics (attempt {attempt + 1}): {e}")
+        except urllib.error.URLError as e:
+            logger.debug(f"URL error fetching metrics (attempt {attempt + 1}): {e}")
+        except json.JSONDecodeError as e:
+            logger.debug(f"JSON decode error (attempt {attempt + 1}): {e}")
+        except Exception as e:
+            logger.debug(f"Error fetching metrics (attempt {attempt + 1}): {e}")
+
+        # Wait before retry (except on last attempt)
+        if attempt < Config.SUMMARY_RETRY_ATTEMPTS - 1:
+            time.sleep(Config.SUMMARY_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],
+    cost: Optional[float],
+    cost_status: str,
+    trace_id: Optional[str] = None,
+) -> None:
+    """
+    Print single-line JSON output summary.
+
+    Args:
+        model: The model name used
+        provider: The provider name (e.g., ANTHROPIC, BEDROCK)
+        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: Total cost if available
+        cost_status: Status of cost retrieval
+        trace_id: Optional trace ID
+    """
+    summary = {
+        "model": model,
+        "provider": provider,
+        "durationSeconds": round(duration_seconds, 3),
+        "inputTokenCount": input_token_count,
+        "outputTokenCount": output_token_count,
+        "totalTokenCount": total_token_count,
+        "cost": cost,
+        "costStatus": cost_status,
+    }
+
+    if trace_id:
+        summary["traceId"] = trace_id
+
+    print(json.dumps(summary, separators=(",", ":")))
+
+
+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],
+    cost_status: str,
+    trace_id: Optional[str] = None,
+) -> None:
+    """
+    Print professional human-readable output summary.
+
+    NO EMOJIS - professional format only.
+
+    Args:
+        model: The model name used
+        provider: The provider name (e.g., ANTHROPIC, BEDROCK)
+        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: Total cost if available
+        cost_status: Status of cost retrieval
+        trace_id: Optional trace ID
+    """
+    separator = "=" * 60
+
+    lines = [
+        separator,
+        "REVENIUM USAGE SUMMARY",
+        separator,
+        f"Model: {model}",
+        f"Provider: {provider}",
+        f"Duration: {duration_seconds:.2f}s",
+        "",
+        "Token Usage:",
+        f"  Input Tokens:  {input_token_count if input_token_count is not None else 'N/A'}",
+        f"  Output Tokens: {output_token_count if output_token_count is not None else 'N/A'}",
+        f"  Total Tokens:  {total_token_count if total_token_count is not None else 'N/A'}",
+        "",
+    ]
+
+    # Add cost information
+    if cost is not None:
+        lines.append(f"Cost: ${cost:.6f}")
+    else:
+        lines.append(f"Cost: {cost_status}")
+
+    # Add trace ID if available
+    if trace_id:
+        lines.append("")
+        lines.append(f"Trace ID: {trace_id}")
+
+    lines.append(separator)
+
+    print("\n".join(lines))
+
+
+def print_usage_summary(
+    model: str,
+    provider: str,
+    request_duration: float,
+    input_token_count: Optional[int],
+    output_token_count: Optional[int],
+    total_token_count: Optional[int],
+    transaction_id: str,
+    trace_id: Optional[str] = None,
+    revenium_api_key: Optional[str] = None,
+) -> None:
+    """
+    Main entry point for printing usage summary.
+
+    This is a fire-and-forget function - it never raises exceptions to the caller.
+
+    Args:
+        model: The model name used
+        provider: The provider name (e.g., ANTHROPIC, BEDROCK)
+        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: The transaction ID for fetching cost
+        trace_id: Optional trace ID
+        revenium_api_key: The Revenium API key for fetching cost
+    """
+    try:
+        # Check if summary printing is enabled
+        summary_format = get_print_summary_config()
+        if summary_format is False:
+            return
+
+        # Convert duration from milliseconds to seconds
+        duration_seconds = request_duration / 1000.0
+
+        # Determine cost status and fetch cost if possible
+        cost: Optional[float] = None
+        cost_status: str = "unavailable"
+        team_id = get_team_id()
+
+        if revenium_api_key and team_id:
+            metrics = fetch_completion_metrics(transaction_id, revenium_api_key)
+            if metrics and metrics.total_cost is not None:
+                cost = metrics.total_cost
+                cost_status = "available"
+            else:
+                cost_status = "Pending (aggregating... check Revenium dashboard)"
+        elif not team_id:
+            cost_status = "Add REVENIUM_TEAM_ID to see pricing"
+
+        # Print in the appropriate format
+        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,
+                cost_status=cost_status,
+                trace_id=trace_id,
+            )
+        else:
+            # Default to human format
+            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,
+                cost_status=cost_status,
+                trace_id=trace_id,
+            )
+
+    except Exception as e:
+        # Fire-and-forget: never fail the main API call
+        logger.debug(f"Failed to print summary: {e}")