"""The bare resolver. Matches a raw value to a canonical id via the strategy chain (exact → normalized → fuzzy → no_match), and — when given a `CanonicalStore` — enriches the result with the matched canonical's metadata, parent edges, model-specific lineage fields, and quantized-chain root collapse. The enrichment matches the HTTP API's response shape exactly. Callers using the resolver standalone get the same `ResolutionResult` they'd get back from `POST /api/v1/resolve`.""" from __future__ import annotations from pathlib import Path from typing import Callable, Optional from eval_entity_resolver.alias_store import AliasStore from eval_entity_resolver.canonical_store import CanonicalStore, _hf_repo_id_of from eval_entity_resolver.models import ( HfIdHit, ResolutionResult, ResolverConfig, looks_like_hf_id, ) from eval_entity_resolver.strategies.exact import exact_match from eval_entity_resolver.strategies.normalized import normalized_match from eval_entity_resolver.strategies.fuzzy import fuzzy_match # Confidence assigned to normalized-match results. Below 1.0 (exact) and # above _STEM_CONFIDENCE (0.90, fuzzy) so the provenance is clear in the # resolution log. _NORMALIZED_CONFIDENCE = 0.95 class Resolver: def __init__( self, store: AliasStore, config: Optional[ResolverConfig] = None, canonical_store: Optional[CanonicalStore] = None, hf_id_checker: Optional[Callable[[str], Optional[HfIdHit]]] = None, ) -> None: """`store` is required (alias matching is the resolver's core job). `canonical_store` is optional — when provided, results are enriched with parent / lineage / metadata fields. Without it, only the basic match fields (canonical_id, strategy, confidence) are populated. `hf_id_checker` is optional — when provided, HF-shaped model raw values are checked against it as a resolution step: a verbatim (case-insensitive) confirmation outranks any alias match, and a normalized (separator-collapse) confirmation ranks between the normalized-alias and fuzzy steps. Without it, the chain is identical to the classic exact → normalized → fuzzy order.""" self.store = store self.config = config or ResolverConfig() self.canonical_store = canonical_store self.hf_id_checker = hf_id_checker @classmethod def from_parquet( cls, path: str | Path, config: Optional[ResolverConfig] = None, ) -> "Resolver": """Load both alias and canonical stores from a parquet directory (e.g. `./fixtures/`) and return a fully-enriching resolver. This is the recommended convenience for callers who want the same response shape as the HTTP API.""" return cls( AliasStore.from_parquet(path), config=config, canonical_store=CanonicalStore.from_parquet(path), ) @classmethod def from_hf( cls, repo_id: str, config: Optional[ResolverConfig] = None, ) -> "Resolver": """Load both stores from a HF Dataset repo and return a fully-enriching resolver.""" return cls( AliasStore.from_hf(repo_id), config=config, canonical_store=CanonicalStore.from_hf(repo_id), ) def resolve( self, raw_value: str, entity_type: str, source_config: Optional[str] = None, mode: str = "resolve", check_hf: bool = True, ) -> ResolutionResult: """`check_hf=False` disables the injected HF id check for this call. Internal inference resolves (e.g. tier-3 stem probing) use it so a batch of stem candidates never draws live-lookup budget.""" # 1. Exact canonical_id = exact_match(raw_value, entity_type, source_config, self.store) # 2. HF id check (models only, HF-shaped only, checker injected). # Called at most once per resolve; the hit is reused by the # normalized-tier step below. hf_hit = ( self._maybe_check_hf_id(raw_value, entity_type, canonical_id) if check_hf else None ) if hf_hit is not None and hf_hit.verbatim: if canonical_id is not None and canonical_id == hf_hit.hf_id: # Agreement — the registry result is richer; stamp the # runtime attestation onto it. result = self._enrich( raw_value, entity_type, source_config, canonical_id, "exact", 1.0 ) return self._stamp_hf_attestation(result, hf_hit) # Disagreement or registry miss — the verbatim HF id wins over # any alias mapping (owner decision: a string that IS a real HF # repo id resolves to itself). return self._hf_checker_result( raw_value, entity_type, source_config, hf_hit, "exact", 1.0 ) if canonical_id is not None: result = self._enrich( raw_value, entity_type, source_config, canonical_id, "exact", 1.0 ) if hf_hit is not None and canonical_id == hf_hit.hf_id: result = self._stamp_hf_attestation(result, hf_hit) return result # 3. Normalized alias (confidence 0.95 — only return if above # threshold). A curated normalized alias outranks a merely # separator-collapsed HF index hit. if _NORMALIZED_CONFIDENCE >= self.config.threshold: canonical_id = normalized_match(raw_value, entity_type, self.store, source_config) if canonical_id is not None: result = self._enrich( raw_value, entity_type, source_config, canonical_id, "normalized", _NORMALIZED_CONFIDENCE, ) if hf_hit is not None and canonical_id == hf_hit.hf_id: result = self._stamp_hf_attestation(result, hf_hit) return result # 4. Normalized-tier HF id check hit (separator-collapse match). if hf_hit is not None and not hf_hit.verbatim: return self._hf_checker_result( raw_value, entity_type, source_config, hf_hit, "normalized", _NORMALIZED_CONFIDENCE, ) # Exact-only mode stops here: no fuzzy inference. if mode == "exact": return ResolutionResult( raw_value=raw_value, entity_type=entity_type, source_config=source_config, canonical_id=None, strategy="no_match", confidence=0.0, ) # 5. Fuzzy — thread the store-backed curated org map (incl. orgs.yaml # alias tier) into the org-agreement guard so org-equivalent namespaces # (AlephAlpha/aleph-alpha, MiniMaxAI/minimax, Alibaba-NLP/alibaba) fold # and match. canonical_id, confidence, inferred_platform = fuzzy_match( raw_value, entity_type, self.config.threshold, self.store, source_config, org_dev_map=self._org_fold_map(), known_orgs=self._known_orgs(), ) if canonical_id is not None: result = self._enrich( raw_value, entity_type, source_config, canonical_id, "fuzzy", confidence, ) # Thread the captured inference_platform onto the result. This is # the per-run platform read off an EXPLICIT host token in the raw # id (a `together/`-prefix or `-bedrock`-suffix), which WINS — an # explicit host token in the id is the strongest per-run platform # fact. Only set it when a token was actually present (None # otherwise), so non-host ids leave the field untouched. if inferred_platform is not None: result.inference_platform = inferred_platform return result # 6. No match return ResolutionResult( raw_value=raw_value, entity_type=entity_type, source_config=source_config, canonical_id=None, strategy="no_match", confidence=0.0, ) # ------------------------------------------------------------------ # HF id check (injected) # ------------------------------------------------------------------ def _maybe_check_hf_id( self, raw_value: str, entity_type: str, exact_canonical_id: Optional[str], ) -> Optional[HfIdHit]: """Run the injected HF id checker when it can change the outcome. Skipped when: no checker, non-model, not HF-shaped, or the exact alias hit is byte-equal to the raw value AND that canonical row either is already HF-attested (oracle `resolution_source == "hf"` or hub-stats-confirmed — the checker could only agree) or carries a curated/models.dev provenance claim (`models_dev`/`NA`/`curated` — a registered off-HF canonical like an OpenRouter or closed-API id, where a live check would just 404 on every resolve). A case-only or name-inferred (tier-3) agreement still runs the checker, because the checker is what recovers HF-true casing for lowercased drafts.""" if self.hf_id_checker is None or entity_type != "model": return None if not looks_like_hf_id(raw_value): return None if exact_canonical_id is not None and exact_canonical_id == raw_value: if self._hf_check_skippable(exact_canonical_id): return None return self.hf_id_checker(raw_value) def _hf_check_skippable(self, canonical_id: str) -> bool: if self.canonical_store is None: return False ent = self.canonical_store.lookup("model", canonical_id) if ent is None: return False if _hf_repo_id_of(ent, canonical_id) is not None: return True src = ent.get("resolution_source") return isinstance(src, str) and src in ("models_dev", "NA", "curated") def _hf_checker_result( self, raw_value: str, entity_type: str, source_config: Optional[str], hit: HfIdHit, strategy: str, confidence: float, ) -> ResolutionResult: """Build the result for a winning HF id check hit. A registered canonical gets the full registry enrichment; an unregistered one gets a bare result flagged `hf_attested_unregistered` so the service layer can auto-create it (write mode) or serve it without minting (read-only).""" known = ( self.canonical_store is not None and self.canonical_store.lookup("model", hit.hf_id) is not None ) if known: result = self._enrich( raw_value, entity_type, source_config, hit.hf_id, strategy, confidence ) return self._stamp_hf_attestation(result, hit) return ResolutionResult( raw_value=raw_value, entity_type=entity_type, source_config=source_config, canonical_id=hit.hf_id, strategy=strategy, confidence=confidence, resolution_source=hit.source, ancestry=[], resolution_detail={"granularity": None, "hf_repo_id": hit.hf_id}, hf_attestation=hit.hf_id, hf_attested_unregistered=True, ) @staticmethod def _stamp_hf_attestation(result: ResolutionResult, hit: HfIdHit) -> ResolutionResult: """The checker just runtime-attested the repo id, so surface it in `resolution_detail.hf_repo_id` even when the canonical row itself carries no HF provenance. Only stamped when the response canonical IS the attested id (root-collapse may have moved it).""" if result.canonical_id == hit.hf_id: result.hf_attestation = hit.hf_id if isinstance(result.resolution_detail, dict): result.resolution_detail["hf_repo_id"] = hit.hf_id return result def _org_fold_map(self) -> dict: """The org-fold map threaded into the fuzzy org-agreement guard, built ONCE from BOTH stores and cached: - canonical_store.org_dev_map: `_ORG_ALIASES` ∪ every canonical_orgs `id`/`hf_org` (lowercased -> curated id); - the org ALIAS rows in the alias table (`iter_alias_pairs("org")`), which carry the curated orgs.yaml alias tier (`AI2`->`allenai`, `ai21labs`->`ai21`) that is NOT a canonical_orgs column. Alias rows win on key collision (they ARE the curated seed). Keyed lowercase; `_fold_org` separator-strips on lookup so case/separator variants fold too.""" cached = getattr(self, "_org_fold_map_cache", None) if cached is not None: return cached from eval_entity_resolver.fold import _ORG_ALIASES m: dict = dict(self.canonical_store.org_dev_map) if self.canonical_store else dict(_ORG_ALIASES) if self.store is not None: for raw, cid in self.store.iter_alias_pairs("org"): m[raw.lower()] = cid self._org_fold_map_cache = m return m def _known_orgs(self) -> frozenset: """The set of canonical_orgs ids (a developer that has a real row). The fuzzy org-agreement guard never separator-strip-merges two prefixes that resolve to DIFFERENT ids in this set — so distinct uploaders kept apart in canonical_orgs (the orgs_distinct_allowlist contract) stay apart at resolve time too. Cached; empty when no canonical_store is attached.""" cached = getattr(self, "_known_orgs_cache", None) if cached is not None: return cached ids: set[str] = set() if self.canonical_store is not None: df = self.canonical_store._tables.get("org") if df is not None and not df.empty and "id" in df.columns: ids = {str(i) for i in df["id"] if isinstance(i, str)} self._known_orgs_cache = frozenset(ids) return self._known_orgs_cache # ------------------------------------------------------------------ # Enrichment (no-op when no canonical_store is attached) # ------------------------------------------------------------------ def build_result( self, raw_value: str, entity_type: str, source_config: Optional[str], canonical_id: str, strategy: str, confidence: float, ) -> ResolutionResult: """Construct an enriched `ResolutionResult` for a canonical_id the caller already knows — useful for callers that bypass the strategy chain (e.g. an alias-table cache hit, an auto-created draft) but want the same rich response shape. Identical to the enrichment that happens inside `resolve()`.""" return self._enrich(raw_value, entity_type, source_config, canonical_id, strategy, confidence) def _enrich( self, raw_value: str, entity_type: str, source_config: Optional[str], matched_canonical_id: str, strategy: str, confidence: float, ) -> ResolutionResult: """Look up the matched canonical's row and populate the rich response fields. When no canonical_store is attached, the rich fields stay None and the result has just the basic match info.""" if self.canonical_store is None: return ResolutionResult( raw_value=raw_value, entity_type=entity_type, source_config=source_config, canonical_id=matched_canonical_id, strategy=strategy, confidence=confidence, ) cs = self.canonical_store matched_entity = cs.lookup(entity_type, matched_canonical_id) review_status = (matched_entity or {}).get("review_status") if matched_entity else None if entity_type == "model": fields = cs.model_metadata_fields(matched_canonical_id, matched_entity) # If the response collapses to a different canonical (root), # surface THAT canonical's review_status — keeps the response # internally consistent. if fields["canonical_id"] != matched_canonical_id: root_entity = cs.lookup("model", fields["canonical_id"]) if root_entity: review_status = root_entity.get("review_status") or review_status return ResolutionResult( raw_value=raw_value, entity_type=entity_type, source_config=source_config, canonical_id=fields["canonical_id"], strategy=strategy, confidence=confidence, review_status=review_status, parent_canonical_id=cs.parent_canonical_id("model", matched_entity), resolved_leaf_id=fields["resolved_leaf_id"], root_model_id=fields["root_model_id"], lineage_origin_org_id=fields["lineage_origin_org_id"], # Extended lineage / provenance fields. None-safe .get so a # store predating these keys still works. model_group_id=fields.get("model_group_id"), model_family_id=fields.get("model_family_id"), lineage_origin_model_id=fields.get("lineage_origin_model_id"), lineage_origin_model_org_id=fields.get("lineage_origin_model_org_id"), inference_platform=fields.get("inference_platform"), resolution_source=fields.get("resolution_source"), resolution_granularity=fields.get("resolution_granularity"), parents=fields["parents"], open_weights=fields["open_weights"], release_date=fields["release_date"], params_billions=fields["params_billions"], ancestry=cs.compute_ancestry("model", fields["canonical_id"], matched_entity), resolution_detail=cs.resolution_detail( "model", fields["canonical_id"], matched_entity=matched_entity ), ) # Benchmark: fill in hierarchy-alignment fields (family_key, # category) by walking canonical_families. composite_keys stays # empty here — see CanonicalStore.benchmark_family_enrichment for # why composite computation belongs in the producer. if entity_type == "benchmark": fam = cs.benchmark_family_enrichment(matched_canonical_id) return ResolutionResult( raw_value=raw_value, entity_type=entity_type, source_config=source_config, canonical_id=matched_canonical_id, strategy=strategy, confidence=confidence, review_status=review_status, parent_canonical_id=cs.parent_canonical_id(entity_type, matched_entity), family_key=fam["family_key"], category=fam["category"], composite_keys=fam["composite_keys"], ancestry=cs.compute_ancestry("benchmark", matched_canonical_id, matched_entity), resolution_detail=cs.resolution_detail( "benchmark", matched_canonical_id, raw_value=raw_value, matched_entity=matched_entity, ), ) # Other non-model types (metric, harness, org, family, composite): # parent_canonical_id + review_status, plus ancestry/detail (family # carries a composite parent; the rest are roots with empty detail). return ResolutionResult( raw_value=raw_value, entity_type=entity_type, source_config=source_config, canonical_id=matched_canonical_id, strategy=strategy, confidence=confidence, review_status=review_status, parent_canonical_id=cs.parent_canonical_id(entity_type, matched_entity), ancestry=cs.compute_ancestry(entity_type, matched_canonical_id, matched_entity), resolution_detail=cs.resolution_detail( entity_type, matched_canonical_id, raw_value=raw_value, matched_entity=matched_entity, ), )