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

revenium_middleware/_core/decorators.py

@@ -0,0 +1,202 @@
+"""
+Decorators for selective metering in Revenium middleware.
+
+This module provides decorators that allow users to selectively meter
+specific functions/endpoints instead of automatically metering all API calls.
+"""
+
+import functools
+import asyncio
+from typing import Optional, Dict, Any, Callable, TypeVar
+
+from .context import (
+    set_decorated_context,
+    clear_decorated_context,
+    set_injected_metadata,
+    clear_injected_metadata,
+)
+
+# Type variables for generic decorator support
+F = TypeVar('F', bound=Callable[..., Any])
+
+
+def revenium_meter(
+    operation_type: Optional[str] = None,
+    metadata: Optional[Dict[str, Any]] = None,
+    name: Optional[str] = None
+) -> Callable[[F], F]:
+    """
+    Decorator to mark a function for selective metering.
+    
+    When selective metering is enabled in the middleware configuration,
+    only functions decorated with @revenium_meter will have their API
+    calls metered. When selective metering is disabled (default), this decorator
+    has no effect and all API calls are metered automatically.
+    
+    Args:
+        operation_type: Optional override for operation type (e.g., 'CHAT', 'EMBED')
+        metadata: Optional default metadata to include in metering calls
+        name: Optional friendly name for the function (for logging/debugging)
+    
+    Returns:
+        Decorated function that sets context when called
+    
+    Example:
+        @revenium_meter(metadata={'task_type': 'chat_analysis'})
+        def analyze_text(text: str) -> str:
+            response = client.chat.completions.create(
+                model="gpt-4o-mini",
+                messages=[{"role": "user", "content": text}]
+            )
+            return response.choices[0].message.content
+    """
+    def decorator(func: F) -> F:
+        # Determine if function is async
+        is_async = asyncio.iscoroutinefunction(func)
+        
+        if is_async:
+            @functools.wraps(func)
+            async def async_wrapper(*args, **kwargs):
+                # Prepare metadata
+                func_metadata = {
+                    'function_name': name or func.__name__,
+                    'is_decorated': True,
+                }
+                if operation_type:
+                    func_metadata['operation_type'] = operation_type
+                if metadata:
+                    func_metadata.update(metadata)
+                
+                # Set context before calling function
+                set_decorated_context(True, func_metadata)
+                try:
+                    result = await func(*args, **kwargs)
+                    return result
+                finally:
+                    # Clear context after function completes
+                    clear_decorated_context()
+            
+            return async_wrapper  # type: ignore
+        else:
+            @functools.wraps(func)
+            def sync_wrapper(*args, **kwargs):
+                # Prepare metadata
+                func_metadata = {
+                    'function_name': name or func.__name__,
+                    'is_decorated': True,
+                }
+                if operation_type:
+                    func_metadata['operation_type'] = operation_type
+                if metadata:
+                    func_metadata.update(metadata)
+                
+                # Set context before calling function
+                set_decorated_context(True, func_metadata)
+                try:
+                    result = func(*args, **kwargs)
+                    return result
+                finally:
+                    # Clear context after function completes
+                    clear_decorated_context()
+            
+            return sync_wrapper  # type: ignore
+    
+    return decorator
+
+
+# Alias for convenience
+track_usage = revenium_meter
+
+
+def revenium_metadata(**metadata_kwargs) -> Callable[[F], F]:
+    """
+    Decorator to inject metadata into all API calls within a function's scope.
+
+    This decorator automatically injects metadata into every API call made
+    within the decorated function, eliminating the need to pass usage_metadata to
+    each individual call. API-level metadata (passed directly to API calls) takes
+    precedence over decorator-injected metadata.
+
+    This decorator works independently of selective metering and can be used:
+    - With automatic metering (default behavior)
+    - With selective metering (when @revenium_meter is also used)
+    - Alongside @revenium_meter on the same function
+
+    Args:
+        **metadata_kwargs: Arbitrary keyword arguments that will be injected as metadata.
+            Common fields include:
+            - trace_id: Unique identifier for a conversation or session
+            - task_type: Classification of the AI operation
+            - organization_id: Customer or department ID
+            - subscription_id: Reference to a billing plan
+            - product_id: Your product or feature making the AI call
+            - agent: Identifier for the specific AI agent
+            - response_quality_score: Quality metric (0-1)
+            - subscriber: Nested object with user information
+
+    Returns:
+        Decorated function that injects metadata into API calls
+
+    Example:
+        @revenium_metadata(org_id="acme", task_type="analysis")
+        def analyze_documents(docs):
+            # All API calls here automatically get the metadata
+            response = client.chat.completions.create(
+                model="gpt-4o",
+                messages=[{"role": "user", "content": "Analyze this"}]
+            )
+            return response
+
+        # Can be combined with @revenium_meter for selective metering
+        @revenium_meter()
+        @revenium_metadata(org_id="acme", trace_id="session-123")
+        def metered_analysis():
+            response = client.chat.completions.create(...)
+            return response
+
+        # API-level metadata overrides decorator metadata
+        @revenium_metadata(task_type="default")
+        def mixed_metadata():
+            # Uses decorator metadata: {"task_type": "default"}
+            response1 = client.chat.completions.create(...)
+
+            # API-level overrides: {"task_type": "special"}
+            response2 = client.chat.completions.create(
+                ...,
+                usage_metadata={"task_type": "special"}
+            )
+            return response1, response2
+    """
+    def decorator(func: F) -> F:
+        # Determine if function is async
+        is_async = asyncio.iscoroutinefunction(func)
+
+        if is_async:
+            @functools.wraps(func)
+            async def async_wrapper(*args, **kwargs):
+                # Set injected metadata context before calling function
+                set_injected_metadata(metadata_kwargs)
+                try:
+                    result = await func(*args, **kwargs)
+                    return result
+                finally:
+                    # Clear injected metadata context after function completes
+                    clear_injected_metadata()
+
+            return async_wrapper  # type: ignore
+        else:
+            @functools.wraps(func)
+            def sync_wrapper(*args, **kwargs):
+                # Set injected metadata context before calling function
+                set_injected_metadata(metadata_kwargs)
+                try:
+                    result = func(*args, **kwargs)
+                    return result
+                finally:
+                    # Clear injected metadata context after function completes
+                    clear_injected_metadata()
+
+            return sync_wrapper  # type: ignore
+
+    return decorator
+

