SPB Git

spb/ai-atlas

Public
25.4 KB · 369 lines python
1"""FactWriter — turns `Facts` into versioned database state.23Temporal rules (per entity × property):4  * no current claim          → insert current claim, set attribute, NEW_* event for new entities5  * same value                → confirm (observed_at bumped), nothing else6  * different, source ≥ tier  → supersede (valid_to = observed), insert new current, update attribute, CHANGE event for material props7  * different, worse source   → store as `conflicting`, flag confidence, review-queue item — never overwrite8Prices and benchmark results have their own append-only tables with the same close/open semantics.9"""10from __future__ import annotations1112import hashlib13import json14import logging15from datetime import UTC, datetime16from typing import Any1718from sqlalchemy.ext.asyncio import AsyncConnection1920from aiatlas.db import execute, fetch_one, jsonb21from aiatlas.ids import new_id, normalize_alias22from aiatlas.sdk.facts import (23    EVENT_CATEGORY_BY_TYPE,24    MATERIAL_PROPERTIES,25    NOISY_PREFIXES,26    EntityRef,27    Facts,28    PriceObs,29    ResultObs,30)31from aiatlas.sdk.resolution import Resolver3233log = logging.getLogger(__name__)3435TIER_CONFIDENCE = {1: "high", 2: "medium", 3: "low", 4: "low"}36SOFT_PROPERTIES = {"description", "summary", "tagline", "availability_note", "abstract", "notes", "training_data_notes", "safety_notes", "hardware_requirements"}37NEW_IMPORTANCE = {"model": 3, "company": 2, "provider": 2, "paper": 1, "dataset": 1, "benchmark": 2, "framework": 1, "hardware": 2,38                  "tool": 1, "repository": 0, "release": 2, "regulation": 2, "incident": 2, "organization": 1, "researcher": 0}394041def _norm_value(v: Any) -> Any:42    if isinstance(v, datetime):43        return v.astimezone(UTC).isoformat(timespec="seconds")44    if isinstance(v, dict):45        return {k: _norm_value(x) for k, x in v.items()}46    if isinstance(v, (set, tuple, list)):47        items = [_norm_value(x) for x in v]48        return sorted(items) if all(isinstance(x, (str, int, float)) and not isinstance(x, bool) for x in items) else items49    if isinstance(v, float) and v.is_integer() and abs(v) < 1e15:50        return int(v)51    return v525354def _same(a: Any, b: Any) -> bool:55    return json.dumps(_norm_value(a), sort_keys=True, default=str) == json.dumps(_norm_value(b), sort_keys=True, default=str)565758class WriteStats:59    def __init__(self) -> None:60        self.entities_created = 061        self.entities_updated = 062        self.claims = 063        self.relations = 064        self.events = 065        self.prices = 066        self.results = 067        self.conflicts = 06869    def as_dict(self) -> dict[str, int]:70        return dict(self.__dict__)717273class FactWriter:74    def __init__(self, conn: AsyncConnection, *, source_id: str | None, snapshot_id: str | None, source_url: str | None,75                 tier: int = 2, connector_name: str | None = None, extractor: str = "deterministic", extractor_version: str = "1",76                 observed_at: datetime | None = None):77        self.conn = conn78        self.source_id = source_id79        self.snapshot_id = snapshot_id80        self.source_url = source_url81        self.tier = tier82        self.connector_name = connector_name83        self.extractor = extractor84        self.extractor_version = extractor_version85        self.observed_at = observed_at or datetime.now(UTC)86        self.resolver = Resolver(conn, snapshot_id=snapshot_id, source_tier=tier)87        self.stats = WriteStats()8889    # ---------------------------------------------------------------------------------------------- entry point90    async def write(self, facts: Facts) -> WriteStats:91        new_before = len(self.resolver.created)92        for ref in facts.entities:93            await self.resolver.resolve(ref)94            for prop, value in ref.attributes.items():95                await self.write_claim(ref, prop, value)96        for c in facts.claims:97            await self.write_claim(c.entity, c.property, c.value, unit=c.unit, confidence=c.confidence, observed_at=c.observed_at,98                                   effective_at=c.effective_at, source_url=c.source_url)99        for r in facts.relations:100            await self.write_relation(r.subject, r.predicate, r.object, r.attributes, confidence=r.confidence, source_url=r.source_url)101        for p in facts.prices:102            await self.write_price(p)103        for res in facts.results:104            await self.write_result(res)105        for e in facts.events:106            eid = await self.resolver.resolve(e.entity) if e.entity else None107            await self.emit_event(e.event_type, e.category, e.summary, entity_id=eid, old_value=e.old_value, new_value=e.new_value,108                                  importance=e.importance, effective_at=e.effective_at, dedupe_key=e.dedupe_key, source_url=e.source_url, meta=e.meta)109        # NEW_* events for every entity created in this write110        for eid in self.resolver.created[new_before:]:111            await self._new_entity_event(eid)112        self.stats.entities_created = len(self.resolver.created)113        self.stats.entities_updated = len(self.resolver.updated - set(self.resolver.created))114        return self.stats115116    # ---------------------------------------------------------------------------------------------- claims117    async def write_claim(self, ref: EntityRef, prop: str, value: Any, *, unit: str | None = None, confidence: str | None = None,118                          observed_at: datetime | None = None, effective_at: datetime | None = None, source_url: str | None = None) -> None:119        if value is None or value == "" or value == [] or value == {}:120            return121        eid = await self.resolver.resolve(ref)122        assert eid123        value = _norm_value(value)124        observed = observed_at or self.observed_at125        conf = confidence or TIER_CONFIDENCE.get(self.tier, "medium")126        url = source_url or self.source_url127        noisy = prop.startswith(NOISY_PREFIXES)128        current = await fetch_one(self.conn, """select id, value, tier, source_url, observed_at from claims129                                                where entity_id = :e and property = :p and status = 'current' order by valid_from desc limit 1""",130                                  e=eid, p=prop)131        is_new_entity = eid in self.resolver.created132        if current is None:133            await self._insert_claim(eid, prop, value, unit, conf, "current", observed, effective_at, url)134            await self._set_attribute(eid, prop, value, unit, conf, url, observed)135            return136        if _same(current["value"], value):137            await execute(self.conn, "update claims set observed_at = greatest(observed_at, :o) where id = :id", o=observed, id=current["id"])138            if noisy:139                await self._set_attribute(eid, prop, value, unit, conf, url, observed)140            return141        same_source = bool(url and current["source_url"] == url)142        if prop in SOFT_PROPERTIES and not same_source:143            # soft text (descriptions, notes): first statement wins until *its own* source changes; never an event, never a conflict144            return145        if self.tier <= (current["tier"] or 2) or same_source:146            # supersede147            await execute(self.conn, "update claims set status = 'superseded', valid_to = :o where id = :id", o=observed, id=current["id"])148            await self._insert_claim(eid, prop, value, unit, conf, "current", observed, effective_at, url)149            await self._set_attribute(eid, prop, value, unit, conf, url, observed)150            if not noisy and not is_new_entity and prop not in SOFT_PROPERTIES:151                await self._property_change_event(eid, prop, current["value"], value, observed, effective_at, url)152        else:153            await self._insert_claim(eid, prop, value, unit, conf, "conflicting", observed, effective_at, url)154            await execute(self.conn, "update claims set confidence = 'conflicted' where id = :id", id=current["id"])155            await execute(self.conn, """update entities set quality = quality || jsonb_build_object('conflicts', coalesce((quality->>'conflicts')::int, 0) + 1)156                                        where id = :id""", id=eid)157            self.stats.conflicts += 1158            dedupe = f"conflict:{eid}:{prop}:{hashlib.sha1(json.dumps(value, sort_keys=True, default=str).encode()).hexdigest()[:10]}"159            await execute(self.conn, """insert into review_queue (id, kind, entity_ids, payload, reason, dedupe_key)160                                        values (:id, 'conflict', :ids, cast(:p as jsonb), :r, :d) on conflict (dedupe_key) do nothing""",161                          id=new_id("review"), ids=[eid], p=jsonb({"property": prop, "current": current["value"], "current_source": current["source_url"],162                                                                  "claimed": value, "claimed_source": url, "tier": self.tier}),163                          r=f"source disagrees on {prop}", d=dedupe)164165    async def _insert_claim(self, eid: str, prop: str, value: Any, unit: str | None, conf: str, status: str, observed: datetime,166                            effective_at: datetime | None, url: str | None) -> None:167        text_val = value if isinstance(value, str) else None168        num_val = float(value) if isinstance(value, (int, float)) and not isinstance(value, bool) else None169        await execute(self.conn, """insert into claims (id, entity_id, property, value, value_text, value_num, unit, source_id, snapshot_id, source_url, tier,170                                    confidence, status, extractor, extractor_version, observed_at, effective_at, valid_from)171                                    values (:id, :e, :p, cast(:v as jsonb), :vt, :vn, :u, :src, :snap, :url, :tier, :conf, :status, :ex, :exv, :o, :eff, :vf)""",172                      id=new_id("claim"), e=eid, p=prop, v=jsonb(value), vt=text_val[:2000] if text_val else None, vn=num_val, u=unit, src=self.source_id,173                      snap=self.snapshot_id, url=url, tier=self.tier, conf=conf, status=status, ex=self.extractor, exv=self.extractor_version,174                      o=observed, eff=effective_at, vf=effective_at or observed)175        self.stats.claims += 1176177    async def _set_attribute(self, eid: str, prop: str, value: Any, unit: str | None, conf: str, url: str | None, observed: datetime) -> None:178        prov = {"source_id": self.source_id, "snapshot_id": self.snapshot_id, "url": url, "observed_at": observed.isoformat(timespec="seconds"),179                "tier": self.tier, "confidence": conf, "extractor": self.extractor}180        if unit:181            prov["unit"] = unit182        await execute(self.conn, """update entities set attributes = attributes || jsonb_build_object(cast(:p as text), cast(:v as jsonb)),183                                    provenance = provenance || jsonb_build_object(cast(:p as text), cast(:prov as jsonb)), last_seen_at = greatest(last_seen_at, :o)184                                    where id = :id""", p=prop, v=jsonb(value), prov=jsonb(prov), o=observed, id=eid)185        if prop == "status" and isinstance(value, str):186            await execute(self.conn, "update entities set status = :s where id = :id", s=value[:40], id=eid)187        if prop == "description" and isinstance(value, str):188            await execute(self.conn, "update entities set description = :d where id = :id", d=value[:4000], id=eid)189190    async def _property_change_event(self, eid: str, prop: str, old: Any, new: Any, observed: datetime, effective_at: datetime | None,191                                     url: str | None) -> None:192        event_type, importance = MATERIAL_PROPERTIES.get(prop, ("PROPERTY_CHANGED", 0))193        row = await fetch_one(self.conn, "select canonical_name, entity_type from entities where id = :id", id=eid)194        name = row["canonical_name"] if row else eid195        category = EVENT_CATEGORY_BY_TYPE.get(row["entity_type"] if row else "", "update")196        if prop == "status" and str(new).lower() in ("deprecated", "retired", "discontinued"):197            importance = 3198        summary = f"{name}: {prop.replace('_', ' ')} changed from {_short(old)} to {_short(new)}"199        dedupe = f"{event_type}:{eid}:{prop}:{hashlib.sha1(json.dumps([old, new], sort_keys=True, default=str).encode()).hexdigest()[:12]}"200        await self.emit_event(event_type, category, summary, entity_id=eid, old_value=old, new_value=new, importance=importance,201                              effective_at=effective_at, dedupe_key=dedupe, source_url=url, meta={"property": prop}, observed_at=observed)202203    # ---------------------------------------------------------------------------------------------- relations204    async def write_relation(self, subject: EntityRef, predicate: str, obj: EntityRef, attributes: dict[str, Any] | None = None, *,205                             confidence: str | None = None, source_url: str | None = None) -> None:206        sid = await self.resolver.resolve(subject)207        oid = await self.resolver.resolve(obj)208        if not sid or not oid or sid == oid:209            return210        conf = confidence or TIER_CONFIDENCE.get(self.tier, "medium")211        existing = await fetch_one(self.conn, "select id, attributes from relations where subject_id = :s and predicate = :p and object_id = :o and valid_to is null",212                                   s=sid, p=predicate, o=oid)213        if existing:214            if attributes and not _same(existing["attributes"], {**existing["attributes"], **attributes}):215                await execute(self.conn, "update relations set attributes = attributes || cast(:a as jsonb), observed_at = :o where id = :id",216                              a=jsonb(attributes), o=self.observed_at, id=existing["id"])217            else:218                await execute(self.conn, "update relations set observed_at = :o where id = :id", o=self.observed_at, id=existing["id"])219            return220        await execute(self.conn, """insert into relations (id, subject_id, predicate, object_id, attributes, source_id, snapshot_id, source_url, tier, confidence, observed_at, valid_from)221                                    values (:id, :s, :p, :o, cast(:a as jsonb), :src, :snap, :url, :tier, :conf, :obs, :obs)""",222                      id=new_id("relation"), s=sid, p=predicate, o=oid, a=jsonb(attributes or {}), src=self.source_id, snap=self.snapshot_id,223                      url=source_url or self.source_url, tier=self.tier, conf=conf, obs=self.observed_at)224        self.stats.relations += 1225226    # ---------------------------------------------------------------------------------------------- prices227    async def write_price(self, p: PriceObs) -> None:228        mid = await self.resolver.resolve(p.model)229        pid = await self.resolver.resolve(p.provider)230        if not mid or not pid:231            return232        if all(v is None for v in p.price_tuple()[:-1]):233            return234        url = p.source_url or self.source_url235        current = await fetch_one(self.conn, """select * from prices where model_id = :m and provider_id = :p and coalesce(provider_model_id,'') = :pm and valid_to is null""",236                                  m=mid, p=pid, pm=p.provider_model_id or "")237        new_tuple = p.price_tuple()238        if current:239            cur_tuple = (current["input_per_mtok"], current["output_per_mtok"], current["cached_input_per_mtok"], current["cache_write_per_mtok"],240                         current["batch_input_per_mtok"], current["batch_output_per_mtok"], current["per_image"], current["per_request"], current["currency"])241            if _same(cur_tuple, new_tuple) and (p.context_length in (None, current["context_length"])):242                await execute(self.conn, "update prices set observed_at = :o where id = :id", o=self.observed_at, id=current["id"])243                return244            await execute(self.conn, "update prices set valid_to = :o where id = :id", o=self.observed_at, id=current["id"])245        await execute(self.conn, """insert into prices (id, model_id, provider_id, provider_model_id, input_per_mtok, output_per_mtok, cached_input_per_mtok, cache_write_per_mtok,246                                    batch_input_per_mtok, batch_output_per_mtok, per_image, per_request, currency, context_length, max_output_tokens, features,247                                    observed_at, valid_from, source_id, snapshot_id, source_url, tier, meta)248                                    values (:id, :m, :p, :pm, :i, :o, :ci, :cw, :bi, :bo, :img, :req, :cur, :ctx, :mo, cast(:f as jsonb), :obs, :obs, :src, :snap, :url, :tier, cast(:meta as jsonb))""",249                      id=new_id("price"), m=mid, p=pid, pm=p.provider_model_id, i=p.input_per_mtok, o=p.output_per_mtok, ci=p.cached_input_per_mtok,250                      cw=p.cache_write_per_mtok, bi=p.batch_input_per_mtok, bo=p.batch_output_per_mtok, img=p.per_image, req=p.per_request, cur=p.currency,251                      ctx=p.context_length, mo=p.max_output_tokens, f=jsonb(p.features), obs=self.observed_at, src=self.source_id, snap=self.snapshot_id,252                      url=url, tier=self.tier, meta=jsonb(p.meta))253        self.stats.prices += 1254        model = await fetch_one(self.conn, "select canonical_name from entities where id = :id", id=mid)255        provider = await fetch_one(self.conn, "select canonical_name from entities where id = :id", id=pid)256        mname = model["canonical_name"] if model else mid257        pname = provider["canonical_name"] if provider else pid258        if current:259            old = {"input_per_mtok": current["input_per_mtok"], "output_per_mtok": current["output_per_mtok"]}260            new = {"input_per_mtok": p.input_per_mtok, "output_per_mtok": p.output_per_mtok}261            summary = f"{pname} changed pricing for {mname}: {_fmt_price(old)} → {_fmt_price(new)}"262            dedupe = f"PRICE_CHANGED:{mid}:{pid}:{p.provider_model_id or ''}:{hashlib.sha1(json.dumps([old, new], sort_keys=True, default=str).encode()).hexdigest()[:12]}"263            await self.emit_event("PRICE_CHANGED", "price", summary, entity_id=mid, old_value=old, new_value=new, importance=2, dedupe_key=dedupe,264                                  source_url=url, meta={"provider_id": pid, "provider": pname})265        else:266            new = {"input_per_mtok": p.input_per_mtok, "output_per_mtok": p.output_per_mtok}267            summary = f"{pname} lists {mname} at {_fmt_price(new)}"268            dedupe = f"PROVIDER_LISTED:{mid}:{pid}:{p.provider_model_id or ''}"269            await self.emit_event("PROVIDER_LISTED", "provider", summary, entity_id=mid, new_value=new, importance=1, dedupe_key=dedupe, source_url=url,270                                  meta={"provider_id": pid, "provider": pname})271        await self.write_relation(p.model, "available_through", p.provider, {"provider_model_id": p.provider_model_id}, source_url=url)272273    # ---------------------------------------------------------------------------------------------- benchmark results274    async def write_result(self, r: ResultObs) -> None:275        mid = await self.resolver.resolve(r.model)276        bid = await self.resolver.resolve(r.benchmark)277        if not mid or not bid:278            return279        url = r.source_url or self.source_url280        cfg_hash = hashlib.sha1(json.dumps(r.config, sort_keys=True, default=str).encode()).hexdigest()[:12]281        dedupe = f"{mid}:{bid}:{cfg_hash}:{r.metric or ''}"282        existing = await fetch_one(self.conn, "select id, score from benchmark_results where dedupe_key = :d", d=dedupe)283        conf = r.confidence or TIER_CONFIDENCE.get(self.tier, "medium")284        if existing:285            if abs((existing["score"] or 0) - r.score) > 1e-9:286                await execute(self.conn, "update benchmark_results set valid_to = :o, dedupe_key = dedupe_key || ':' || :suffix where id = :id",287                              o=self.observed_at, suffix=new_id("result")[-10:], id=existing["id"])288                model = await fetch_one(self.conn, "select canonical_name from entities where id = :id", id=mid)289                bench = await fetch_one(self.conn, "select canonical_name from entities where id = :id", id=bid)290                await self.emit_event("BENCHMARK_UPDATED", "benchmark",291                                      f"{model['canonical_name'] if model else mid} on {bench['canonical_name'] if bench else bid}: {existing['score']:g} → {r.score:g}",292                                      entity_id=mid, old_value=existing["score"], new_value=r.score, importance=1,293                                      dedupe_key=f"BENCHMARK_UPDATED:{dedupe}:{r.score:g}", source_url=url, meta={"benchmark_id": bid})294            else:295                await execute(self.conn, "update benchmark_results set observed_at = :o where id = :id", o=self.observed_at, id=existing["id"])296                return297        await execute(self.conn, """insert into benchmark_results (id, model_id, benchmark_id, score, metric, unit, higher_is_better, config, evaluated_at, observed_at,298                                    source_id, snapshot_id, source_url, tier, confidence, dedupe_key)299                                    values (:id, :m, :b, :s, :metric, :unit, :hib, cast(:cfg as jsonb), :ev, :obs, :src, :snap, :url, :tier, :conf, :d)""",300                      id=new_id("result"), m=mid, b=bid, s=r.score, metric=r.metric, unit=r.unit, hib=r.higher_is_better, cfg=jsonb(r.config),301                      ev=r.evaluated_at, obs=self.observed_at, src=self.source_id, snap=self.snapshot_id, url=url, tier=self.tier, conf=conf, d=dedupe)302        self.stats.results += 1303        if not existing:304            model = await fetch_one(self.conn, "select canonical_name from entities where id = :id", id=mid)305            bench = await fetch_one(self.conn, "select canonical_name from entities where id = :id", id=bid)306            await self.emit_event("BENCHMARK_RESULT", "benchmark",307                                  f"{model['canonical_name'] if model else mid} scores {r.score:g}{r.unit or ''} on {bench['canonical_name'] if bench else bid}",308                                  entity_id=mid, new_value=r.score, importance=1, dedupe_key=f"BENCHMARK_RESULT:{dedupe}", source_url=url,309                                  meta={"benchmark_id": bid, "metric": r.metric})310        await self.write_relation(r.model, "evaluated_on", r.benchmark, source_url=url)311312    # ---------------------------------------------------------------------------------------------- events313    async def emit_event(self, event_type: str, category: str, summary: str, *, entity_id: str | None = None, old_value: Any = None,314                         new_value: Any = None, importance: int = 2, effective_at: datetime | None = None, dedupe_key: str | None = None,315                         source_url: str | None = None, meta: dict[str, Any] | None = None, observed_at: datetime | None = None) -> None:316        dedupe = dedupe_key or f"{event_type}:{entity_id or ''}:{normalize_alias(summary)[:120]}"317        await execute(self.conn, """insert into change_events (id, entity_id, event_type, category, property, old_value, new_value, summary, importance, observed_at,318                                    effective_at, source_id, snapshot_id, source_url, connector_name, dedupe_key, meta)319                                    values (:id, :e, :t, :c, :p, cast(:o as jsonb), cast(:n as jsonb), :s, :imp, :obs, :eff, :src, :snap, :url, :conn, :d, cast(:meta as jsonb))320                                    on conflict (dedupe_key) do nothing""",321                      id=new_id("change_event"), e=entity_id, t=event_type, c=category, p=(meta or {}).get("property"), o=jsonb(old_value) if old_value is not None else None,322                      n=jsonb(new_value) if new_value is not None else None, s=summary[:500], imp=max(0, min(3, importance)), obs=observed_at or self.observed_at,323                      eff=effective_at, src=self.source_id, snap=self.snapshot_id, url=source_url or self.source_url, conn=self.connector_name, d=dedupe[:400],324                      meta=jsonb(meta or {}))325        self.stats.events += 1326327    async def _new_entity_event(self, eid: str) -> None:328        row = await fetch_one(self.conn, "select canonical_name, entity_type, attributes, organization_id from entities where id = :id", id=eid)329        if not row:330            return331        etype = row["entity_type"]332        if etype in ("researcher", "country", "license", "quantization"):333            return334        importance = NEW_IMPORTANCE.get(etype, 1)335        if self.tier > 2:336            importance = max(0, importance - 1)337        org = None338        if row["organization_id"]:339            o = await fetch_one(self.conn, "select canonical_name from entities where id = :id", id=row["organization_id"])340            org = o["canonical_name"] if o else None341        label = etype.replace("_", " ")342        summary = f"New {label}: {row['canonical_name']}" + (f" ({org})" if org else "")343        effective = None344        rel = row["attributes"].get("release_date") or row["attributes"].get("published_at")345        if isinstance(rel, str):346            from aiatlas.sdk.extract.dates import parse_datetime347348            effective = parse_datetime(rel)349        await self.emit_event(f"NEW_{etype.upper()}", EVENT_CATEGORY_BY_TYPE.get(etype, "update"), summary, entity_id=eid, importance=importance,350                              dedupe_key=f"NEW_{etype.upper()}:{eid}", effective_at=effective)351352353def _short(v: Any) -> str:354    s = json.dumps(v, default=str, ensure_ascii=False) if not isinstance(v, str) else v355    return s if len(s) <= 60 else s[:57] + "…"356357358def _fmt_price(p: dict[str, Any]) -> str:359    i, o = p.get("input_per_mtok"), p.get("output_per_mtok")360    parts = []361    if i is not None:362        parts.append(f"${i:g} in")363    if o is not None:364        parts.append(f"${o:g} out")365    return " / ".join(parts) + " per 1M tokens" if parts else "n/a"366367368__all__ = ["FactWriter", "WriteStats"]369
Files · 0f064299de6124ef553ef1bf8a851ae63aeae373
Menu

Theme