auditi 0.1.0__py3-none-any.whl

auditi/providers/google.py

@@ -0,0 +1,182 @@
+"""
+Google Gemini provider implementation for usage extraction and cost calculation.
+"""
+
+from typing import Optional, Any, Dict, Tuple
+from .base import BaseProvider
+
+
+def _coerce_int(value: Any) -> Optional[int]:
+    """Helper to safely convert values to int."""
+    if value is None:
+        return None
+    try:
+        return int(value)
+    except (TypeError, ValueError):
+        return None
+
+
+class GoogleProvider(BaseProvider):
+    """Provider implementation for Google Gemini models."""
+
+    @property
+    def name(self) -> str:
+        return "google"
+
+    @property
+    def model_pricing(self) -> Dict[str, Tuple[float, float]]:
+        """
+        Google Gemini model pricing per 1M tokens (input, output) in USD.
+        Updated as of January 2025.
+
+        Note: Pricing may vary by context window size and features.
+        """
+        return {
+            # Gemini 2.0 family
+            "gemini-2.0-flash-exp": (0.00, 0.00),  # Free during preview
+            "gemini-2.0-flash": (0.10, 0.40),
+            "gemini-2.0-flash-thinking-exp": (0.00, 0.00),  # Free during preview
+            # Gemini 1.5 family
+            "gemini-1.5-pro": (1.25, 5.00),
+            "gemini-1.5-pro-001": (1.25, 5.00),
+            "gemini-1.5-pro-002": (1.25, 5.00),
+            "gemini-1.5-pro-latest": (1.25, 5.00),
+            "gemini-1.5-flash": (0.075, 0.30),
+            "gemini-1.5-flash-001": (0.075, 0.30),
+            "gemini-1.5-flash-002": (0.075, 0.30),
+            "gemini-1.5-flash-latest": (0.075, 0.30),
+            "gemini-1.5-flash-8b": (0.0375, 0.15),
+            # Gemini 1.0 family (legacy)
+            "gemini-1.0-pro": (0.50, 1.50),
+            "gemini-1.0-pro-001": (0.50, 1.50),
+            "gemini-1.0-pro-latest": (0.50, 1.50),
+            "gemini-pro": (0.50, 1.50),
+            # Model variants with features
+            "gemini-1.5-pro-exp-0827": (1.25, 5.00),
+            "gemini-exp-1206": (0.00, 0.00),  # Experimental, free
+        }
+
+    def get_default_pricing(self) -> Tuple[float, float]:
+        """Conservative default for unknown Google models."""
+        return (1.25, 5.00)  # Similar to Gemini 1.5 Pro
+
+    def get_model_prefixes(self) -> list[str]:
+        return ["gemini-", "gemini"]
+
+    def extract_usage(self, usage: Any) -> Tuple[Optional[int], Optional[int], Optional[int]]:
+        """
+        Extract usage from Google Gemini response.
+
+        Google structure (note the different field names):
+        {
+            "usageMetadata": {
+                "promptTokenCount": 100,
+                "candidatesTokenCount": 50,
+                "totalTokenCount": 150
+            }
+        }
+
+        OR older structure:
+        {
+            "usage_metadata": {
+                "prompt_token_count": 100,
+                "candidates_token_count": 50,
+                "total_token_count": 150
+            }
+        }
+        """
+        if usage is None:
+            return None, None, None
+
+        input_tokens = None
+        output_tokens = None
+        total_tokens = None
+
+        if isinstance(usage, dict):
+            # Try camelCase (newer API)
+            input_tokens = _coerce_int(usage.get("promptTokenCount"))
+            output_tokens = _coerce_int(usage.get("candidatesTokenCount"))
+            total_tokens = _coerce_int(usage.get("totalTokenCount"))
+
+            # Fallback to snake_case (older API)
+            if input_tokens is None:
+                input_tokens = _coerce_int(usage.get("prompt_token_count"))
+            if output_tokens is None:
+                output_tokens = _coerce_int(usage.get("candidates_token_count"))
+            if total_tokens is None:
+                total_tokens = _coerce_int(usage.get("total_token_count"))
+        else:
+            # Handle object attributes (camelCase)
+            input_tokens = _coerce_int(getattr(usage, "promptTokenCount", None))
+            output_tokens = _coerce_int(getattr(usage, "candidatesTokenCount", None))
+            total_tokens = _coerce_int(getattr(usage, "totalTokenCount", None))
+
+            # Fallback to snake_case attributes
+            if input_tokens is None:
+                input_tokens = _coerce_int(getattr(usage, "prompt_token_count", None))
+            if output_tokens is None:
+                output_tokens = _coerce_int(getattr(usage, "candidates_token_count", None))
+            if total_tokens is None:
+                total_tokens = _coerce_int(getattr(usage, "total_token_count", None))
+
+        # Calculate total if not provided
+        if total_tokens is None and (input_tokens is not None or output_tokens is not None):
+            total_tokens = (input_tokens or 0) + (output_tokens or 0)
+
+        return input_tokens, output_tokens, total_tokens
+
+    def extract_model(self, response: Any) -> Optional[str]:
+        """Extract model name from Google response."""
+        if response is None:
+            return None
+
+        # Try dict access
+        if isinstance(response, dict):
+            # Try common locations
+            model = response.get("model") or response.get("modelVersion")
+            if model:
+                return model
+
+        # Try object attribute
+        if hasattr(response, "model"):
+            return str(response.model)
+        if hasattr(response, "model_version"):
+            return str(response.model_version)
+
+        return None
+
+    def matches_response(self, response: Any) -> bool:
+        """
+        Detect Google Gemini responses by structure.
+
+        Gemini responses typically have:
+        - 'candidates' array (not 'choices')
+        - 'usageMetadata' or 'usage_metadata' (not 'usage')
+        - 'promptTokenCount' style fields
+        """
+        if response is None:
+            return False
+
+        # Check for Gemini-specific structure
+        if isinstance(response, dict):
+            has_candidates = "candidates" in response
+            has_gemini_usage = "usageMetadata" in response or "usage_metadata" in response
+
+            # Check for Gemini-specific token counting fields
+            usage_metadata = response.get("usageMetadata") or response.get("usage_metadata")
+            if isinstance(usage_metadata, dict):
+                has_prompt_token_count = (
+                    "promptTokenCount" in usage_metadata or "prompt_token_count" in usage_metadata
+                )
+                if has_prompt_token_count:
+                    return True
+
+            if has_candidates or has_gemini_usage:
+                return True
+        elif hasattr(response, "candidates"):
+            return True
+        elif hasattr(response, "usage_metadata") or hasattr(response, "usageMetadata"):
+            return True
+
+        # Fallback to model prefix matching
+        return super().matches_response(response)