revenium_middleware/_core/metering.py

@@ -0,0 +1,207 @@
+import os
+import time
+import logging
+import asyncio
+import threading
+import atexit
+import signal
+from typing import Literal, Awaitable, Any, Optional, Callable
+from revenium_metering import ReveniumMetering
+
+# Get the logger that was configured in __init__.py
+logger = logging.getLogger("revenium_middleware")
+
+# Define a StopReason literal type for strict typing of stop_reason
+StopReason = Literal["END", "END_SEQUENCE", "TIMEOUT", "TOKEN_LIMIT", "COST_LIMIT", "COMPLETION_LIMIT", "ERROR"]
+
+api_key = os.environ.get("REVENIUM_METERING_API_KEY") or "DUMMY_API_KEY"
+client = ReveniumMetering(api_key=api_key)
+
+# Keep track of active metering threads
+active_threads = []
+shutdown_event = threading.Event()
+
+def handle_exit(signum=None, frame=None):
+    # Check if shutdown is already initiated to prevent redundant logging/actions
+    if shutdown_event.is_set():
+        return
+
+    logger.debug("Shutdown initiated, waiting for metering calls to complete...")
+    shutdown_event.set()
+
+    # Give threads a chance to notice the shutdown event
+    # Use a small delay, but avoid blocking excessively if called from signal handler
+    try:
+        time.sleep(0.1)
+    except InterruptedError:
+        # Handle potential interruption if called from a signal handler during sleep
+        logger.debug("Sleep interrupted during shutdown.")
+        # Ensure the event is still set
+        shutdown_event.set()
+
+
+    # Iterate over a copy of the list to avoid modification issues during iteration
+    threads_to_join = list(active_threads)
+    for thread in threads_to_join:
+        if thread.is_alive():
+            logger.debug(f"Waiting for metering thread {thread.name} to finish...")
+            thread.join(timeout=5.0) # Wait up to 5 seconds for the thread
+            if thread.is_alive():
+                logger.warning(f"Metering thread {thread.name} did not complete in time.")
+            else:
+                logger.debug(f"Metering thread {thread.name} finished.")
+        # Clean up thread reference if it's already finished or after joining
+        if thread in active_threads:
+             try:
+                 active_threads.remove(thread)
+             except ValueError:
+                 # Thread might have been removed by itself in the finally block
+                 pass
+
+
+    logger.debug("Shutdown complete")
+    
+    # If called from a signal handler, exit the program
+    if signum is not None:
+        logger.debug(f"Exiting due to signal {signum}")
+        os._exit(0)
+
+# Always register atexit handler, works in any thread
+atexit.register(handle_exit)
+
+# Only register signal handlers if in the main thread
+if threading.current_thread() is threading.main_thread():
+    try:
+        signal.signal(signal.SIGINT, handle_exit)
+        signal.signal(signal.SIGTERM, handle_exit)
+        logger.debug("SIGINT and SIGTERM handlers registered.")
+    except ValueError as e:
+        # This can happen in environments where signal handling is restricted (e.g., mod_wsgi)
+        logger.warning(f"Could not register signal handlers: {e}. Shutdown will rely on atexit.")
+else:
+    logger.debug("Not running in main thread, skipping signal handler registration. Shutdown will rely on atexit.")
+
+
+class MeteringThread(threading.Thread):
+    def __init__(self, coro, *args, **kwargs):
+        # Default to non-daemon threads so atexit handlers wait for them
+        daemon = kwargs.pop('daemon', False)
+        super().__init__(*args, **kwargs)
+        self.coro = coro
+        self.daemon = daemon # Store daemon status
+        self.error = None
+        self.loop = None
+        # Assign a more descriptive name if not provided
+        if self.name is None:
+             self.name = f"MeteringThread-{id(self)}"
+
+
+    def run(self):
+        # Check shutdown event *before* starting the loop
+        if shutdown_event.is_set():
+            logger.debug(f"Metering thread {self.name} not starting due to shutdown.")
+            # Ensure thread is removed from active_threads if it was added
+            if self in active_threads:
+                try:
+                    active_threads.remove(self)
+                except ValueError:
+                    pass # Should not happen if logic is correct, but handle defensively
+            return
+
+        try:
+            # Create and set a new event loop for this thread
+            self.loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(self.loop)
+            logger.debug(f"Metering thread {self.name} started with loop {id(self.loop)}")
+            try:
+                # Run the coroutine until it completes
+                self.loop.run_until_complete(self.coro)
+            finally:
+                # Ensure async generators are properly shut down
+                logger.debug(f"Shutting down async generators for loop {id(self.loop)} in thread {self.name}")
+                self.loop.run_until_complete(self.loop.shutdown_asyncgens())
+                # Close the event loop
+                logger.debug(f"Closing event loop {id(self.loop)} in thread {self.name}")
+                self.loop.close()
+                logger.debug(f"Event loop {id(self.loop)} closed in thread {self.name}")
+        except Exception as e:
+            # Log errors unless it's during shutdown
+            if not shutdown_event.is_set():
+                self.error = e
+                # Use exc_info=True to include traceback in the log
+                logger.warning(f"Error in metering thread {self.name}: {str(e)}", exc_info=True)
+            else:
+                logger.debug(f"Exception ignored in metering thread {self.name} during shutdown: {str(e)}")
+        finally:
+            # Ensure the thread is removed from the active list upon completion or error
+            if self in active_threads:
+                try:
+                    active_threads.remove(self)
+                    logger.debug(f"Removed thread {self.name} from active list.")
+                except ValueError:
+                     # Can happen if handle_exit removes it first during shutdown join timeout
+                     logger.debug(f"Thread {self.name} already removed from active list.")
+            else:
+                 # This case might indicate the thread wasn't properly added or removed elsewhere
+                 logger.warning(f"Thread {self.name} finished but was not found in active_threads list.")
+
+
+def run_async_in_thread(coroutine_or_func):
+    """
+    Helper function to run an async coroutine or a regular function in a background thread
+    with better handling of interpreter shutdown.
+
+    Args:
+        coroutine_or_func: Either an awaitable coroutine or a regular function
+
+    Returns:
+        Optional[threading.Thread]: The thread running the task, or None if shutdown initiated.
+    """
+    if shutdown_event.is_set():
+        logger.warning("Not starting new metering thread during shutdown")
+        return None
+
+    # Check if we received a coroutine or a regular function
+    if asyncio.iscoroutine(coroutine_or_func):
+        # It's a coroutine, use it directly
+        coro = coroutine_or_func
+    elif callable(coroutine_or_func):
+         # It's a callable (sync function), wrap it
+         async def wrapper():
+             # Check shutdown again before potentially long-running sync call
+             if shutdown_event.is_set():
+                 logger.debug("Skipping sync function execution due to shutdown.")
+                 return None # Or raise an exception? Returning None seems safer.
+             try:
+                 return coroutine_or_func()
+             except Exception as e:
+                 logger.warning(f"Exception in wrapped sync function: {e}", exc_info=True)
+                 # Propagate or handle error as needed
+                 raise # Re-raise the exception to be caught by MeteringThread run method
+         coro = wrapper()
+    else:
+         # If it's neither a coroutine nor callable, it's likely an error or unexpected input
+         logger.error(f"Invalid type passed to run_async_in_thread: {type(coroutine_or_func)}. Expected coroutine or callable.")
+         # Decide how to handle this: return None, raise TypeError, etc.
+         # Returning None might be safest to avoid crashing the caller unexpectedly.
+         return None
+
+
+    # Create and start the thread
+    # Pass daemon=False explicitly if that's the desired default
+    thread = MeteringThread(coro, daemon=False)
+    active_threads.append(thread)
+    logger.debug(f"Starting and adding thread {thread.name} to active list (now {len(active_threads)} threads).")
+    try:
+        thread.start()
+    except RuntimeError as e:
+        logger.error(f"Failed to start thread {thread.name}: {e}", exc_info=True)
+        # Clean up: remove the thread we failed to start
+        if thread in active_threads:
+            try:
+                active_threads.remove(thread)
+            except ValueError:
+                pass # Should be there, but handle defensively
+        return None # Indicate failure to start
+
+    return thread

