HTML 77.2%
TypeScript 10.5%
Python 9.6%
JavaScript 2.5%
1"""Resolver guards against the local dev database (rolled back)."""2from __future__ import annotations34import uuid5from collections.abc import AsyncIterator67import pytest8from sqlalchemy.ext.asyncio import AsyncConnection910from aiatlas import db11from aiatlas.db import execute, fetch_one12from aiatlas.sdk.facts import EntityRef, Facts13from aiatlas.sdk.resolution import Resolver, compatible_types14from aiatlas.sdk.writer import FactWriter15from aiatlas.services.merge import merge_entities, record_decision161718@pytest.fixture19async def conn() -> AsyncIterator[AsyncConnection]:20 async with db.engine().connect() as c:21 trans = await c.begin()22 try:23 yield c24 finally:25 await trans.rollback()26 await db.dispose()272829def _tag() -> str:30 return uuid.uuid4().hex[:8]313233async def _write(conn: AsyncConnection, facts: Facts, *, tier: int = 1) -> FactWriter:34 w = FactWriter(conn, source_id=None, snapshot_id=None, source_url="https://example.com/x", tier=tier, connector_name="test", run_id="run_test", source_key="")35 await w.write(facts)36 return w373839def test_compatible_types() -> None:40 assert set(compatible_types("model")) == {"model", "artifact"} and compatible_types("paper") == ("paper",)414243async def test_alias_collision_requires_variant_key(conn: AsyncConnection) -> None:44 t = _tag()45 f = Facts()46 a = f.entity("model", f"Qwen{t}-8B")47 await _write(conn, f)48 r = Resolver(conn, source_tier=2)49 # "Qwen{t} 38B" normalises to the same alias key but is another size → must not resolve onto the 8B model50 other = await r.resolve(EntityRef("model", f"Qwen{t} 38B"), create=False)51 assert other is None52 same = await r.resolve(EntityRef("model", f"Qwen{t} 8B"), create=False)53 assert same == a.id545556async def test_keep_separate_decision_blocks_alias_and_merge(conn: AsyncConnection) -> None:57 t = _tag()58 f = Facts()59 org1 = f.entity("company", f"Org One {t}")60 org2 = f.entity("company", f"Org Two {t}")61 a = f.entity("model", f"Nimbus {t} Pro", organization=org1, identifiers={"hf_repo": f"one/nimbus-{t}"}, aliases=[f"Nimbus {t}"])62 b = f.entity("model", f"Nimbus {t} b", organization=org2, aliases=[f"Nimbus {t}"])63 await _write(conn, f)64 await record_decision(conn, a.id, b.id, "keep_separate", actor="test", note="different vendors")65 r = Resolver(conn, source_tier=2)66 # ambiguous alias with both candidates kept separate → no guess; with the organisation of B → B; the pair is never merged automatically67 assert await r.resolve(EntityRef("model", f"Nimbus {t}"), create=False) is None68 assert await r.resolve(EntityRef("model", f"Nimbus {t}", organization=EntityRef("company", f"Org One {t}", id=org1.id)), create=False) == a.id69 assert await r.resolve(EntityRef("model", f"Nimbus {t}", organization=EntityRef("company", f"Org Two {t}", id=org2.id)), create=False) == b.id70 with pytest.raises(ValueError):71 await merge_entities(conn, a.id, b.id)727374async def test_model_ref_resolves_to_artifact_row(conn: AsyncConnection) -> None:75 t = _tag()76 f = Facts()77 art = f.entity("artifact", f"unsloth/Zeta-{t}-GGUF", identifiers={"hf_repo": f"unsloth/Zeta-{t}-GGUF"})78 await _write(conn, f)79 r = Resolver(conn, source_tier=2)80 assert await r.resolve(EntityRef("model", f"Zeta-{t}-GGUF", identifiers={"hf_repo": f"unsloth/Zeta-{t}-GGUF"}), create=False) == art.id818283async def test_resolve_variant_only_for_evaluator_refs(conn: AsyncConnection) -> None:84 t = _tag()85 f = Facts()86 org = f.entity("company", f"Zorg {t}")87 base = f.entity("model", f"Zeta {t}", organization=org, identifiers={"artificial_analysis": f"zeta-{t}"})88 await _write(conn, f, tier=2)89 r = Resolver(conn, source_tier=2)90 folded = await r.resolve_variant(EntityRef("model", f"zeta-{t}-xhigh"))91 assert folded == (base.id, {"reasoning_effort": "xhigh"})92 assert await r.resolve_variant(EntityRef("model", f"zeta-{t}")) is None # not a variant93 assert await r.resolve_variant(EntityRef("model", f"omega-{t}-high")) is None # base unknown94 # a ref with its own hub identifier is a real model even if its name ends in "-thinking": created, never folded95 real = await r.resolve(EntityRef("model", f"zeta-{t}-thinking", identifiers={"hf_repo": f"zorg/Zeta-{t}-Thinking"}))96 assert real != base.id97 # a tier-1 resolver never folds either98 r1 = Resolver(conn, source_tier=1)99 official = await r1.resolve(EntityRef("model", f"zeta-{t}-high"), create=False)100 assert official is None101 # an evaluator-only ref folds and records the fold102 r2 = Resolver(conn, source_tier=2)103 eid = await r2.resolve(EntityRef("model", f"zeta-{t}-high", identifiers={"artificial_analysis": f"zeta-{t}-high"}))104 assert eid == base.id and r2.folded[f"model:artificial_analysis=zeta-{t}-high"] == (base.id, {"reasoning_effort": "high"})105106107async def test_first_seen_hint_and_touch(conn: AsyncConnection) -> None:108 from datetime import UTC, datetime109110 t = _tag()111 r = Resolver(conn, source_tier=1)112 eid = await r.resolve(EntityRef("model", f"Zeta {t}", first_seen_hint=datetime(2023, 5, 1, tzinfo=UTC)))113 row = await fetch_one(conn, "select first_seen_at from entities where id = :id", id=eid)114 assert row["first_seen_at"] == datetime(2023, 5, 1, tzinfo=UTC)115 r2 = Resolver(conn, source_tier=1)116 await r2.resolve(EntityRef("model", f"Zeta {t}", first_seen_hint=datetime(2022, 1, 1, tzinfo=UTC)))117 row = await fetch_one(conn, "select first_seen_at from entities where id = :id", id=eid)118 assert row["first_seen_at"] == datetime(2022, 1, 1, tzinfo=UTC)119 # a future hint never moves first_seen forward120 await execute(conn, "update entities set first_seen_at = :d where id = :id", d=datetime(2022, 1, 1, tzinfo=UTC), id=eid)121 await Resolver(conn).resolve(EntityRef("model", f"Zeta {t}", first_seen_hint=datetime(2030, 1, 1, tzinfo=UTC)))122 row = await fetch_one(conn, "select first_seen_at from entities where id = :id", id=eid)123 assert row["first_seen_at"] == datetime(2022, 1, 1, tzinfo=UTC)124