superpenguin 0.1.0__py3-none-any.whl

superpenguin/__init__.py

@@ -0,0 +1,79 @@
+"""SuperPenguin Python SDK — AI cost management, attribution, and spend tracking."""
+
+from __future__ import annotations
+
+import os
+from typing import Any
+
+from superpenguin._client import SuperPenguinClient
+from superpenguin._litellm import patch_litellm
+from superpenguin._trace import trace
+from superpenguin._wrap import wrap
+
+__all__ = ["init", "wrap", "trace", "patch_litellm", "get_client", "flush"]
+__version__ = "0.1.0"
+
+_client: SuperPenguinClient | None = None
+
+_DEFAULT_BASE_URL = "https://api.carrotlabs.ai"
+
+
+def init(
+    *,
+    api_key: str | None = None,
+    base_url: str | None = None,
+    flush_interval: float = 5.0,
+    batch_size: int = 50,
+) -> None:
+    """Initialize the SuperPenguin SDK.
+
+    Must be called before ``wrap()`` or ``@trace`` capture any data.
+    Alternatively, set the ``SP_API_KEY`` environment variable for
+    auto-initialization on first use.
+
+    Args:
+        api_key: Your SuperPenguin API key (``sp_...``).  Falls back to
+            ``SP_API_KEY`` env var.
+        base_url: Override the API endpoint.  Falls back to
+            ``SP_BASE_URL`` env var, then ``https://api.carrotlabs.ai``.
+        flush_interval: Seconds between background flushes (default 5).
+        batch_size: Max events per batch POST (default 50).
+    """
+    global _client
+
+    if api_key is None:
+        api_key = os.environ.get("SP_API_KEY")
+    if api_key is None:
+        raise ValueError(
+            "api_key is required — pass it to sp.init() "
+            "or set the SP_API_KEY environment variable"
+        )
+    if base_url is None:
+        base_url = os.environ.get("SP_BASE_URL", _DEFAULT_BASE_URL)
+
+    _client = SuperPenguinClient(
+        api_key=api_key,
+        base_url=base_url,
+        flush_interval=flush_interval,
+        batch_size=batch_size,
+    )
+
+
+def get_client() -> SuperPenguinClient:
+    """Return the global client, auto-initializing if ``SP_API_KEY`` is set."""
+    global _client
+    if _client is None:
+        api_key = os.environ.get("SP_API_KEY")
+        if api_key:
+            init(api_key=api_key)
+        else:
+            raise RuntimeError(
+                "sp.init() has not been called and SP_API_KEY is not set"
+            )
+    return _client
+
+
+def flush() -> None:
+    """Flush any pending events to the server immediately."""
+    if _client is not None:
+        _client.flush()

superpenguin/_client.py

