1"""Shared types and helpers for model-fetching scripts."""
2
3from __future__ import annotations
4
5import json
6import re
7import sys
8from dataclasses import asdict, dataclass
9from typing import TextIO
10
11
12@dataclass(frozen=True)
13class Model:
14 """Normalized model metadata used to update client configs."""
15
16 id: str
17 name: str
18 context_length: int
19 max_completion_tokens: int
20 supports_reasoning: bool
21 supports_image: bool
22
23
24# Known brand names that need non-trivial capitalisation.
25# Lowercase lookup → correct display form.
26_BRANDS: dict[str, str] = {
27 "deepseek": "DeepSeek",
28 "glm": "GLM",
29 "gpt": "GPT",
30 "gemma": "Gemma",
31 "kimi": "Kimi",
32 "devstral": "Devstral",
33 "llama": "Llama",
34 "mimo": "MiMo",
35 "minimax": "MiniMax",
36 "qwen": "Qwen",
37 "qwen3": "Qwen3",
38 "claude": "Claude",
39 "gemini": "Gemini",
40 "google": "Google",
41 "nvidia": "NVIDIA",
42 "nemotron": "Nemotron",
43}
44
45# Pure-alpha tokens that should be all-caps.
46_ACRONYMS: set[str] = {"oss", "ar"}
47
48# Alpha+digit tokens that should be all-caps.
49_MIXED_ACRONYMS: set[str] = {"fp8", "int4", "nvfp4"}
50
51
52def _split_brand_version(token: str) -> list[str]:
53 """Split a brand-version compound like 'Qwen3.6' into ['Qwen', '3.6'].
54
55 Only splits when the leading letters match a known brand **and** the
56 version contains a decimal point. This avoids splitting 'Qwen3'
57 (no decimal) which should stay together in forms like 'Qwen3 Coder'.
58 """
59 m = re.fullmatch(r"([a-zA-Z]{2,})(\d+\.\d+)", token)
60 if m and m.group(1).lower() in _BRANDS:
61 return [m.group(1), m.group(2)]
62 return [token]
63
64
65def _clean_token(tok: str) -> str:
66 low = tok.lower()
67
68 if low in _BRANDS:
69 return _BRANDS[low]
70 if low in _ACRONYMS:
71 return tok.upper()
72 if low in _MIXED_ACRONYMS:
73 return tok.upper()
74 if tok.isalpha() and tok.isupper() and len(tok) >= 2:
75 return tok
76 if re.fullmatch(r"\d+(?:\.\d+)?", tok):
77 return tok
78
79 # '26b' → '26B', '128e' → '128E'
80 m = re.fullmatch(r"(\d+)([a-zA-Z]+)", tok)
81 if m:
82 return m.group(1) + m.group(2).upper()
83
84 # 'v4' → 'V4', 'k2.5' → 'K2.5', 'a4b' → 'A4B', 'a35b' → 'A35B'
85 m = re.fullmatch(r"([a-zA-Z]+)(\d+(?:\.\d+)?)([a-zA-Z]*)", tok)
86 if m:
87 prefix = m.group(1).capitalize()
88 nums = m.group(2)
89 suffix = m.group(3).upper() if m.group(3) else ""
90 return prefix + nums + suffix
91
92 if tok.isalpha():
93 return tok.capitalize()
94
95 return tok
96
97
98def clean_name(raw: str) -> str:
99 """Normalise a display name from an upstream catalog.
100
101 - Strip redundant vendor prefixes ('DeepSeek: DeepSeek V4 Pro' → 'DeepSeek V4 Pro')
102 - Remove parentheses ('(fast)' → 'fast')
103 - Replace hyphens with spaces
104 - Split brand-version compounds ('Qwen3.6' → 'Qwen 3.6')
105 - Normalise capitalisation for known brands and acronyms
106 """
107 if ":" in raw:
108 raw = raw.split(":")[-1].strip()
109 raw = re.sub(r"\(([^)]+)\)", r"\1", raw)
110 raw = raw.replace("-", " ")
111
112 tokens: list[str] = []
113 for tok in raw.split():
114 tokens.extend(_split_brand_version(tok))
115
116 return " ".join(_clean_token(t) for t in tokens)
117
118
119def emit_tsv(models: list[Model], stream: TextIO) -> None:
120 """Write models as TSV: id, name, ctx, max_tokens, reasoning, image."""
121 for m in models:
122 reasoning = "reasoning" if m.supports_reasoning else ""
123 image = "image" if m.supports_image else ""
124 stream.write(
125 f"{m.id}\t{m.name}\t{m.context_length}\t{m.max_completion_tokens}"
126 f"\t{reasoning}\t{image}\n"
127 )
128
129
130def emit_json(models: list[Model], stream: TextIO) -> None:
131 """Write models as a JSON array (indent=2, trailing newline)."""
132 json.dump([asdict(m) for m in models], stream, indent=2)
133 stream.write("\n")