auditi/providers/openai.py

@@ -0,0 +1,147 @@
+"""
+OpenAI provider implementation for usage extraction and cost calculation.
+"""
+
+from typing import Optional, Any, Dict, Tuple
+from .base import BaseProvider
+
+
+def _coerce_int(value: Any) -> Optional[int]:
+    """Helper to safely convert values to int."""
+    if value is None:
+        return None
+    try:
+        return int(value)
+    except (TypeError, ValueError):
+        return None
+
+
+class OpenAIProvider(BaseProvider):
+    """Provider implementation for OpenAI models."""
+
+    @property
+    def name(self) -> str:
+        return "openai"
+
+    @property
+    def model_pricing(self) -> Dict[str, Tuple[float, float]]:
+        """
+        OpenAI model pricing per 1M tokens (input, output) in USD.
+        Updated as of January 2025.
+        """
+        return {
+            # GPT-4o family
+            "gpt-4o": (2.50, 10.00),
+            "gpt-4o-2024-11-20": (2.50, 10.00),
+            "gpt-4o-2024-08-06": (2.50, 10.00),
+            "gpt-4o-2024-05-13": (5.00, 15.00),
+            "gpt-4o-mini": (0.15, 0.60),
+            "gpt-4o-mini-2024-07-18": (0.15, 0.60),
+            # GPT-4 Turbo
+            "gpt-4-turbo": (10.00, 30.00),
+            "gpt-4-turbo-2024-04-09": (10.00, 30.00),
+            "gpt-4-turbo-preview": (10.00, 30.00),
+            "gpt-4-0125-preview": (10.00, 30.00),
+            "gpt-4-1106-preview": (10.00, 30.00),
+            # GPT-4 base
+            "gpt-4": (30.00, 60.00),
+            "gpt-4-0613": (30.00, 60.00),
+            "gpt-4-32k": (60.00, 120.00),
+            "gpt-4-32k-0613": (60.00, 120.00),
+            # GPT-3.5 Turbo
+            "gpt-3.5-turbo": (0.50, 1.50),
+            "gpt-3.5-turbo-0125": (0.50, 1.50),
+            "gpt-3.5-turbo-1106": (1.00, 2.00),
+            "gpt-3.5-turbo-16k": (3.00, 4.00),
+            # o1 models
+            "o1-preview": (15.00, 60.00),
+            "o1-preview-2024-09-12": (15.00, 60.00),
+            "o1-mini": (3.00, 12.00),
+            "o1-mini-2024-09-12": (3.00, 12.00),
+            "o1": (15.00, 60.00),
+        }
+
+    def get_default_pricing(self) -> Tuple[float, float]:
+        """Conservative default for unknown OpenAI models."""
+        return (10.00, 30.00)  # Similar to GPT-4 Turbo
+
+    def get_model_prefixes(self) -> list[str]:
+        return ["gpt-", "o1-", "o1"]
+
+    def extract_usage(self, usage: Any) -> Tuple[Optional[int], Optional[int], Optional[int]]:
+        """
+        Extract usage from OpenAI response.
+
+        OpenAI structure:
+        {
+            "usage": {
+                "prompt_tokens": 100,
+                "completion_tokens": 50,
+                "total_tokens": 150
+            }
+        }
+        """
+        if usage is None:
+            return None, None, None
+
+        input_tokens = None
+        output_tokens = None
+        total_tokens = None
+
+        if isinstance(usage, dict):
+            input_tokens = _coerce_int(usage.get("prompt_tokens"))
+            output_tokens = _coerce_int(usage.get("completion_tokens"))
+            total_tokens = _coerce_int(usage.get("total_tokens"))
+        else:
+            # Handle object attributes
+            input_tokens = _coerce_int(getattr(usage, "prompt_tokens", None))
+            output_tokens = _coerce_int(getattr(usage, "completion_tokens", None))
+            total_tokens = _coerce_int(getattr(usage, "total_tokens", None))
+
+        # Calculate total if not provided
+        if total_tokens is None and (input_tokens is not None or output_tokens is not None):
+            total_tokens = (input_tokens or 0) + (output_tokens or 0)
+
+        return input_tokens, output_tokens, total_tokens
+
+    def extract_model(self, response: Any) -> Optional[str]:
+        """Extract model name from OpenAI response."""
+        if response is None:
+            return None
+
+        # Try dict access
+        if isinstance(response, dict):
+            return response.get("model")
+
+        # Try object attribute
+        if hasattr(response, "model"):
+            return str(response.model)
+
+        return None
+
+    def matches_response(self, response: Any) -> bool:
+        """
+        Detect OpenAI responses by structure.
+
+        OpenAI responses typically have:
+        - 'choices' array with 'message' or 'text'
+        - 'usage' with 'prompt_tokens' and 'completion_tokens'
+        """
+        if response is None:
+            return False
+
+        # Check for OpenAI-specific structure
+        if isinstance(response, dict):
+            has_choices = "choices" in response
+            has_openai_usage = (
+                "usage" in response
+                and isinstance(response.get("usage"), dict)
+                and "prompt_tokens" in response.get("usage", {})
+            )
+            if has_choices or has_openai_usage:
+                return True
+        elif hasattr(response, "choices") and hasattr(response, "usage"):
+            return True
+
+        # Fallback to model prefix matching
+        return super().matches_response(response)

