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

frugal_metrics/__init__.py

@@ -0,0 +1,13 @@
+"""frugal-metrics: call-site instrumentation for AI API calls."""
+
+from importlib.metadata import PackageNotFoundError, version
+
+from frugal_metrics.wrappers.openai import wrap_openai
+from frugal_metrics.wrappers.anthropic import wrap_anthropic
+from frugal_metrics.wrappers.bedrock import wrap_bedrock
+
+__all__ = ["wrap_openai", "wrap_anthropic", "wrap_bedrock"]
+try:
+    __version__ = version("frugal-sdk-python")
+except PackageNotFoundError:
+    __version__ = "0+unknown"

frugal_metrics/block_hash_lru.py

@@ -0,0 +1,75 @@
+"""Per-call-site LRU storing the previous call's block hashes.
+
+Used by the prompt analyzer to derive cache-prefix-stable signals without
+leaving raw hashes anywhere outside the process. Default capacity matches the
+existing call-site cardinality cap so a busy process can't blow memory by
+tracking thousands of distinct sites.
+"""
+
+from __future__ import annotations
+
+import threading
+from collections import OrderedDict
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class BlockHashes:
+    system_hash: str | None = None
+    user_hash: str | None = None
+    tools_hash: str | None = None
+    schema_hash: str | None = None
+
+
+class BlockHashLRU:
+    def __init__(self, capacity: int = 500) -> None:
+        self._capacity = capacity
+        self._map: "OrderedDict[str, BlockHashes]" = OrderedDict()
+        self._lock = threading.Lock()
+
+    def get(self, call_site_id: str) -> BlockHashes | None:
+        with self._lock:
+            entry = self._map.get(call_site_id)
+            if entry is None:
+                return None
+            self._map.move_to_end(call_site_id)
+            return entry
+
+    def set(self, call_site_id: str, hashes: BlockHashes) -> None:
+        with self._lock:
+            if call_site_id in self._map:
+                self._map.move_to_end(call_site_id)
+                self._map[call_site_id] = hashes
+                return
+            self._map[call_site_id] = hashes
+            if self._capacity > 0:
+                while len(self._map) > self._capacity:
+                    self._map.popitem(last=False)
+
+    def size(self) -> int:
+        with self._lock:
+            return len(self._map)
+
+    def clear(self) -> None:
+        with self._lock:
+            self._map.clear()
+
+
+_shared: BlockHashLRU | None = None
+_shared_lock = threading.Lock()
+
+
+def get_shared_block_hash_lru(capacity: int | None = None) -> BlockHashLRU:
+    global _shared
+    if _shared is not None:
+        return _shared
+    with _shared_lock:
+        if _shared is None:
+            _shared = BlockHashLRU(capacity if capacity is not None else 500)
+        return _shared
+
+
+def reset_shared_block_hash_lru_for_tests() -> None:
+    global _shared
+    with _shared_lock:
+        _shared = None

frugal_metrics/caller.py

