synckey 0.1.0__py3-none-any.whl

synckey/state.py

@@ -0,0 +1,152 @@
+"""Durable key health state.
+
+In-memory dict is the hot path (zero I/O on reads). SQLite is the durability
+layer: written immediately on every health transition so a restart never forgets
+which keys are cooling or dead.
+
+Health transitions:
+    LIVE    -> COOLING  on 429 (auto-recovers after cooldown_until)
+    LIVE    -> DEAD     on 401/403 or explicit quota exhaustion
+    COOLING -> LIVE     automatically when cooldown_until passes (no write needed)
+    DEAD    -> LIVE     only via `synckey key enable` (user action)
+"""
+
+from __future__ import annotations
+
+import time
+from dataclasses import dataclass
+from enum import IntEnum
+
+import sqlite3
+from pathlib import Path
+
+
+class Health(IntEnum):
+    LIVE    = 0
+    COOLING = 1
+    DEAD    = 2
+
+
+@dataclass
+class KeyState:
+    health: Health = Health.LIVE
+    cooldown_until: float = 0.0
+    dead_reason: str = ""
+    rpm_observed: float | None = None
+    tpm_observed: float | None = None
+
+    def is_available(self) -> bool:
+        if self.health == Health.DEAD:
+            return False
+        if self.health == Health.COOLING and time.time() < self.cooldown_until:
+            return False
+        return True
+
+    def cooldown_remaining(self) -> float:
+        if self.health == Health.COOLING:
+            return max(0.0, self.cooldown_until - time.time())
+        return 0.0
+
+
+_CREATE = """
+CREATE TABLE IF NOT EXISTS key_state (
+    key_id          INTEGER PRIMARY KEY,
+    health          INTEGER NOT NULL DEFAULT 0,
+    cooldown_until  REAL    NOT NULL DEFAULT 0,
+    dead_reason     TEXT    NOT NULL DEFAULT '',
+    rpm_observed    REAL,
+    tpm_observed    REAL,
+    updated_at      REAL    NOT NULL DEFAULT 0
+);
+"""
+
+_UPSERT = """
+INSERT INTO key_state(key_id, health, cooldown_until, dead_reason, rpm_observed, tpm_observed, updated_at)
+VALUES (?, ?, ?, ?, ?, ?, ?)
+ON CONFLICT(key_id) DO UPDATE SET
+    health         = excluded.health,
+    cooldown_until = excluded.cooldown_until,
+    dead_reason    = excluded.dead_reason,
+    rpm_observed   = excluded.rpm_observed,
+    tpm_observed   = excluded.tpm_observed,
+    updated_at     = excluded.updated_at
+"""
+
+
+class StateStore:
+    """In-memory key health store backed by SQLite for durability."""
+
+    def __init__(self, db_path: Path):
+        self._mem: dict[int, KeyState] = {}
+        self._conn = sqlite3.connect(str(db_path), check_same_thread=False)
+        self._conn.execute("PRAGMA journal_mode=WAL")
+        self._conn.execute("PRAGMA synchronous=NORMAL")
+        self._conn.execute(_CREATE)
+        self._conn.commit()
+        self._load()
+
+    def _load(self) -> None:
+        now = time.time()
+        for row in self._conn.execute("SELECT * FROM key_state").fetchall():
+            key_id, health, cooldown_until, dead_reason, rpm_obs, tpm_obs, _ = row
+            h = Health(health)
+            # Expired cooldowns auto-recover to LIVE on load.
+            if h == Health.COOLING and cooldown_until <= now:
+                h = Health.LIVE
+                cooldown_until = 0.0
+            self._mem[key_id] = KeyState(
+                health=h,
+                cooldown_until=cooldown_until,
+                dead_reason=dead_reason or "",
+                rpm_observed=rpm_obs,
+                tpm_observed=tpm_obs,
+            )
+
+    def _persist(self, key_id: int) -> None:
+        st = self._mem[key_id]
+        self._conn.execute(
+            _UPSERT,
+            (
+                key_id,
+                int(st.health),
+                st.cooldown_until,
+                st.dead_reason,
+                st.rpm_observed,
+                st.tpm_observed,
+                time.time(),
+            ),
+        )
+        self._conn.commit()
+
+    def get(self, key_id: int) -> KeyState:
+        return self._mem.get(key_id, KeyState())
+
+    def set_cooling(self, key_id: int, until: float, rpm_observed: float | None = None) -> None:
+        st = self._mem.get(key_id, KeyState())
+        st.health = Health.COOLING
+        st.cooldown_until = until
+        if rpm_observed is not None:
+            st.rpm_observed = rpm_observed
+        self._mem[key_id] = st
+        self._persist(key_id)
+
+    def set_dead(self, key_id: int, reason: str) -> None:
+        st = self._mem.get(key_id, KeyState())
+        st.health = Health.DEAD
+        st.dead_reason = reason
+        self._mem[key_id] = st
+        self._persist(key_id)
+
+    def set_live(self, key_id: int) -> None:
+        st = self._mem.get(key_id, KeyState())
+        if st.health != Health.LIVE:
+            st.health = Health.LIVE
+            st.cooldown_until = 0.0
+            self._mem[key_id] = st
+            self._persist(key_id)
+
+    def all(self) -> dict[int, KeyState]:
+        return dict(self._mem)
+
+    def close(self) -> None:
+        self._conn.close()

