fenra 0.1.0__py3-none-any.whl

fenra/__init__.py

@@ -0,0 +1,187 @@
+"""Fenra: Privacy-first AI cost tracking library."""
+
+import os
+from typing import Any
+
+from fenra._context import clear_context, get_context, set_context
+from fenra._core import enqueue_transaction, flush, init, shutdown
+
+__version__ = "0.1.0"
+
+# Check environment variable for disabling auto-tracking
+_auto_track_disabled = os.environ.get("FENRA_DISABLE_AUTO_TRACK", "").lower() in (
+    "1",
+    "true",
+    "yes",
+)
+
+
+def disable_auto_track() -> None:
+    """
+    Disable auto-tracking/patching of AI provider SDKs.
+
+    This function must be called BEFORE importing fenra if you want to prevent
+    automatic patching. Alternatively, set the environment variable
+    FENRA_DISABLE_AUTO_TRACK=1 before importing fenra.
+
+    When auto-tracking is disabled, you can still use fenra.track() for manual
+    tracking.
+    """
+    global _auto_track_disabled
+    _auto_track_disabled = True
+
+
+def track(
+    provider: str,
+    model: str,
+    input_tokens: int,
+    output_tokens: int,
+    *,
+    total_tokens: int | None = None,
+    reasoning_tokens: int | None = None,
+    cached_tokens: int | None = None,
+    **context_overrides: Any,
+) -> None:
+    """
+    Manually track an AI call.
+
+    Merges context_overrides with the current context from set_context().
+
+    Args:
+        provider: Provider name (e.g., "openai", "anthropic", "custom")
+        model: Model identifier (e.g., "gpt-4o", "claude-3-opus")
+        input_tokens: Number of input/prompt tokens
+        output_tokens: Number of output/completion tokens
+        total_tokens: Total tokens (calculated if not provided)
+        reasoning_tokens: Reasoning tokens (for o1/o3 models)
+        cached_tokens: Cached tokens (if applicable)
+        **context_overrides: Additional context to merge with current context
+    """
+    from fenra._context import get_context
+
+    # Merge context_overrides with current context
+    context = get_context()
+    context.update(context_overrides)
+
+    # Calculate total_tokens if not provided
+    if total_tokens is None:
+        total_tokens = input_tokens + output_tokens
+
+    # Build metrics
+    metrics: dict[str, Any] = {
+        "input_tokens": input_tokens,
+        "output_tokens": output_tokens,
+        "total_tokens": total_tokens,
+    }
+
+    if reasoning_tokens is not None:
+        metrics["reasoning_tokens"] = reasoning_tokens
+
+    if cached_tokens is not None:
+        metrics["cached_tokens"] = cached_tokens
+
+    # Build transaction
+    transaction = {
+        "provider": provider,
+        "model": model,
+        "usage": [
+            {
+                "type": "tokens",
+                "metrics": metrics,
+            }
+        ],
+        "context": context,
+    }
+
+    enqueue_transaction(transaction)
+
+
+# Auto-patch OpenAI when imported (if OpenAI is installed)
+def _auto_patch_openai() -> None:
+    """Automatically patch OpenAI if it's installed."""
+    if _auto_track_disabled:
+        return
+    try:
+        from fenra.integrations.openai import (
+            patch_openai,
+            patch_openai_async,
+            patch_openai_beta_parse,
+            patch_openai_beta_parse_async,
+            patch_openai_images,
+            patch_openai_images_async,
+            patch_openai_responses,
+            patch_openai_responses_async,
+        )
+
+        patch_openai()
+        patch_openai_async()
+        patch_openai_responses()
+        patch_openai_responses_async()
+        patch_openai_images()
+        patch_openai_images_async()
+        patch_openai_beta_parse()
+        patch_openai_beta_parse_async()
+    except ImportError:
+        # OpenAI not installed, that's fine
+        pass
+
+
+# Auto-patch Gemini when imported (if Gemini is installed)
+def _auto_patch_gemini() -> None:
+    """Automatically patch Gemini if it's installed."""
+    if _auto_track_disabled:
+        return
+    try:
+        from fenra.integrations.gemini import (
+            patch_gemini,
+            patch_gemini_async,
+            patch_gemini_stream,
+            patch_gemini_stream_async,
+        )
+
+        patch_gemini()
+        patch_gemini_async()
+        patch_gemini_stream()
+        patch_gemini_stream_async()
+    except ImportError:
+        # Gemini not installed, that's fine
+        pass
+
+
+# Auto-patch Anthropic when imported (if Anthropic is installed)
+def _auto_patch_anthropic() -> None:
+    """Automatically patch Anthropic if it's installed."""
+    if _auto_track_disabled:
+        return
+    try:
+        from fenra.integrations.anthropic import (
+            patch_anthropic,
+            patch_anthropic_async,
+            patch_anthropic_stream,
+            patch_anthropic_stream_async,
+        )
+
+        patch_anthropic()
+        patch_anthropic_async()
+        patch_anthropic_stream()
+        patch_anthropic_stream_async()
+    except ImportError:
+        # Anthropic not installed, that's fine
+        pass
+
+
+# Auto-patch on import
+_auto_patch_openai()
+_auto_patch_gemini()
+_auto_patch_anthropic()
+
+__all__ = [
+    "init",
+    "shutdown",
+    "flush",
+    "track",
+    "set_context",  # Primary context API - merges into existing
+    "get_context",
+    "clear_context",
+    "disable_auto_track",
+]

