SPB Git forge

spb/ai-atlas

Public
41commits 1branches 0releases
4.6 MBsize
maindefault branch
12 days agolast push
HTML 77.2% TypeScript 10.5% Python 9.6% JavaScript 2.5%
20.5 KB · 388 lines python
Raw Blame History
1"""Shared model-identity helper for third-party sources (leaderboards, aggregators, provider price lists).23Those sources name models with *API ids* (`anthropic/claude-3-7-sonnet-20250219`, `gemini/gemini-2.5-pro-preview-05-06`,4`openrouter/x-ai/grok-4`, `gpt-4o-2024-08-06`, `Qwen/Qwen2.5-Coder-32B-Instruct`) or free text. This module turns such an id into:56  * the developer organisation (registry `org_ref`) when the id carries a first-party prefix or family word — never a guess beyond7    the deterministic tables below;8  * resolver-friendly aliases (the raw id, the id without provider prefix, the id without evaluator effort suffixes);9  * identifiers, **only when the caller marks the id as trusted** (= the string really is the vendor's own API id, e.g. the `--model`10    argument aider passed to the API). Free-text tags never become identifiers: the resolver refuses to merge an alias match when the11    entity already carries a *different* value for the same scheme, so a wrong `anthropic_model_id` would split an entity instead of12    linking it.13  * the evaluation-effort configuration encoded in the id or its label ("(xhigh)", "-thinking-64k", "(Non-reasoning)") — effort variants14    are result *configurations*, never model entities.15"""16from __future__ import annotations1718import re19from dataclasses import dataclass, field20from typing import Any2122from aiatlas.ontology.models import OFFICIAL_THINKING_RELEASES, analyze_model_name, family_release_hint23from aiatlas.registry import org_by_hf, org_ref, organizations24from aiatlas.sdk.facts import EntityRef, Facts2526# provider / gateway prefixes used by litellm, aider, OpenRouter, LiveBench… → (identifier scheme for the remainder, developer org key)27# scheme None = the prefix is a gateway or an OpenAI-compatible endpoint, the remainder decides.28PROVIDER_PREFIXES: dict[str, tuple[str | None, str | None]] = {29    "anthropic": ("anthropic_model_id", "anthropic"),30    "gemini": ("gemini_model_id", "google"),31    "google": ("gemini_model_id", "google"),32    "vertex_ai": (None, "google"),33    "deepseek": ("deepseek_model_id", "deepseek"),34    "mistral": ("mistral_model_id", "mistral"),35    "mistralai": ("mistral_model_id", "mistral"),36    "xai": ("xai_model_id", "xai"),37    "x-ai": ("xai_model_id", "xai"),38    "cohere": ("cohere_model_id", "cohere"),39    "groq": ("groq_model_id", None),40    "together_ai": ("together_ai_model_slug", None),41    "fireworks_ai": ("fireworks_model_id", None),42    "openai": (None, None),          # litellm/aider route *any* OpenAI-compatible endpoint through `openai/…`43    "azure": (None, None),44    "bedrock": (None, None),45    "nvidia_nim": (None, None),46    "openrouter": (None, None),47}48# OpenRouter vendor slug → registry organisation (shared with the OpenRouter connector)49VENDOR_ORG: dict[str, str] = {50    "meta": "meta-ai", "meta-llama": "meta-ai", "liquid": "liquid-ai", "x-ai": "xai", "mistralai": "mistral", "z-ai": "zhipu", "thudm": "zhipu",51    "moonshotai": "moonshot", "bytedance": "bytedance", "bytedance-seed": "bytedance", "ibm-granite": "ibm", "google": "google",52    "amazon": "amazon", "qwen": "qwen", "deepseek": "deepseek", "anthropic": "anthropic", "openai": "openai", "cohere": "cohere",53    "nvidia": "nvidia", "microsoft": "microsoft", "perplexity": "perplexity", "minimax": "minimax", "tencent": "tencent", "baidu": "baidu",54    "ai21": "ai21", "nousresearch": "nous-research", "openrouter": "openrouter", "all-hands": "all-hands-ai", "stepfun": "stepfun",55    "stepfun-ai": "stepfun", "upstage": "upstage", "kwaipilot": "kwaipilot", "thinking-machines": "thinking-machines", "thinkingmachines": "thinking-machines",56    "allenai": "allenai", "eleutherai": "eleutherai", "internlm": "internlm",57}58# first-party naming patterns of bare API ids → (identifier scheme or None, developer org key). Ordered: specific families first.59_BARE: list[tuple[re.Pattern[str], str | None, str | None]] = [60    (re.compile(r"^claude[-_ ]"), "anthropic_model_id", "anthropic"),61    (re.compile(r"^(chatgpt-|gpt-(?!oss)|o[1-9](-|$)|codex|davinci|text-embedding-|dall-e|sora|tts-1|computer-use-preview)"), "openai_model_id", "openai"),62    (re.compile(r"^gpt-oss"), None, "openai"),63    (re.compile(r"^(gemini-|imagen-|veo-|lyria-|gemini-embedding)"), "gemini_model_id", "google"),64    (re.compile(r"^(gemma-|paligemma|medgemma|shieldgemma)"), None, "google"),65    (re.compile(r"^grok-"), "xai_model_id", "xai"),66    (re.compile(r"^deepseek-"), "deepseek_model_id", "deepseek"),67    (re.compile(r"^(mistral-|mixtral-|ministral-|codestral|magistral|devstral|pixtral|voxtral|open-mistral|open-mixtral|mistral-embed)"), "mistral_model_id", "mistral"),68    (re.compile(r"^(command-|command$|c4ai-|embed-(english|multilingual|v)|rerank-)"), "cohere_model_id", "cohere"),69    (re.compile(r"^aya-"), None, "cohere"),70    (re.compile(r"^(hermes)"), None, "nous-research"),71    (re.compile(r"^(llama[- ]?\d[\d.]*[- ].*nemotron|nemotron|nvidia-)"), None, "nvidia"),72    (re.compile(r"^(meta-)?llama-\d"), None, "meta-ai"),73    (re.compile(r"^(qwen|qwq|qvq)"), None, "qwen"),74    (re.compile(r"^kimi-"), None, "moonshot"),75    (re.compile(r"^(glm-|chatglm)"), None, "zhipu"),76    (re.compile(r"^minimax-"), None, "minimax"),77    (re.compile(r"^nova-"), None, "amazon"),78    (re.compile(r"^phi-"), None, "microsoft"),79    (re.compile(r"^sonar"), None, "perplexity"),80    (re.compile(r"^jamba"), None, "ai21"),81    (re.compile(r"^granite-"), None, "ibm"),82    (re.compile(r"^hunyuan"), None, "tencent"),83    (re.compile(r"^ernie"), None, "baidu"),84    (re.compile(r"^(doubao|seed-)"), None, "bytedance"),85    (re.compile(r"^lfm"), None, "liquid-ai"),86    (re.compile(r"^step-?\d"), None, "stepfun"),87    (re.compile(r"^solar-"), None, "upstage"),88    (re.compile(r"^(kat-|kwaipilot)"), None, "kwaipilot"),89    (re.compile(r"^dbrx"), None, "databricks"),90    (re.compile(r"^(olmo|molmo|tulu)"), None, "allenai"),91    (re.compile(r"^(smollm|smolvlm)"), None, "huggingface"),92]93_FIREWORKS_PATH = re.compile(r"^accounts/fireworks/models/")94_DATE_IN_ID = re.compile(r"(20\d{2}-?(0[1-9]|1[0-2])-?(0[1-9]|[12]\d|3[01]))|(?<![0-9])(0[1-9]|1[0-2])(0[1-9]|[12]\d|3[01])(?![0-9])|(exp|preview)-\d{2}-\d{2}")95# compound evaluator suffixes that the ontology table does not know (LiveBench): "-thinking-auto" = adaptive thinking budget96_EXTRA_SUFFIXES: dict[str, dict[str, str]] = {"thinking-auto": {"reasoning": "on", "thinking_budget": "auto"}, "auto": {"thinking_budget": "auto"},97                                              "effort": {}}98_EFFORT_WORDS = re.compile(r"^(x-?high|extra-high|high|medium|low|minimal|max|effort|thinking|non-?reasoning|no-?think(ing)?|reasoning|auto|\d+k|default|\(.*\))$", re.IGNORECASE)99_LABEL_EFFORT = re.compile(r"^(?P<eff>x-?high|extra-high|high|medium|low|minimal|max)(?:\s+effort)?(?:\s+with\s+fallback)?$", re.IGNORECASE)100_LABEL_BUDGET = re.compile(r"^(?P<k>\d+)k\s*(think(ing)?|reasoning)?(\s*tokens)?$", re.IGNORECASE)101_LABEL_OFF = re.compile(r"^(non-?reasoning|no-?\s?think(ing)?|thinking\s*off|non-?thinking|without\s+thinking)$", re.IGNORECASE)102_LABEL_ON = re.compile(r"^(thinking|reasoning|think|thinking\s*on|reasoner)$", re.IGNORECASE)103_LABEL_DEFAULT = re.compile(r"^(default(\s+think(ing)?)?)$", re.IGNORECASE)104_LABEL_ADAPTIVE = re.compile(r"^(adaptive(\s+(thinking|reasoning))?|thinking\s+auto|auto)$", re.IGNORECASE)105_PAREN = re.compile(r"\s*\(([^()]*)\)\s*$")106107108@dataclass109class ModelIdentity:110    raw: str111    base_id: str                                   # id without provider prefix, variant suffix and effort suffixes112    scheme: str | None = None                      # identifier scheme the id belongs to (when recognised)113    identifiers: dict[str, str] = field(default_factory=dict)114    aliases: list[str] = field(default_factory=list)115    org_key: str | None = None                     # registry organisation key of the developer116    effort: dict[str, str] = field(default_factory=dict)117    provider_prefix: str | None = None118    pinned: bool = False                           # the id names a dated snapshot (stable), not a rolling alias119120    @property121    def is_effort_variant(self) -> bool:122        return bool(self.effort)123124125# ---------------------------------------------------------------------------------------------- effort labels / suffixes126def effort_from_label(label: str | None) -> dict[str, str] | None:127    """'(xhigh)' → {reasoning_effort: xhigh} · '(Non-reasoning)' → {reasoning: off} · '(32k thinking tokens)' → {reasoning: on,128    thinking_budget: 32k} · '(default think)' → {thinking_budget: default}. None when the label is not an evaluation setting."""129    if not label:130        return None131    s = label.strip().strip("()").strip().replace("_", "-")132    if not s:133        return None134    m = _LABEL_EFFORT.match(s)135    if m:136        eff = m.group("eff").lower().replace("x-high", "xhigh").replace("extra-high", "xhigh")137        out = {"reasoning_effort": eff}138        if "fallback" in s.lower():139            out["effort_fallback"] = "on"140        return out141    m = _LABEL_BUDGET.match(s)142    if m:143        return {"reasoning": "on", "thinking_budget": f"{m.group('k')}k"}144    if _LABEL_OFF.match(s):145        return {"reasoning": "off"}146    if _LABEL_ON.match(s):147        return {"reasoning": "on"}148    if _LABEL_ADAPTIVE.match(s):149        return {"reasoning": "adaptive"}150    if _LABEL_DEFAULT.match(s):151        return {"thinking_budget": "default"}152    return None153154155def split_effort_label(name: str | None) -> tuple[str, dict[str, str]]:156    """'Claude Opus 5 (xhigh)' → ('Claude Opus 5', {reasoning_effort: xhigh}); 'GPT-5.6 Sol xHigh Effort' → ('GPT-5.6 Sol', {reasoning_effort:157    xhigh}); 'Mistral Small 4 Non-reasoning' → reasoning off. Labels that are not settings ("(0324)", "(Jul)", "Thinking" alone) stay in the name."""158    if not name:159        return "", {}160    n = name.strip()161    m = _PAREN.search(n)162    if m:163        eff = effort_from_label(m.group(1))164        if eff is not None:165            return n[: m.start()].strip(), eff166    m = _TRAILING_EFFORT.search(n)167    if m:168        eff = m.group("eff").lower().replace("x-high", "xhigh").replace("extra-high", "xhigh")169        return n[: m.start()].strip(), {"reasoning_effort": eff}170    m = _TRAILING_OFF.search(n)171    if m:172        return n[: m.start()].strip(), {"reasoning": "off"}173    return n, {}174175176# Suffixes that also name size tiers or official products ("mistral-medium", "sonar-reasoning"): they count as an evaluator setting only177# when the remaining stem still carries a version/size digit ("o3-mini-high", "deepseek-v3-1-reasoning"). ("max", "fast", "instant" are178# not effort suffixes at all in the ontology — Qwen3-Max, Grok 4.1 Fast and Claude Instant are model tiers.)179_AMBIGUOUS_SUFFIXES = {"medium", "low", "high", "reasoning", "think"}180_TRAILING_EFFORT = re.compile(r"\s+(?P<eff>x-?high|extra-high|high|medium|low|minimal|max)\s+effort$", re.IGNORECASE)181_TRAILING_OFF = re.compile(r"\s+(non-?reasoning|no-?think(ing)?|non-?thinking)$", re.IGNORECASE)182183184def strip_effort(model_id: str) -> tuple[str, dict[str, str]]:185    """Remove every trailing evaluator effort suffix of an id ('…-thinking-64k-high-effort' → '…', {reasoning: on, thinking_budget: 64k,186    reasoning_effort: high}). Official '-thinking' releases (ontology list) are left alone."""187    low = re.sub(r"[\s_]+", "-", model_id.strip().lower())188    effort: dict[str, str] = {}189    for _ in range(4):190        if low in OFFICIAL_THINKING_RELEASES:191            break192        a = analyze_model_name(low)193        if a.is_effort_variant and a.effort_suffix:194            stem = low[: -(len(a.effort_suffix) + 1)]195            if a.effort_suffix in _AMBIGUOUS_SUFFIXES and not re.search(r"\d", stem):196                break197            effort = {**a.effort, **effort}198            low = stem199            continue200        hit = next((s for s in sorted(_EXTRA_SUFFIXES, key=len, reverse=True) if low.endswith("-" + s)), None)201        if hit:202            effort = {**_EXTRA_SUFFIXES[hit], **effort}203            low = low[: -(len(hit) + 1)]204            continue205        break206    # keep the original casing of the surviving prefix207    base = model_id.strip()[: len(low)] if re.sub(r"[\s_]+", "-", model_id.strip().lower()).startswith(low) else low208    return base.rstrip("-_ "), effort209210211def strip_effort_words(display: str, *, only_if: bool = True) -> str:212    """Trailing effort words of a display name ('Claude 4.6 Opus Thinking High Effort' → 'Claude 4.6 Opus'). Applied only when the213    caller knows (from the id) that the row is an effort variant — 'Kimi K2 Thinking' and 'Sonar Reasoning' are real releases."""214    if not only_if:215        return display.strip()216    words = display.strip().split()217    while len(words) > 1 and _EFFORT_WORDS.match(words[-1]):218        words.pop()219    return " ".join(words)220221222def is_pinned(model_id: str) -> bool:223    return bool(_DATE_IN_ID.search(model_id))224225226# ---------------------------------------------------------------------------------------------- identity227def org_key_for_vendor(vendor: str | None) -> str | None:228    if not vendor:229        return None230    v = vendor.strip().lower()231    key = VENDOR_ORG.get(v)232    if key and key in organizations():233        return key234    known = org_by_hf(v)235    if known:236        return known["key"]237    if v in organizations():238        return v239    return None240241242def org_key_for_bare_id(model_id: str) -> tuple[str | None, str | None]:243    low = model_id.strip().lower()244    for pat, scheme, org in _BARE:245        if pat.search(low):246            return scheme, (org if org in organizations() else None)247    return None, None248249250def model_identity(api_id: str | None, *, trusted: bool = False) -> ModelIdentity | None:251    """Analyse an API id. `trusted=True` when the string is the vendor's own API id (adds the vendor identifier when recognised)."""252    if not api_id or not isinstance(api_id, str):253        return None254    raw = api_id.strip()255    if not raw or " " in raw:256        return None257    ident = ModelIdentity(raw=raw, base_id=raw)258    rest = raw259    prefix, _, tail = raw.partition("/")260    plow = prefix.lower()261    scheme: str | None = None262    org_key: str | None = None263    if tail and plow in PROVIDER_PREFIXES:264        ident.provider_prefix = plow265        scheme, org_key = PROVIDER_PREFIXES[plow]266        rest = tail267        if plow == "openrouter":268            vendor, _, slug = rest.partition("/")269            if slug:270                base_slug = slug.split(":", 1)[0]271                ident.aliases += [f"{vendor}/{base_slug}", base_slug]272                if trusted:273                    ident.identifiers["openrouter"] = f"{vendor}/{base_slug}"274                org_key = org_key_for_vendor(vendor)275                rest = base_slug276        elif plow == "fireworks_ai":277            rest = _FIREWORKS_PATH.sub("", rest)278            if trusted and "/" not in rest:279                ident.identifiers["fireworks_model_id"] = f"fireworks/{rest}"280            scheme = None281        elif plow in ("nvidia_nim",):282            vendor, _, slug = rest.partition("/")283            if slug:284                org_key = org_key_for_vendor(vendor)285                rest = slug286        elif plow == "openai":287            # `openai/Qwen/Qwen2.5-Coder-32B-Instruct` (OpenAI-compatible endpoint): the remainder decides288            bare_scheme, bare_org = org_key_for_bare_id(rest.split("/")[-1])289            if "/" in rest:290                hf_org = rest.split("/")[0]291                known = org_by_hf(hf_org)292                if known:293                    org_key = known["key"]294                    if trusted:295                        ident.identifiers["hf_repo"] = rest296                    ident.aliases.append(rest)297                rest = rest.split("/")[-1]298            else:299                scheme, org_key = bare_scheme, bare_org300    elif tail and "/" in raw and org_by_hf(prefix):301        # `Qwen/Qwen2.5-Coder-32B-Instruct` — a Hugging Face repository id302        org_key = org_by_hf(prefix)["key"]  # type: ignore[index]303        if trusted:304            ident.identifiers["hf_repo"] = raw305        ident.aliases.append(raw)306        rest = tail307    if org_key is None:308        bare_scheme, bare_org = org_key_for_bare_id(rest)309        scheme = scheme or bare_scheme310        org_key = bare_org311    base, effort = strip_effort(rest)312    ident.base_id = base or rest313    ident.effort = effort314    ident.scheme = scheme315    ident.org_key = org_key if org_key in organizations() else None316    ident.pinned = is_pinned(ident.base_id)317    if scheme and trusted and "/" not in ident.base_id and not ident.identifiers:318        ident.identifiers[scheme] = ident.base_id319    for a in (raw, rest, ident.base_id):320        if a and a not in ident.aliases:321            ident.aliases.append(a)322    return ident323324325def family_ref(name: str | None, org: EntityRef | None) -> EntityRef | None:326    """`model_family` hint for a model name ("Claude Opus 5" → Claude; "Qwen3.6-35B-A3B" → Qwen3.6). None when no family is detected."""327    if not name:328        return None329    # size tokens first ("Qwen3-0.6B" would otherwise read as version 3.0): 0.6B / 35B-A3B / 17B-128E330    bare = re.sub(r"[-_ ]?\d+(?:\.\d+)?\s?[bmt](?:[-_]?a\d+(?:\.\d+)?[bmt])?(?:[-_]?\d+e)?(?![a-z0-9])", "", name.split("/")[-1], flags=re.IGNORECASE)331    label = family_release_hint(bare or name)332    if not label:333        return None334    return EntityRef(entity_type="model_family", name=label, organization=org, identity_confidence="medium",335                     identifiers={"family_key": re.sub(r"[^a-z0-9.]+", "-", label.lower()).strip("-") + (f"@{org.slug_hint}" if org and org.slug_hint else "")})336337338def org_ref_in(facts: Facts, key: str | None) -> EntityRef | None:339    """Registry organisation appended once to `facts.entities`."""340    if not key or key not in organizations():341        return None342    for e in facts.entities:343        if e.identifiers.get("registry_org") == key:344            return e345    ref = org_ref(key)346    facts.entities.append(ref)347    return ref348349350def model_ref_from_api_id(facts: Facts, api_id: str | None, *, name: str | None = None, trusted: bool = False, identity_confidence: str = "medium",351                          organization: EntityRef | None = None, extra_aliases: list[str] | None = None, family: bool = True,352                          attributes: dict[str, Any] | None = None) -> tuple[EntityRef, dict[str, str]]:353    """Model EntityRef for a third-party API id (deduplicated inside `facts`), plus the effort configuration stripped from the id/label.354    `name` is the source's display label (effort parentheticals are folded into the configuration); without it the base id is the name."""355    ident = model_identity(api_id, trusted=trusted)356    label, label_effort = split_effort_label(name) if name else ("", {})357    effort = {**(ident.effort if ident else {}), **label_effort}358    display = label or (ident.base_id if ident else (api_id or "").strip())359    if ident and ident.is_effort_variant and label and not label_effort:360        display = strip_effort_words(label, only_if=True)361    # the first-party pattern org (what the lab connectors use) wins over the board's own creator label, so that the resolver362    # disambiguates aliases with the same organisation the official source attached363    org = (org_ref_in(facts, ident.org_key) if ident else None) or organization364    identifiers = dict(ident.identifiers) if ident else {}365    aliases = [a for a in dict.fromkeys([*(ident.aliases if ident else []), *(extra_aliases or [])]) if a and a != display]366    # one ref per identity inside the document: same identifier, or same display name (a board's rows for one model share its label)367    for e in facts.entities:368        if e.entity_type != "model":369            continue370        if (identifiers and any(e.identifiers.get(k) == v for k, v in identifiers.items())) or e.name.lower() == display.lower():371            for k, v in identifiers.items():372                e.identifiers.setdefault(k, v)373            for a in aliases:374                if a not in e.aliases and a != e.name:375                    e.aliases.append(a)376            if e.organization is None and org is not None:377                e.organization = org378            if e.family is None and family:379                e.family = family_ref(display, org)380            return e, effort381    ref = facts.entity("model", display[:200], identifiers=identifiers, aliases=aliases, organization=org, identity_confidence=identity_confidence,382                       family=family_ref(display, org) if family else None, attributes=dict(attributes or {}))383    return ref, effort384385386__all__ = ["PROVIDER_PREFIXES", "VENDOR_ORG", "ModelIdentity", "effort_from_label", "family_ref", "is_pinned", "model_identity", "model_ref_from_api_id",387           "org_key_for_bare_id", "org_key_for_vendor", "org_ref_in", "split_effort_label", "strip_effort", "strip_effort_words"]388