revenium_middleware/_core/prompt_extraction.py

@@ -0,0 +1,55 @@
+"""
+Shared prompt extraction utilities.
+
+This module provides prompt extraction functions shared across all provider
+middlewares. Provider-specific prompt_extractor modules import from here.
+"""
+
+import logging
+from typing import Dict, Any, Optional
+
+from .config import Config
+
+logger = logging.getLogger("revenium_middleware.extension")
+
+
+def extract_streaming_response_content(
+    accumulated_content: str,
+    prompts_truncated: bool = False,
+    max_prompt_length: Optional[int] = None,
+) -> Dict[str, Any]:
+    """
+    Extract output response content from accumulated streaming chunks.
+
+    Args:
+        accumulated_content: Accumulated text from streaming response
+        prompts_truncated: Whether prompts were already truncated (from request)
+        max_prompt_length: Maximum prompt length (defaults to Config.MAX_PROMPT_LENGTH)
+
+    Returns:
+        Dict containing:
+            - outputResponse: String or None (assistant response content)
+            - promptsTruncated: Boolean (True if any field was truncated)
+    """
+    if max_prompt_length is None:
+        max_prompt_length = Config.MAX_PROMPT_LENGTH
+
+    output_response = accumulated_content if accumulated_content else None
+    was_truncated = prompts_truncated
+
+    # Apply truncation - keep total at max_prompt_length
+    if output_response and len(output_response) > max_prompt_length:
+        marker = "...[TRUNCATED]"
+        marker_len = len(marker)
+        truncate_at = max_prompt_length - marker_len
+        output_response = output_response[:truncate_at] + marker
+        was_truncated = True
+        logger.debug(
+            f"Streaming output response truncated to "
+            f"{max_prompt_length} characters"
+        )
+
+    return {
+        'outputResponse': output_response,
+        'promptsTruncated': was_truncated
+    }