@@ -0,0 +1,142 @@
+"""Background event submitter — batches and POSTs cost events to the API."""
+
+from __future__ import annotations
+
+import atexit
+import json
+import logging
+import queue
+import threading
+import time
+from typing import Any
+from urllib.error import URLError
+from urllib.request import Request, urlopen
+
+logger = logging.getLogger("superpenguin")
+
+_SENTINEL = object()
+
+
+class SuperPenguinClient:
+    """Batched HTTP client that queues cost events and flushes them
+    to the server in a background daemon thread.
+
+    ``flush()`` stops the worker, drains remaining items, sends them
+    synchronously, and restarts the worker so nothing is lost.
+    """
+
+    def __init__(
+        self,
+        api_key: str,
+        base_url: str = "https://api.carrotlabs.ai",
+        flush_interval: float = 5.0,
+        batch_size: int = 50,
+        max_queue: int = 1000,
+    ) -> None:
+        self._api_key = api_key
+        self._base_url = base_url.rstrip("/")
+        self._queue: queue.Queue[dict[str, Any] | object] = queue.Queue(
+            maxsize=max_queue
+        )
+        self._flush_interval = flush_interval
+        self._batch_size = batch_size
+        self._shutdown = threading.Event()
+        self._lifecycle_lock = threading.Lock()
+        self._thread = self._start_worker()
+        atexit.register(self.shutdown)
+
+    def submit(self, event: dict[str, Any]) -> None:
+        """Enqueue a cost event for background submission. Non-blocking."""
+        try:
+            self._queue.put_nowait(event)
+        except queue.Full:
+            logger.warning("superpenguin: event queue full, dropping event")
+
+    def flush(self) -> None:
+        """Stop worker, drain queue, send everything, restart worker."""
+        with self._lifecycle_lock:
+            self._stop_worker()
+            self._drain_and_send()
+            self._thread = self._start_worker()
+
+    def shutdown(self) -> None:
+        """Final flush on interpreter exit. Does not restart the worker."""
+        with self._lifecycle_lock:
+            self._stop_worker()
+            self._drain_and_send()
+
+    # -- internals ------------------------------------------------------------
+
+    def _start_worker(self) -> threading.Thread:
+        self._shutdown.clear()
+        t = threading.Thread(target=self._run, daemon=True, name="sp-flush")
+        t.start()
+        return t
+
+    def _stop_worker(self) -> None:
+        self._shutdown.set()
+        try:
+            self._queue.put_nowait(_SENTINEL)
+        except queue.Full:
+            pass
+        self._thread.join(timeout=10.0)
+
+    def _drain_and_send(self) -> None:
+        batch: list[dict[str, Any]] = []
+        while True:
+            try:
+                item = self._queue.get_nowait()
+                if item is _SENTINEL:
+                    continue
+                batch.append(item)  # type: ignore[arg-type]
+            except queue.Empty:
+                break
+        if batch:
+            self._send(batch)
+
+    def _run(self) -> None:
+        while not self._shutdown.is_set():
+            batch: list[dict[str, Any]] = []
+            deadline = time.monotonic() + self._flush_interval
+            while time.monotonic() < deadline and len(batch) < self._batch_size:
+                if self._shutdown.is_set():
+                    break
+                timeout = min(0.5, max(0.05, deadline - time.monotonic()))
+                try:
+                    item = self._queue.get(timeout=timeout)
+                    if item is _SENTINEL:
+                        break
+                    batch.append(item)  # type: ignore[arg-type]
+                except queue.Empty:
+                    continue
+            if batch:
+                self._send(batch)
+
+    def _send(self, events: list[dict[str, Any]], *, _retries: int = 2) -> None:
+        url = f"{self._base_url}/api/sdk/ingest"
+        body = json.dumps({"events": events}).encode()
+
+        for attempt in range(_retries + 1):
+            req = Request(url, data=body, method="POST")
+            req.add_header("Content-Type", "application/json")
+            req.add_header("Authorization", f"Bearer {self._api_key}")
+            try:
+                with urlopen(req, timeout=10) as resp:
+                    resp.read()
+                logger.debug("superpenguin: sent %d event(s)", len(events))
+                return
+            except (URLError, OSError) as exc:
+                if attempt < _retries:
+                    time.sleep(0.3 * (attempt + 1))
+                    continue
+                logger.warning(
+                    "superpenguin: failed to send %d event(s) after %d attempts: %s",
+                    len(events),
+                    _retries + 1,
+                    exc,
+                )
+            except Exception:
+                logger.warning(
+                    "superpenguin: unexpected error sending events", exc_info=True
+                )
+                return

superpenguin/_context.py

@@ -0,0 +1,7 @@
+from __future__ import annotations
+
+import contextvars
+
+current_trace_id: contextvars.ContextVar[str | None] = contextvars.ContextVar(
+    "sp_trace_id", default=None
+)

superpenguin/_litellm.py