auditi/providers/registry.py

@@ -0,0 +1,166 @@
+"""
+Provider registry for auto-detecting and managing LLM providers.
+"""
+
+from typing import Optional, Any, List
+from .base import BaseProvider
+from .openai import OpenAIProvider
+from .anthropic import AnthropicProvider
+from .google import GoogleProvider
+
+
+class ProviderRegistry:
+    """
+    Central registry for LLM providers.
+
+    Handles:
+    - Provider registration
+    - Auto-detection from model names
+    - Auto-detection from response structures
+    - Fallback to generic provider
+    """
+
+    def __init__(self):
+        self._providers: List[BaseProvider] = []
+        self._default_provider: Optional[BaseProvider] = None
+
+        # Register built-in providers
+        self.register(OpenAIProvider())
+        self.register(AnthropicProvider())
+        self.register(GoogleProvider())
+
+        # Set OpenAI as default fallback (most common format)
+        self._default_provider = self.get_provider("openai")
+
+    def register(self, provider: BaseProvider) -> None:
+        """Register a new provider."""
+        self._providers.append(provider)
+
+    def get_provider(self, name: str) -> Optional[BaseProvider]:
+        """Get provider by name."""
+        for provider in self._providers:
+            if provider.name == name:
+                return provider
+        return None
+
+    def detect_from_model(self, model: Optional[str]) -> Optional[BaseProvider]:
+        """
+        Detect provider from model name.
+
+        Args:
+            model: Model name string (e.g., "gpt-4", "claude-3-opus-20240229")
+
+        Returns:
+            Matching provider, or None if no match
+
+        Example:
+            >>> registry.detect_from_model("gpt-4o")
+            OpenAIProvider()
+            >>> registry.detect_from_model("claude-3-5-sonnet-20241022")
+            AnthropicProvider()
+        """
+        if not model:
+            return None
+
+        for provider in self._providers:
+            if provider.matches_model(model):
+                return provider
+
+        return None
+
+    def detect_from_response(self, response: Any) -> Optional[BaseProvider]:
+        """
+        Detect provider from response structure.
+
+        This is useful when the model name is not known but we have
+        the API response object.
+
+        Args:
+            response: Raw API response object or dict
+
+        Returns:
+            Matching provider, or None if no match
+
+        Example:
+            >>> response = {"choices": [...], "usage": {"prompt_tokens": 10}}
+            >>> registry.detect_from_response(response)
+            OpenAIProvider()
+        """
+        if response is None:
+            return None
+
+        # Try to extract model first (fastest method)
+        for provider in self._providers:
+            model = provider.extract_model(response)
+            if model and provider.matches_model(model):
+                return provider
+
+        # Fall back to structure detection
+        for provider in self._providers:
+            if provider.matches_response(response):
+                return provider
+
+        return None
+
+    def get_provider_or_default(
+        self, model: Optional[str] = None, response: Any = None
+    ) -> BaseProvider:
+        """
+        Get provider with fallback to default.
+
+        Tries in order:
+        1. Detect from model name
+        2. Detect from response structure
+        3. Use default provider
+
+        Args:
+            model: Optional model name
+            response: Optional API response
+
+        Returns:
+            Provider (never None, returns default if detection fails)
+        """
+        # Try model detection first
+        if model:
+            provider = self.detect_from_model(model)
+            if provider:
+                return provider
+
+        # Try response detection
+        if response:
+            provider = self.detect_from_response(response)
+            if provider:
+                return provider
+
+        # Fallback to default
+        return self._default_provider or self._providers[0]
+
+
+# Global singleton registry
+_registry = ProviderRegistry()
+
+
+def get_registry() -> ProviderRegistry:
+    """Get the global provider registry."""
+    return _registry
+
+
+def detect_provider(model: Optional[str] = None, response: Any = None) -> BaseProvider:
+    """
+    Detect provider from model name or response.
+
+    This is the main entry point for provider detection.
+
+    Args:
+        model: Optional model name
+        response: Optional API response
+
+    Returns:
+        Detected provider (falls back to default if detection fails)
+
+    Example:
+        >>> provider = detect_provider(model="gpt-4o")
+        >>> input_tokens, output_tokens, total = provider.extract_usage(usage)
+        >>> cost = provider.calculate_cost(model, input_tokens, output_tokens)
+    """
+    return _registry.get_provider_or_default(model=model, response=response)

