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%
6.0 KB · 124 lines python
Raw Blame History
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