packages/eval-entity-resolver/src/eval_entity_resolver/resolver.py

"""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,
            ),
        )