@@ -0,0 +1,278 @@
+"""patch_litellm() — automatic cost tracking for litellm.completion / acompletion."""
+
+from __future__ import annotations
+
+import functools
+import logging
+import time
+from typing import Any
+
+from superpenguin._pricing import calculate_cost_micros
+from superpenguin._wrap import (
+    _AsyncStreamProxy,
+    _OpenAIAccumulator,
+    _SyncStreamProxy,
+    _normalize_openai,
+    _extract_call_metadata,
+)
+
+logger = logging.getLogger("superpenguin")
+
+_patched = False
+
+
+def patch_litellm(
+    *,
+    name: str | None = None,
+    metadata: dict[str, Any] | None = None,
+    tags: list[str] | None = None,
+) -> None:
+    """Monkey-patch ``litellm.completion`` and ``litellm.acompletion``.
+
+    After calling this, every ``litellm.completion()`` and
+    ``litellm.acompletion()`` call is automatically tracked.
+
+    Args:
+        name: Override the default event name (``"chat.completions"``).
+        metadata: Default metadata for every litellm call.
+        tags: Tags attached to every event.
+
+    Usage::
+
+        import superpenguin as sp
+        import litellm
+
+        sp.init(api_key="sp_...")
+        sp.patch_litellm()
+
+        response = litellm.completion(
+            model="openai/gpt-4o",
+            messages=[{"role": "user", "content": "Hello"}],
+        )
+    """
+    global _patched
+    if _patched:
+        return
+
+    try:
+        import litellm  # noqa: F811
+    except ImportError as exc:
+        raise ImportError(
+            "litellm is not installed. Install it with: pip install litellm"
+        ) from exc
+
+    extra: dict[str, Any] = {}
+    if name:
+        extra["name"] = name
+    if metadata:
+        extra["metadata"] = metadata
+    if tags:
+        extra["tags"] = tags
+
+    _do_patch(litellm, "completion", is_async=False, extra=extra)
+    _do_patch(litellm, "acompletion", is_async=True, extra=extra)
+    _patched = True
+
+
+def _do_patch(
+    module: Any, attr: str, *, is_async: bool, extra: dict[str, Any],
+) -> None:
+    original = getattr(module, attr)
+
+    if is_async:
+
+        @functools.wraps(original)
+        async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+            return await _track_async(original, extra, args, kwargs)
+
+        setattr(module, attr, async_wrapper)
+    else:
+
+        @functools.wraps(original)
+        def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+            return _track_sync(original, extra, args, kwargs)
+
+        setattr(module, attr, sync_wrapper)
+
+
+# ---------------------------------------------------------------------------
+# Sync / async tracking
+# ---------------------------------------------------------------------------
+
+
+def _track_sync(
+    fn: Any,
+    extra: dict[str, Any],
+    args: tuple[Any, ...],
+    kwargs: dict[str, Any],
+) -> Any:
+    started_at = time.time()
+    is_stream = kwargs.get("stream", False)
+    call_meta = _capture_litellm_input(args, kwargs)
+
+    if is_stream:
+        kwargs = dict(kwargs)
+        opts = dict(kwargs.get("stream_options") or {})
+        opts["include_usage"] = True
+        kwargs["stream_options"] = opts
+
+    try:
+        result = fn(*args, **kwargs)
+    except Exception as exc:
+        _submit_event(
+            extra, call_meta, {}, started_at, time.time(),
+            status_code=getattr(exc, "status_code", 500),
+            streaming=False,
+        )
+        raise
+
+    if is_stream:
+        acc = _OpenAIAccumulator()
+
+        def on_done() -> None:
+            _submit_event(
+                extra, call_meta, acc.result(),
+                started_at, time.time(), status_code=200, streaming=True,
+            )
+
+        return _SyncStreamProxy(result, acc, on_done)
+
+    output = _normalize_openai(result)
+    _submit_event(
+        extra, call_meta, output,
+        started_at, time.time(), status_code=200, streaming=False,
+    )
+    return result
+
+
+async def _track_async(
+    fn: Any,
+    extra: dict[str, Any],
+    args: tuple[Any, ...],
+    kwargs: dict[str, Any],
+) -> Any:
+    started_at = time.time()
+    is_stream = kwargs.get("stream", False)
+    call_meta = _capture_litellm_input(args, kwargs)
+
+    if is_stream:
+        kwargs = dict(kwargs)
+        opts = dict(kwargs.get("stream_options") or {})
+        opts["include_usage"] = True
+        kwargs["stream_options"] = opts
+
+    try:
+        result = await fn(*args, **kwargs)
+    except Exception as exc:
+        _submit_event(
+            extra, call_meta, {}, started_at, time.time(),
+            status_code=getattr(exc, "status_code", 500),
+            streaming=False,
+        )
+        raise
+
+    if is_stream:
+        acc = _OpenAIAccumulator()
+
+        def on_done() -> None:
+            _submit_event(
+                extra, call_meta, acc.result(),
+                started_at, time.time(), status_code=200, streaming=True,
+            )
+
+        return _AsyncStreamProxy(result, acc, on_done)
+
+    output = _normalize_openai(result)
+    _submit_event(
+        extra, call_meta, output,
+        started_at, time.time(), status_code=200, streaming=False,
+    )
+    return result
+
+
+# ---------------------------------------------------------------------------
+# litellm-specific input capture
+# ---------------------------------------------------------------------------
+
+_POS_PARAMS = ("model", "messages")
+
+
+def _capture_litellm_input(
+    args: tuple[Any, ...], kwargs: dict[str, Any],
+) -> dict[str, Any]:
+    captured: dict[str, Any] = {}
+    for i, param_name in enumerate(_POS_PARAMS):
+        if param_name in kwargs:
+            captured[param_name] = kwargs[param_name]
+        elif i < len(args):
+            captured[param_name] = args[i]
+
+    meta_raw = kwargs.get("metadata") or {}
+    if isinstance(meta_raw, dict):
+        call_meta = _extract_call_metadata({"extra_body": {"metadata": meta_raw}})
+        captured.update(call_meta)
+
+    return captured
+
+
+# ---------------------------------------------------------------------------
+# Event submission
+# ---------------------------------------------------------------------------
+
+
+def _submit_event(
+    extra: dict[str, Any],
+    call_meta: dict[str, Any],
+    output_data: dict[str, Any],
+    started_at: float,
+    ended_at: float,
+    *,
+    status_code: int,
+    streaming: bool,
+) -> None:
+    from superpenguin import get_client
+
+    usage = output_data.get("usage", {})
+    input_tokens = usage.get("input_tokens", 0)
+    output_tokens = usage.get("output_tokens", 0)
+    cached_tokens = usage.get("cached_tokens", 0)
+    model = output_data.get("model") or call_meta.get("model", "")
+
+    cost_micros = calculate_cost_micros(
+        model, input_tokens, output_tokens, cached_tokens,
+    )
+
+    merged_meta = dict(extra.get("metadata") or {})
+    merged_meta.update(
+        {k: v for k, v in call_meta.items() if k not in ("model", "messages")}
+    )
+
+    event: dict[str, Any] = {
+        "provider": "litellm",
+        "model": model,
+        "input_tokens": input_tokens,
+        "output_tokens": output_tokens,
+        "cached_tokens": cached_tokens,
+        "cost_usd_micros": cost_micros,
+        "latency_ms": round((ended_at - started_at) * 1000),
+        "status_code": status_code,
+        "streaming": streaming,
+        "has_tools": output_data.get("has_tools", False),
+        "has_vision": False,
+    }
+
+    for key in (
+        "customer_id", "feature", "team", "environment",
+        "prompt_key", "prompt_version",
+    ):
+        val = merged_meta.get(key)
+        if val is not None:
+            event[key] = str(val) if key == "prompt_version" else val
+
+    if merged_meta.get("custom_tags"):
+        event["custom_tags"] = merged_meta["custom_tags"]
+
+    try:
+        get_client().submit(event)
+    except Exception:
+        logger.debug("superpenguin: failed to submit litellm event", exc_info=True)

