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%
18.3 KB · 267 lines python
Raw Blame History
1"""FactWriter rules against the local dev database, inside a transaction that is always rolled back (nothing persists)."""2from __future__ import annotations34import uuid5from collections.abc import AsyncIterator6from datetime import UTC, datetime, timedelta78import pytest9from sqlalchemy.ext.asyncio import AsyncConnection1011from aiatlas import db12from aiatlas.db import fetch_all, fetch_one13from aiatlas.sdk.facts import Claim, EntityRef, Event, Facts, PriceObs, ResultObs, Target, facts_from_json, facts_to_json14from aiatlas.sdk.writer import FactWriter1516pytestmark = pytest.mark.usefixtures("conn")171819@pytest.fixture20async def conn() -> AsyncIterator[AsyncConnection]:21    async with db.engine().connect() as c:22        trans = await c.begin()23        try:24            yield c25        finally:26            await trans.rollback()27    await db.dispose()282930def _tag() -> str:31    return uuid.uuid4().hex[:8]323334def _writer(conn: AsyncConnection, *, tier: int = 1, extractor: str = "deterministic", url: str = "https://example.com/doc", run_id: str | None = "run_test",35            observed_at: datetime | None = None, source_key: str | None = None, is_first_run: bool = False) -> FactWriter:36    return FactWriter(conn, source_id=None, snapshot_id=None, source_url=url, tier=tier, connector_name="test", extractor=extractor,37                      extractor_version="t", observed_at=observed_at, run_id=run_id, source_key=source_key or "", is_first_run=is_first_run)383940# ---------------------------------------------------------------------------------------------- pure: Facts JSON round-trip41def test_facts_json_round_trip() -> None:42    f = Facts()43    org = f.entity("company", "Zorg Labs", identifiers={"domain": "zorg.example"})44    m = f.entity("model", "Zeta 9", identifiers={"hf_repo": "zorg/Zeta-9"}, organization=org, attributes={"license": "Apache 2.0", "released": datetime(2026, 1, 2, tzinfo=UTC)},45                 first_seen_hint=datetime(2025, 12, 31, tzinfo=UTC), family=EntityRef("model_family", "Zeta"), artifact_kind=None)46    f.claim(m, "context_length", 128000, unit="tokens", observed_at=datetime(2026, 1, 3, tzinfo=UTC))47    f.relate(org, "develops", m, attributes={"role": "developer"})48    f.event("RELEASE", "model", "Zeta 9 released", entity=m, effective_at=datetime(2026, 1, 2, tzinfo=UTC), importance=3)49    f.price(model=m, provider=f.entity("provider", "Zorg Cloud"), input_per_mtok=1.5, output_per_mtok=6.0, provider_model_id="zeta-9")50    f.result(model=m, benchmark=f.entity("benchmark", "GPQA Diamond"), score=71.2, metric="accuracy", config={"variant": "Diamond"}, evaluated_at=datetime(2026, 1, 5, tzinfo=UTC))51    f.follow("https://zorg.example/more", doc_type="page", key="more", meta={"x": 1})52    f.document_entity = m53    f.document_title = "Zeta 9"54    data = facts_to_json(f)55    back = facts_from_json(data)56    assert [e.name for e in back.entities] == [e.name for e in f.entities]57    assert back.entities[1].first_seen_hint == datetime(2025, 12, 31, tzinfo=UTC) and back.entities[1].family and back.entities[1].family.name == "Zeta"58    assert back.entities[1].attributes["released"] == datetime(2026, 1, 2, tzinfo=UTC) and back.entities[1].organization.identifiers == {"domain": "zorg.example"}59    assert back.claims[0].observed_at == datetime(2026, 1, 3, tzinfo=UTC) and back.claims[0].entity.key() == m.key()60    assert back.relations[0].attributes == {"role": "developer"} and back.events[0].effective_at == datetime(2026, 1, 2, tzinfo=UTC)61    assert back.prices[0].input_per_mtok == 1.5 and back.results[0].evaluated_at == datetime(2026, 1, 5, tzinfo=UTC) and back.results[0].config == {"variant": "Diamond"}62    assert back.targets[0].key == "more" and back.document_entity.key() == m.key() and back.document_title == "Zeta 9"63    assert facts_to_json(back) == data646566# ---------------------------------------------------------------------------------------------- taxonomy at write time67async def test_license_normalised_with_raw_and_mapping(conn: AsyncConnection) -> None:68    t = _tag()69    f = Facts()70    m = f.entity("model", f"Zeta {t}", identifiers={"hf_repo": f"zorg/zeta-{t}"}, attributes={"license": "Apache 2.0", "openness": "restricted", "modalities": ["Text", "pdf"]})71    w = _writer(conn)72    await w.write(f)73    claims = {r["property"]: r for r in await fetch_all(conn, "select * from claims where entity_id = :e and status = 'current'", e=m.id)}74    assert claims["license"]["value"] == "Apache-2.0" and claims["license"]["value_raw"] == "Apache 2.0" and claims["license"]["run_id"] == "run_test"75    assert claims["license_key"]["value"] == "Apache-2.0"76    assert claims["openness"]["value"] == "restricted-weights" and claims["openness"]["value_raw"] == "restricted"77    assert claims["modalities"]["value"] == ["document", "text"]78    ent = await fetch_one(conn, "select attributes from entities where id = :e", e=m.id)79    assert ent["attributes"]["license"] == "Apache-2.0" and ent["attributes"]["license_raw"] == "Apache 2.0" and ent["attributes"]["license_key"] == "Apache-2.0"80    assert ent["attributes"]["modalities"] == ["document", "text"] and ent["attributes"]["modalities_raw"] == "Text, pdf"81    mapping = await fetch_one(conn, "select canonical from taxonomy_mappings where domain = 'license' and raw = 'Apache 2.0'")82    assert mapping and mapping["canonical"] == "Apache-2.0"83    # a second source spelling the same licence differently confirms, never a LICENSE_CHANGED event84    f2 = Facts()85    m2 = f2.entity("model", f"Zeta {t}", identifiers={"hf_repo": f"zorg/zeta-{t}"}, attributes={"license": "apache-2.0"})86    await _writer(conn, tier=2, url="https://other.example/x").write(f2)87    assert m2.id == m.id88    n = await fetch_one(conn, "select count(*) as n from claims where entity_id = :e and property = 'license' and status = 'current'", e=m.id)89    assert n["n"] == 190    ev = await fetch_all(conn, "select event_type from change_events where entity_id = :e", e=m.id)91    assert {e["event_type"] for e in ev} == {"NEW_MODEL"}929394async def test_unknown_taxonomy_value_kept(conn: AsyncConnection) -> None:95    t = _tag()96    f = Facts()97    m = f.entity("model", f"Zeta {t}", attributes={"license": f"custom-{t}"})98    await _writer(conn).write(f)99    c = await fetch_one(conn, "select value, value_raw from claims where entity_id = :e and property = 'license'", e=m.id)100    assert c["value"] == f"custom-{t}" and c["value_raw"] is None101    mp = await fetch_one(conn, "select canonical from taxonomy_mappings where domain = 'license' and raw = :r", r=f"custom-{t}")102    assert mp and mp["canonical"] is None103104105# ---------------------------------------------------------------------------------------------- same-source loophole & supersede106async def test_llm_never_supersedes_deterministic_same_source(conn: AsyncConnection) -> None:107    t = _tag()108    url = f"https://docs.example/{t}"109    f = Facts()110    m = f.entity("model", f"Zeta {t}", identifiers={"hf_repo": f"zorg/zeta-{t}"}, attributes={"context_length": 100_000})111    await _writer(conn, tier=1, url=url).write(f)112    # LLM extraction from the same URL, one tier lower113    f2 = Facts()114    m2 = f2.entity("model", f"Zeta {t}", identifiers={"hf_repo": f"zorg/zeta-{t}"}, attributes={"context_length": 200_000})115    llm = _writer(conn, tier=2, extractor="llm", url=url)116    await llm.write(f2)117    rows = await fetch_all(conn, "select value, status, extractor from claims where entity_id = :e and property = 'context_length' order by value_num", e=m.id)118    assert [(r["value"], r["status"]) for r in rows] == [(100_000, "current"), (200_000, "conflicting")]119    assert llm.stats.conflicts == 1120    assert await fetch_one(conn, "select 1 from review_queue where kind = 'conflict' and :e = any(entity_ids)", e=m.id)121    # the deterministic extractor correcting itself from the same URL supersedes → CONTEXT_CHANGED with deterministic importance122    f3 = Facts()123    f3.entity("model", f"Zeta {t}", identifiers={"hf_repo": f"zorg/zeta-{t}"}, attributes={"context_length": 1_000_000})124    await _writer(conn, tier=1, url=url).write(f3)125    cur = await fetch_one(conn, "select value from claims where entity_id = :e and property = 'context_length' and status = 'current'", e=m.id)126    assert cur["value"] == 1_000_000 and m2.id == m.id127    ev = await fetch_one(conn, "select importance, is_backfill, run_id, recorded_at from change_events where entity_id = :e and event_type = 'CONTEXT_CHANGED'", e=m.id)128    assert ev and ev["importance"] == 3 and ev["is_backfill"] is False and ev["run_id"] == "run_test" and ev["recorded_at"] is not None129130131async def test_derived_writer_conflicts_without_review_items(conn: AsyncConnection) -> None:132    t = _tag()133    f = Facts()134    m = f.entity("model", f"Zeta {t}", attributes={"openness": "open-weights"})135    await _writer(conn, tier=1).write(f)136    f2 = Facts()137    f2.claim(EntityRef("model", f"Zeta {t}", id=m.id), "openness", "restricted-weights")138    w = _writer(conn, tier=2, extractor="derived", url=None)139    await w.write(f2)140    assert w.stats.conflicts == 1141    assert (await fetch_one(conn, "select value from claims where entity_id = :e and property = 'openness' and status = 'current'", e=m.id))["value"] == "open-weights"142    assert not await fetch_one(conn, "select 1 from review_queue where kind = 'conflict' and :e = any(entity_ids)", e=m.id)143144145# ---------------------------------------------------------------------------------------------- results: comparability, one current row per key, bounds146async def test_results_one_current_per_config_key(conn: AsyncConnection) -> None:147    t = _tag()148    f = Facts()149    m = f.entity("model", f"Zeta {t}", identifiers={"hf_repo": f"zorg/zeta-{t}"})150    b = f.entity("benchmark", f"LiveBench {t}", identifiers={"registry_benchmark": f"livebench-{t}"})151    f.result(model=m, benchmark=b, score=61.0, metric="global_average", unit="%", config={"release": "2026-05-01", "livebench_model_id": "zeta"})152    await _writer(conn, tier=2, source_key="livebench.ai").write(f)153    first = await fetch_one(conn, "select * from benchmark_results where model_id = :m", m=m.id)154    assert first["config_key"] and first["trust_level"] == "official-benchmark" and first["run_group"] == "2026-05-01" and first["is_current"] is True155    assert first["run_id"] == "run_test" and first["extractor"] == "deterministic"156    f2 = Facts()157    m2 = f2.entity("model", f"Zeta {t}", identifiers={"hf_repo": f"zorg/zeta-{t}"})158    b2 = f2.entity("benchmark", f"LiveBench {t}", identifiers={"registry_benchmark": f"livebench-{t}"})159    f2.result(model=m2, benchmark=b2, score=64.0, metric="global_average", unit="%", config={"release": "2026-06-25", "livebench_model_id": "zeta"})160    await _writer(conn, tier=2, source_key="livebench.ai", observed_at=datetime.now(UTC) + timedelta(seconds=1)).write(f2)161    rows = await fetch_all(conn, "select score, is_current, valid_to, config_key, run_group from benchmark_results where model_id = :m order by observed_at", m=m.id)162    assert [(r["score"], r["is_current"], r["valid_to"] is None) for r in rows] == [(61.0, False, False), (64.0, True, True)]163    assert rows[0]["config_key"] == rows[1]["config_key"]164    # same run group, different condition (reasoning effort) → both stay current, same config_key165    f3 = Facts()166    m3 = f3.entity("model", f"Zeta {t}", identifiers={"hf_repo": f"zorg/zeta-{t}"})167    b3 = f3.entity("benchmark", f"LiveBench {t}", identifiers={"registry_benchmark": f"livebench-{t}"})168    f3.result(model=m3, benchmark=b3, score=66.0, metric="global_average", unit="%", config={"release": "2026-06-25", "livebench_model_id": "zeta", "reasoning_effort": "high"})169    await _writer(conn, tier=2, source_key="livebench.ai", observed_at=datetime.now(UTC) + timedelta(seconds=2)).write(f3)170    cur = await fetch_all(conn, "select score from benchmark_results where model_id = :m and is_current order by score", m=m.id)171    assert [r["score"] for r in cur] == [64.0, 66.0]172173174async def test_result_out_of_bounds_flagged(conn: AsyncConnection) -> None:175    t = _tag()176    f = Facts()177    m = f.entity("model", f"Zeta {t}")178    b = f.entity("benchmark", f"Bench {t}")179    f.result(model=m, benchmark=b, score=140.0, metric="accuracy", unit="%")180    await _writer(conn, tier=2).write(f)181    r = await fetch_one(conn, "select confidence from benchmark_results where model_id = :m", m=m.id)182    assert r["confidence"] == "low"183    a = await fetch_one(conn, "select check_name, severity, status from anomalies where entity_id = :m", m=m.id)184    assert a and a["check_name"] == "score_above_max" and a["severity"] == "critical" and a["status"] == "open"185186187async def test_effort_variant_result_config_augmented(conn: AsyncConnection) -> None:188    """An evaluator row named '<model>-high' lands on the canonical model with reasoning_effort in the config."""189    t = _tag()190    base = Facts()191    m = base.entity("model", f"Zeta {t}", identifiers={"artificial_analysis": f"zeta-{t}"})192    await _writer(conn, tier=2, source_key="artificialanalysis.ai").write(base)193    f = Facts()194    v = f.entity("model", f"zeta-{t}-high", identifiers={"artificial_analysis": f"zeta-{t}-high"})195    b = f.entity("benchmark", f"GPQA {t}", identifiers={"registry_benchmark": f"gpqa-{t}"})196    f.result(model=v, benchmark=b, score=80.0, metric="accuracy", unit="%", config={"evaluator": "Artificial Analysis", "index_version": "4.3", "aa_slug": f"zeta-{t}-high"})197    w = _writer(conn, tier=2, source_key="artificialanalysis.ai")198    await w.write(f)199    assert v.id == m.id and w.stats.folded_variants == 1200    r = await fetch_one(conn, "select config, trust_level from benchmark_results where model_id = :m", m=m.id)201    assert r["config"]["reasoning_effort"] == "high" and r["config"]["aa_variant_slug"] == f"zeta-{t}-high" and r["trust_level"] == "independent-evaluator"202    ids = await fetch_all(conn, "select value from entity_identifiers where entity_id = :m and scheme = 'artificial_analysis' order by value", m=m.id)203    assert [i["value"] for i in ids] == [f"zeta-{t}", f"zeta-{t}-high"]204    assert not await fetch_one(conn, "select 1 from entities where canonical_name = :n", n=f"zeta-{t}-high")205206207# ---------------------------------------------------------------------------------------------- events: backfill, group key, NEW_* importance208async def test_new_entity_events_backfill_and_grouping(conn: AsyncConnection) -> None:209    t = _tag()210    f = Facts()211    m = f.entity("model", f"Zeta {t}", attributes={"release_date": "2024-01-15"})212    await _writer(conn, tier=1).write(f)213    ev = await fetch_one(conn, "select is_backfill, group_key, effective_at, importance from change_events where entity_id = :e and event_type = 'NEW_MODEL'", e=m.id)214    assert ev["is_backfill"] is True and ev["group_key"] == f"release:{m.id}:2024-01" and ev["importance"] == 2215    # first run of a connector → everything is backfill even without dates216    f2 = Facts()217    m2 = f2.entity("model", f"Zeta {t} b")218    await _writer(conn, tier=1, is_first_run=True).write(f2)219    ev2 = await fetch_one(conn, "select is_backfill from change_events where entity_id = :e", e=m2.id)220    assert ev2["is_backfill"] is True221    # artifacts and families never make importance-3 news222    f3 = Facts()223    art = f3.entity("artifact", f"unsloth/Zeta-{t}-GGUF", identifiers={"hf_repo": f"unsloth/Zeta-{t}-GGUF"}, canonical=EntityRef("model", f"Zeta {t}", id=m.id), artifact_kind="quantization")224    fam = f3.entity("model_family", f"Zeta family {t}")225    await _writer(conn, tier=2).write(f3)226    imps = {r["event_type"]: r["importance"] for r in await fetch_all(conn, "select event_type, importance from change_events where entity_id in (:a, :f)", a=art.id, f=fam.id)}227    assert imps == {"NEW_ARTIFACT": 0, "NEW_MODEL_FAMILY": 1}228    row = await fetch_one(conn, "select canonical_id, artifact_kind from entities where id = :a", a=art.id)229    assert row["canonical_id"] == m.id and row["artifact_kind"] == "quantization"230    assert await fetch_one(conn, "select 1 from relations where subject_id = :a and predicate = 'artifact_of' and object_id = :m and valid_to is null", a=art.id, m=m.id)231232233async def test_family_hint_and_first_seen_hint(conn: AsyncConnection) -> None:234    t = _tag()235    f = Facts()236    m = f.entity("model", f"Llama {t} 9B", family=EntityRef("model_family", f"Llama {t}"), first_seen_hint=datetime(2024, 3, 1, tzinfo=UTC), identity_confidence="medium")237    await _writer(conn, tier=1).write(f)238    row = await fetch_one(conn, "select family_id, first_seen_at, identity_confidence from entities where id = :m", m=m.id)239    assert row["family_id"] == m.family.id and row["first_seen_at"] == datetime(2024, 3, 1, tzinfo=UTC) and row["identity_confidence"] == "medium"240    assert await fetch_one(conn, "select 1 from relations where subject_id = :m and predicate = 'member_of_family' and object_id = :f", m=m.id, f=m.family.id)241    fam = await fetch_one(conn, "select entity_type, slug from entities where id = :f", f=m.family.id)242    assert fam["entity_type"] == "model_family"243244245async def test_price_change_importance_and_run_id(conn: AsyncConnection) -> None:246    t = _tag()247    f = Facts()248    m = f.entity("model", f"Zeta {t}", identifiers={"hf_repo": f"zorg/zeta-{t}"})249    p = f.entity("provider", f"Zorg Cloud {t}")250    f.price(model=m, provider=p, provider_model_id="zeta", input_per_mtok=10.0, output_per_mtok=30.0)251    await _writer(conn, tier=1).write(f)252    f2 = Facts()253    m2 = f2.entity("model", f"Zeta {t}", identifiers={"hf_repo": f"zorg/zeta-{t}"})254    p2 = f2.entity("provider", f"Zorg Cloud {t}")255    f2.price(model=m2, provider=p2, provider_model_id="zeta", input_per_mtok=4.0, output_per_mtok=30.0)256    await _writer(conn, tier=1, run_id="run_2", observed_at=datetime.now(UTC) + timedelta(seconds=1)).write(f2)257    ev = await fetch_one(conn, "select importance, run_id from change_events where entity_id = :m and event_type = 'PRICE_CHANGED'", m=m.id)258    assert ev["importance"] == 3 and ev["run_id"] == "run_2"     # -60 %259    prices = await fetch_all(conn, "select input_per_mtok, run_id, valid_to from prices where model_id = :m order by valid_from", m=m.id)260    assert [(r["input_per_mtok"], r["run_id"], r["valid_to"] is None) for r in prices] == [(10.0, "run_test", False), (4.0, "run_2", True)]261262263async def test_target_and_unused_imports_keep_dataclasses_stable() -> None:264    assert Target(url="https://x").doc_type == "page" and Claim(EntityRef("model", "x"), "p", 1).unit is None265    assert Event("RELEASE", "model", "s").importance == 2 and PriceObs(EntityRef("model", "m"), EntityRef("provider", "p")).currency == "USD"266    assert ResultObs(EntityRef("model", "m"), EntityRef("benchmark", "b"), 1.0).trust_level is None267