synckey/tiers.py

@@ -0,0 +1,157 @@
+"""Model tier catalog, pricing, and tier-floored fallback logic.
+
+Tiers:
+    FRONTIER (4)  best quality, most expensive (Opus, GPT-5, o3)
+    HIGH     (3)  capable, premium (Sonnet, GPT-4o, o1)
+    MID      (2)  fast, affordable (Haiku, GPT-4o-mini, Flash, 70b)
+    LOW      (1)  budget, high-volume (8b, Gemma, Qwen)
+
+A request for a frontier model will never fall back below FRONTIER unless the
+caller explicitly sets X-Quality-Floor: high (or lower). This is the Opus ->
+GPT-5 guarantee, never Opus -> oss-llama.
+"""
+
+from __future__ import annotations
+
+import re
+
+FRONTIER = 4
+HIGH     = 3
+MID      = 2
+LOW      = 1
+
+TIER_NAMES = {FRONTIER: "frontier", HIGH: "high", MID: "mid", LOW: "low"}
+TIER_BY_NAME = {v: k for k, v in TIER_NAMES.items()}
+
+# (pattern, tier) -- checked in order; more-specific patterns first.
+# gpt-4o-mini must come before gpt-4o to avoid premature match.
+_PATTERNS: list[tuple[re.Pattern, int]] = [
+    (re.compile(r"claude-opus",         re.I), FRONTIER),
+    (re.compile(r"gpt-5",               re.I), FRONTIER),
+    (re.compile(r"^o3",                 re.I), FRONTIER),
+    (re.compile(r"gemini-2\.5-pro",     re.I), FRONTIER),
+    (re.compile(r"grok-3",              re.I), FRONTIER),
+    (re.compile(r"magistral",           re.I), FRONTIER),
+
+    (re.compile(r"claude-sonnet",       re.I), HIGH),
+    (re.compile(r"gpt-4\.1(?!-mini)",   re.I), HIGH),
+    (re.compile(r"gpt-4o(?!-mini)",     re.I), HIGH),
+    (re.compile(r"^o1",                 re.I), HIGH),
+    (re.compile(r"gemini-2\.0-pro",     re.I), HIGH),
+    (re.compile(r"gemini-1\.5-pro",     re.I), HIGH),
+    (re.compile(r"mistral-large",       re.I), HIGH),
+    (re.compile(r"grok-2",              re.I), HIGH),
+    (re.compile(r"deepseek-r1",         re.I), HIGH),
+
+    (re.compile(r"claude-haiku",        re.I), MID),
+    (re.compile(r"gpt-4o-mini",         re.I), MID),
+    (re.compile(r"gpt-4\.1-mini",       re.I), MID),
+    (re.compile(r"gemini-2\.0-flash",   re.I), MID),
+    (re.compile(r"gemini-flash",        re.I), MID),
+    (re.compile(r"llama-3\.3-70b",      re.I), MID),
+    (re.compile(r"llama-3\.1-70b",      re.I), MID),
+    (re.compile(r"llama-3-70b",         re.I), MID),
+    (re.compile(r"mistral-medium",      re.I), MID),
+    (re.compile(r"codestral",           re.I), MID),
+    (re.compile(r"deepseek-chat",       re.I), MID),
+    (re.compile(r"command-r-plus",      re.I), MID),
+    (re.compile(r"command-r(?!-plus)",  re.I), MID),
+    (re.compile(r"qwen-72b",            re.I), MID),
+    (re.compile(r"mixtral",             re.I), MID),
+    (re.compile(r"llama3\.1-70b",       re.I), MID),
+
+    (re.compile(r"gemma",               re.I), LOW),
+    (re.compile(r"llama-3\.1-8b",       re.I), LOW),
+    (re.compile(r"llama-3-8b",          re.I), LOW),
+    (re.compile(r"llama3\.1-8b",        re.I), LOW),
+    (re.compile(r"qwen(?!-72b)",        re.I), LOW),
+    (re.compile(r"mistral-7b",          re.I), LOW),
+    (re.compile(r"mistral-0",           re.I), LOW),
+    (re.compile(r"command-light",       re.I), LOW),
+    (re.compile(r"gemini-1\.5-flash-8", re.I), LOW),
+    (re.compile(r"phi-3",               re.I), LOW),
+]
+
+# (input $/M tokens, output $/M tokens) -- best-effort, users can override
+_PRICES: list[tuple[re.Pattern, tuple[float, float]]] = [
+    (re.compile(r"claude-opus-4-8",     re.I), (15.00, 75.00)),
+    (re.compile(r"claude-opus",         re.I), (15.00, 75.00)),
+    (re.compile(r"claude-sonnet-4-6",   re.I), ( 3.00, 15.00)),
+    (re.compile(r"claude-sonnet",       re.I), ( 3.00, 15.00)),
+    (re.compile(r"claude-haiku-4-5",    re.I), ( 0.80,  4.00)),
+    (re.compile(r"claude-haiku",        re.I), ( 0.80,  4.00)),
+    (re.compile(r"gpt-5",               re.I), (10.00, 40.00)),
+    (re.compile(r"gpt-4\.1-mini",       re.I), ( 0.40,  1.60)),
+    (re.compile(r"gpt-4\.1",            re.I), ( 2.00,  8.00)),
+    (re.compile(r"gpt-4o-mini",         re.I), ( 0.15,  0.60)),
+    (re.compile(r"gpt-4o",              re.I), ( 5.00, 15.00)),
+    (re.compile(r"^o3",                 re.I), (10.00, 40.00)),
+    (re.compile(r"^o1-mini",            re.I), ( 3.00, 12.00)),
+    (re.compile(r"^o1",                 re.I), (15.00, 60.00)),
+    (re.compile(r"gemini-2\.5-pro",     re.I), ( 3.50, 10.50)),
+    (re.compile(r"gemini-2\.0-flash",   re.I), ( 0.10,  0.40)),
+    (re.compile(r"gemini-flash",        re.I), ( 0.10,  0.40)),
+    (re.compile(r"gemini-1\.5-pro",     re.I), ( 3.50, 10.50)),
+    (re.compile(r"gemini-1\.5-flash",   re.I), ( 0.075, 0.30)),
+    (re.compile(r"llama-3\.3-70b",      re.I), ( 0.59,  0.79)),
+    (re.compile(r"llama-3\.1-70b",      re.I), ( 0.52,  0.75)),
+    (re.compile(r"llama-3\.1-8b",       re.I), ( 0.05,  0.08)),
+    (re.compile(r"llama3\.1-70b",       re.I), ( 0.52,  0.75)),
+    (re.compile(r"llama3\.1-8b",        re.I), ( 0.05,  0.08)),
+    (re.compile(r"mistral-large",       re.I), ( 3.00,  9.00)),
+    (re.compile(r"mistral-medium",      re.I), ( 2.70,  8.10)),
+    (re.compile(r"codestral",           re.I), ( 0.30,  0.90)),
+    (re.compile(r"deepseek-r1",         re.I), ( 0.55,  2.19)),
+    (re.compile(r"deepseek-chat",       re.I), ( 0.27,  1.10)),
+    (re.compile(r"command-r-plus",      re.I), ( 2.50, 10.00)),
+    (re.compile(r"command-r",           re.I), ( 0.15,  0.60)),
+    (re.compile(r"grok-3",              re.I), ( 3.00, 15.00)),
+    (re.compile(r"grok-2",              re.I), ( 2.00, 10.00)),
+    (re.compile(r"mixtral",             re.I), ( 0.24,  0.24)),
+    (re.compile(r"gemma",               re.I), ( 0.10,  0.10)),
+]
+
+# User-configurable overrides; set by load_overrides().
+_tier_overrides: dict[str, int] = {}
+_price_overrides: dict[str, tuple[float, float]] = {}
+
+
+def load_overrides(tiers: dict[str, str], prices: dict[str, list]) -> None:
+    """Load user-defined tier and price overrides from config."""
+    for model, tier_name in tiers.items():
+        t = TIER_BY_NAME.get(tier_name.lower())
+        if t is not None:
+            _tier_overrides[model.lower()] = t
+    for model, pair in prices.items():
+        if len(pair) == 2:
+            _price_overrides[model.lower()] = (float(pair[0]), float(pair[1]))
+
+
+def model_tier(model: str) -> int | None:
+    """Return the capability tier for a model name, or None if unknown."""
+    key = model.lower()
+    if key in _tier_overrides:
+        return _tier_overrides[key]
+    for rx, tier in _PATTERNS:
+        if rx.search(model):
+            return tier
+    return None
+
+
+def model_price(model: str) -> tuple[float, float] | None:
+    """Return (input, output) price in $/M tokens, or None if unknown."""
+    key = model.lower()
+    if key in _price_overrides:
+        return _price_overrides[key]
+    for rx, price in _PRICES:
+        if rx.search(model):
+            return price
+    return None
+
+
+def cost_usd(model: str, prompt_tokens: int, completion_tokens: int) -> float | None:
+    """Compute cost in USD for a completed request, or None if price unknown."""
+    price = model_price(model)
+    if price is None:
+        return None
+    return (prompt_tokens * price[0] + completion_tokens * price[1]) / 1_000_000

