EventHubAiRouter/scripts/gen-litellm-config.py

#!/usr/bin/env python3
"""Generate litellm_config.yaml from PRIMARY_PROVIDER + config/*.yaml."""
from __future__ import annotations

import os
import sys
from pathlib import Path

import yaml

ROOT = Path(__file__).resolve().parent.parent
PROFILE = os.environ.get("PRIMARY_PROVIDER", "hybrid").strip().lower()


def load(path: Path) -> dict:
    with path.open(encoding="utf-8") as f:
        return yaml.safe_load(f) or {}


def api_key_ref(name: str) -> str:
    return f"os.environ/{name}"


def _apply_thinking(params: dict, spec: dict) -> None:
    """Qwen3 defaults to thinking; empty content + hang unless disabled."""
    model = str(params.get("model") or spec.get("model") or "").lower()
    force = bool(spec.get("disable_thinking"))
    auto = ("qwen3" in model) and not spec.get("enable_thinking")
    if force or auto:
        params["extra_body"] = {
            "chat_template_kwargs": {"enable_thinking": False},
            "enable_thinking": False,
        }


def lane_entry(name: str, spec: dict, rpm: int | None) -> dict:
    params: dict = {
        "model": spec["model"],
        "api_key": api_key_ref(spec["api_key"]),
        # Per-deployment timeout: litellm_settings.request_timeout often shows as
        # Deployment Info timeout: None on stream/tool hangs.
        "timeout": int(os.environ.get("LITELLM_MODEL_TIMEOUT", "180")),
    }
    if rpm:
        params["rpm"] = rpm
    _apply_thinking(params, spec)
    return {"model_name": name, "litellm_params": params}


def fixed_entry(name: str, spec: dict) -> dict:
    params: dict = {
        "model": spec["model"],
        "api_key": api_key_ref(spec["api_key"]),
        "timeout": int(os.environ.get("LITELLM_MODEL_TIMEOUT", "180")),
    }
    for key in ("ssl_verify", "max_tokens", "temperature"):
        if key in spec:
            params[key] = spec[key]
    _apply_thinking(params, spec)
    return {"model_name": name, "litellm_params": params}


def smart_router_internal(rules: dict) -> dict:
    litellm_rules = rules.get("litellm", {})
    return {
        "model_name": "smart-router-internal",
        "litellm_params": {
            "model": "auto_router/complexity_router",
            "drop_params": True,
            "complexity_router_default_model": "a-medium-ops",
            "complexity_router_config": {
                "tiers": {
                    "SIMPLE": "a-simple",
                    "MEDIUM": "a-medium-ops",
                    "MEDIUM_CODE": "a-medium-code",
                    "COMPLEX": "a-complex",
                    "REASONING": "a-reasoning",
                },
                "classifier_fallback": "heuristic",
                "keyword_tier_rules": litellm_rules.get("keyword_tier_rules", []),
                "custom_technical_keywords": litellm_rules.get(
                    "custom_technical_keywords", []
                ),
                "token_thresholds": {"simple": 20, "complex": 500},
                "tier_boundaries": {
                    "simple_medium": 0.18,
                    "medium_complex": 0.38,
                    "complex_reasoning": 0.62,
                },
                "session_affinity": True,
                "session_affinity_ttl_seconds": 1800,
            },
        },
    }


def main() -> int:
    providers = load(ROOT / "config/providers.yaml")
    profiles = providers.get("profiles", {})
    if PROFILE not in profiles:
        print(f"ERROR: unknown PRIMARY_PROVIDER={PROFILE!r}", file=sys.stderr)
        print(f"Available: {', '.join(sorted(profiles))}", file=sys.stderr)
        return 1

    profile = profiles[PROFILE]
    base = load(ROOT / "config/litellm.base.yaml")
    matrix = load(ROOT / "config/model_matrix.yaml")
    rules = load(ROOT / "config/routing_rules.yaml")
    meta = matrix.get("models", {})

    model_list: list[dict] = []

    # Classifier first
    fixed = providers.get("fixed_models", {})
    if "gigachat-classifier" in fixed:
        model_list.append(fixed_entry("gigachat-classifier", fixed["gigachat-classifier"]))

    # Lane models from active profile
    lanes = profile.get("lanes", {})
    for name, spec in lanes.items():
        rpm = meta.get(name, {}).get("rpm")
        model_list.append(lane_entry(name, spec, rpm))

    # Fixed models (Claude / OR fallbacks / GigaChat / optional free tiers)
    for name in (
        "novita-planner",
        "novita-verifier",
        "claude-haiku-planner",
        "claude-sonnet-verifier",
        "groq-llama-8b",
        "groq-qwen-coder",
        "gemini-flash",
        "grok-3",
        "gigachat-pro",
    ):
        if name in fixed:
            model_list.append(fixed_entry(name, fixed[name]))

    model_list.append(smart_router_internal(rules))

    # smart-router alias (Zed default)
    sr = profile.get("smart_router", {})
    if sr:
        sr_params: dict = {
            "model": sr["model"],
            "api_key": api_key_ref(sr["api_key"]),
        }
        _apply_thinking(sr_params, sr)
        model_list.append({"model_name": "smart-router", "litellm_params": sr_params})

    router = base.setdefault("router_settings", {})
    router["fallbacks"] = [
        {k: v for k, v in row.items()}
        for row in _fallback_list(profile.get("fallbacks", {}))
    ]
    router["default_fallbacks"] = profile.get(
        "default_fallbacks", ["a-medium-code", "groq-qwen-coder"]
    )

    out_cfg = {**base, "model_list": model_list}
    out_path = ROOT / "litellm_config.yaml"
    header = (
        f"# LiteLLM — generated for PRIMARY_PROVIDER={PROFILE}\n"
        f"# Profile: {profile.get('label', PROFILE)}\n"
        f"# Regenerate: PRIMARY_PROVIDER={PROFILE} bash scripts/gen-litellm-config.py\n\n"
    )
    body = yaml.dump(out_cfg, allow_unicode=True, sort_keys=False)
    out_path.write_text(header + body, encoding="utf-8")
    print(f"Wrote {out_path} (profile={PROFILE}, {len(model_list)} models)")
    return 0


def _fallback_list(fallbacks: dict) -> list[dict]:
    return [{model: targets} for model, targets in fallbacks.items()]


if __name__ == "__main__":
    sys.exit(main())