superpenguin/_pricing.py

@@ -0,0 +1,77 @@
+"""Built-in cost estimation for known LLM models.
+
+Prices are in USD per 1M tokens. Cost is returned in USD micros
+(1 USD = 1,000,000 micros) for lossless integer arithmetic.
+"""
+
+from __future__ import annotations
+
+import re
+
+_MODEL_PRICING: dict[str, dict[str, float]] = {
+    # OpenAI
+    "gpt-5.4":          {"input": 2.5,   "output": 15.0,  "cache_read": 0.25},
+    "gpt-5.4-mini":     {"input": 0.75,  "output": 4.5,   "cache_read": 0.075},
+    "gpt-5.4-nano":     {"input": 0.2,   "output": 1.25,  "cache_read": 0.02},
+    "gpt-5":            {"input": 1.25,  "output": 10.0,  "cache_read": 0.125},
+    "gpt-5-mini":       {"input": 0.25,  "output": 2.0,   "cache_read": 0.025},
+    "gpt-5-nano":       {"input": 0.05,  "output": 0.4,   "cache_read": 0.005},
+    "gpt-4o":           {"input": 2.5,   "output": 10.0},
+    "gpt-4o-mini":      {"input": 0.15,  "output": 0.6},
+    "gpt-4.1":          {"input": 2.0,   "output": 8.0},
+    "gpt-4.1-mini":     {"input": 0.4,   "output": 1.6},
+    "gpt-4.1-nano":     {"input": 0.1,   "output": 0.4},
+    "o3":               {"input": 2.0,   "output": 8.0},
+    "o4-mini":          {"input": 1.1,   "output": 4.4},
+    # Anthropic
+    "claude-opus-4":     {"input": 15.0, "output": 75.0,  "cache_read": 1.5},
+    "claude-opus-4-6":   {"input": 5.0,  "output": 25.0,  "cache_read": 0.5},
+    "claude-sonnet-4":   {"input": 3.0,  "output": 15.0,  "cache_read": 0.3},
+    "claude-sonnet-4-5": {"input": 3.0,  "output": 15.0,  "cache_read": 0.3},
+    "claude-sonnet-4-6": {"input": 3.0,  "output": 15.0,  "cache_read": 0.3},
+    "claude-haiku-4-5":  {"input": 1.0,  "output": 5.0,   "cache_read": 0.1},
+    # Google Gemini
+    "gemini-2.5-pro":        {"input": 1.25,  "output": 10.0,  "cache_read": 0.3125},
+    "gemini-2.5-flash":      {"input": 0.15,  "output": 0.6,   "cache_read": 0.0375},
+    "gemini-2.0-pro":        {"input": 1.25,  "output": 10.0,  "cache_read": 0.3125},
+    "gemini-2.0-flash":      {"input": 0.1,   "output": 0.4,   "cache_read": 0.025},
+    "gemini-2.0-flash-lite": {"input": 0.075, "output": 0.3},
+    # xAI / Grok
+    "grok-3":      {"input": 3.0,  "output": 15.0},
+    "grok-3-mini": {"input": 0.3,  "output": 0.5},
+    "grok-3-fast": {"input": 5.0,  "output": 25.0},
+}
+
+_DATE_SUFFIX = re.compile(r"-\d{4}-\d{2}-\d{2}$|-\d{8}$")
+
+
+def _normalize_model(raw: str) -> str:
+    # Strip provider prefix (e.g. "openai/gpt-4o" → "gpt-4o")
+    if "/" in raw:
+        raw = raw.rsplit("/", 1)[-1]
+    return _DATE_SUFFIX.sub("", raw)
+
+
+def calculate_cost_micros(
+    model: str,
+    input_tokens: int,
+    output_tokens: int,
+    cached_tokens: int = 0,
+) -> int:
+    """Calculate request cost in USD micros (1 USD = 1,000,000 micros).
+
+    Returns 0 if the model is unknown.
+    """
+    normalized = _normalize_model(model)
+    pricing = _MODEL_PRICING.get(normalized)
+    if pricing is None:
+        return 0
+
+    uncached_input = max(0, input_tokens - cached_tokens)
+    input_cost = (uncached_input / 1_000_000) * pricing["input"]
+    output_cost = (output_tokens / 1_000_000) * pricing["output"]
+    cache_cost = (cached_tokens / 1_000_000) * pricing.get(
+        "cache_read", pricing["input"] * 0.5
+    )
+
+    return round((input_cost + output_cost + cache_cost) * 1_000_000)