From 164f0f9fe9a69d929c42f54d2e1b16109adc0119 Mon Sep 17 00:00:00 2001 From: Paul Furgale Date: Mon, 7 Sep 2026 15:13:09 +0000 Subject: [PATCH] feat(llm): config-driven provider declaration and catalog drift check An llm_config.yaml alias can now declare its logical provider and model compat group. Opaque enterprise/gateway model ids resolve to no provider on their own, so without a declaration they carry no identity and replay compatibility fails closed. The declaration is the config's answer: - ModelConfig gains typed optional 'provider' and 'compat_group' fields (extra=allow passthrough unchanged); the registry schema documents them. - get_llm_client resolves a declaring alias through the contracts parser and attaches a ProviderIdentity to the client as metadata only: request parameters are untouched, and undeclared aliases behave exactly as before. - nooa.unifiedllm.declaration translates the declaration into process-lifetime compat-group registration and reasoning-capability overrides, applied lazily on first use of the alias. Declared values override the built-in catalogs; a provider that contradicts what the model string itself resolves raises rather than guessing, and a group name owned by a different provider is rejected. The declared group is authoritative for the declaring alias's opaque replay key, so other groups claiming the same model cannot silently win. - scripts/refresh_model_id_corpora.py rebuilds the openrouter and nvidia_gateway corpus sections from the live catalogs, prints the resolution/misattribution summary, and supports --check to exit 1 on drift without writing. Offline it exits non-zero with a clear message and never truncates the fixture; the provenance comment is excluded from drift so its fetch date does not trip the check. Signed-off-by: Paul Furgale --- scripts/refresh_model_id_corpora.py | 271 ++++++++++++ src/nooa/config/model_config.py | 8 + src/nooa/unifiedllm/__init__.py | 3 + src/nooa/unifiedllm/declaration.py | 273 ++++++++++++ src/nooa/unifiedllm/registry.py | 16 + src/nooa/unifiedllm/unifiedllm.py | 6 + tests/unifiedllm/test_declaration.py | 630 +++++++++++++++++++++++++++ 7 files changed, 1207 insertions(+) create mode 100755 scripts/refresh_model_id_corpora.py create mode 100644 src/nooa/unifiedllm/declaration.py create mode 100644 tests/unifiedllm/test_declaration.py diff --git a/scripts/refresh_model_id_corpora.py b/scripts/refresh_model_id_corpora.py new file mode 100755 index 000000000..07edb7001 --- /dev/null +++ b/scripts/refresh_model_id_corpora.py @@ -0,0 +1,271 @@ +#!/usr/bin/env python +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Refresh the model-id corpora fixture from the live model catalogs. + +Fetches the public OpenRouter catalog (``/v1/models``) and, when +``NVIDIA_INFERENCE_API_KEY`` is set, the NVIDIA inference gateway catalog, +rebuilds the ``openrouter`` / ``nvidia_gateway`` sections of +``tests/unifiedllm/fixtures/model_id_corpora.json``, and prints the +resolution/misattribution summary for the fresh catalog. + +Offline-safe: with no network it exits non-zero with a clear message and +never truncates the existing fixture (a failed fetch is never written). + +Usage:: + + uv run python scripts/refresh_model_id_corpora.py # refresh + uv run python scripts/refresh_model_id_corpora.py --check # report only + +``--check`` exits 1 on drift without writing. +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +import urllib.error +import urllib.request +from datetime import date +from pathlib import Path + +REPO = Path(__file__).resolve().parent.parent +FIXTURE = REPO / "tests/unifiedllm/fixtures/model_id_corpora.json" + +OPENROUTER_URL = "https://openrouter.ai/api/v1/models" +NVIDIA_GATEWAY_URL = "https://inference-api.nvidia.com/v1/models" + +#: Sections not fetched live (azure/vertex_ai/bedrock come from the transport +#: library's bundled catalog, not a public endpoint): preserved verbatim on +#: rewrite, in fixture key order. +PRESERVED_SECTIONS = ("azure", "vertex_ai", "bedrock") + + +def fetch_catalog(url: str, headers: dict[str, str] | None = None) -> dict: + req = urllib.request.Request(url, headers=headers or {}) + with urllib.request.urlopen(req, timeout=30) as resp: + return json.loads(resp.read().decode("utf-8")) + + +def openrouter_entries(payload: dict) -> list[dict]: + """Catalog section entries: ground truth = leading id segment.""" + entries = [] + for model in payload["data"]: + model_id = model["id"] + segments = model_id.split("/") + vendor = segments[0].lstrip("~") if len(segments) > 1 else None + entries.append({"id": model_id, "vendor": vendor}) + return entries + + +def nvidia_gateway_entries(payload: dict) -> list[dict]: + """Catalog section entries. + + Ground truth = the middle vendor segment of ``nvidia//`` + spellings; every other spelling the gateway serves (us/..., gcp/..., + nvcf/..., bare ids) has no segment the catalog itself vouches for, so + vendor stays null and conformance treats it as unverifiable. + """ + entries = [] + for model in payload["data"]: + model_id = model["id"] + segments = model_id.split("/") + vendor = segments[1] if len(segments) == 3 and segments[0] == "nvidia" else None + entries.append({"id": model_id, "vendor": vendor}) + return entries + + +def comment(fetch_date: str, openrouter_count: int, nvidia_count: int | None) -> str: + """Provenance note for the fixture's ``_comment`` member.""" + nvidia = ( + f"nvidia_gateway: {nvidia_count} ids from the NVIDIA inference gateway /v1/models; " + "ground truth = middle vendor segment for nvidia// spellings. " + if nvidia_count is not None + else "nvidia_gateway: not refreshed (NVIDIA_INFERENCE_API_KEY unset); " + "previous section preserved. " + ) + return ( + "Model-id corpora for parse_model_string conformance. " + f"openrouter: {openrouter_count} ids from the public OpenRouter catalog " + f"(fetched {fetch_date}); ground truth = leading id segment. " + + nvidia + + "azure/vertex_ai/bedrock: deployment-prefixed ids from the catalog bundled with the " + "transport library; ground truth = the logical provider of the served model (bare-id " + "provider, or the Bedrock vendor.model head). Transport labels (azure/bedrock/...) are " + "never treated as logical providers. vendor=null means the catalog cannot verify the " + "logical provider." + ) + + +#: Catalog vendor label -> canonical logical provider. Mirrors the table in +#: tests/unifiedllm/test_model_id_corpora.py: these are spelling variants of +#: the SAME logical provider ("z-ai" vs "glm", "moonshotai" vs "kimi"), not +#: misattributions, so the summary must not count them as such. Unknown +#: labels pass through verbatim and usually fail closed (unverifiable). +_CANON = { + "openai": "openai", + "anthropic": "anthropic", + "~anthropic": "anthropic", + "google": "google", + "meta": "meta", + "meta-llama": "meta", + "mistral": "mistral", + "mistralai": "mistral", + "x-ai": "xai", + "~x-ai": "xai", + "xai": "xai", + "nvidia": "nvidia", + "deepseek": "deepseek", + "deepseek-ai": "deepseek", + "~deepseek": "deepseek", + "qwen": "qwen", + "z-ai": "glm", + "~z-ai": "glm", + "zai": "glm", + "zai-org": "glm", + "moonshot": "kimi", + "moonshotai": "kimi", + "minimaxai": "minimax", + "microsoft": "microsoft", +} + + +def _canon(label: str) -> str: + return _CANON.get(label.lower().strip(), label.lower().strip()) + + +def summarize(name: str, entries: list[dict]) -> None: + """Print the resolution/misattribution summary for one catalog section. + + Uses the same canonicalization as the conformance test, so the numbers + here are the numbers the test would measure. + """ + from nooa.unifiedllm.contracts import parse_model_string + + misattributed: list[tuple[str, str, str]] = [] + resolved = 0 + with_truth = 0 + for entry in entries: + vendor = entry["vendor"] + if not vendor: + continue # no ground truth for this spelling + truth = _canon(vendor) + with_truth += 1 + parsed = parse_model_string(entry["id"]) + if parsed.provider is not None and parsed.provider != truth: + misattributed.append((entry["id"], vendor, parsed.provider)) + elif parsed.provider is not None: + resolved += 1 + rate = resolved / with_truth if with_truth else 0.0 + print(f"{name}: {len(entries)} ids, {with_truth} with catalog ground truth") + print(f" resolution rate: {rate:.3f} ({resolved}/{with_truth})") + print(f" misattributed: {len(misattributed)}") + for model_id, vendor, got in misattributed[:20]: + print(f" {model_id}: catalog says {vendor}, parser says {got}") + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--check", + action="store_true", + help="only report drift; exit 1 if the fixture differs from the live catalogs, " + "without writing", + ) + parser.add_argument( + "--openrouter-url", + default=OPENROUTER_URL, + help="override the OpenRouter catalog URL", + ) + args = parser.parse_args(argv) + + try: + print(f"Fetching OpenRouter catalog: {args.openrouter_url}") + openrouter_section = openrouter_entries(fetch_catalog(args.openrouter_url)) + except (urllib.error.URLError, OSError, ValueError, KeyError) as exc: + print( + f"error: could not fetch the OpenRouter catalog ({args.openrouter_url}): {exc}\n" + "The refresh needs network access; the fixture was left untouched.", + file=sys.stderr, + ) + return 1 + if not openrouter_section: + print( + "error: the OpenRouter catalog returned no models; refusing to write.", file=sys.stderr + ) + return 1 + + nvidia_section: list[dict] | None = None + api_key = os.getenv("NVIDIA_INFERENCE_API_KEY") + if api_key: + try: + print(f"Fetching NVIDIA gateway catalog: {NVIDIA_GATEWAY_URL}") + nvidia_section = nvidia_gateway_entries( + fetch_catalog(NVIDIA_GATEWAY_URL, headers={"Authorization": f"Bearer {api_key}"}) + ) + except (urllib.error.URLError, OSError, ValueError, KeyError) as exc: + print( + f"error: could not fetch the NVIDIA gateway catalog ({NVIDIA_GATEWAY_URL}): " + f"{exc}\nThe fixture was left untouched.", + file=sys.stderr, + ) + return 1 + else: + print( + "NVIDIA_INFERENCE_API_KEY is unset: keeping the fixture's nvidia_gateway section " + "and its ground truth as-is." + ) + + for name, section in (("openrouter", openrouter_section), ("nvidia_gateway", nvidia_section)): + if section is not None: + summarize(name, section) + + current = json.loads(FIXTURE.read_text()) + # Keys keep the fixture's original order so a refresh diffs only what + # actually changed. Sections not fetched live (their source is the + # transport library's bundled catalog, not a public endpoint) are + # preserved verbatim. + fresh: dict = { + "_comment": comment( + date.today().isoformat(), + len(openrouter_section), + len(nvidia_section) if nvidia_section is not None else None, + ), + "openrouter": openrouter_section, + "nvidia_gateway": nvidia_section + if nvidia_section is not None + else current["nvidia_gateway"], + **{name: current[name] for name in PRESERVED_SECTIONS}, + } + + # Drift is about corpus content, not the fixture's provenance note: the + # fetch date in _comment changes daily, so counting it would make --check + # fail against an otherwise-identical corpus. + drifted = [ + key + for key in ("openrouter", "nvidia_gateway", *PRESERVED_SECTIONS) + if fresh[key] != current.get(key) + ] + if args.check: + if drifted: + print( + "Drift detected in corpus section(s): " + + ", ".join(drifted) + + " (rerun without --check to refresh the fixture).", + file=sys.stderr, + ) + return 1 + print("No drift: the fixture's corpus already matches the live catalogs.") + return 0 + + # Same layout the fixture has always used (indent=1, one compact entry per + # model) so the diff stays reviewable. + FIXTURE.write_text(json.dumps(fresh, indent=1) + "\n") + print(f"Wrote {FIXTURE}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/nooa/config/model_config.py b/src/nooa/config/model_config.py index c081837e8..8c3396b56 100644 --- a/src/nooa/config/model_config.py +++ b/src/nooa/config/model_config.py @@ -39,6 +39,14 @@ class ModelConfig(BaseModel): # Name of the env var holding the API key (NOT the key itself). api_key_env: str | None = None client_type: str | None = None + # Canonical logical provider for ids whose routing string resolves to no + # provider (opaque enterprise/gateway ids). ``None`` means undeclared — + # identity then stays unset and consumption-time behavior is unchanged. + provider: str | None = None + # Name of the model compat group the declared model belongs to. Gives an + # otherwise-opaque id an opaque replay boundary derived from the declared + # group; requires ``provider`` (or a model string that resolves one). + compat_group: str | None = None context_window: int | None = None max_tokens: int | None = None temperature: float | None = None diff --git a/src/nooa/unifiedllm/__init__.py b/src/nooa/unifiedllm/__init__.py index e457479e3..3e527a2d4 100644 --- a/src/nooa/unifiedllm/__init__.py +++ b/src/nooa/unifiedllm/__init__.py @@ -17,6 +17,7 @@ register_compat_group, register_reasoning_capabilities, ) +from nooa.unifiedllm.declaration import apply_alias_declaration from nooa.unifiedllm.fake import FakeLLMClient from nooa.unifiedllm.http_config import HttpConfig from nooa.unifiedllm.registry import ( @@ -63,6 +64,8 @@ "parse_model_string", "register_compat_group", "register_reasoning_capabilities", + # Config-driven declarations (registry edge) + "apply_alias_declaration", # Core classes "UnifiedLLM", "CompletionClient", diff --git a/src/nooa/unifiedllm/declaration.py b/src/nooa/unifiedllm/declaration.py new file mode 100644 index 000000000..fef7920f3 --- /dev/null +++ b/src/nooa/unifiedllm/declaration.py @@ -0,0 +1,273 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Config-driven provider declarations at the registry edge. + +A registry alias can declare what its model id cannot prove on its own: + +- ``provider`` — the canonical logical provider ("openai", "glm", ...), for + opaque enterprise/gateway model ids whose routing string resolves to no + provider; +- ``compat_group`` — the model compat group the id belongs to, giving an + opaque id an opaque replay boundary derived from the declared group + instead of none at all. + +:func:`apply_alias_declaration` is called lazily by +:func:`nooa.unifiedllm.registry.get_llm_client` on first use of a declaring +alias. It translates the declaration into process-lifetime registrations +that override the built-in catalogs (compat-group membership via +:func:`register_compat_group`, reasoning capabilities via +:func:`register_reasoning_capabilities`) and returns the +:class:`ProviderIdentity` the registry attaches to the client as metadata — +never as a request parameter. + +Fail-closed rules: + +- No declaration → no identity; consumption-time behavior is unchanged. +- ``compat_group`` without a resolvable provider (declared or parsed) raises + ``ValueError`` — groups are provider-scoped, so a group-only declaration + would be a guess. +- A declared provider that contradicts what :func:`parse_model_string` + resolves from the model string raises ``ValueError`` — replay artifacts + must never be routed to a provider the config contradicts. +- No ``capabilities`` declaration leaves the catalog answer as-is: unknown + providers still look up to ``None``. A ``capabilities``-only declaration + registers an override (when a provider is resolvable) but carries no + identity. +""" + +from __future__ import annotations + +import json +import logging +import threading +from collections.abc import Mapping +from typing import Any + +from nooa.unifiedllm.contracts import ( + _COMPAT_GROUPS, + _PROVIDER_ALIASES, + ModelCompatGroup, + ProviderIdentity, + ReasoningCapabilities, + derive_opaque_replay_key, + parse_model_string, + register_compat_group, + register_reasoning_capabilities, +) + +__all__ = ["apply_alias_declaration", "canonical_provider"] + +logger = logging.getLogger(__name__) + +# Serializes the registrations and the once-per-declaration bookkeeping. +# Reentrant so nested application (a declaration whose registration path +# re-enters the module) cannot self-deadlock. +_declaration_lock = threading.RLock() + +#: alias -> declaration signature already applied this process. Registrations +#: are process-lifetime: re-applying an identical declaration is a no-op, and a +#: *changed* declaration (the registry was reloaded with different values) +#: re-applies and overrides what the earlier one registered. +_applied: dict[str, tuple[Any, ...]] = {} + + +def canonical_provider(provider: str) -> str: + """Canonicalize a declared provider through the contracts alias map. + + ``derive_opaque_replay_key`` canonicalizes the provider with the same + map, so going through it here is what keeps a declared group registration + and the derived key from disagreeing about scope (a "zai" declaration + registering under "zai" while the key derives under "glm" would never + match). + """ + key = provider.strip().lower() + return _PROVIDER_ALIASES.get(key, key) + + +def _register_declared_group(*, provider: str, group_name: str, model: str) -> None: + """Register the declared *model* into *group_name* under *provider*. + + A declaration may create a new group or extend a same-provider one (an + opaque enterprise id joining a hand-verified group). Extending keeps the + group's existing members — they were live-verified, and the declaration + asserts one more member, not a replacement. A same-name group owned by a + *different* provider is rejected: accepting it would move that provider's + replay boundary, which a single alias's declaration must not do. + """ + with _declaration_lock: + existing = _COMPAT_GROUPS.get(group_name) + if existing is None: + register_compat_group( + ModelCompatGroup(name=group_name, provider=provider, models=frozenset({model})) + ) + elif existing.provider.lower() != provider.lower(): + raise ValueError( + f"Compat group {group_name!r} is already declared for provider " + f"{existing.provider!r}, but this declaration says {provider!r}; " + "declare a different group name." + ) + else: + register_compat_group( + existing.model_copy(update={"models": existing.models | frozenset({model})}) + ) + # The registration is process-wide, so a model that other same- + # provider groups already claim now sits in several groups; lookups + # outside this alias resolve deterministically (alphabetically first + # group name), which may not be this one. Surface it rather than let + # the declaration appear to take effect when it did not. + conflicts = sorted( + group.name + for group in _COMPAT_GROUPS.values() + if group.name != group_name + and group.provider.lower() == provider.lower() + and model.lower() in {m.lower() for m in group.models} + ) + if conflicts: + logger.warning( + "Model %r is now a member of compat groups %r and %r; catalog " + "lookups resolve to the alphabetically first group name.", + model, + group_name, + conflicts, + ) + + +def _apply_capability_override(alias: str, provider: str, config: Mapping[str, Any]) -> None: + """Register the alias's declared ``capabilities`` for *provider*. + + ``capabilities`` is an optional registry mapping shaped like + :class:`ReasoningCapabilities` (``capture_kinds``, ``effort_map``, + ``replay_field``, ...). It is validated, never trusted: an invalid + declaration is warned about and dropped, because a wrong *capability* + profile degrades reasoning handling without breaking the request path — + unlike a wrong *identity*, which must fail loudly. + """ + caps = config.get("capabilities") + if not caps: + return + if not isinstance(caps, Mapping): + logger.warning( + "Model %r has an invalid capabilities declaration: expected a mapping, got %s; ignoring it.", + alias, + type(caps).__name__, + ) + return + try: + register_reasoning_capabilities(provider, ReasoningCapabilities.model_validate(dict(caps))) + except Exception as exc: + logger.warning( + "Model %r has an invalid capabilities declaration (%s); ignoring it.", alias, exc + ) + + +def apply_alias_declaration( + alias: str, + config: Mapping[str, Any], + *, + api_style: str, + transport: str, +) -> ProviderIdentity | None: + """Translate *alias*'s declared provider/compat_group into registrations + plus a :class:`ProviderIdentity`, or return ``None`` when nothing is + declared. + + The registration side effects (compat-group membership, capability + overrides) happen once per distinct declaration for the process lifetime; + the returned identity is a pure value and is rebuilt on every call. + + Raises: + ValueError: the declaration is unusable — a ``compat_group`` with no + resolvable provider, a provider that contradicts the model + string, a group name owned by another provider, or a non-string + declaration value. + """ + for field, value in ( + ("provider", config.get("provider")), + ("compat_group", config.get("compat_group")), + ): + if value is not None and not isinstance(value, str): + raise ValueError( + f"Model {alias!r} declares {field}={value!r}: expected a string or null." + ) + declared_provider = ( + config["provider"].strip() if isinstance(config.get("provider"), str) else None + ) + declared_group = ( + config["compat_group"].strip() if isinstance(config.get("compat_group"), str) else None + ) + caps = config.get("capabilities") + if not declared_provider and not declared_group and not caps: + return None + + model_name = config.get("model_name") + model_string = model_name if isinstance(model_name, str) and model_name else alias + parsed = parse_model_string(model_string) + + if declared_provider: + provider = canonical_provider(declared_provider) + elif parsed.provider is not None: + # No provider declared: the parser resolves one, so the declaration + # can be honored without guessing. + provider = parsed.provider + elif declared_group: + raise ValueError( + f"Model {alias!r} declares compat_group {declared_group!r} without a provider, " + f"and model string {model_string!r} resolves to no provider either; declare " + "'provider' alongside 'compat_group'." + ) + else: + # Capabilities-only declaration for an id with no resolvable + # provider: there is nothing to register them under. + logger.warning( + "Model %r declares capabilities without a resolvable provider; ignoring the " + "declaration.", + alias, + ) + return None + if parsed.provider is not None and parsed.provider != provider: + raise ValueError( + f"Model {alias!r} declares provider {declared_provider!r}, but model string " + f"{model_string!r} resolves to {parsed.provider!r}; fix the declaration." + ) + + caps_key = json.dumps(caps, sort_keys=True, default=repr) if caps else None + signature = (model_string, declared_provider, declared_group, caps_key) + with _declaration_lock: + if _applied.get(alias) != signature: + if declared_group: + _register_declared_group( + provider=provider, group_name=declared_group, model=parsed.model + ) + _apply_capability_override(alias, provider, config) + _applied[alias] = signature + + # Capabilities-only declarations register overrides but carry no + # identity: provider/compat_group is what an identity declaration is. + if not declared_provider and not declared_group: + return None + + # The declared group is authoritative for THIS alias's replay boundary: + # deriving against a mapping that contains exactly the declaration keeps + # other groups (built-in or registered by other aliases) that also claim + # the model from silently winning the lookup. + declared_groups = ( + { + declared_group: ModelCompatGroup( + name=declared_group, provider=provider, models=frozenset({parsed.model}) + ) + } + if declared_group + else None + ) + return ProviderIdentity( + provider=provider, + api_style=api_style, + model=parsed.model, + transport=transport, + opaque_replay_key=derive_opaque_replay_key( + provider=provider, + api_style=api_style, + model=parsed.model, + compat_groups=declared_groups, + ), + ) diff --git a/src/nooa/unifiedllm/registry.py b/src/nooa/unifiedllm/registry.py index 7ac724715..5713d49c2 100644 --- a/src/nooa/unifiedllm/registry.py +++ b/src/nooa/unifiedllm/registry.py @@ -38,6 +38,10 @@ top_p: 1.0 # optional max_tokens: 4096 # optional drop_params: true # optional, defaults to true + provider: openai # optional: logical provider for ids + # the model string cannot resolve + compat_group: openai-gpt-5 # optional: declared compat group + # (opaque replay boundary) Set a model to ``null`` in a later layer to remove it. """ @@ -54,6 +58,8 @@ import yaml +from nooa.unifiedllm.declaration import apply_alias_declaration + if TYPE_CHECKING: from nooa.unifiedllm import UnifiedLLM @@ -412,4 +418,14 @@ def get_llm_client(name: str, *, client_type: str | None = None, **overrides) -> client_type = client_type or config.get("client_type", "completion") client = ResponsesClient(**params) if client_type == "responses" else CompletionClient(**params) client._registry_config = config # For context_window lookup + + # An alias may declare its logical provider and compat group — the two + # things an opaque enterprise/gateway model id cannot prove on its own. + # Applied lazily (first use of the alias) and attaches identity as + # metadata only: the request path is untouched. Aliases that declare + # nothing stay exactly as they were — no identity attached, and + # consumers keep failing closed at consumption time. + client.provider_identity = apply_alias_declaration( + name, config, api_style=client_type or "completion", transport="litellm" + ) return client diff --git a/src/nooa/unifiedllm/unifiedllm.py b/src/nooa/unifiedllm/unifiedllm.py index 569c6b589..ff9b58279 100644 --- a/src/nooa/unifiedllm/unifiedllm.py +++ b/src/nooa/unifiedllm/unifiedllm.py @@ -17,6 +17,7 @@ import litellm from pydantic import BaseModel, RootModel +from .contracts import ProviderIdentity from .http_config import HttpConfig from .retry import EmptyContentError, sync_retry, with_retry from .retry_config import RetryConfig @@ -1152,11 +1153,16 @@ def _update_token_calibration( class UnifiedLLM(ABC): _registry_config: dict[str, Any] | None + #: Logical provider identity for models the registry declares one for. + #: Attached by get_llm_client() as metadata (never a request param); + #: None means undeclared, and consumers must fail closed on that. + provider_identity: ProviderIdentity | None def __init__(self, model: str, **config): self.model = model self.config = config self._registry_config = None + self.provider_identity = None # Cache control injection — shared by CompletionClient and ResponsesClient self.cache_control_injection_points: list[dict[str, Any]] = ( DEFAULT_CACHE_CONTROL_INJECTION_POINTS diff --git a/tests/unifiedllm/test_declaration.py b/tests/unifiedllm/test_declaration.py new file mode 100644 index 000000000..2845802f9 --- /dev/null +++ b/tests/unifiedllm/test_declaration.py @@ -0,0 +1,630 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Tests for config-driven provider declarations at the registry edge. + +An alias can declare ``provider`` / ``compat_group`` (and optional +``capabilities``) in ``llm_config.yaml``. These tests pin the lazy +registration flow, the fail-closed paths, and that declaration is metadata +only — the request path and undeclared aliases are untouched. +""" + +from __future__ import annotations + +import importlib.util +import sys +import textwrap +from collections.abc import Iterator +from pathlib import Path + +import pytest + +from nooa.unifiedllm import ( + CompletionClient, + ResponsesClient, + get_llm_client, + reload_registry, +) +from nooa.unifiedllm import contracts as contracts +from nooa.unifiedllm.declaration import apply_alias_declaration + +REPO = Path(__file__).resolve().parent.parent.parent +SCRIPT = REPO / "scripts/refresh_model_id_corpora.py" +FIXTURE = REPO / "tests/unifiedllm/fixtures/model_id_corpora.json" + + +@pytest.fixture() +def scratch_config(tmp_path, monkeypatch) -> Path: + """Isolated registry pointing at one YAML file; restored afterwards. + + User/project dirs point at an empty temp dir (the real user's + ``~/.config/nooa/llm_config.yaml`` must not leak in) and bundled-default + entry-points are stubbed empty, mirroring test_model_registry.py. + """ + user = tmp_path / "user" + user.mkdir() + monkeypatch.setenv("NEMO_OO_USER_DIR", str(user)) + monkeypatch.setenv("NEMO_OO_PROJECT_DIR", str(tmp_path / "proj")) + monkeypatch.delenv("NEMO_OO_LLM_CONFIG", raising=False) + monkeypatch.setattr("nooa.llm_config.bundled_config_paths", lambda: []) + monkeypatch.chdir(tmp_path) + cfg = tmp_path / "llm_config.yaml" + cfg.write_text("") + yield cfg + reload_registry() # reset the registry for the next test + + +def write_models(cfg: Path, body: str) -> None: + cfg.write_text(textwrap.dedent(body)) + + +@pytest.fixture(autouse=True) +def _restore_contract_registries() -> Iterator[None]: + """Declarations register into process-lifetime catalogs; restore them.""" + yield + contracts._COMPAT_GROUPS.clear() + contracts._COMPAT_GROUPS.update({g.name: g for g in contracts._DEFAULT_COMPAT_GROUPS}) + contracts._REASONING_CAPABILITIES.clear() + contracts._REASONING_CAPABILITIES.update(dict(contracts.DEFAULT_REASONING_CAPABILITIES)) + from nooa.unifiedllm import declaration as declaration + + declaration._applied.clear() + + +# --- ModelConfig typed fields ---------------------------------------------- + + +class TestModelConfigDeclarationFields: + def test_provider_and_compat_group_are_typed_fields(self): + from nooa.config import ModelConfig + + mc = ModelConfig.from_registry( + "opaque-enterprise", {"provider": "openai", "compat_group": "openai-gpt-5"} + ) + assert mc.provider == "openai" + assert mc.compat_group == "openai-gpt-5" + + def test_declaration_fields_default_to_none(self): + from nooa.config import ModelConfig + + mc = ModelConfig.from_registry("plain", {"model_name": "gpt-5.6"}) + assert mc.provider is None + assert mc.compat_group is None + + def test_extra_allow_keeps_litellm_passthrough(self): + from nooa.config import ModelConfig + + mc = ModelConfig.from_registry( + "a", {"provider": "openai", "num_retries": 7, "extra_body": {"trace": True}} + ) + assert mc.num_retries == 7 # type: ignore[attr-defined] + assert mc.extra_body == {"trace": True} # type: ignore[attr-defined] + + def test_non_string_declaration_rejected(self): + from pydantic import ValidationError + + from nooa.config import ModelConfig + + with pytest.raises(ValidationError): + ModelConfig.from_registry("bad", {"provider": 42}) + + +# --- Direct declaration behavior ------------------------------------------- + + +class TestApplyAliasDeclaration: + def test_no_declaration_returns_none(self): + assert ( + apply_alias_declaration( + "some-unknown-model-xyz", {}, api_style="responses", transport="litellm" + ) + is None + ) + + def test_declared_group_gives_opaque_id_a_key(self): + identity = apply_alias_declaration( + "enterprise-internal", + {"provider": "openai", "compat_group": "openai-gpt-5"}, + api_style="responses", + transport="litellm", + ) + assert identity is not None + assert identity.provider == "openai" + assert identity.model == "enterprise-internal" + # Key matches the one derived straight from the declared group. + expected = contracts.derive_opaque_replay_key( + provider="openai", + api_style="responses", + model="enterprise-internal", + compat_groups={ + "openai-gpt-5": contracts.ModelCompatGroup( + name="openai-gpt-5", + provider="openai", + models=frozenset({"enterprise-internal"}), + ) + }, + ) + assert identity.opaque_replay_key == expected + assert identity.opaque_replay_key is not None + + def test_group_only_declaration_uses_parsed_provider(self): + # No `provider` declared: the model string itself must resolve one, + # which is what makes a compat_group-only declaration non-speculative. + identity = apply_alias_declaration( + "alias", + {"model_name": "gpt-5.6-sol", "compat_group": "openai-gpt-5"}, + api_style="responses", + transport="litellm", + ) + assert identity is not None + assert identity.provider == "openai" + assert identity.model == "gpt-5.6-sol" + + def test_group_only_with_unresolvable_model_raises(self): + with pytest.raises(ValueError, match="without a provider"): + apply_alias_declaration( + "enterprise-internal", + {"compat_group": "my-group"}, + api_style="responses", + transport="litellm", + ) + + def test_declared_provider_must_not_contradict_the_model_string(self): + # The model string resolves to a provider the declaration disagrees + # with; trusting the declaration would route replay artifacts across + # providers, so it is rejected instead. + with pytest.raises(ValueError, match="fix the declaration"): + apply_alias_declaration( + "alias", + { + "model_name": "claude-opus-4-5", + "provider": "openai", + "compat_group": "openai-gpt-5", + }, + api_style="responses", + transport="litellm", + ) + + def test_alias_provider_is_canonicalized_like_the_key_derivation(self): + identity = apply_alias_declaration( + "a", + {"provider": "ZAI", "compat_group": "glm-family"}, + api_style="responses", + transport="litellm", + ) + assert identity is not None + assert identity.provider == "glm" + # The group must have been registered under the canonical provider so + # the case-insensitive catalog lookup in derive_opaque_replay_key + # agrees about the scope. + group = contracts.compat_group_for("glm", "a") + assert group is not None and group.name == "glm-family" + + def test_registration_is_process_lifetime_and_declared_group_overrides_catalog(self): + # "gpt-5.6" is in the built-in openai-gpt-5 group; the declaration + # routes it to its own group, and the derived key reflects that. + identity = apply_alias_declaration( + "alias", + {"model_name": "gpt-5.6", "provider": "openai", "compat_group": "my-gpt5"}, + api_style="responses", + transport="litellm", + ) + assert identity is not None + assert identity.opaque_replay_key is not None + group = contracts.compat_group_for("openai", "gpt-5.6") + assert group is not None + # Both the built-in and the declared group now claim the model; + # the declaration wins for THIS alias, and catalog lookups resolve + # deterministically (alphabetically first). + assert group.name == "my-gpt5" + + def test_existing_group_members_are_extended_not_replaced(self): + contracts.register_compat_group( + contracts.ModelCompatGroup( + name="verified-group", provider="openai", models=frozenset({"gpt-5.6"}) + ) + ) + apply_alias_declaration( + "alias", + { + "model_name": "opaque-enterprise", + "provider": "openai", + "compat_group": "verified-group", + }, + api_style="responses", + transport="litellm", + ) + group = contracts.compat_group_for("openai", "opaque-enterprise") + assert group is not None + assert {"gpt-5.6", "opaque-enterprise"} <= group.models + + def test_group_name_owned_by_other_provider_is_rejected(self): + with pytest.raises(ValueError, match="already declared for provider"): + apply_alias_declaration( + "alias", + { + "model_name": "claude-sonnet-4-5", + "provider": "anthropic", + "compat_group": "openai-gpt-5", + }, + api_style="responses", + transport="litellm", + ) + + def test_reapplication_with_same_declaration_is_a_noop(self): + cfg = {"provider": "openai", "compat_group": "gpt-family"} + first = apply_alias_declaration("a", cfg, api_style="responses", transport="litellm") + second = apply_alias_declaration("a", cfg, api_style="responses", transport="litellm") + assert first == second + from nooa.unifiedllm import declaration as declaration + + assert declaration._applied["a"][1:] == ("openai", "gpt-family", None) + + def test_changed_declaration_reapplies_and_overrides(self): + apply_alias_declaration( + "a", + {"provider": "openai", "compat_group": "group-one"}, + api_style="responses", + transport="litellm", + ) + # The registry was reloaded with different values for the same alias. + identity = apply_alias_declaration( + "a", + {"provider": "openai", "compat_group": "group-two"}, + api_style="responses", + transport="litellm", + ) + assert identity is not None + # The new declaration is registered and is authoritative for THIS + # alias's key derivation. (The earlier group-one registration is + # process-lifetime too, so the catalog now holds both; a fresh + # process resolves only group-two.) + assert "group-two" in contracts._COMPAT_GROUPS + assert "a" in {m.lower() for m in contracts._COMPAT_GROUPS["group-two"].models} + expected = contracts.derive_opaque_replay_key( + provider="openai", + api_style="responses", + model="a", + compat_groups={ + "group-two": contracts.ModelCompatGroup( + name="group-two", provider="openai", models=frozenset({"a"}) + ) + }, + ) + assert identity.opaque_replay_key == expected + + def test_capabilities_override_registered(self): + caps = { + "capture_kinds": ["text"], + "native_replay_kinds": [], + "effort_map": {"medium": "banana"}, + "replay_field": "thoughts", + } + apply_alias_declaration( + "alias", + {"provider": "openai", "compat_group": "openai-gpt-5", "capabilities": caps}, + api_style="responses", + transport="litellm", + ) + got = contracts.get_reasoning_capabilities("openai") + assert got is not None + assert "banana" in got.capture_kinds or got.effort_map["medium"] == "banana" + + def test_invalid_capabilities_dropped_with_warning(self, caplog): + with caplog.at_level("WARNING"): + apply_alias_declaration( + "alias", + { + "provider": "openai", + "compat_group": "openai-gpt-5", + "capabilities": {"bogus": True}, + }, + api_style="responses", + transport="litellm", + ) + assert any("capabilities" in r.message for r in caplog.records) + # Catalog answer unchanged. + assert ( + contracts.get_reasoning_capabilities("openai") + == contracts.DEFAULT_REASONING_CAPABILITIES["openai"] + ) + + def test_none_declaration_fields_mean_undeclared(self): + identity = apply_alias_declaration( + "a", + {"provider": None, "compat_group": None}, + api_style="responses", + transport="litellm", + ) + assert identity is None + + def test_capabilities_only_registers_override_but_no_identity(self): + caps = { + "capture_kinds": ["text"], + "native_replay_kinds": [], + "effort_map": {"medium": None}, + } + identity = apply_alias_declaration( + "a", + {"model_name": "gpt-5.6-sol", "capabilities": caps}, + api_style="responses", + transport="litellm", + ) + # No provider/compat_group declared: the override applies, but the + # alias has no identity declaration, so none is attached. + assert identity is None + # The declared profile replaces the catalog's wholesale. + got = contracts.get_reasoning_capabilities("openai") + assert got is not None + assert got.effort_map == {"medium": None} + assert "text" in got.capture_kinds + + def test_capabilities_only_without_resolvable_provider_is_dropped(self, caplog): + with caplog.at_level("WARNING"): + identity = apply_alias_declaration( + "opaque-enterprise", + {"capabilities": {"capture_kinds": ["text"], "effort_map": {}}}, + api_style="responses", + transport="litellm", + ) + assert identity is None + assert any("without a resolvable provider" in r.message for r in caplog.records) + + +# --- Registry end-to-end (get_llm_client) ---------------------------------- + + +class TestRegistryDeclaration: + def test_opaque_alias_derives_group_key(self, scratch_config): + write_models( + scratch_config, + """\ + models: + enterprise-gpt: + model_name: enterprise-internal-b9f2 + api_base: https://gw.example.com/v1 + api_key_env: MY_GATEWAY_KEY + provider: openai + compat_group: openai-gpt-5 + """, + ) + reload_registry(scratch_config) + client = get_llm_client("enterprise-gpt") + assert client.model == "enterprise-internal-b9f2" + identity = client.provider_identity + assert identity is not None + assert identity.provider == "openai" + assert identity.model == "enterprise-internal-b9f2" + assert identity.transport == "litellm" + # Derived from the DECLARED provider+group, not from the opaque id. + assert identity.opaque_replay_key is not None + + def test_same_alias_without_declaration_has_no_identity(self, scratch_config): + write_models( + scratch_config, + """\ + models: + enterprise-gpt: + model_name: enterprise-internal-b9f2 + """, + ) + reload_registry(scratch_config) + client = get_llm_client("enterprise-gpt") + assert client.provider_identity is None + # And no speculative group was registered for the opaque id. + assert contracts.compat_group_for("openai", "enterprise-internal-b9f2") is None + + def test_client_type_declares_api_style(self, scratch_config): + write_models( + scratch_config, + """\ + models: + enterprise-gpt: + model_name: enterprise-internal-b9f2 + client_type: responses + provider: openai + compat_group: openai-gpt-5 + """, + ) + reload_registry(scratch_config) + client = get_llm_client("enterprise-gpt") + assert isinstance(client, ResponsesClient) + assert client.provider_identity is not None + assert client.provider_identity.api_style == "responses" + + def test_declared_group_registered_for_process_lifetime(self, scratch_config): + write_models( + scratch_config, + """\ + models: + enterprise-gpt: + model_name: enterprise-internal-b9f2 + provider: openai + compat_group: openai-gpt-5 + """, + ) + reload_registry(scratch_config) + get_llm_client("enterprise-gpt") + # The declaration registered the opaque id into the named group, so + # even catalog lookups outside this alias now resolve it. + group = contracts.compat_group_for("openai", "enterprise-internal-b9f2") + assert group is not None + assert group.name == "openai-gpt-5" + + def test_declaration_is_metadata_only(self, scratch_config, monkeypatch): + """Identity is attached as an attribute, never as a request param.""" + captured: dict = {} + + class RecordingClient(CompletionClient): + def __init__(self, **kwargs): + captured.update(kwargs) + super().__init__(**kwargs) + + # get_llm_client imports the client classes from the package at call + # time, so patching the package attribute is enough. + monkeypatch.setattr("nooa.unifiedllm.ResponsesClient", RecordingClient) + write_models( + scratch_config, + """\ + models: + enterprise-gpt: + model_name: enterprise-internal-b9f2 + client_type: responses + provider: openai + compat_group: openai-gpt-5 + """, + ) + reload_registry(scratch_config) + client = get_llm_client("enterprise-gpt") + assert isinstance(client, RecordingClient) + assert client.provider_identity is not None + assert "provider_identity" not in captured + assert "provider" not in captured + assert "compat_group" not in captured + + def test_undeclared_aliases_keep_exact_previous_config(self, scratch_config): + write_models( + scratch_config, + """\ + models: + my-alias: + model_name: openai/my-org/my-model + temperature: 0.3 + """, + ) + reload_registry(scratch_config) + client = get_llm_client("my-alias") + assert client.provider_identity is None + assert client.model == "openai/my-org/my-model" + assert client.config["temperature"] == 0.3 + + def test_invalid_declaration_fails_loudly_at_client_construction(self, scratch_config): + """A contradictory declaration is a config error, not a silent no-op.""" + write_models( + scratch_config, + """\ + models: + wrong: + model_name: claude-opus-4-5 + provider: openai + compat_group: openai-gpt-5 + """, + ) + reload_registry(scratch_config) + with pytest.raises(ValueError, match="fix the declaration"): + get_llm_client("wrong") + + def test_non_string_declaration_rejected_at_client_construction(self, scratch_config): + write_models( + scratch_config, + """\ + models: + wrong: + model_name: enterprise-internal + provider: 42 + """, + ) + reload_registry(scratch_config) + with pytest.raises(ValueError, match="expected a string"): + get_llm_client("wrong") + + +# --- Refresh script --check mode ------------------------------------------- + + +@pytest.fixture(scope="module") +def script(): + """Load refresh_model_id_corpora.py, which lives in scripts/ (not importable).""" + spec = importlib.util.spec_from_file_location("_refresh_model_id_corpora", SCRIPT) + assert spec and spec.loader + module = importlib.util.module_from_spec(spec) + sys.modules["_refresh_model_id_corpora"] = module + spec.loader.exec_module(module) + return module + + +class TestCorpusDriftLogic: + def test_openrouter_entries_from_synthetic_payload(self, script): + payload = {"data": [{"id": "openai/gpt-6-astra"}, {"id": "~z-ai/glm-latest"}]} + assert script.openrouter_entries(payload) == [ + {"id": "openai/gpt-6-astra", "vendor": "openai"}, + {"id": "~z-ai/glm-latest", "vendor": "z-ai"}, + ] + + def test_nvidia_gateway_entries_from_synthetic_payload(self, script): + payload = { + "data": [ + {"id": "nvidia/zai-org/glm-5.3"}, + {"id": "gcp/google/gemini-omni-flash-preview"}, + {"id": "nvidia/nvidia/cosmos3-nano-reasoner"}, + ] + } + assert script.nvidia_gateway_entries(payload) == [ + {"id": "nvidia/zai-org/glm-5.3", "vendor": "zai-org"}, + {"id": "gcp/google/gemini-omni-flash-preview", "vendor": None}, + {"id": "nvidia/nvidia/cosmos3-nano-reasoner", "vendor": "nvidia"}, + ] + + def test_offline_fetch_failure_exits_nonzero_without_touching_fixture( + self, script, monkeypatch, capsys + ): + """No network → non-zero exit, fixture untouched, clear message.""" + before = FIXTURE.read_text() + monkeypatch.setattr( + script.urllib.request, + "urlopen", + lambda *a, **k: (_ for _ in ()).throw(OSError("no network")), + ) + assert script.main([]) == 1 + err = capsys.readouterr().err + assert "needs network access" in err + assert FIXTURE.read_text() == before + + def test_check_mode_reports_drift_without_writing(self, script, monkeypatch, tmp_path, capsys): + """Drifted catalog + --check → exit 1, fixture untouched.""" + before = FIXTURE.read_text() + fixture_copy = tmp_path / "model_id_corpora.json" + fixture_copy.write_text(before) + monkeypatch.setattr(script, "FIXTURE", fixture_copy) + live = {"data": [{"id": "openai/gpt-6-astra"}]} # one model: not the corpus + monkeypatch.setattr(script, "fetch_catalog", lambda *a, **k: live) + monkeypatch.delenv("NVIDIA_INFERENCE_API_KEY", raising=False) + assert script.main(["--check"]) == 1 + assert "Drift detected" in capsys.readouterr().err + assert fixture_copy.read_text() == before + + def test_check_mode_green_when_no_drift(self, script, monkeypatch, tmp_path, capsys): + """--check against the current fixture's own openrouter section → 0.""" + import json + + current = json.loads(FIXTURE.read_text()) + fixture_copy = tmp_path / "model_id_corpora.json" + fixture_copy.write_text(json.dumps(current, indent=1) + "\n") + monkeypatch.setattr(script, "FIXTURE", fixture_copy) + payload = {"data": [{"id": e["id"]} for e in current["openrouter"]]} + monkeypatch.setattr(script, "fetch_catalog", lambda *a, **k: payload) + monkeypatch.delenv("NVIDIA_INFERENCE_API_KEY", raising=False) + assert script.main(["--check"]) == 0 + out = capsys.readouterr().out + assert "No drift" in out + # _comment is intentionally excluded from drift: it carries the fetch + # date, which changes daily. + assert fixture_copy.read_text() == json.dumps(current, indent=1) + "\n" + + def test_refresh_mode_rewrites_live_sections_only(self, script, monkeypatch, tmp_path): + import json + + current = json.loads(FIXTURE.read_text()) + fixture_copy = tmp_path / "model_id_corpora.json" + fixture_copy.write_text(FIXTURE.read_text()) + monkeypatch.setattr(script, "FIXTURE", fixture_copy) + payload = {"data": [{"id": "openai/gpt-6-astra"}, {"id": "z-ai/glm-5.3"}]} + monkeypatch.setattr(script, "fetch_catalog", lambda *a, **k: payload) + monkeypatch.delenv("NVIDIA_INFERENCE_API_KEY", raising=False) + assert script.main([]) == 0 + fresh = json.loads(fixture_copy.read_text()) + assert [e["id"] for e in fresh["openrouter"]] == ["openai/gpt-6-astra", "z-ai/glm-5.3"] + # Non-live sections and the un-refreshed nvidia section are preserved. + assert fresh["azure"] == current["azure"] + assert fresh["vertex_ai"] == current["vertex_ai"] + assert fresh["bedrock"] == current["bedrock"] + assert fresh["nvidia_gateway"] == current["nvidia_gateway"]