revenium_middleware/_core/subscriber.py

@@ -0,0 +1,51 @@
+"""
+Shared subscriber extraction from usage metadata.
+
+This module provides the canonical implementation of subscriber extraction
+shared across all provider middlewares.
+"""
+
+
+def extract_subscriber_from_metadata(usage_metadata: dict) -> dict:
+    """
+    Extract subscriber from usage_metadata.
+
+    Supports both nested and flat patterns:
+    - Nested: usage_metadata["subscriber"]["id"], ["email"], ["credential"]
+    - Flat: usage_metadata["subscriber_id"], ["subscriber_email"], etc.
+
+    Args:
+        usage_metadata: Dictionary containing usage metadata
+
+    Returns:
+        Dictionary containing subscriber information (id, email, credential)
+    """
+    subscriber = {}
+
+    # Pattern 1: Nested subscriber object (OpenAI, Ollama, Anthropic non-Bedrock, LiteLLM)
+    if "subscriber" in usage_metadata and isinstance(usage_metadata["subscriber"], dict):
+        nested = usage_metadata["subscriber"]
+
+        if nested.get("id"):
+            subscriber["id"] = nested["id"]
+        if nested.get("email"):
+            subscriber["email"] = nested["email"]
+        if nested.get("credential") and isinstance(nested["credential"], dict):
+            subscriber["credential"] = {
+                "name": nested["credential"].get("name"),
+                "value": nested["credential"].get("value"),
+            }
+
+    # Pattern 2: Flat field pattern (Anthropic Bedrock legacy)
+    if not subscriber:
+        if usage_metadata.get("subscriber_id"):
+            subscriber["id"] = usage_metadata["subscriber_id"]
+        if usage_metadata.get("subscriber_email"):
+            subscriber["email"] = usage_metadata["subscriber_email"]
+        if usage_metadata.get("subscriber_credential_name"):
+            subscriber["credential"] = {
+                "name": usage_metadata["subscriber_credential_name"],
+                "value": usage_metadata.get("subscriber_credential"),
+            }
+
+    return subscriber