@@ -0,0 +1,213 @@
+"""Stack walking to resolve the call site of a wrapped AI call.
+
+Algorithm: starting from the frame above the wrapper, walk upward until we find
+a frame whose source file is not (a) inside our own package, (b) inside the
+wrapped SDK's package, or (c) under any FRUGAL_PATH_BLOCK_LIST prefix. Return
+the file relative to FRUGAL_REPO_ROOT plus the function name.
+
+Performance: realpath results and full caller resolutions are cached so the
+stack walk and path resolution only happen once per unique call site. Subsequent
+calls from the same site are a dict lookup.
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+import threading
+from dataclasses import dataclass
+
+from frugal_metrics.semconv import (
+    CALLER_FILE_EXTERNAL,
+    CALLER_FILE_UNKNOWN,
+    CALLER_OVERFLOW,
+)
+
+# One entry per distinct (filename, function) pair that ever calls a wrapped AI
+# method. Even a very large codebase has <10k unique sites.
+_CACHE_MAX_SIZE = 10_000
+
+_realpath_cache: dict[str, str] = {}
+_realpath_lock = threading.Lock()
+
+# Keyed on (raw_filename, function_name) — the skip_dirs, block_list, and
+# repo_root are stable for the lifetime of a process (they come from env vars
+# loaded once at config time), so we don't need them in the key.
+_caller_cache: dict[tuple[str, str], CallerInfo | None] = {}
+_caller_cache_lock = threading.Lock()
+
+# Cardinality cap state — first-come-first-served. Once `_seen_call_sites`
+# hits the configured max, subsequent unseen pairs overflow (partial if the
+# file is already in `_seen_files`, full otherwise).
+_seen_call_sites: set[tuple[str, str]] = set()
+_seen_files: set[str] = set()
+_cap_state_lock = threading.Lock()
+_cap_hit_notified = False
+_cap_hit_pending = False
+
+
+@dataclass(frozen=True)
+class CallerInfo:
+    file: str
+    function: str
+
+
+def resolve_caller(
+    *,
+    skip_package_dirs: tuple[str, ...],
+    block_list: tuple[str, ...],
+    repo_root: str,
+    max_call_sites: int = 0,
+) -> CallerInfo:
+    """Walk the stack for the first frame outside our wrapper / wrapped SDK / block list.
+
+    Results are cached by (raw_filename, function_name) so the expensive path
+    resolution only runs once per call site.
+    """
+    frame = sys._getframe(1)
+    # Skip the wrapper frame that called us.
+    if frame is not None:
+        frame = frame.f_back
+
+    while frame is not None:
+        raw_filename = frame.f_code.co_filename
+        function = frame.f_code.co_name
+
+        # Fast path: check the caller cache before doing any path work.
+        cache_key = (raw_filename, function)
+        cached = _caller_cache.get(cache_key)
+        if cached is not None:
+            return cached
+
+        abs_filename = _cached_realpath(raw_filename) if raw_filename else ""
+
+        if _is_in_any(abs_filename, skip_package_dirs):
+            # Cache this as a "skip" so we don't realpath again next time.
+            _cache_put(cache_key, None)
+            frame = frame.f_back
+            continue
+
+        rel = _repo_relative(abs_filename, repo_root)
+        if rel is not None and _is_blocked(rel, block_list):
+            _cache_put(cache_key, None)
+            frame = frame.f_back
+            continue
+
+        resolved_file = CALLER_FILE_EXTERNAL if rel is None else rel
+        result = _apply_cap(resolved_file, function, max_call_sites)
+
+        _cache_put(cache_key, result)
+        return result
+
+    return CallerInfo(file=CALLER_FILE_UNKNOWN, function="<unknown>")
+
+
+def _apply_cap(file: str, function: str, max_call_sites: int) -> CallerInfo:
+    global _cap_hit_notified, _cap_hit_pending
+    if max_call_sites <= 0:
+        return CallerInfo(file=file, function=function)
+    key = (file, function)
+    # Fast path — already known, no lock needed for read (sets are thread-safe
+    # for membership in CPython; we take the lock only on mutation).
+    if key in _seen_call_sites:
+        return CallerInfo(file=file, function=function)
+    with _cap_state_lock:
+        # Re-check under the lock.
+        if key in _seen_call_sites:
+            return CallerInfo(file=file, function=function)
+        if len(_seen_call_sites) < max_call_sites:
+            _seen_call_sites.add(key)
+            _seen_files.add(file)
+            return CallerInfo(file=file, function=function)
+        # At cap.
+        if not _cap_hit_notified:
+            _cap_hit_notified = True
+            _cap_hit_pending = True
+        if file in _seen_files:
+            return CallerInfo(file=file, function=CALLER_OVERFLOW)
+        return CallerInfo(file=CALLER_OVERFLOW, function=CALLER_OVERFLOW)
+
+
+def consume_cap_hit_event() -> bool:
+    """Returns True once per process when the cap is first crossed.
+
+    The instrument layer calls this after each resolve; on True it emits a
+    single `frugal.gen_ai.internal_errors{stage=cardinality_cap_hit}`.
+    """
+    global _cap_hit_pending
+    with _cap_state_lock:
+        if _cap_hit_pending:
+            _cap_hit_pending = False
+            return True
+        return False
+
+
+def _cached_realpath(path: str) -> str:
+    cached = _realpath_cache.get(path)
+    if cached is not None:
+        return cached
+    resolved = os.path.realpath(path)
+    with _realpath_lock:
+        if len(_realpath_cache) < _CACHE_MAX_SIZE:
+            _realpath_cache[path] = resolved
+    return resolved
+
+
+def _cache_put(key: tuple[str, str], value: CallerInfo | None) -> None:
+    with _caller_cache_lock:
+        if len(_caller_cache) < _CACHE_MAX_SIZE:
+            _caller_cache[key] = value
+
+
+def _is_in_any(abs_filename: str, dirs: tuple[str, ...]) -> bool:
+    if not abs_filename:
+        return False
+    for d in dirs:
+        if not d:
+            continue
+        if abs_filename == d or abs_filename.startswith(d.rstrip(os.sep) + os.sep):
+            return True
+    return False
+
+
+def _repo_relative(abs_filename: str, repo_root: str) -> str | None:
+    if not abs_filename:
+        return None
+    root = repo_root.rstrip(os.sep)
+    if abs_filename == root:
+        return ""
+    prefix = root + os.sep
+    if abs_filename.startswith(prefix):
+        return abs_filename[len(prefix):]
+    return None
+
+
+def _is_blocked(rel_path: str, block_list: tuple[str, ...]) -> bool:
+    for entry in block_list:
+        if not entry:
+            continue
+        if rel_path == entry or rel_path.startswith(entry.rstrip("/") + "/"):
+            return True
+    return False
+
+
+def package_dir_for(module: object) -> str:
+    """Return the absolute directory of a module, suitable for skip_package_dirs."""
+    module_file = getattr(module, "__file__", None)
+    if not module_file:
+        return ""
+    return os.path.realpath(os.path.dirname(module_file))
+
+
+def reset_caches_for_tests() -> None:
+    """Clear all caches — test-only hook."""
+    global _cap_hit_notified, _cap_hit_pending
+    with _realpath_lock:
+        _realpath_cache.clear()
+    with _caller_cache_lock:
+        _caller_cache.clear()
+    with _cap_state_lock:
+        _seen_call_sites.clear()
+        _seen_files.clear()
+        _cap_hit_notified = False
+        _cap_hit_pending = False

frugal_metrics/config.py

@@ -0,0 +1,231 @@
+"""Environment-variable configuration loading.
+
+All configuration comes from env vars — no init() API. If any required variable
+is missing, load_config() returns None and the wrappers no-op (returning the
+client unwrapped). A single warning is logged on the first no-op.
+
+Minimum required production config (two env vars):
+    FRUGAL_API_KEY     — per-customer bearer token (Frugal mints these per tenant)
+    FRUGAL_PROJECT_ID  — the Frugal project this deployment belongs to. Stats are
+                         attributed per project (file/function is unique within a
+                         project), so the same repo deployed across microservices
+                         no longer collides. Frugal shows this value in the UI.
+
+Everything else has sensible defaults:
+    - Endpoint defaults to https://metrics.frugal.co/opentelemetry/v1/metrics
+    - REPO_ROOT defaults to os.getcwd() (the strip prefix for caller_file)
+    - CUSTOMER_ID is optional; when unset, the tenant label (auto-injected
+      server-side by VMAuth from the bearer token) is the canonical customer
+      identifier and the frugal.customer_id label is omitted.
+    - COMPONENTS is optional; a comma-separated list (e.g. "ai-handler,reco").
+      When set, every metric is ALSO emitted once per component (tagged with
+      frugal.component) on top of the base project-level series, so usage can be
+      viewed for the whole project (component absent) or split by component.
+
+The older env vars (FRUGAL_OTEL_EXPORTER_OTLP_ENDPOINT,
+FRUGAL_OTEL_EXPORTER_OTLP_HEADERS, FRUGAL_REPO_ROOT, FRUGAL_CUSTOMER_ID,
+plus the standard OTel ones) all still work — explicit values override the
+defaults so existing customers don't need to change anything.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import threading
+from dataclasses import dataclass, field
+
+logger = logging.getLogger("frugal_metrics")
+
+_ENV_API_KEY = "FRUGAL_API_KEY"
+_ENV_CUSTOMER_ID = "FRUGAL_CUSTOMER_ID"
+_ENV_PROJECT_ID = "FRUGAL_PROJECT_ID"
+_ENV_COMPONENTS = "FRUGAL_COMPONENTS"
+_ENV_REPO_ROOT = "FRUGAL_REPO_ROOT"
+_ENV_PATH_BLOCK_LIST = "FRUGAL_PATH_BLOCK_LIST"
+_ENV_DISABLED = "FRUGAL_DISABLED"
+
+_ENV_FRUGAL_OTLP_ENDPOINT = "FRUGAL_OTEL_EXPORTER_OTLP_ENDPOINT"
+_ENV_FRUGAL_OTLP_HEADERS = "FRUGAL_OTEL_EXPORTER_OTLP_HEADERS"
+_ENV_OTLP_ENDPOINT = "OTEL_EXPORTER_OTLP_ENDPOINT"
+_ENV_OTLP_HEADERS = "OTEL_EXPORTER_OTLP_HEADERS"
+_ENV_EXPORT_INTERVAL_MS = "FRUGAL_METRIC_EXPORT_INTERVAL_MS"
+_ENV_MAX_CALL_SITES = "FRUGAL_MAX_CALL_SITES"
+_ENV_PROMPT_ANALYSIS_SAMPLE_RATE = "FRUGAL_PROMPT_ANALYSIS_SAMPLE_RATE"
+
+_DEFAULT_OTLP_ENDPOINT = "https://metrics.frugal.co/opentelemetry/v1/metrics"
+_DEFAULT_MAX_CALL_SITES = 500
+_DEFAULT_PROMPT_ANALYSIS_SAMPLE_RATE = 0.1
+
+_warned_once = False
+_warned_lock = threading.Lock()
+
+
+@dataclass(frozen=True)
+class Config:
+    # Optional now — when None the tenant label (auto-injected by VMAuth)
+    # identifies the customer.
+    customer_id: str | None
+    project_id: str
+    repo_root: str
+    # Optional logical components; metrics fan out one extra series per entry.
+    components: tuple[str, ...] = field(default_factory=tuple)
+    path_block_list: tuple[str, ...] = field(default_factory=tuple)
+    otlp_endpoint: str = _DEFAULT_OTLP_ENDPOINT
+    otlp_headers: str | None = None
+    export_interval_ms: int = 60_000
+    # Max distinct (file, function) combinations to attribute before overflow.
+    # 0 disables the cap.
+    max_call_sites: int = _DEFAULT_MAX_CALL_SITES
+    # Fraction of analyzed calls to emit prompt-analysis metrics for, 0..1.
+    prompt_analysis_sample_rate: float = _DEFAULT_PROMPT_ANALYSIS_SAMPLE_RATE
+
+
+def _truthy(value: str | None) -> bool:
+    return value is not None and value.strip().lower() in {"1", "true", "yes", "on"}
+
+
+def _warn_once(message: str) -> None:
+    global _warned_once
+    with _warned_lock:
+        if _warned_once:
+            return
+        _warned_once = True
+    logger.warning("frugal-metrics disabled: %s", message)
+
+
+def load_config() -> Config | None:
+    """Read env vars and return a Config, or None if disabled or misconfigured.
+
+    Side effect: logs a single warning the first time we return None due to
+    missing required vars. Repeated calls stay silent.
+    """
+    if _truthy(os.environ.get(_ENV_DISABLED)):
+        return None
+
+    customer_id = os.environ.get(_ENV_CUSTOMER_ID) or None
+    project_id = os.environ.get(_ENV_PROJECT_ID)
+    repo_root_env = os.environ.get(_ENV_REPO_ROOT)
+    api_key = os.environ.get(_ENV_API_KEY)
+
+    # Explicit headers override anything we'd build from the API key — keeps
+    # existing deployments that set FRUGAL_OTEL_EXPORTER_OTLP_HEADERS directly
+    # working unchanged.
+    explicit_headers = (
+        os.environ.get(_ENV_FRUGAL_OTLP_HEADERS)
+        or os.environ.get(_ENV_OTLP_HEADERS)
+    )
+    if explicit_headers:
+        otlp_headers: str | None = explicit_headers
+    elif api_key:
+        otlp_headers = f"Authorization=Bearer {api_key}"
+    else:
+        otlp_headers = None
+
+    # Required for any useful run: PROJECT_ID (the attribution key) and SOME auth
+    # (either FRUGAL_API_KEY or an explicit OTLP_HEADERS containing the bearer).
+    missing = []
+    if not project_id:
+        missing.append(_ENV_PROJECT_ID)
+    if not otlp_headers:
+        missing.append(f"{_ENV_API_KEY} or {_ENV_FRUGAL_OTLP_HEADERS}")
+    if missing:
+        _warn_once(f"missing required env vars: {', '.join(missing)}")
+        return None
+
+    # Normalize repo_root to a real absolute path so caller-file relativization
+    # is consistent even when the process sees symlinked mounts (e.g. macOS
+    # /var -> /private/var). Default to os.getcwd() — works for the standard
+    # container pattern where WORKDIR is set to the source-code root.
+    repo_root_raw = repo_root_env or os.getcwd()
+    repo_root_abs = os.path.realpath(repo_root_raw)
+
+    block_list_raw = os.environ.get(_ENV_PATH_BLOCK_LIST, "")
+    block_list = tuple(
+        _normalize_block_entry(entry)
+        for entry in block_list_raw.split(",")
+        if entry.strip()
+    )
+
+    # Optional comma-separated logical components (e.g. "ai-handler,reco"). Each
+    # becomes an extra per-component series on top of the base project series.
+    components_raw = os.environ.get(_ENV_COMPONENTS, "")
+    components = tuple(
+        entry.strip() for entry in components_raw.split(",") if entry.strip()
+    )
+
+    otlp_endpoint = (
+        os.environ.get(_ENV_FRUGAL_OTLP_ENDPOINT)
+        or os.environ.get(_ENV_OTLP_ENDPOINT)
+        or _DEFAULT_OTLP_ENDPOINT
+    )
+
+    export_interval_ms = _parse_int(
+        os.environ.get(_ENV_EXPORT_INTERVAL_MS), default=60_000
+    )
+    max_call_sites = _parse_int(
+        os.environ.get(_ENV_MAX_CALL_SITES),
+        default=_DEFAULT_MAX_CALL_SITES,
+    )
+    if max_call_sites < 0:
+        max_call_sites = _DEFAULT_MAX_CALL_SITES
+
+    sample_rate = _parse_float(
+        os.environ.get(_ENV_PROMPT_ANALYSIS_SAMPLE_RATE),
+        default=_DEFAULT_PROMPT_ANALYSIS_SAMPLE_RATE,
+    )
+    if sample_rate < 0 or sample_rate > 1:
+        sample_rate = _DEFAULT_PROMPT_ANALYSIS_SAMPLE_RATE
+
+    return Config(
+        customer_id=customer_id,
+        project_id=project_id,  # type: ignore[arg-type]
+        repo_root=repo_root_abs,
+        components=components,
+        path_block_list=block_list,
+        otlp_endpoint=otlp_endpoint,
+        otlp_headers=otlp_headers,
+        export_interval_ms=export_interval_ms,
+        max_call_sites=max_call_sites,
+        prompt_analysis_sample_rate=sample_rate,
+    )
+
+
+def _normalize_block_entry(entry: str) -> str:
+    stripped = entry.strip().strip("/")
+    return stripped
+
+
+def _parse_int(raw: str | None, *, default: int) -> int:
+    if raw is None or not raw.strip():
+        return default
+    try:
+        return int(raw)
+    except ValueError:
+        logger.warning(
+            "frugal-metrics: could not parse int env var value %r, using default %d",
+            raw,
+            default,
+        )
+        return default
+
+
+def _parse_float(raw: str | None, *, default: float) -> float:
+    if raw is None or not raw.strip():
+        return default
+    try:
+        return float(raw)
+    except ValueError:
+        logger.warning(
+            "frugal-metrics: could not parse float env var value %r, using default %s",
+            raw,
+            default,
+        )
+        return default
+
+
+def reset_warning_state_for_tests() -> None:
+    """Test-only hook to re-arm the one-shot warning."""
+    global _warned_once
+    with _warned_lock:
+        _warned_once = False