219 lines
7.9 KiB
Python
219 lines
7.9 KiB
Python
"""Snapshot collector — turns live Router state into the discovery data model.
|
|
|
|
Reads only structures the Router already maintains for routing:
|
|
* ``state.usage_counts`` — live per-endpoint/model active connections
|
|
* ``config.endpoints`` / ``llama_server_endpoints`` / ``api_keys`` / limits
|
|
* ``backends.probe`` model discovery + endpoint health
|
|
* ``cache.LLMCache.stats()`` — global cache hit rate
|
|
* the token time-series DB — 24h token totals and per-model ``last_used``
|
|
|
|
Fields the Router does not currently instrument (per-model latency percentiles,
|
|
request/error counts, per-model cache hit rate, token averages) are emitted as
|
|
``None`` rather than fabricated. See the plugin docs / spec §4 for the rationale.
|
|
"""
|
|
import time
|
|
from datetime import datetime, timezone
|
|
from typing import Optional
|
|
|
|
from config import Config
|
|
from state import usage_counts
|
|
from cache import get_llm_cache
|
|
from db import get_db
|
|
from routing import get_max_connections
|
|
from backends.probe import fetch, _endpoint_health
|
|
from backends.normalize import (
|
|
is_ext_openai_endpoint,
|
|
_normalize_llama_model_name,
|
|
_extract_llama_quant,
|
|
)
|
|
|
|
_DAY_SECONDS = 86400
|
|
|
|
|
|
def _iso_now() -> str:
|
|
return datetime.now(tz=timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
|
|
|
|
|
|
def _iso_from_unix(ts: Optional[int]) -> Optional[str]:
|
|
if not ts:
|
|
return None
|
|
return datetime.fromtimestamp(ts, tz=timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
|
|
|
|
|
|
def _endpoint_type(cfg: Config, ep: str) -> str:
|
|
"""Best-effort classification from config membership + URL heuristics."""
|
|
low = ep.lower()
|
|
if "googleapis" in low or "gemini" in low:
|
|
return "gemini"
|
|
if "anthropic" in low:
|
|
return "anthropic"
|
|
if "cohere" in low:
|
|
return "cohere"
|
|
if "cerebras" in low:
|
|
return "cerebras"
|
|
if "inceptionlabs" in low or "inception" in low:
|
|
return "inception labs"
|
|
if ep in cfg.llama_server_endpoints:
|
|
# llama_server_endpoints covers both llama.cpp and vLLM (both OpenAI-
|
|
# compatible). vLLM commonly serves on :8000; treat the rest as llama.cpp.
|
|
return "vllm" if ":8000" in ep else "llama_cpp"
|
|
if is_ext_openai_endpoint(ep):
|
|
return "openai"
|
|
return "ollama"
|
|
|
|
|
|
def _auth_method(cfg: Config, ep: str) -> str:
|
|
return "bearer" if cfg.api_keys.get(ep) else "none"
|
|
|
|
|
|
def _status_label(health: dict) -> str:
|
|
# _endpoint_health only distinguishes ok/error; map onto the spec's enum.
|
|
return "healthy" if health.get("status") == "ok" else "unhealthy"
|
|
|
|
|
|
async def _collect_endpoints(cfg: Config) -> tuple[list[dict], int]:
|
|
"""Build the endpoint registry and the global active-connection count."""
|
|
seen: list[str] = []
|
|
for ep in list(cfg.endpoints) + list(cfg.llama_server_endpoints):
|
|
if ep not in seen:
|
|
seen.append(ep)
|
|
|
|
endpoints: list[dict] = []
|
|
total_active = 0
|
|
for ep in seen:
|
|
concurrent = sum(usage_counts.get(ep, {}).values())
|
|
total_active += concurrent
|
|
try:
|
|
health = await _endpoint_health(ep, timeout=5)
|
|
except Exception:
|
|
health = {"status": "error"}
|
|
endpoints.append({
|
|
"url": ep,
|
|
"type": _endpoint_type(cfg, ep),
|
|
"status": _status_label(health),
|
|
"concurrent_connections": concurrent,
|
|
"max_concurrent_connections": get_max_connections(ep),
|
|
"tls_enabled": ep.lower().startswith("https://"),
|
|
"auth_method": _auth_method(cfg, ep),
|
|
})
|
|
return endpoints, total_active
|
|
|
|
|
|
def _ollama_detail_index(details: list[dict]) -> dict:
|
|
"""Index Ollama /api/tags entries by model name for size/quant lookup."""
|
|
index: dict[str, dict] = {}
|
|
for item in details or []:
|
|
name = item.get("name") or item.get("model")
|
|
if name:
|
|
index[name] = item
|
|
return index
|
|
|
|
|
|
async def _collect_models(cfg: Config, last_used: dict) -> list[dict]:
|
|
models: list[dict] = []
|
|
for ep in list(cfg.endpoints) + list(cfg.llama_server_endpoints):
|
|
is_llama = ep in cfg.llama_server_endpoints
|
|
try:
|
|
available = await fetch.available_models(ep, cfg.api_keys.get(ep))
|
|
except Exception:
|
|
available = set()
|
|
if not available:
|
|
continue
|
|
|
|
# Ollama exposes weight size + quantization via /api/tags details.
|
|
detail_index: dict = {}
|
|
if not is_llama and not is_ext_openai_endpoint(ep):
|
|
try:
|
|
details = await fetch.endpoint_details(ep, "/api/tags", "models")
|
|
detail_index = _ollama_detail_index(details)
|
|
except Exception:
|
|
detail_index = {}
|
|
|
|
for raw_name in sorted(available):
|
|
if is_llama:
|
|
name = _normalize_llama_model_name(raw_name)
|
|
quant = _extract_llama_quant(raw_name) or "none"
|
|
version = raw_name.split(":", 1)[1] if ":" in raw_name else "latest"
|
|
size_gb = None
|
|
else:
|
|
name = raw_name
|
|
version = raw_name.split(":", 1)[1] if ":" in raw_name else "latest"
|
|
meta = detail_index.get(raw_name, {})
|
|
det = meta.get("details", {}) if isinstance(meta, dict) else {}
|
|
quant = det.get("quantization_level") or "none"
|
|
size_bytes = meta.get("size") if isinstance(meta, dict) else None
|
|
size_gb = round(size_bytes / 1e9, 3) if size_bytes else None
|
|
|
|
models.append({
|
|
"name": name,
|
|
"version": version,
|
|
"quantization": quant,
|
|
"size_gb": size_gb,
|
|
"endpoint": ep,
|
|
"last_used": _iso_from_unix(last_used.get((ep, raw_name)) or last_used.get((ep, name))),
|
|
# Un-instrumented today — emitted as null rather than fabricated.
|
|
"request_count_24h": None,
|
|
"avg_latency_ms": None,
|
|
"p99_latency_ms": None,
|
|
"error_count_24h": None,
|
|
"cache_hit_rate": None,
|
|
"avg_input_tokens": None,
|
|
"avg_output_tokens": None,
|
|
})
|
|
return models
|
|
|
|
|
|
async def _collect_telemetry(cfg: Config, active_connections: int, started_at: float) -> dict:
|
|
cache = get_llm_cache()
|
|
if cache is not None:
|
|
cstats = cache.stats()
|
|
total_cache_hits = cstats.get("hits", 0)
|
|
cache_hit_rate = cstats.get("hit_rate", 0.0)
|
|
else:
|
|
total_cache_hits = 0
|
|
cache_hit_rate = 0.0
|
|
|
|
cutoff = int(time.time()) - _DAY_SECONDS
|
|
try:
|
|
tokens_in, tokens_out = await get_db().get_token_totals_since(cutoff)
|
|
except Exception:
|
|
tokens_in, tokens_out = 0, 0
|
|
|
|
uptime_hours = round((time.monotonic() - started_at) / 3600.0, 3)
|
|
|
|
return {
|
|
# Request/error totals are un-instrumented today (null, not fabricated).
|
|
"total_requests_24h": None,
|
|
"total_cache_hits": total_cache_hits,
|
|
"cache_hit_rate": cache_hit_rate,
|
|
"error_rate": None,
|
|
# Router tracks tokens, not raw bytes; surfaced under the byte fields as
|
|
# the closest available signal (24h windowed token counts).
|
|
"total_bytes_in_24h": tokens_in,
|
|
"total_bytes_out_24h": tokens_out,
|
|
"active_connections": active_connections,
|
|
"uptime_hours": uptime_hours,
|
|
}
|
|
|
|
|
|
async def collect_snapshot(cfg: Config, started_at: float) -> dict:
|
|
"""Assemble the full discovery snapshot body (without signature).
|
|
|
|
``started_at`` is a ``time.monotonic()`` reading taken at plugin start,
|
|
used to derive ``uptime_hours``.
|
|
"""
|
|
try:
|
|
last_used = await get_db().get_last_used_map()
|
|
except Exception:
|
|
last_used = {}
|
|
|
|
endpoints, active_connections = await _collect_endpoints(cfg)
|
|
models = await _collect_models(cfg, last_used)
|
|
telemetry = await _collect_telemetry(cfg, active_connections, started_at)
|
|
|
|
return {
|
|
"timestamp": _iso_now(),
|
|
"models": models,
|
|
"endpoints": endpoints,
|
|
"telemetry": telemetry,
|
|
}
|