214 lines
7.4 KiB
Python
214 lines
7.4 KiB
Python
#!/usr/bin/env python3
|
|
"""HTTPS client for LiteLLM / OpenAI-compatible model gateways. No agent CLI, no SDK.
|
|
|
|
Built on urllib so the program has zero third-party HTTP dependencies.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
import random
|
|
import threading
|
|
import time
|
|
import urllib.error
|
|
import urllib.request
|
|
from dataclasses import dataclass, field
|
|
|
|
from .config import CFG
|
|
from .log import log
|
|
|
|
PRICES = {
|
|
"aistudio/gemini-flash-lite-latest": (0.3, 2.5),
|
|
"vertex/gemini-flash-lite-latest": (0.3, 2.5),
|
|
"gemini-flash-lite-latest": (0.3, 2.5),
|
|
"google/gemini-3.5-flash-lite": (0.3, 2.5),
|
|
"aistudio/gemini-flash-latest": (1.5, 9.0),
|
|
"vertex/gemini-flash-latest": (1.5, 9.0),
|
|
"aistudio/gemini-pro-latest": (2.5, 12.0),
|
|
"codex/gpt-sol-latest": (2.0, 8.0),
|
|
"openai/gpt-sol-latest": (2.0, 8.0),
|
|
"codex/gpt-luna-latest": (1.5, 6.0),
|
|
"xai/grok-latest": (2.0, 10.0),
|
|
"anthropic/claude-haiku-4.5": (1.0, 5.0),
|
|
"anthropic/claude-sonnet-5": (2.0, 10.0),
|
|
"anthropic/claude-sonnet-5.5": (2.0, 10.0),
|
|
"anthropic/claude-sonnet-4.6": (3.0, 15.0),
|
|
"anthropic/claude-opus-5": (5.0, 25.0),
|
|
}
|
|
|
|
_sem = threading.BoundedSemaphore(CFG.concurrency)
|
|
_NO_TEMPERATURE: set[str] = {
|
|
"anthropic/claude-sonnet-5",
|
|
"anthropic/claude-sonnet-5.5",
|
|
"anthropic/claude-opus-5",
|
|
"codex/gpt-sol-latest",
|
|
"codex/gpt-luna-latest",
|
|
}
|
|
_lock = threading.Lock()
|
|
_key: list[str] = []
|
|
|
|
|
|
def _auth() -> str:
|
|
with _lock:
|
|
if not _key:
|
|
_key.append(CFG.api_key())
|
|
return _key[0]
|
|
|
|
|
|
@dataclass
|
|
class Usage:
|
|
calls: int = 0
|
|
prompt_tokens: int = 0
|
|
completion_tokens: int = 0
|
|
retries: int = 0
|
|
failures: int = 0
|
|
cost_usd: float = 0.0
|
|
latencies: list[float] = field(default_factory=list)
|
|
_lk: threading.Lock = field(default_factory=threading.Lock, repr=False)
|
|
|
|
def add(self, model, pt, ct, latency, retries=0):
|
|
with self._lk:
|
|
self.calls += 1
|
|
self.prompt_tokens += pt
|
|
self.completion_tokens += ct
|
|
self.retries += retries
|
|
if latency:
|
|
self.latencies.append(latency)
|
|
pin, pout = PRICES.get(model, (0.0, 0.0))
|
|
self.cost_usd += pt / 1e6 * pin + ct / 1e6 * pout
|
|
|
|
def summary(self) -> dict:
|
|
lat = sorted(self.latencies)
|
|
return {
|
|
"calls": self.calls,
|
|
"prompt_tokens": self.prompt_tokens,
|
|
"completion_tokens": self.completion_tokens,
|
|
"cost_usd": round(self.cost_usd, 4),
|
|
"retries": self.retries,
|
|
"failures": self.failures,
|
|
"latency_avg": round(sum(lat) / len(lat), 2) if lat else 0,
|
|
"latency_max": round(lat[-1], 2) if lat else 0,
|
|
}
|
|
|
|
|
|
def _cache_path(model: str, payload: str):
|
|
h = hashlib.sha256(f"{model}\x00{payload}".encode()).hexdigest()[:32]
|
|
d = CFG.cache_dir / "llm" / model.replace("/", "_")
|
|
d.mkdir(parents=True, exist_ok=True)
|
|
return d / f"{h}.json"
|
|
|
|
|
|
def chat(model: str, prompt: str, *, system: str | None = None,
|
|
max_tokens: int = 8000, temperature: float = 0.0,
|
|
usage: Usage | None = None) -> str:
|
|
msgs = ([{"role": "system", "content": system}] if system else []) + \
|
|
[{"role": "user", "content": prompt}]
|
|
|
|
def payload_of() -> str:
|
|
body = {"model": model, "messages": msgs, "max_tokens": max_tokens}
|
|
if model not in _NO_TEMPERATURE:
|
|
body["temperature"] = temperature
|
|
return json.dumps(body, ensure_ascii=False, sort_keys=True)
|
|
|
|
payload = payload_of()
|
|
cpath = _cache_path(model, payload) if CFG.cache_llm else None
|
|
if cpath is not None and cpath.exists():
|
|
try:
|
|
c = json.loads(cpath.read_text())
|
|
if usage:
|
|
usage.add(model, 0, 0, 0)
|
|
return c["text"]
|
|
except Exception:
|
|
pass
|
|
|
|
retries = 0
|
|
last = None
|
|
for attempt in range(CFG.max_retries):
|
|
with _sem:
|
|
try:
|
|
req = urllib.request.Request(
|
|
CFG.endpoint, data=payload.encode(),
|
|
headers={"Content-Type": "application/json",
|
|
"Authorization": f"Bearer {_auth()}"})
|
|
t0 = time.time()
|
|
with urllib.request.urlopen(req, timeout=CFG.http_timeout) as r:
|
|
d = json.loads(r.read())
|
|
latency = time.time() - t0
|
|
except urllib.error.HTTPError as e:
|
|
detail = ""
|
|
try:
|
|
detail = e.read().decode("utf-8", "ignore")
|
|
except Exception:
|
|
pass
|
|
last = f"HTTP{e.code}: {detail[:200]}"
|
|
if e.code == 400 and "temperature" in detail.lower():
|
|
with _lock:
|
|
_NO_TEMPERATURE.add(model)
|
|
payload = payload_of()
|
|
if CFG.cache_llm:
|
|
cpath = _cache_path(model, payload)
|
|
retries += 1
|
|
continue
|
|
if e.code not in (408, 429, 500, 502, 503, 504):
|
|
if usage:
|
|
usage.failures += 1
|
|
raise RuntimeError(f"{model} {last}") from e
|
|
retries += 1
|
|
delay = min(2 ** attempt + random.uniform(0, 1.5), 45)
|
|
log(f"llm retry {attempt+1}/{CFG.max_retries} {model} {last} "
|
|
f"sleep {delay:.0f}s", level="WARN")
|
|
time.sleep(delay)
|
|
continue
|
|
except Exception as e: # network / SSL / timeout
|
|
last = repr(e)
|
|
retries += 1
|
|
time.sleep(min(2 ** attempt + random.uniform(0, 1.5), 45))
|
|
continue
|
|
|
|
text = (d["choices"][0]["message"].get("content") or "").strip()
|
|
u = d.get("usage", {})
|
|
pt, ct = u.get("prompt_tokens", 0), u.get("completion_tokens", 0)
|
|
if not text:
|
|
last = "empty completion"
|
|
retries += 1
|
|
time.sleep(1.5 * (attempt + 1))
|
|
continue
|
|
if cpath is not None:
|
|
cpath.write_text(json.dumps({"text": text}, ensure_ascii=False))
|
|
if usage:
|
|
usage.add(model, pt, ct, latency, retries=retries)
|
|
return text
|
|
|
|
if usage:
|
|
usage.failures += 1
|
|
raise RuntimeError(f"{model} failed after {CFG.max_retries} attempts: {last}")
|
|
|
|
|
|
def chat_many(model: str, prompts: list[str], **kw) -> list[str]:
|
|
"""Parallel map preserving order. Concurrency is bounded by the module
|
|
semaphore, so handing it an 87-chunk transcript is safe."""
|
|
import concurrent.futures as cf
|
|
|
|
out: list[str] = [""] * len(prompts)
|
|
errors: list[tuple[int, BaseException]] = []
|
|
with cf.ThreadPoolExecutor(max_workers=CFG.concurrency) as ex:
|
|
futs = {ex.submit(chat, model, p, **kw): i for i, p in enumerate(prompts)}
|
|
for f in cf.as_completed(futs):
|
|
i = futs[f]
|
|
try:
|
|
out[i] = f.result()
|
|
except BaseException as e: # noqa: BLE001 - recorded, decided by caller
|
|
errors.append((i, e))
|
|
if errors:
|
|
for i, e in errors:
|
|
log(f"chunk {i} failed: {e}", level="ERROR")
|
|
raise PartialFailure(out, [i for i, _ in errors])
|
|
return out
|
|
|
|
|
|
class PartialFailure(Exception):
|
|
def __init__(self, outputs: list[str], failed: list[int]):
|
|
super().__init__(f"{len(failed)} of {len(outputs)} chunks failed")
|
|
self.outputs = outputs
|
|
self.failed = failed
|