synckey/usage.py

@@ -0,0 +1,99 @@
+"""Batched off-thread usage writer.
+
+The gateway calls record() once per request. Non-blocking: the row queues and
+a background thread flushes in batches. One fsync per batch instead of per request.
+"""
+
+from __future__ import annotations
+
+import queue
+import sqlite3
+import threading
+import time
+from pathlib import Path
+
+from .db import USAGE_INSERT, tune
+
+Row = tuple
+
+
+class UsageWriter:
+    def __init__(self, path: Path, batch_size: int = 200, flush_interval: float = 0.5):
+        self.path = Path(path)
+        self.batch_size = batch_size
+        self.flush_interval = flush_interval
+        self.queue: queue.Queue[Row] = queue.Queue(maxsize=20000)
+        self._stop = threading.Event()
+        self._thread: threading.Thread | None = None
+
+    def start(self) -> None:
+        self._thread = threading.Thread(target=self._run, name="synckey-usage", daemon=True)
+        self._thread.start()
+
+    def record(
+        self,
+        *,
+        provider: str,
+        model: str,
+        key_id: int | None,
+        prompt_tokens: int = 0,
+        completion_tokens: int = 0,
+        total_tokens: int = 0,
+        cost_usd: float | None = None,
+        status_code: int = 0,
+        latency_ms: int = 0,
+        stream: bool = False,
+        tier: int | None = None,
+        error: str | None = None,
+    ) -> None:
+        row = (
+            time.time(),
+            provider,
+            model,
+            key_id,
+            prompt_tokens,
+            completion_tokens,
+            total_tokens,
+            cost_usd,
+            status_code,
+            latency_ms,
+            1 if stream else 0,
+            tier,
+            error,
+        )
+        try:
+            self.queue.put_nowait(row)
+        except queue.Full:
+            pass  # drop metrics under extreme backpressure
+
+    def stop(self, timeout: float = 5.0) -> None:
+        self._stop.set()
+        if self._thread:
+            self._thread.join(timeout=timeout)
+
+    def _run(self) -> None:
+        conn = sqlite3.connect(str(self.path), check_same_thread=False)
+        tune(conn)
+        try:
+            while True:
+                batch = self._collect()
+                if batch:
+                    conn.executemany(USAGE_INSERT, batch)
+                    conn.commit()
+                elif self._stop.is_set() and self.queue.empty():
+                    return
+        finally:
+            conn.close()
+
+    def _collect(self) -> list[Row]:
+        batch: list[Row] = []
+        try:
+            batch.append(self.queue.get(timeout=self.flush_interval))
+        except queue.Empty:
+            return batch
+        while len(batch) < self.batch_size:
+            try:
+                batch.append(self.queue.get_nowait())
+            except queue.Empty:
+                break
+        return batch