devplacepy/devplacepy/services/openai_gateway/usage.py

# retoor <retoor@molodetz.nl>

from __future__ import annotations

import json
import logging
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from typing import Any, Optional

import httpx

from devplacepy.database import db, get_table
from devplacepy.services.openai_gateway import config

logger = logging.getLogger(__name__)

GATEWAY_LEDGER = "gateway_usage_ledger"
GATEWAY_CONCURRENCY = "gateway_concurrency_samples"
PER_MILLION = 1_000_000


def _now() -> datetime:
    return datetime.now(timezone.utc)


def _iso(moment: datetime) -> str:
    return moment.isoformat()


@dataclass(frozen=True)
class Pricing:
    chat_cache_hit_per_m: float
    chat_cache_miss_per_m: float
    chat_output_per_m: float
    vision_input_per_m: float
    vision_output_per_m: float


def pricing_from_cfg(cfg: dict) -> Pricing:
    return Pricing(
        chat_cache_hit_per_m=float(cfg.get("gateway_price_cache_hit_per_m", config.PRICE_CACHE_HIT_PER_M_DEFAULT)),
        chat_cache_miss_per_m=float(cfg.get("gateway_price_cache_miss_per_m", config.PRICE_CACHE_MISS_PER_M_DEFAULT)),
        chat_output_per_m=float(cfg.get("gateway_price_output_per_m", config.PRICE_OUTPUT_PER_M_DEFAULT)),
        vision_input_per_m=float(cfg.get("gateway_vision_price_input_per_m", config.VISION_PRICE_INPUT_PER_M_DEFAULT)),
        vision_output_per_m=float(cfg.get("gateway_vision_price_output_per_m", config.VISION_PRICE_OUTPUT_PER_M_DEFAULT)),
    )


def parse_context_map(raw: Any) -> dict[str, int]:
    if isinstance(raw, dict):
        return {str(k): int(v) for k, v in raw.items()}
    if not raw:
        return dict(config.MODEL_CONTEXT_MAP_DEFAULT)
    try:
        loaded = json.loads(raw)
        if isinstance(loaded, dict):
            return {str(k): int(v) for k, v in loaded.items()}
    except (ValueError, TypeError):
        logger.warning("Invalid gateway_model_context_map, using defaults")
    return dict(config.MODEL_CONTEXT_MAP_DEFAULT)


def normalize_usage(usage: Optional[dict]) -> dict:
    usage = usage or {}
    prompt = int(usage.get("prompt_tokens", 0) or 0)
    completion = int(usage.get("completion_tokens", 0) or 0)
    total = int(usage.get("total_tokens", prompt + completion) or 0)

    hit = usage.get("prompt_cache_hit_tokens")
    if hit is None:
        details = usage.get("prompt_tokens_details") or {}
        hit = details.get("cached_tokens", 0)
    hit = int(hit or 0)

    miss = usage.get("prompt_cache_miss_tokens")
    if miss is None:
        miss = max(prompt - hit, 0)
    miss = int(miss or 0)

    completion_details = usage.get("completion_tokens_details") or {}
    reasoning = int(completion_details.get("reasoning_tokens", 0) or 0)
    return {
        "prompt": prompt,
        "completion": completion,
        "total": total,
        "cache_hit": hit,
        "cache_miss": miss,
        "reasoning": reasoning,
    }


def compute_cost(usage: dict, norm: dict, pricing: Pricing, backend: str) -> tuple[float, float, float, bool]:
    native = usage.get("cost") if isinstance(usage, dict) else None
    if isinstance(native, (int, float)) and not isinstance(native, bool):
        total = float(native)
        denom = norm["prompt"] + norm["completion"]
        input_cost = total * norm["prompt"] / denom if denom > 0 else 0.0
        return total, input_cost, total - input_cost, True
    if backend == "vision":
        input_cost = norm["prompt"] / PER_MILLION * pricing.vision_input_per_m
        output_cost = norm["completion"] / PER_MILLION * pricing.vision_output_per_m
        return input_cost + output_cost, input_cost, output_cost, False
    input_cost = (norm["cache_hit"] / PER_MILLION * pricing.chat_cache_hit_per_m
                  + norm["cache_miss"] / PER_MILLION * pricing.chat_cache_miss_per_m)
    output_cost = norm["completion"] / PER_MILLION * pricing.chat_output_per_m
    return input_cost + output_cost, input_cost, output_cost, False


def context_utilization(total_tokens: int, model: str, context_map: dict) -> tuple[Optional[int], Optional[float]]:
    window = context_map.get(model)
    if not window or window <= 0:
        return None, None
    return int(window), round(total_tokens / window, 4)