auditi/transport.py

@@ -0,0 +1,78 @@
+"""
+Transport layer for sending traces to the Auditi platform.
+
+Provides different transport implementations for various use cases:
+- SyncHttpTransport: Synchronous HTTP transport (default)
+- DebugTransport: Debug transport that prints to console
+"""
+
+import abc
+import httpx
+from typing import Any, Dict
+import logging
+
+logger = logging.getLogger("auditi")
+
+
+class BaseTransport(abc.ABC):
+    """
+    Abstract base class for transport implementations.
+    Subclass this to create custom transport mechanisms.
+    """
+
+    @abc.abstractmethod
+    def send_trace(self, trace_data: Dict[str, Any]) -> None:
+        """
+        Send trace data to the Auditi platform.
+
+        Args:
+            trace_data: Serialized trace data as a dictionary
+        """
+        pass
+
+
+class SyncHttpTransport(BaseTransport):
+    """
+    Synchronous HTTP transport using httpx.
+    Sends traces immediately via HTTP POST.
+    """
+
+    def __init__(self, base_url: str, api_key: str = None):
+        """
+        Initialize the transport.
+
+        Args:
+            base_url: Base URL of the Auditi API
+            api_key: API key for authentication (optional)
+        """
+        self.base_url = base_url.rstrip("/")
+        self.api_key = api_key
+        self.headers = {
+            "Content-Type": "application/json",
+        }
+        if api_key:
+            self.headers["Authorization"] = f"Bearer {api_key}"
+
+    def send_trace(self, trace_data: Dict[str, Any]) -> None:
+        """Send trace data via HTTP POST."""
+        url = f"{self.base_url}/api/v1/ingest"
+        try:
+            with httpx.Client() as client:
+                response = client.post(url, json=trace_data, headers=self.headers, timeout=5.0)
+                response.raise_for_status()
+                logger.debug(f"Trace sent successfully: {trace_data.get('id')}")
+        except Exception as e:
+            logger.error(f"Failed to send trace to Auditi: {e}")
+
+
+class DebugTransport(BaseTransport):
+    """
+    Debug transport that prints trace data to console.
+    Useful for local development and testing.
+    """
+
+    def send_trace(self, trace_data: Dict[str, Any]) -> None:
+        """Print trace data to console."""
+        trace_id = trace_data.get("id", "unknown")
+        spans_count = len(trace_data.get("spans", []))
+        print(f"[Auditi] Trace captured: {trace_id} ({spans_count} spans)")