fenra/_context.py

@@ -0,0 +1,42 @@
+"""Context management using contextvars for async-safe context propagation."""
+
+from contextvars import ContextVar
+from typing import Any
+
+_fenra_context: ContextVar[dict[str, Any]] = ContextVar("fenra_context", default={})
+
+
+def set_context(**kwargs: Any) -> None:
+    """
+    Merge new metadata into the current context.
+
+    This function is additive. Each call merges new keys into the existing
+    context dictionary. Later calls with the same key will overwrite.
+
+    Example middleware stacking:
+        # Auth middleware
+        fenra.set_context(billable_customer_id="acme-corp", user_id="u_123")
+
+        # Feature middleware (adds to existing context)
+        fenra.set_context(feature="chat-assistant", environment="production")
+
+        # Final context: {billable_customer_id, user_id, feature, environment}
+    """
+    current = _fenra_context.get().copy()
+    current.update(kwargs)
+    _fenra_context.set(current)
+
+
+def get_context() -> dict[str, Any]:
+    """
+    Get the full merged context dictionary.
+
+    Returns a copy to prevent external mutation.
+    Called by auto-instrumentation to include in logging payload.
+    """
+    return _fenra_context.get().copy()
+
+
+def clear_context() -> None:
+    """Reset context to empty dict. Useful at request boundaries."""
+    _fenra_context.set({})

fenra/_core.py