def extract_params(body: Any) -> dict:
    if not isinstance(body, dict):
        return {"requested_model": "", "stream_requested": False, "temperature": None,
                "top_p": None, "max_tokens": None, "has_tools": False}
    temperature = body.get("temperature")
    top_p = body.get("top_p")
    max_tokens = body.get("max_tokens")
    if max_tokens is None:
        max_tokens = body.get("max_completion_tokens")
    return {
        "requested_model": body.get("model") or "",
        "stream_requested": bool(body.get("stream")),
        "temperature": float(temperature) if isinstance(temperature, (int, float)) and not isinstance(temperature, bool) else None,
        "top_p": float(top_p) if isinstance(top_p, (int, float)) and not isinstance(top_p, bool) else None,
        "max_tokens": int(max_tokens) if isinstance(max_tokens, (int, float)) and not isinstance(max_tokens, bool) else None,
        "has_tools": bool(body.get("tools") or body.get("functions")),
    }


def classify_error(status_code: int, exc: Optional[Exception] = None, message: str = "") -> str:
    if exc is not None:
        if isinstance(exc, httpx.TimeoutException):
            return "timeout"
        return "gateway"
    if status_code == 429:
        return "rate_limit"
    if status_code in (401, 403):
        return "auth"
    if status_code == 404:
        return "model_not_found"
    if status_code in (400, 422):
        lowered = (message or "").lower()
        if "context" in lowered or "maximum" in lowered or "too long" in lowered:
            return "context_length"
        return "bad_request"
    if status_code and status_code >= 500:
        return "upstream_error"
    return "gateway"


class GatewayUsageLedger:
    def record(self, raw: dict, pricing: Pricing, context_map: dict) -> None:
        try:
            usage = raw.get("usage") or {}
            norm = normalize_usage(usage)
            cost_usd, input_cost, output_cost, native = compute_cost(usage, norm, pricing, raw["backend"])
            window, util = context_utilization(norm["total"], raw.get("model") or "", context_map)
            upstream_ms = float(raw.get("upstream_latency_ms") or 0)
            completion = norm["completion"]
            tps = completion / (upstream_ms / 1000.0) if upstream_ms > 0 and completion else 0.0
            row = {
                "created_at": _iso(_now()),
                "owner_kind": raw.get("owner_kind") or "unknown",
                "owner_id": raw.get("owner_id") or "unknown",
                "backend": raw["backend"],
                "endpoint": raw.get("endpoint") or "",
                "requested_model": raw.get("requested_model") or "",
                "model": raw.get("model") or "",
                "status_code": int(raw.get("status_code") or 0),
                "success": 1 if raw.get("success") else 0,
                "error_category": raw.get("error_category"),
                "upstream_latency_ms": upstream_ms,
                "gateway_overhead_ms": float(raw.get("gateway_overhead_ms") or 0),
                "queue_wait_ms": float(raw.get("queue_wait_ms") or 0),
                "connect_ms": float(raw.get("connect_ms") or 0),
                "total_latency_ms": float(raw.get("total_latency_ms") or 0),
                "prompt_tokens": norm["prompt"],
                "completion_tokens": norm["completion"],
                "cache_hit_tokens": norm["cache_hit"],
                "cache_miss_tokens": norm["cache_miss"],
                "reasoning_tokens": norm["reasoning"],
                "total_tokens": norm["total"],
                "tokens_per_second": round(tps, 3),
                "context_window": window,
                "context_utilization": util,
                "cost_usd": round(cost_usd, 8),
                "input_cost_usd": round(input_cost, 8),
                "output_cost_usd": round(output_cost, 8),
                "native_cost": 1 if native else 0,
                "stream_requested": 1 if raw.get("stream_requested") else 0,
                "temperature": raw.get("temperature"),
                "top_p": raw.get("top_p"),
                "max_tokens": raw.get("max_tokens"),
                "has_tools": 1 if raw.get("has_tools") else 0,
                "retries_attempted": int(raw.get("retries_attempted") or 0),
                "retry_succeeded": 1 if raw.get("retry_succeeded") else 0,
                "circuit_open": 1 if raw.get("circuit_open") else 0,
                "user_agent": (raw.get("user_agent") or "")[:300],
            }
            get_table(GATEWAY_LEDGER).insert(row)
        except Exception as exc:
            logger.warning("gateway usage record failed: %s", exc)

    def sample_concurrency(self, in_flight: int) -> None:
        try:
            get_table(GATEWAY_CONCURRENCY).insert({
                "created_at": _iso(_now()),
                "in_flight": int(in_flight),
            })
        except Exception as exc:
            logger.warning("gateway concurrency sample failed: %s", exc)

    def prune(self, older_than_hours: int) -> tuple[int, int]:
        cutoff = _iso(_now() - timedelta(hours=max(1, older_than_hours)))
        ledger_removed = 0
        samples_removed = 0
        if GATEWAY_LEDGER in db.tables:
            ledger_removed = int(get_table(GATEWAY_LEDGER).delete(created_at={"<": cutoff}))
        if GATEWAY_CONCURRENCY in db.tables:
            samples_removed = int(get_table(GATEWAY_CONCURRENCY).delete(created_at={"<": cutoff}))
        return ledger_removed, samples_removed