auditi/types/__init__.py

@@ -0,0 +1,12 @@
+"""
+Type definitions for Auditi SDK.
+"""
+
+from .api_types import TraceInput, SpanInput, EvaluationResult, TraceResponse
+
+__all__ = [
+    "TraceInput",
+    "SpanInput", 
+    "EvaluationResult",
+    "TraceResponse",
+]

auditi/types/api_types.py

@@ -0,0 +1,107 @@
+"""
+Pydantic models for Auditi SDK API types.
+"""
+
+from pydantic import BaseModel, Field, field_validator
+from typing import List, Optional, Dict, Any
+from datetime import datetime
+from uuid import UUID
+
+
+class SpanInput(BaseModel):
+    """
+    Represents a single operation (tool call, LLM call) within a trace.
+
+    UPDATED: Added processing_time field for performance tracking.
+    """
+
+    id: UUID
+    trace_id: UUID
+    parent_id: Optional[UUID] = None
+    name: str
+    span_type: str  # "tool", "llm", "retrieval", etc.
+    start_time: datetime
+    end_time: Optional[datetime] = None
+    processing_time: Optional[float] = None  # NEW: Duration in seconds
+
+    # Input/Output
+    inputs: Optional[Dict[str, Any]] = None
+    outputs: Optional[str] = None
+
+    # LLM specific
+    model: Optional[str] = None
+    input_tokens: Optional[int] = None
+    output_tokens: Optional[int] = None
+    tokens: Optional[int] = None
+    cost: Optional[float] = None
+
+    # Status
+    status: Optional[str] = None  # "ok", "error"
+    error: Optional[str] = None
+
+    # Metadata
+    metadata: Optional[Dict[str, Any]] = Field(default_factory=dict)
+
+
+class TraceInput(BaseModel):
+    """
+    Represents a complete agent interaction with user.
+    """
+
+    id: UUID
+    user_id: Optional[str] = None
+    conversation_id: Optional[str] = None
+
+    # Timing
+    start_time: datetime
+    end_time: Optional[datetime] = None
+
+    # Content
+    # Content
+    name: str
+    user_input: str = ""
+    assistant_output: Optional[str] = None
+
+    # Metrics
+    total_tokens: Optional[int] = None
+    cost: Optional[float] = None
+
+    # Evaluation
+    status: Optional[str] = None  # "pass", "fail", "review", "pending"
+    score: Optional[float] = None
+    failure_mode: Optional[str] = None
+    eval_reason: Optional[str] = None
+
+    # Relations
+    spans: List[SpanInput] = Field(default_factory=list)
+
+    # Metadata
+    tags: List[str] = Field(default_factory=list)
+    metadata: Optional[Dict[str, Any]] = Field(default_factory=dict)
+    error: Optional[str] = None
+
+    @field_validator("user_input", mode="before")
+    def normalize_user_input(cls, v):
+        return v or ""
+
+
+class EvaluationResult(BaseModel):
+    """
+    Result of evaluating a trace.
+    Contains pass/fail status, score, and optional failure information.
+    """
+
+    status: str  # pass, fail
+    score: float
+    reason: Optional[str] = None
+    failure_mode: Optional[str] = None
+    recommended_action: Optional[str] = None
+
+
+class TraceResponse(BaseModel):
+    """
+    Response from the Auditi API after ingesting a trace.
+    """
+
+    success: bool
+    count: int