@@ -0,0 +1,229 @@
+"""Core Fenra SDK functionality: background worker, queue, and configuration."""
+
+import logging
+import queue
+import threading
+import time
+from dataclasses import dataclass, field
+from typing import Any
+
+import requests
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class FenraConfig:
+    """Configuration for Fenra SDK."""
+
+    api_key: str
+    base_url: str = "https://ingest.fenra.io"
+    content_logging: bool = False  # Privacy-first default
+    enabled: bool = True
+    flush_interval: float = 1.0
+    batch_size: int = 10
+    queue_size: int = 10000
+
+    def __post_init__(self) -> None:
+        """Validate configuration."""
+        if not self.api_key:
+            raise ValueError("api_key is required")
+
+
+# Global configuration singleton
+_config: FenraConfig | None = None
+_transaction_queue: queue.Queue[dict[str, Any]] | None = None
+_worker_thread: threading.Thread | None = None
+_worker_stop_event: threading.Event | None = None
+
+
+def _get_config() -> FenraConfig:
+    """Get the global configuration, raising if not initialized."""
+    if _config is None:
+        raise RuntimeError("Fenra not initialized. Call fenra.init() first.")
+    return _config
+
+
+def _get_queue() -> queue.Queue[dict[str, Any]]:
+    """Get the global transaction queue, raising if not initialized."""
+    if _transaction_queue is None:
+        raise RuntimeError("Fenra not initialized. Call fenra.init() first.")
+    return _transaction_queue
+
+
+def _background_worker() -> None:
+    """Background worker thread that processes transactions from the queue."""
+    batch: list[dict[str, Any]] = []
+    last_flush = time.time()
+
+    while True:
+        try:
+            # Check if we should stop
+            if _worker_stop_event and _worker_stop_event.is_set():
+                # Flush any remaining items before stopping
+                if batch:
+                    _flush_batch(batch)
+                break
+
+            config = _get_config()
+            if not config.enabled:
+                time.sleep(0.1)
+                continue
+
+            # Try to get an item from the queue with timeout
+            try:
+                transaction = _get_queue().get(timeout=config.flush_interval)
+                batch.append(transaction)
+            except queue.Empty:
+                # Timeout - flush if we have items and enough time has passed
+                if batch and (time.time() - last_flush) >= config.flush_interval:
+                    _flush_batch(batch)
+                    batch = []
+                    last_flush = time.time()
+                continue
+
+            # Check if we should flush due to batch size
+            if len(batch) >= config.batch_size:
+                _flush_batch(batch)
+                batch = []
+                last_flush = time.time()
+
+        except Exception as e:
+            # Never crash the host application
+            logger.error(f"Error in Fenra background worker: {e}", exc_info=True)
+            time.sleep(0.1)
+
+
+def _flush_batch(batch: list[dict[str, Any]]) -> None:
+    """Flush a batch of transactions to the Fenra API."""
+    if not batch:
+        return
+
+    try:
+        config = _get_config()
+
+        # Prepare payload - single transaction or bulk
+        if len(batch) == 1:
+            payload = batch[0]
+        else:
+            payload = {"transactions": batch}
+
+        # Send to API
+        response = requests.post(
+            f"{config.base_url}/usage/transactions",
+            json=payload,
+            headers={"X-Api-Key": config.api_key},
+            timeout=10.0,
+        )
+
+        # Log errors but don't raise
+        if response.status_code >= 400:
+            logger.warning(
+                f"Fenra API error: {response.status_code} - {response.text}"
+            )
+        else:
+            logger.debug(f"Fenra: Successfully sent {len(batch)} transaction(s)")
+
+    except Exception as e:
+        # Never crash the host application
+        logger.error(f"Error sending transactions to Fenra API: {e}", exc_info=True)
+
+
+def init(
+    api_key: str,
+    *,
+    base_url: str = "https://ingest.fenra.io",
+    content_logging: bool = False,
+    enabled: bool = True,
+    flush_interval: float = 1.0,
+    batch_size: int = 10,
+    queue_size: int = 10000,
+) -> None:
+    """
+    Initialize Fenra SDK.
+
+    Args:
+        api_key: Your Fenra API key
+        base_url: Base URL for Fenra API (default: https://ingest.fenra.io)
+        content_logging: Whether to log prompt content (default: False, privacy-first)
+        enabled: Whether tracking is enabled (default: True)
+        flush_interval: Seconds between automatic flushes (default: 1.0)
+        batch_size: Maximum transactions per batch (default: 10)
+        queue_size: Maximum queue size before dropping transactions (default: 10000)
+    """
+    global _config, _transaction_queue, _worker_thread, _worker_stop_event
+
+    # Stop existing worker if any
+    shutdown()
+
+    # Create configuration
+    _config = FenraConfig(
+        api_key=api_key,
+        base_url=base_url,
+        content_logging=content_logging,
+        enabled=enabled,
+        flush_interval=flush_interval,
+        batch_size=batch_size,
+        queue_size=queue_size,
+    )
+
+    _transaction_queue = queue.Queue(maxsize=_config.queue_size)
+
+    # Create and start worker thread
+    _worker_stop_event = threading.Event()
+    _worker_thread = threading.Thread(target=_background_worker, daemon=True)
+    _worker_thread.start()
+
+    logger.info("Fenra SDK initialized")
+
+
+def shutdown() -> None:
+    """Shutdown Fenra SDK and flush remaining transactions."""
+    global _worker_thread, _worker_stop_event
+
+    if _worker_thread is not None and _worker_thread.is_alive():
+        if _worker_stop_event:
+            _worker_stop_event.set()
+        _worker_thread.join(timeout=5.0)
+        _worker_thread = None
+        _worker_stop_event = None
+
+    logger.info("Fenra SDK shut down")
+
+
+def flush() -> None:
+    """Manually flush all pending transactions."""
+    if _transaction_queue is None:
+        return
+
+    batch: list[dict[str, Any]] = []
+    # Drain the queue
+    while True:
+        try:
+            transaction = _transaction_queue.get_nowait()
+            batch.append(transaction)
+        except queue.Empty:
+            break
+
+    if batch:
+        _flush_batch(batch)
+
+
+def enqueue_transaction(transaction: dict[str, Any]) -> None:
+    """
+    Enqueue a transaction for background processing.
+
+    This function is safe to call from any thread. Failures are logged
+    but never raise exceptions.
+    """
+    try:
+        config = _get_config()
+        if not config.enabled:
+            return
+
+        _get_queue().put_nowait(transaction)
+    except queue.Full:
+        logger.warning("Fenra transaction queue is full, dropping transaction")
+    except Exception as e:
+        # Never crash the host application
+        logger.error(f"Error enqueueing Fenra transaction: {e}", exc_info=True)

fenra/integrations/__init__.py

@@ -0,0 +1 @@
+"""Fenra integrations for various AI providers."""