SDK hardening: SSRF guard with per-hop redirect validation, run_id on every fact, same-source loophole closed, taxonomy at write time, result comparability, event semantics, hierarchy hints, quarantine
- fetch: validate_destination() on every request and redirect hop (max 5, manual loop): non-http(s), localhost/.local/.internal, RFC1918,
loopback, link-local (metadata), CGNAT, IPv6 loopback/link-local/ULA, unspecified → FetchError("blocked destination")
- facts: facts_to_json / facts_from_json (tagged datetimes) for quarantined runs
- writer: run_id/source_key/is_first_run; a claim supersedes only when tier <= current or same source AND same extractor (LLM never
overwrites deterministic from the same URL); licence/openness/modalities/status/kind/org_kind normalised via the ontology with value_raw
+ taxonomy_mappings, same assertion in another spelling re-encoded in place (no event); license_key claim; results get config_key,
trust_level, variant, run_group, extractor, effort config for variant refs, one current row per (model, benchmark, metric, config_key),
out-of-range scores → confidence low + anomaly; events get recorded_at, is_backfill (classify_backfill), group_key, deterministic
importance; family/canonical hints materialised (family_id, canonical_id, artifact_kind, member_of_family, artifact_of); NEW_ARTIFACT 0,
NEW_MODEL_FAMILY 1; conflicting claims recorded once; derived writers never open review items and mark their events backfill
- resolver: model/artifact compatible lookups; keep_separate decisions honoured; digit-collapsing aliases require the same variant_key;
first_seen_hint → first_seen_at = min(now, hint); resolve_variant() folds evaluator-only effort variants onto the canonical model
- connector: run_id/source_key/first-run flag into the writer; --url reprocess rebuilds Targets from documents.meta._target (key, doc_type,
meta, needs_llm, entity persisted at first fetch); quarantine for tier ≥ 2 connectors (rolling median baseline of the last 5 full runs:
new entities > max(50, 1.3×), prices > 3×, entities < 0.8× / results < 0.5× on full extraction) → quarantined_runs + review item, nothing written
- services/events: classify_backfill, group_key_for, importance_for · services/anomalies: record / run_checks / list_anomalies
- merge: modes merge|alias|variant|family_member, keep_separate refusal, org-type and model/artifact pairs, dedupe/config_key recompute with
collision handling, one current result per key, embeddings/sources/domains/review_queue/family_id/canonical_id re-pointed,
resolution_decisions + admin_audit_log rows
- handlers: LLM extraction writer carries the snapshot's run_id
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
13 changed files +1,757 −133
modified
src/aiatlas/sdk/connector.py
+200 −17
@@ -11,25 +11,29 @@ from __future__ import annotations | ||
| 11 | 11 | import asyncio |
| 12 | 12 | import json |
| 13 | 13 | import logging |
| 14 | +import statistics | |
| 14 | 15 | import time |
| 15 | 16 | from dataclasses import dataclass, field |
| 16 | 17 | from datetime import UTC, datetime, timedelta |
| 17 | 18 | from typing import Any |
| 18 | 19 | |
| 19 | −from aiatlas.db import execute, fetch_one, jsonb, transaction | |
| 20 | +from aiatlas.db import connection, execute, fetch_one, jsonb, transaction | |
| 20 | 21 | from aiatlas.ids import new_id |
| 21 | 22 | from aiatlas.sdk import archive |
| 22 | 23 | from aiatlas.sdk.extract.feeds import FeedItem, parse_feed |
| 23 | 24 | from aiatlas.sdk.extract.html import HtmlDoc, parse_html |
| 24 | 25 | from aiatlas.sdk.extract.markdown import MarkdownDoc, parse_markdown |
| 25 | −from aiatlas.sdk.facts import Facts, Target | |
| 26 | +from aiatlas.sdk.facts import EntityRef, Facts, Target, facts_to_json | |
| 26 | 27 | from aiatlas.sdk.fetch import BlockedError, Fetcher, FetchError, FetchResult, NotModified, canonicalize_url, file_result |
| 28 | +from aiatlas.sdk.resolution import Resolver | |
| 27 | 29 | from aiatlas.sdk.writer import FactWriter |
| 28 | 30 | |
| 29 | 31 | log = logging.getLogger(__name__) |
| 30 | 32 | |
| 31 | 33 | CIRCUIT_FAILURES = 3 |
| 32 | 34 | CIRCUIT_COOLDOWN_S = 1800 |
| 35 | +BASELINE_RUNS = 5 # rolling window (median) of full-extraction runs used by the quarantine detector | |
| 36 | +QUARANTINE_MIN_NEW_ENTITIES = 50 # a run may always create up to this many new entities without tripping the detector | |
| 33 | 37 | |
| 34 | 38 | |
| 35 | 39 | class ConnectorError(Exception): |
@@ -77,6 +81,17 @@ class RunStats: | ||
| 77 | 81 | meta: dict[str, Any] = field(default_factory=dict) |
| 78 | 82 | |
| 79 | 83 | |
| 84 | +@dataclass | |
| 85 | +class Buffered: | |
| 86 | + """Facts extracted from one document, held until the end of the run when the connector is quarantine-capable.""" | |
| 87 | + target: Target | |
| 88 | + doc_id: str | |
| 89 | + snapshot_id: str | |
| 90 | + source_url: str | |
| 91 | + fetched_at: datetime | |
| 92 | + facts: Facts | |
| 93 | + | |
| 94 | + | |
| 80 | 95 | @dataclass |
| 81 | 96 | class RunContext: |
| 82 | 97 | run_id: str |
@@ -89,6 +104,11 @@ class RunContext: | ||
| 89 | 104 | max_targets: int = 2000 |
| 90 | 105 | stats: RunStats = field(default_factory=RunStats) |
| 91 | 106 | seen_urls: set[str] = field(default_factory=set) |
| 107 | + is_first_run: bool = False # no earlier successful run → every event of this run is backfill | |
| 108 | + source_key: str | None = None | |
| 109 | + buffered: list[Buffered] = field(default_factory=list) # quarantine mode: facts held until the run-end anomaly check | |
| 110 | + quarantined: bool = False | |
| 111 | + quarantine_reason: str | None = None | |
| 92 | 112 | log: logging.LoggerAdapter[logging.Logger] = field(init=False) |
| 93 | 113 | |
| 94 | 114 | def __post_init__(self) -> None: |
@@ -114,11 +134,21 @@ class BaseConnector: | ||
| 114 | 134 | concurrency: int = 3 |
| 115 | 135 | default_doc_type: str = "page" |
| 116 | 136 | needs_llm: bool = False # queue documents for LLM extraction after deterministic extraction |
| 137 | + # Hold every fact until the end of the run and compare the run's counts with the connector's rolling baseline before writing; | |
| 138 | + # a collapse or an explosion parks the whole run in `quarantined_runs` for review. None = default (tier ≥ 2 hub/leaderboard/registry | |
| 139 | + # connectors are quarantine-capable; tier-1 lab documentation is not — it is the primary source). | |
| 140 | + quarantine: bool | None = None | |
| 117 | 141 | |
| 118 | 142 | def __init__(self, config: dict[str, Any] | None = None): |
| 119 | 143 | self.config = config or {} |
| 120 | 144 | self.source_id: str | None = None |
| 121 | 145 | |
| 146 | + @property | |
| 147 | + def quarantine_enabled(self) -> bool: | |
| 148 | + if self.quarantine is not None: | |
| 149 | + return self.quarantine | |
| 150 | + return self.tier >= 2 | |
| 151 | + | |
| 122 | 152 | # ------------------------------------------------------------------------------------------ to implement |
| 123 | 153 | async def discover(self, ctx: RunContext) -> list[Target]: |
| 124 | 154 | """Return the targets for this run (seed pages, feeds, sitemaps, org pages…).""" |
@@ -165,10 +195,14 @@ class BaseConnector: | ||
| 165 | 195 | ctx = RunContext(run_id=run_id, connector=self, fetcher=fetcher, started_at=started, force=force, reprocess=reprocess, |
| 166 | 196 | file_overrides=file_overrides or {}, max_targets=max_targets or int(self.config.get("max_targets", 2000))) |
| 167 | 197 | async with transaction() as conn: |
| 168 | − state = await fetch_one(conn, "select source_id, enabled, circuit_open_until from connectors where name = :n", n=self.name) | |
| 198 | + state = await fetch_one(conn, """select c.source_id, c.enabled, c.circuit_open_until, s.key as source_key, | |
| 199 | + exists(select 1 from connector_runs r where r.connector_name = c.name and r.status in ('success','unchanged','suspect')) as has_success | |
| 200 | + from connectors c left join sources s on s.id = c.source_id where c.name = :n""", n=self.name) | |
| 169 | 201 | if state is None: |
| 170 | 202 | raise ConnectorError(f"connector {self.name!r} not registered (run `aia seed`)") |
| 171 | 203 | self.source_id = state["source_id"] |
| 204 | + ctx.source_key = state["source_key"] or self.source_key or None | |
| 205 | + ctx.is_first_run = not state["has_success"] | |
| 172 | 206 | if not state["enabled"] and not force: |
| 173 | 207 | ctx.log.info("connector disabled, skipping") |
| 174 | 208 | await self._record(conn, ctx, "skipped", error="disabled") |
@@ -184,21 +218,23 @@ class BaseConnector: | ||
| 184 | 218 | try: |
| 185 | 219 | targets = await self.discover(ctx) |
| 186 | 220 | if only_urls: |
| 187 | − wanted = {canonicalize_url(u) for u in only_urls} | |
| 188 | − targets = [t for t in targets if canonicalize_url(t.url) in wanted] or [Target(url=u) for u in only_urls] | |
| 221 | + targets = await self._select_targets(targets, only_urls) | |
| 189 | 222 | ctx.stats.docs_discovered = len(targets) |
| 190 | 223 | await self._process_all(ctx, targets) |
| 191 | − status = self._final_status(ctx) | |
| 224 | + if ctx.buffered: | |
| 225 | + await self._flush_or_quarantine(ctx) | |
| 226 | + status = "quarantined" if ctx.quarantined else self._final_status(ctx) | |
| 192 | 227 | async with transaction() as conn: |
| 193 | 228 | await self._record(conn, ctx, status) |
| 194 | 229 | changed = ctx.stats.docs_changed > 0 |
| 195 | − await execute(conn, """update connectors set consecutive_failures = 0, circuit_open_until = null, last_success_at = now(), | |
| 196 | − last_change_at = case when cast(:changed as boolean) then now() else last_change_at end, | |
| 230 | + await execute(conn, """update connectors set consecutive_failures = 0, circuit_open_until = null, | |
| 231 | + last_success_at = case when cast(:q as boolean) then last_success_at else now() end, | |
| 232 | + last_change_at = case when cast(:changed as boolean) and not cast(:q as boolean) then now() else last_change_at end, | |
| 197 | 233 | consecutive_unchanged = case when cast(:changed as boolean) then 0 else consecutive_unchanged + 1 end, |
| 198 | 234 | interval_seconds = cast(:interval as integer), next_run_at = now() + make_interval(secs => cast(:interval as double precision)), |
| 199 | 235 | last_duration_ms = :d, health = :health, updated_at = now() where name = :n""", |
| 200 | − changed=changed, interval=await self._next_interval(conn, changed), d=int((time.perf_counter() - t0) * 1000), | |
| 201 | − health="degraded" if status == "suspect" else "ok", n=self.name) | |
| 236 | + changed=changed, q=ctx.quarantined, interval=await self._next_interval(conn, changed), d=int((time.perf_counter() - t0) * 1000), | |
| 237 | + health="degraded" if status in ("suspect", "quarantined") else "ok", n=self.name) | |
| 202 | 238 | if status == "suspect": |
| 203 | 239 | await execute(conn, """insert into review_queue (id, kind, entity_ids, payload, reason, dedupe_key) |
| 204 | 240 | values (:id, 'parser_breakage', '{}', cast(:p as jsonb), :r, :d) on conflict (dedupe_key) do nothing""", |
@@ -223,11 +259,58 @@ class BaseConnector: | ||
| 223 | 259 | raise |
| 224 | 260 | return ctx |
| 225 | 261 | |
| 262 | + async def _select_targets(self, targets: list[Target], only_urls: list[str]) -> list[Target]: | |
| 263 | + """`--url` / single-document reprocess: keep the discovered targets for those URLs; for URLs `discover()` no longer lists, rebuild | |
| 264 | + the Target from what the first fetch persisted in `documents.meta` (key, doc_type, meta, needs_llm, entity) so extractors that | |
| 265 | + branch on `target.key`/`target.meta` behave exactly as they did originally.""" | |
| 266 | + wanted = {canonicalize_url(u) for u in only_urls} | |
| 267 | + selected = [t for t in targets if canonicalize_url(t.url) in wanted] | |
| 268 | + missing = wanted - {canonicalize_url(t.url) for t in selected} | |
| 269 | + if missing: | |
| 270 | + async with transaction() as conn: | |
| 271 | + for url in sorted(missing): | |
| 272 | + doc = await fetch_one(conn, "select url, doc_type, meta, needs_llm, entity_id, priority from documents where url = :u", u=url) | |
| 273 | + if not doc: | |
| 274 | + selected.append(Target(url=url)) | |
| 275 | + continue | |
| 276 | + meta = dict(doc["meta"] or {}) | |
| 277 | + stored = meta.pop("_target", {}) if isinstance(meta.get("_target"), dict) else {} | |
| 278 | + entity = None | |
| 279 | + if stored.get("entity"): | |
| 280 | + try: | |
| 281 | + from aiatlas.sdk.facts import _ref_from # noqa: PLC0415 | |
| 282 | + | |
| 283 | + entity = _ref_from(stored["entity"]) | |
| 284 | + except Exception: # noqa: BLE001 | |
| 285 | + entity = None | |
| 286 | + elif doc["entity_id"]: | |
| 287 | + row = await fetch_one(conn, "select entity_type, canonical_name from entities where id = :id", id=doc["entity_id"]) | |
| 288 | + if row: | |
| 289 | + entity = EntityRef(entity_type=row["entity_type"], name=row["canonical_name"], id=doc["entity_id"]) | |
| 290 | + selected.append(Target(url=doc["url"], doc_type=stored.get("doc_type") or doc["doc_type"] or self.default_doc_type, entity=entity, | |
| 291 | + meta=stored.get("meta") if isinstance(stored.get("meta"), dict) else meta, key=stored.get("key"), | |
| 292 | + needs_llm=bool(stored.get("needs_llm", doc["needs_llm"])), priority=int(doc["priority"] or 2))) | |
| 293 | + return selected | |
| 294 | + | |
| 295 | + @staticmethod | |
| 296 | + def _target_meta(target: Target, doc_type: str) -> dict[str, Any]: | |
| 297 | + """What `documents.meta` remembers about the target that produced the document (rebuilt by `_select_targets`).""" | |
| 298 | + stored: dict[str, Any] = {"key": target.key, "doc_type": doc_type, "meta": _jsonable(target.meta), "needs_llm": bool(target.needs_llm)} | |
| 299 | + if target.entity is not None: | |
| 300 | + from aiatlas.sdk.facts import _encode # noqa: PLC0415 | |
| 301 | + | |
| 302 | + stored["entity"] = _encode(target.entity) | |
| 303 | + return {**_jsonable(target.meta), "_target": stored} | |
| 304 | + | |
| 305 | + def _full_extraction(self, ctx: RunContext) -> bool: | |
| 306 | + s = ctx.stats | |
| 307 | + return s.docs_fetched > 0 and s.docs_changed >= s.docs_fetched - s.docs_failed and s.docs_unchanged == 0 | |
| 308 | + | |
| 226 | 309 | def _final_status(self, ctx: RunContext) -> str: |
| 227 | 310 | s = ctx.stats |
| 228 | 311 | # Breakage detection only makes sense when every fetched document was (re)extracted: unchanged documents legitimately |
| 229 | 312 | # produce zero records (304 / same hash), so an incremental run is never "suspect". |
| 230 | − full_extraction = s.docs_fetched > 0 and s.docs_changed >= s.docs_fetched - s.docs_failed and s.docs_unchanged == 0 | |
| 313 | + full_extraction = self._full_extraction(ctx) | |
| 231 | 314 | if self.expected_min_records and full_extraction and s.records < self.expected_min_records and not ctx.reprocess: |
| 232 | 315 | return "suspect" |
| 233 | 316 | if s.docs_fetched and s.docs_failed == s.docs_fetched and s.docs_fetched > 0: |
@@ -298,8 +381,11 @@ class BaseConnector: | ||
| 298 | 381 | if doc is None: |
| 299 | 382 | doc_id = new_id("document") |
| 300 | 383 | await execute(conn, """insert into documents (id, source_id, connector_name, url, doc_type, priority, meta) values (:id, :s, :c, :u, :t, :p, cast(:m as jsonb)) |
| 301 | − on conflict (url) do nothing""", id=doc_id, s=self.source_id, c=self.name, u=url, t=doc_type, p=target.priority, m=jsonb(target.meta)) | |
| 384 | + on conflict (url) do nothing""", id=doc_id, s=self.source_id, c=self.name, u=url, t=doc_type, p=target.priority, | |
| 385 | + m=jsonb(self._target_meta(target, doc_type))) | |
| 302 | 386 | doc = await fetch_one(conn, "select * from documents where url = :u", u=url) |
| 387 | + elif (target.key or target.meta) and not ctx.reprocess and (doc["meta"] or {}).get("_target", {}).get("key") != target.key: | |
| 388 | + await execute(conn, "update documents set meta = meta || cast(:m as jsonb) where id = :id", m=jsonb(self._target_meta(target, doc_type)), id=doc["id"]) | |
| 303 | 389 | assert doc |
| 304 | 390 | # ---- reprocess mode: replay latest snapshot without network |
| 305 | 391 | if ctx.reprocess: |
@@ -407,15 +493,26 @@ class BaseConnector: | ||
| 407 | 493 | return [] |
| 408 | 494 | if facts is None: |
| 409 | 495 | return [] |
| 496 | + item = Buffered(target=target, doc_id=doc["id"], snapshot_id=snapshot_id, source_url=res.final_url or res.url, fetched_at=res.fetched_at, facts=facts) | |
| 497 | + ctx.stats.records += len(facts.entities) + len(facts.prices) + len(facts.results) | |
| 498 | + if self.quarantine_enabled and not ctx.reprocess: | |
| 499 | + ctx.buffered.append(item) # written (or quarantined) at the end of the run, see _flush_or_quarantine | |
| 500 | + else: | |
| 501 | + await self._write_buffered(ctx, item) | |
| 502 | + return facts.targets | |
| 503 | + | |
| 504 | + async def _write_buffered(self, ctx: RunContext, item: Buffered) -> None: | |
| 505 | + facts, target, snapshot_id = item.facts, item.target, item.snapshot_id | |
| 410 | 506 | async with transaction() as conn: |
| 411 | − writer = FactWriter(conn, source_id=self.source_id, snapshot_id=snapshot_id, source_url=res.final_url or res.url, tier=self.tier, | |
| 412 | − connector_name=self.name, extractor="deterministic", extractor_version=self.parser_version, observed_at=res.fetched_at) | |
| 507 | + writer = FactWriter(conn, source_id=self.source_id, snapshot_id=snapshot_id, source_url=item.source_url, tier=self.tier, | |
| 508 | + connector_name=self.name, extractor="deterministic", extractor_version=self.parser_version, observed_at=item.fetched_at, | |
| 509 | + run_id=ctx.run_id, source_key=ctx.source_key, is_first_run=ctx.is_first_run) | |
| 413 | 510 | ws = await writer.write(facts) |
| 414 | 511 | main = facts.document_entity or target.entity |
| 415 | 512 | if main and main.id is None: |
| 416 | 513 | await writer.resolver.resolve(main) |
| 417 | 514 | if main and main.id: |
| 418 | − await execute(conn, "update documents set entity_id = coalesce(entity_id, :e), title = coalesce(:t, title) where id = :id", e=main.id, t=facts.document_title, id=doc["id"]) | |
| 515 | + await execute(conn, "update documents set entity_id = coalesce(entity_id, :e), title = coalesce(:t, title) where id = :id", e=main.id, t=facts.document_title, id=item.doc_id) | |
| 419 | 516 | await execute(conn, "update snapshots set processing_status = :st where id = :id", |
| 420 | 517 | st="llm_pending" if (target.needs_llm or self.needs_llm or facts.llm_hint) else "extracted", id=snapshot_id) |
| 421 | 518 | if target.needs_llm or self.needs_llm or facts.llm_hint: |
@@ -430,8 +527,94 @@ class BaseConnector: | ||
| 430 | 527 | s.claims_written += ws.claims |
| 431 | 528 | s.relations_written += ws.relations |
| 432 | 529 | s.events_emitted += ws.events |
| 433 | − s.records += len(facts.entities) + len(facts.prices) + len(facts.results) | |
| 434 | − return facts.targets | |
| 530 | + | |
| 531 | + # ------------------------------------------------------------------------------------------ quarantine | |
| 532 | + async def _observed_counts(self, ctx: RunContext) -> dict[str, int]: | |
| 533 | + """Counts of what the buffered facts would write: entity references, prices, results and *new* entity names (refs that resolve to | |
| 534 | + nothing today). Resolution runs in a transaction that is rolled back so the check leaves no trace.""" | |
| 535 | + refs: dict[str, EntityRef] = {} | |
| 536 | + prices = results = 0 | |
| 537 | + for item in ctx.buffered: | |
| 538 | + prices += len(item.facts.prices) | |
| 539 | + results += len(item.facts.results) | |
| 540 | + for ref in item.facts.entities: | |
| 541 | + refs.setdefault(ref.key(), ref) | |
| 542 | + new_entities = 0 | |
| 543 | + async with connection() as conn: | |
| 544 | + trans = await conn.begin() | |
| 545 | + try: | |
| 546 | + resolver = Resolver(conn, source_tier=self.tier) | |
| 547 | + for ref in refs.values(): | |
| 548 | + probe = EntityRef(entity_type=ref.entity_type, name=ref.name, identifiers=dict(ref.identifiers), aliases=list(ref.aliases), slug_hint=ref.slug_hint) | |
| 549 | + if await resolver.resolve(probe, create=False) is None: | |
| 550 | + new_entities += 1 | |
| 551 | + finally: | |
| 552 | + await trans.rollback() | |
| 553 | + return {"entities": len(refs), "prices": prices, "results": results, "new_entities": new_entities} | |
| 554 | + | |
| 555 | + @staticmethod | |
| 556 | + def _quarantine_reason(observed: dict[str, int], baseline: dict[str, Any] | None, *, full_extraction: bool) -> str | None: | |
| 557 | + """Explosion rules always apply; collapse rules only when every document was re-extracted (an incremental run legitimately yields little).""" | |
| 558 | + if not baseline or not baseline.get("history"): | |
| 559 | + return None | |
| 560 | + med = baseline.get("medians") or {} | |
| 561 | + b_new, b_ent, b_pr, b_res = med.get("new_entities", 0), med.get("entities", 0), med.get("prices", 0), med.get("results", 0) | |
| 562 | + if observed["new_entities"] > max(QUARANTINE_MIN_NEW_ENTITIES, 1.3 * b_new): | |
| 563 | + return f"new entities {observed['new_entities']} > max({QUARANTINE_MIN_NEW_ENTITIES}, 1.3 × baseline {b_new:g})" | |
| 564 | + if b_pr > 0 and observed["prices"] > 3 * b_pr: | |
| 565 | + return f"price rows {observed['prices']} > 3 × baseline {b_pr:g}" | |
| 566 | + if full_extraction: | |
| 567 | + if b_ent > 0 and observed["entities"] < 0.8 * b_ent: | |
| 568 | + return f"entity references {observed['entities']} < 0.8 × baseline {b_ent:g}" | |
| 569 | + if b_res > 0 and observed["results"] < 0.5 * b_res: | |
| 570 | + return f"benchmark results {observed['results']} < 0.5 × baseline {b_res:g}" | |
| 571 | + return None | |
| 572 | + | |
| 573 | + @staticmethod | |
| 574 | + def _next_baseline(baseline: dict[str, Any] | None, observed: dict[str, int]) -> dict[str, Any]: | |
| 575 | + history = list((baseline or {}).get("history") or []) | |
| 576 | + history.append({k: int(observed[k]) for k in ("entities", "prices", "results", "new_entities")}) | |
| 577 | + history = history[-BASELINE_RUNS:] | |
| 578 | + medians = {k: float(statistics.median(h[k] for h in history)) for k in ("entities", "prices", "results", "new_entities")} | |
| 579 | + return {"history": history, "medians": medians, "runs": len(history), "updated_at": datetime.now(UTC).isoformat(timespec="seconds")} | |
| 580 | + | |
| 581 | + async def _flush_or_quarantine(self, ctx: RunContext) -> None: | |
| 582 | + observed = await self._observed_counts(ctx) | |
| 583 | + full = self._full_extraction(ctx) | |
| 584 | + async with transaction() as conn: | |
| 585 | + row = await fetch_one(conn, "select baseline from connectors where name = :n", n=self.name) | |
| 586 | + baseline = (row or {}).get("baseline") | |
| 587 | + reason = None if ctx.force and ctx.file_overrides else self._quarantine_reason(observed, baseline, full_extraction=full) | |
| 588 | + ctx.stats.meta["quarantine_check"] = {"observed": observed, "baseline": (baseline or {}).get("medians"), "full_extraction": full} | |
| 589 | + if reason: | |
| 590 | + ctx.quarantined = True | |
| 591 | + ctx.quarantine_reason = reason | |
| 592 | + ctx.log.warning("run quarantined", extra={"reason": reason, **observed}) | |
| 593 | + payload = [{"doc_id": b.doc_id, "snapshot_id": b.snapshot_id, "source_url": b.source_url, "fetched_at": b.fetched_at.isoformat(), | |
| 594 | + "target": {"url": b.target.url, "doc_type": b.target.doc_type, "key": b.target.key, "needs_llm": b.target.needs_llm, "meta": _jsonable(b.target.meta)}, | |
| 595 | + "facts": facts_to_json(b.facts)} for b in ctx.buffered] | |
| 596 | + qid = f"quar_{new_id('connector_run').split('_', 1)[1]}" | |
| 597 | + async with transaction() as conn: | |
| 598 | + await execute(conn, """insert into quarantined_runs (id, run_id, connector_name, reason, stats, facts) values (:id, :r, :c, :why, cast(:s as jsonb), cast(:f as jsonb))""", | |
| 599 | + id=qid, r=ctx.run_id, c=self.name, why=reason, s=jsonb({"observed": observed, "baseline": baseline, "full_extraction": full, "documents": len(ctx.buffered)}), | |
| 600 | + f=jsonb(payload)) | |
| 601 | + await execute(conn, "update connector_runs set quarantined = true, baseline = cast(:b as jsonb) where id = :id", b=jsonb(baseline), id=ctx.run_id) | |
| 602 | + await execute(conn, "update snapshots set processing_status = 'quarantined' where id = any(cast(:ids as text[]))", ids=[b.snapshot_id for b in ctx.buffered]) | |
| 603 | + await execute(conn, """insert into review_queue (id, kind, entity_ids, payload, reason, dedupe_key) values (:id, 'quarantine', '{}', cast(:p as jsonb), :r, :d) | |
| 604 | + on conflict (dedupe_key) do nothing""", | |
| 605 | + id=new_id("review"), p=jsonb({"connector": self.name, "run_id": ctx.run_id, "quarantine_id": qid, "observed": observed, | |
| 606 | + "baseline": (baseline or {}).get("medians")}), | |
| 607 | + r=f"{self.name}: run held in quarantine — {reason}; nothing written, nothing deleted", d=f"quarantine:{qid}") | |
| 608 | + ctx.buffered.clear() | |
| 609 | + return | |
| 610 | + for item in ctx.buffered: | |
| 611 | + await self._write_buffered(ctx, item) | |
| 612 | + ctx.buffered.clear() | |
| 613 | + if full and not ctx.reprocess and not ctx.file_overrides: | |
| 614 | + async with transaction() as conn: | |
| 615 | + nb = self._next_baseline(baseline, observed) | |
| 616 | + await execute(conn, "update connectors set baseline = cast(:b as jsonb) where name = :n", b=jsonb(nb), n=self.name) | |
| 617 | + await execute(conn, "update connector_runs set baseline = cast(:b as jsonb) where id = :id", b=jsonb(nb), id=ctx.run_id) | |
| 435 | 618 | |
| 436 | 619 | async def _record(self, conn: Any, ctx: RunContext, status: str, *, error: str | None = None) -> None: |
| 437 | 620 | s = ctx.stats |
modified
src/aiatlas/sdk/facts.py
+88 −0
@@ -227,6 +227,92 @@ EVENT_CATEGORY_BY_TYPE: dict[str, str] = { | ||
| 227 | 227 | "mcp_server": "tool", "researcher": "company", "license": "model", "product": "tool", "robot": "hardware", |
| 228 | 228 | } |
| 229 | 229 | |
| 230 | + | |
| 231 | +# ---------------------------------------------------------------------------------------------- JSON round-trip (quarantine) | |
| 232 | +# Quarantined runs hold their Facts in `quarantined_runs.facts` (jsonb) until an operator releases or discards them. Datetimes are | |
| 233 | +# tagged so they come back as real datetime objects (asyncpg needs them, not ISO strings). Entity references are inlined; the | |
| 234 | +# resolver caches by `EntityRef.key()`, so equal refs resolve to the same entity after a round-trip. | |
| 235 | +_DT_TAG = "__datetime__" | |
| 236 | + | |
| 237 | + | |
| 238 | +def _encode(value: Any) -> Any: | |
| 239 | + if isinstance(value, datetime): | |
| 240 | + return {_DT_TAG: value.isoformat()} | |
| 241 | + if isinstance(value, EntityRef): | |
| 242 | + return {"entity_type": value.entity_type, "name": value.name, "identifiers": dict(value.identifiers), "aliases": list(value.aliases), | |
| 243 | + "slug_hint": value.slug_hint, "organization": _encode(value.organization), "description": value.description, "status": value.status, | |
| 244 | + "attributes": _encode(value.attributes), "first_seen_hint": _encode(value.first_seen_hint), "family": _encode(value.family), | |
| 245 | + "canonical": _encode(value.canonical), "artifact_kind": value.artifact_kind, "identity_confidence": value.identity_confidence, "id": value.id} | |
| 246 | + if isinstance(value, dict): | |
| 247 | + return {str(k): _encode(v) for k, v in value.items()} | |
| 248 | + if isinstance(value, (list, tuple, set)): | |
| 249 | + return [_encode(v) for v in value] | |
| 250 | + if isinstance(value, (str, int, float, bool)) or value is None: | |
| 251 | + return value | |
| 252 | + return str(value) | |
| 253 | + | |
| 254 | + | |
| 255 | +def _decode(value: Any) -> Any: | |
| 256 | + if isinstance(value, dict): | |
| 257 | + if set(value) == {_DT_TAG}: | |
| 258 | + return datetime.fromisoformat(value[_DT_TAG]) | |
| 259 | + return {k: _decode(v) for k, v in value.items()} | |
| 260 | + if isinstance(value, list): | |
| 261 | + return [_decode(v) for v in value] | |
| 262 | + return value | |
| 263 | + | |
| 264 | + | |
| 265 | +def _ref_from(d: dict[str, Any] | None) -> EntityRef | None: | |
| 266 | + if not d: | |
| 267 | + return None | |
| 268 | + return EntityRef(entity_type=d["entity_type"], name=d["name"], identifiers=dict(d.get("identifiers") or {}), aliases=list(d.get("aliases") or []), | |
| 269 | + slug_hint=d.get("slug_hint"), organization=_ref_from(d.get("organization")), description=d.get("description"), status=d.get("status"), | |
| 270 | + attributes=_decode(d.get("attributes") or {}), first_seen_hint=_decode(d.get("first_seen_hint")), family=_ref_from(d.get("family")), | |
| 271 | + canonical=_ref_from(d.get("canonical")), artifact_kind=d.get("artifact_kind"), identity_confidence=d.get("identity_confidence"), id=d.get("id")) | |
| 272 | + | |
| 273 | + | |
| 274 | +def facts_to_json(facts: Facts) -> dict[str, Any]: | |
| 275 | + """Serialise a Facts object to a JSON-compatible dict (see `facts_from_json`).""" | |
| 276 | + def obj(o: Any) -> dict[str, Any]: | |
| 277 | + return {k: _encode(v) for k, v in o.__dict__.items()} | |
| 278 | + | |
| 279 | + return { | |
| 280 | + "version": 1, | |
| 281 | + "entities": [_encode(e) for e in facts.entities], | |
| 282 | + "claims": [obj(c) for c in facts.claims], | |
| 283 | + "relations": [obj(r) for r in facts.relations], | |
| 284 | + "events": [obj(e) for e in facts.events], | |
| 285 | + "prices": [obj(p) for p in facts.prices], | |
| 286 | + "results": [obj(r) for r in facts.results], | |
| 287 | + "targets": [obj(t) for t in facts.targets], | |
| 288 | + "document_title": facts.document_title, | |
| 289 | + "document_entity": _encode(facts.document_entity), | |
| 290 | + "llm_hint": facts.llm_hint, | |
| 291 | + } | |
| 292 | + | |
| 293 | + | |
| 294 | +def facts_from_json(data: dict[str, Any]) -> Facts: | |
| 295 | + """Inverse of `facts_to_json`.""" | |
| 296 | + def ref_fields(d: dict[str, Any], *keys: str) -> dict[str, Any]: | |
| 297 | + out = {k: _decode(v) for k, v in d.items() if k not in keys} | |
| 298 | + for k in keys: | |
| 299 | + out[k] = _ref_from(d.get(k)) | |
| 300 | + return out | |
| 301 | + | |
| 302 | + facts = Facts() | |
| 303 | + facts.entities = [r for r in (_ref_from(e) for e in data.get("entities") or []) if r] | |
| 304 | + facts.claims = [Claim(**ref_fields(c, "entity")) for c in data.get("claims") or []] | |
| 305 | + facts.relations = [Relation(**ref_fields(r, "subject", "object")) for r in data.get("relations") or []] | |
| 306 | + facts.events = [Event(**ref_fields(e, "entity")) for e in data.get("events") or []] | |
| 307 | + facts.prices = [PriceObs(**ref_fields(p, "model", "provider")) for p in data.get("prices") or []] | |
| 308 | + facts.results = [ResultObs(**ref_fields(r, "model", "benchmark")) for r in data.get("results") or []] | |
| 309 | + facts.targets = [Target(**ref_fields(t, "entity")) for t in data.get("targets") or []] | |
| 310 | + facts.document_title = data.get("document_title") | |
| 311 | + facts.document_entity = _ref_from(data.get("document_entity")) | |
| 312 | + facts.llm_hint = data.get("llm_hint") | |
| 313 | + return facts | |
| 314 | + | |
| 315 | + | |
| 230 | 316 | __all__ = [ |
| 231 | 317 | "CONFIDENCE", |
| 232 | 318 | "EVENT_CATEGORY_BY_TYPE", |
@@ -240,4 +326,6 @@ __all__ = [ | ||
| 240 | 326 | "Relation", |
| 241 | 327 | "ResultObs", |
| 242 | 328 | "Target", |
| 329 | + "facts_from_json", | |
| 330 | + "facts_to_json", | |
| 243 | 331 | ] |
modified
src/aiatlas/sdk/fetch.py
+187 −24
@@ -7,11 +7,13 @@ from __future__ import annotations | ||
| 7 | 7 | |
| 8 | 8 | import asyncio |
| 9 | 9 | import hashlib |
| 10 | +import ipaddress | |
| 10 | 11 | import logging |
| 12 | +import socket | |
| 11 | 13 | import time |
| 12 | 14 | from dataclasses import dataclass, field |
| 13 | 15 | from datetime import UTC, datetime |
| 14 | −from urllib.parse import urlparse, urlunparse | |
| 16 | +from urllib.parse import urljoin, urlparse, urlunparse | |
| 15 | 17 | from urllib.robotparser import RobotFileParser |
| 16 | 18 | |
| 17 | 19 | import httpx |
@@ -22,6 +24,100 @@ log = logging.getLogger(__name__) | ||
| 22 | 24 | |
| 23 | 25 | TRANSIENT_STATUS = {408, 425, 429, 500, 502, 503, 504} |
| 24 | 26 | BLOCK_STATUS = {401, 403, 451, 999} |
| 27 | +REDIRECT_STATUS = {301, 302, 303, 307, 308} | |
| 28 | +MAX_REDIRECTS = 5 | |
| 29 | + | |
| 30 | +# ---------------------------------------------------------------------------------------------- SSRF guard | |
| 31 | +# Follow-up targets come from crawled content (links in model cards, feeds, READMEs): every destination — and every redirect | |
| 32 | +# hop — is validated before a connection is opened. Nothing internal is ever fetched. | |
| 33 | +_BLOCKED_SUFFIXES = (".local", ".internal", ".localhost", ".localdomain", ".lan", ".home", ".corp", ".intranet") | |
| 34 | +_BLOCKED_HOSTS = {"localhost", "metadata.google.internal", "metadata", "instance-data", "kubernetes.default", "kubernetes.default.svc"} | |
| 35 | +_BLOCKED_NETWORKS = [ipaddress.ip_network(n) for n in ( | |
| 36 | + "0.0.0.0/8", # unspecified / "this network" | |
| 37 | + "10.0.0.0/8", # RFC1918 | |
| 38 | + "100.64.0.0/10", # CGNAT | |
| 39 | + "127.0.0.0/8", # loopback | |
| 40 | + "169.254.0.0/16", # link-local (cloud metadata endpoints 169.254.169.254 …) | |
| 41 | + "172.16.0.0/12", # RFC1918 | |
| 42 | + "192.0.0.0/24", # IETF protocol assignments | |
| 43 | + "192.168.0.0/16", # RFC1918 | |
| 44 | + "198.18.0.0/15", # benchmarking | |
| 45 | + "240.0.0.0/4", # reserved + broadcast | |
| 46 | + "::/128", # unspecified | |
| 47 | + "::1/128", # loopback | |
| 48 | + "fc00::/7", # ULA | |
| 49 | + "fe80::/10", # link-local | |
| 50 | + "::ffff:0:0/96", # IPv4-mapped (checked again as IPv4 below) | |
| 51 | + "64:ff9b::/96", # NAT64 | |
| 52 | +)] | |
| 53 | + | |
| 54 | + | |
| 55 | +class BlockedDestination(ValueError): | |
| 56 | + """The URL points at a private, local or otherwise non-public destination.""" | |
| 57 | + | |
| 58 | + | |
| 59 | +def _ip_blocked(ip: ipaddress.IPv4Address | ipaddress.IPv6Address) -> bool: | |
| 60 | + if isinstance(ip, ipaddress.IPv6Address) and ip.ipv4_mapped is not None: | |
| 61 | + ip = ip.ipv4_mapped | |
| 62 | + if ip.is_private or ip.is_loopback or ip.is_link_local or ip.is_unspecified or ip.is_reserved or ip.is_multicast: | |
| 63 | + return True | |
| 64 | + return any(ip in net for net in _BLOCKED_NETWORKS) | |
| 65 | + | |
| 66 | + | |
| 67 | +def validate_destination(url: str, *, resolved_ips: list[str] | None = None) -> None: | |
| 68 | + """Raise `BlockedDestination` for non-http(s) schemes, local/internal hostnames and private, loopback, link-local, CGNAT, | |
| 69 | + IPv6 ULA/link-local or unspecified addresses. `resolved_ips` lets callers (and tests) inject the DNS answer.""" | |
| 70 | + p = urlparse(url.strip()) | |
| 71 | + if p.scheme.lower() not in ("http", "https"): | |
| 72 | + raise BlockedDestination(f"blocked destination: scheme {p.scheme!r} for {url}") | |
| 73 | + host = (p.hostname or "").strip().lower().rstrip(".") | |
| 74 | + if not host: | |
| 75 | + raise BlockedDestination(f"blocked destination: no host in {url}") | |
| 76 | + if p.username or p.password: | |
| 77 | + raise BlockedDestination(f"blocked destination: credentials in URL {url}") | |
| 78 | + if host in _BLOCKED_HOSTS or host.endswith(_BLOCKED_SUFFIXES) or "." not in host and not _is_ip_literal(host): | |
| 79 | + raise BlockedDestination(f"blocked destination: local host {host!r}") | |
| 80 | + literal = _parse_ip(host) | |
| 81 | + if literal is not None: | |
| 82 | + if _ip_blocked(literal): | |
| 83 | + raise BlockedDestination(f"blocked destination: non-public address {host}") | |
| 84 | + return | |
| 85 | + ips = resolved_ips if resolved_ips is not None else _resolve(host, p.port or (443 if p.scheme.lower() == "https" else 80)) | |
| 86 | + for raw in ips: | |
| 87 | + ip = _parse_ip(raw) | |
| 88 | + if ip is not None and _ip_blocked(ip): | |
| 89 | + raise BlockedDestination(f"blocked destination: {host} resolves to non-public address {raw}") | |
| 90 | + | |
| 91 | + | |
| 92 | +def _is_ip_literal(host: str) -> bool: | |
| 93 | + return _parse_ip(host) is not None | |
| 94 | + | |
| 95 | + | |
| 96 | +def _parse_ip(host: str) -> ipaddress.IPv4Address | ipaddress.IPv6Address | None: | |
| 97 | + try: | |
| 98 | + return ipaddress.ip_address(host.strip("[]").split("%")[0]) | |
| 99 | + except ValueError: | |
| 100 | + return None | |
| 101 | + | |
| 102 | + | |
| 103 | +def _resolve(host: str, port: int) -> list[str]: | |
| 104 | + try: | |
| 105 | + infos = socket.getaddrinfo(host, port, proto=socket.IPPROTO_TCP) | |
| 106 | + except socket.gaierror: | |
| 107 | + return [] # unresolvable hosts fail later with a transport error, not as SSRF | |
| 108 | + return sorted({info[4][0] for info in infos}) | |
| 109 | + | |
| 110 | + | |
| 111 | +async def validate_destination_async(url: str) -> None: | |
| 112 | + """Non-blocking wrapper (DNS resolution runs in a thread).""" | |
| 113 | + p = urlparse(url.strip()) | |
| 114 | + host = (p.hostname or "").strip().lower() | |
| 115 | + if not host or _parse_ip(host) is not None: | |
| 116 | + validate_destination(url, resolved_ips=[]) | |
| 117 | + return | |
| 118 | + validate_destination(url, resolved_ips=[]) # scheme / hostname rules first (no DNS) | |
| 119 | + ips = await asyncio.to_thread(_resolve, host, p.port or (443 if p.scheme.lower() == "https" else 80)) | |
| 120 | + validate_destination(url, resolved_ips=ips) | |
| 25 | 121 | |
| 26 | 122 | |
| 27 | 123 | class FetchError(Exception): |
@@ -138,7 +234,7 @@ class _Robots: | ||
| 138 | 234 | if cached is None or cached[0] < time.monotonic(): |
| 139 | 235 | rp: RobotFileParser | None = RobotFileParser() |
| 140 | 236 | try: |
| 141 | − r = await client.get(f"{key}/robots.txt", timeout=15, follow_redirects=True) | |
| 237 | + r = await client.get(f"{key}/robots.txt", timeout=15, follow_redirects=False) # host already validated by the caller | |
| 142 | 238 | if r.status_code == 200 and len(r.content) < 512 * 1024: |
| 143 | 239 | rp.parse(r.text.splitlines()) # type: ignore[union-attr] |
| 144 | 240 | else: |
@@ -170,7 +266,8 @@ class Fetcher: | ||
| 170 | 266 | self._http2 = http2 |
| 171 | 267 | |
| 172 | 268 | async def __aenter__(self) -> Fetcher: |
| 173 | − self._client = httpx.AsyncClient(headers=self.headers, timeout=httpx.Timeout(self.timeout_s, connect=20), follow_redirects=True, | |
| 269 | + # redirects are followed manually in `get()` so that every hop goes through the SSRF guard | |
| 270 | + self._client = httpx.AsyncClient(headers=self.headers, timeout=httpx.Timeout(self.timeout_s, connect=20), follow_redirects=False, | |
| 174 | 271 | http2=self._http2, limits=httpx.Limits(max_connections=16, max_keepalive_connections=8)) |
| 175 | 272 | return self |
| 176 | 273 | |
@@ -187,9 +284,13 @@ class Fetcher: | ||
| 187 | 284 | |
| 188 | 285 | async def get(self, url: str, *, etag: str | None = None, last_modified: str | None = None, min_bytes: int = 1, |
| 189 | 286 | escalate: bool = False, retries: int = 2, accept: str | None = None, rate_per_min: int | None = None) -> FetchResult: |
| 190 | − """Direct fetch with conditional headers. Raises NotModified (304), FetchError, BlockedError.""" | |
| 287 | + """Direct fetch with conditional headers. Raises NotModified (304), FetchError, BlockedError. | |
| 288 | + Every destination and every redirect hop (max 5) is validated against the SSRF guard before a connection is opened.""" | |
| 191 | 289 | client = self.client |
| 192 | − domain = domain_of(url) | |
| 290 | + try: | |
| 291 | + await validate_destination_async(url) | |
| 292 | + except BlockedDestination as exc: | |
| 293 | + raise FetchError(f"blocked destination: {exc}", url=url) from exc | |
| 193 | 294 | if self.robots and not await _robots.allowed(client, url): |
| 194 | 295 | raise BlockedError(f"robots.txt disallows {url}", status=None, url=url) |
| 195 | 296 | headers: dict[str, str] = {} |
@@ -200,13 +301,34 @@ class Fetcher: | ||
| 200 | 301 | if accept: |
| 201 | 302 | headers["Accept"] = accept |
| 202 | 303 | last_exc: Exception | None = None |
| 304 | + current = url | |
| 305 | + hops = 0 | |
| 203 | 306 | for attempt in range(retries + 1): |
| 204 | − await _limiter.wait(domain, rate_per_min or self.rate_per_min) | |
| 307 | + await _limiter.wait(domain_of(current), rate_per_min or self.rate_per_min) | |
| 205 | 308 | t0 = time.perf_counter() |
| 206 | 309 | try: |
| 207 | − async with client.stream("GET", url, headers=headers) as r: | |
| 310 | + async with client.stream("GET", current, headers=headers) as r: | |
| 208 | 311 | if r.status_code == 304: |
| 209 | 312 | raise NotModified() |
| 313 | + if r.status_code in REDIRECT_STATUS: | |
| 314 | + location = r.headers.get("location") | |
| 315 | + if not location: | |
| 316 | + raise FetchError(f"http {r.status_code} without Location for {current}", status=r.status_code, url=url) | |
| 317 | + hops += 1 | |
| 318 | + if hops > MAX_REDIRECTS: | |
| 319 | + raise FetchError(f"too many redirects (> {MAX_REDIRECTS}) for {url}", status=r.status_code, url=url) | |
| 320 | + nxt = urljoin(current, location) | |
| 321 | + try: | |
| 322 | + await validate_destination_async(nxt) | |
| 323 | + except BlockedDestination as exc: | |
| 324 | + raise FetchError(f"blocked destination: redirect {current} → {nxt}: {exc}", status=r.status_code, url=url) from exc | |
| 325 | + if domain_of(nxt) != domain_of(current): | |
| 326 | + headers.pop("If-None-Match", None) # validators belong to the original resource only | |
| 327 | + headers.pop("If-Modified-Since", None) | |
| 328 | + current = nxt | |
| 329 | + # a redirect is not a retry: loop again without consuming the attempt budget | |
| 330 | + return await self._follow(client, current, headers, url=url, hops=hops, min_bytes=min_bytes, escalate=escalate, | |
| 331 | + retries=retries, rate_per_min=rate_per_min) | |
| 210 | 332 | if r.status_code in TRANSIENT_STATUS and attempt < retries: |
| 211 | 333 | retry_after = r.headers.get("retry-after") |
| 212 | 334 | delay = min(60.0, float(retry_after)) if retry_after and retry_after.isdigit() else 2.0 * (2 ** attempt) |
@@ -219,22 +341,7 @@ class Fetcher: | ||
| 219 | 341 | raise BlockedError(f"http {r.status_code} for {url}", status=r.status_code, url=url) |
| 220 | 342 | if r.status_code >= 400: |
| 221 | 343 | raise FetchError(f"http {r.status_code} for {url}", status=r.status_code, url=url) |
| 222 | − chunks: list[bytes] = [] | |
| 223 | − size = 0 | |
| 224 | − async for chunk in r.aiter_bytes(): | |
| 225 | − size += len(chunk) | |
| 226 | − if size > settings.max_body_bytes: | |
| 227 | − raise FetchError(f"body exceeds {settings.max_body_bytes} bytes", status=r.status_code, url=url) | |
| 228 | − chunks.append(chunk) | |
| 229 | − content = b"".join(chunks) | |
| 230 | − if len(content) < min_bytes: | |
| 231 | − raise FetchError(f"suspiciously short body ({len(content)} bytes) for {url}", status=r.status_code, url=url) | |
| 232 | − res = FetchResult(url=url, final_url=str(r.url), status=r.status_code, headers={k.lower(): v for k, v in r.headers.items()}, | |
| 233 | − content=content, content_type=r.headers.get("content-type", "").lower(), fetched_at=datetime.now(UTC), | |
| 234 | − duration_ms=int((time.perf_counter() - t0) * 1000), transport="direct") | |
| 235 | − if escalate and res.is_html and _looks_like_challenge(content): | |
| 236 | − return await self._escalate(url, reason="anti-bot challenge page") | |
| 237 | − return res | |
| 344 | + return await self._read(r, url=url, final_url=current, t0=t0, min_bytes=min_bytes, escalate=escalate) | |
| 238 | 345 | except NotModified: |
| 239 | 346 | raise |
| 240 | 347 | except (BlockedError, FetchError): |
@@ -247,6 +354,61 @@ class Fetcher: | ||
| 247 | 354 | raise FetchError(f"{exc.__class__.__name__}: {exc}", url=url) from exc |
| 248 | 355 | raise FetchError(f"fetch failed: {last_exc}", url=url) |
| 249 | 356 | |
| 357 | + async def _follow(self, client: httpx.AsyncClient, current: str, headers: dict[str, str], *, url: str, hops: int, min_bytes: int, | |
| 358 | + escalate: bool, retries: int, rate_per_min: int | None) -> FetchResult: | |
| 359 | + """Continue a redirect chain: each hop is validated, capped at MAX_REDIRECTS, rate-limited per domain.""" | |
| 360 | + while True: | |
| 361 | + await _limiter.wait(domain_of(current), rate_per_min or self.rate_per_min) | |
| 362 | + t0 = time.perf_counter() | |
| 363 | + try: | |
| 364 | + async with client.stream("GET", current, headers=headers) as r: | |
| 365 | + if r.status_code in REDIRECT_STATUS: | |
| 366 | + location = r.headers.get("location") | |
| 367 | + if not location: | |
| 368 | + raise FetchError(f"http {r.status_code} without Location for {current}", status=r.status_code, url=url) | |
| 369 | + hops += 1 | |
| 370 | + if hops > MAX_REDIRECTS: | |
| 371 | + raise FetchError(f"too many redirects (> {MAX_REDIRECTS}) for {url}", status=r.status_code, url=url) | |
| 372 | + nxt = urljoin(current, location) | |
| 373 | + try: | |
| 374 | + await validate_destination_async(nxt) | |
| 375 | + except BlockedDestination as exc: | |
| 376 | + raise FetchError(f"blocked destination: redirect {current} → {nxt}: {exc}", status=r.status_code, url=url) from exc | |
| 377 | + if domain_of(nxt) != domain_of(current): | |
| 378 | + headers.pop("If-None-Match", None) | |
| 379 | + headers.pop("If-Modified-Since", None) | |
| 380 | + current = nxt | |
| 381 | + continue | |
| 382 | + if r.status_code == 304: | |
| 383 | + raise NotModified() | |
| 384 | + if r.status_code in BLOCK_STATUS: | |
| 385 | + if escalate: | |
| 386 | + return await self._escalate(url, reason=f"http {r.status_code}") | |
| 387 | + raise BlockedError(f"http {r.status_code} for {current}", status=r.status_code, url=url) | |
| 388 | + if r.status_code >= 400: | |
| 389 | + raise FetchError(f"http {r.status_code} for {current}", status=r.status_code, url=url) | |
| 390 | + return await self._read(r, url=url, final_url=current, t0=t0, min_bytes=min_bytes, escalate=escalate) | |
| 391 | + except (httpx.TimeoutException, httpx.TransportError) as exc: | |
| 392 | + raise FetchError(f"{exc.__class__.__name__}: {exc}", url=url) from exc | |
| 393 | + | |
| 394 | + async def _read(self, r: httpx.Response, *, url: str, final_url: str, t0: float, min_bytes: int, escalate: bool) -> FetchResult: | |
| 395 | + chunks: list[bytes] = [] | |
| 396 | + size = 0 | |
| 397 | + async for chunk in r.aiter_bytes(): | |
| 398 | + size += len(chunk) | |
| 399 | + if size > settings.max_body_bytes: | |
| 400 | + raise FetchError(f"body exceeds {settings.max_body_bytes} bytes", status=r.status_code, url=url) | |
| 401 | + chunks.append(chunk) | |
| 402 | + content = b"".join(chunks) | |
| 403 | + if len(content) < min_bytes: | |
| 404 | + raise FetchError(f"suspiciously short body ({len(content)} bytes) for {url}", status=r.status_code, url=url) | |
| 405 | + res = FetchResult(url=url, final_url=final_url, status=r.status_code, headers={k.lower(): v for k, v in r.headers.items()}, | |
| 406 | + content=content, content_type=r.headers.get("content-type", "").lower(), fetched_at=datetime.now(UTC), | |
| 407 | + duration_ms=int((time.perf_counter() - t0) * 1000), transport="direct") | |
| 408 | + if escalate and res.is_html and _looks_like_challenge(content): | |
| 409 | + return await self._escalate(url, reason="anti-bot challenge page") | |
| 410 | + return res | |
| 411 | + | |
| 250 | 412 | # ---------------------------------------------------------------------------------------------- escalation (optional) |
| 251 | 413 | async def _escalate(self, url: str, *, reason: str) -> FetchResult: |
| 252 | 414 | log.warning("escalating fetch", extra={"url": url, "reason": reason}) |
@@ -326,4 +488,5 @@ def file_result(path: str, *, url: str, content_type: str = "application/octet-s | ||
| 326 | 488 | _limiter = _RateLimiter() |
| 327 | 489 | _robots = _Robots() |
| 328 | 490 | |
| 329 | −__all__ = ["BlockedError", "FetchError", "FetchResult", "Fetcher", "NotModified", "canonicalize_url", "domain_of", "file_result"] | |
| 491 | +__all__ = ["MAX_REDIRECTS", "BlockedDestination", "BlockedError", "FetchError", "FetchResult", "Fetcher", "NotModified", "canonicalize_url", "domain_of", | |
| 492 | + "file_result", "validate_destination", "validate_destination_async"] | |
modified
src/aiatlas/sdk/resolution.py
+124 −11
@@ -3,11 +3,19 @@ | ||
| 3 | 3 | 1. identifiers (scheme, value) → exact entity |
| 4 | 4 | 2. normalized alias within the same entity type (disambiguated by organization when several match) |
| 5 | 5 | 3. slug collision within the same type |
| 6 | − 4. otherwise create, and park ambiguous cases in the review queue as merge candidates | |
| 6 | + 4. evaluation-effort variants of an existing model ("gpt-5-4-mini-medium") fold onto the canonical model — they are a result | |
| 7 | + configuration, never an entity of their own | |
| 8 | + 5. otherwise create, and park ambiguous cases in the review queue as merge candidates | |
| 9 | + | |
| 10 | +Guards: persisted `resolution_decisions` (keep_separate) are honoured — two entities an operator kept apart are never re-fused by an | |
| 11 | +alias match; aliases whose normalisation collapses digits and dots ("Qwen3-8B" ≡ "Qwen 38B") additionally require the same | |
| 12 | +`variant_key`; `model` and `artifact` are compatible types for lookups (an artifact used to be typed model — the same hub repo must | |
| 13 | +keep resolving to the same row). | |
| 7 | 14 | """ |
| 8 | 15 | from __future__ import annotations |
| 9 | 16 | |
| 10 | 17 | import logging |
| 18 | +import re | |
| 11 | 19 | from datetime import UTC, datetime |
| 12 | 20 | from typing import Any |
| 13 | 21 | |
@@ -15,6 +23,7 @@ from sqlalchemy.ext.asyncio import AsyncConnection | ||
| 15 | 23 | |
| 16 | 24 | from aiatlas.db import execute, fetch_all, fetch_one, jsonb |
| 17 | 25 | from aiatlas.ids import ENTITY_TYPES, new_id, normalize_alias, slugify |
| 26 | +from aiatlas.ontology.models import analyze_model_name, base_name, variant_key | |
| 18 | 27 | from aiatlas.sdk.facts import EntityRef |
| 19 | 28 | |
| 20 | 29 | log = logging.getLogger(__name__) |
@@ -22,16 +31,28 @@ log = logging.getLogger(__name__) | ||
| 22 | 31 | # Types whose slug should be prefixed by the organization slug to stay unique and readable (models: `qwen-qwen3-8b` is ugly → |
| 23 | 32 | # we keep the model name and only prefix on collision). |
| 24 | 33 | GENERIC_NAMES = {"model", "models", "api", "pricing", "docs", "blog", "news", "research", "overview"} |
| 34 | +# lookup-compatible type groups: a ref of one type may resolve to a stored entity of another type in the same group | |
| 35 | +_COMPATIBLE: dict[str, tuple[str, ...]] = {"model": ("model", "artifact"), "artifact": ("artifact", "model")} | |
| 36 | +# names whose normalised alias is ambiguous: a digit, a separator, a digit ("Qwen3-8B" / "Qwen 38B" → "qwen38b") | |
| 37 | +_DIGIT_SEP_DIGIT = re.compile(r"\d[.\-\s_]\d") | |
| 38 | +# identifier schemes issued by evaluators: an entity known ONLY through these may be an evaluation configuration rather than a model | |
| 39 | +EVALUATOR_SCHEMES = frozenset({"artificial_analysis", "livebench_model_id", "aider_model", "lmarena"}) | |
| 40 | + | |
| 41 | + | |
| 42 | +def compatible_types(entity_type: str) -> tuple[str, ...]: | |
| 43 | + return _COMPATIBLE.get(entity_type, (entity_type,)) | |
| 25 | 44 | |
| 26 | 45 | |
| 27 | 46 | class Resolver: |
| 28 | − def __init__(self, conn: AsyncConnection, *, snapshot_id: str | None = None, source_tier: int = 2): | |
| 47 | + def __init__(self, conn: AsyncConnection, *, snapshot_id: str | None = None, source_tier: int = 2, variant_index: dict[str, str] | None = None): | |
| 29 | 48 | self.conn = conn |
| 30 | 49 | self.snapshot_id = snapshot_id |
| 31 | 50 | self.tier = source_tier |
| 51 | + self.variant_index = variant_index # optional precomputed variant_key → canonical model id (canonicalization) | |
| 32 | 52 | self._cache: dict[str, str] = {} |
| 33 | 53 | self.created: list[str] = [] |
| 34 | 54 | self.updated: set[str] = set() |
| 55 | + self.folded: dict[str, tuple[str, dict[str, str]]] = {} # ref key → (canonical id, effort config) for folded variants | |
| 35 | 56 | |
| 36 | 57 | async def resolve(self, ref: EntityRef, *, create: bool = True) -> str | None: |
| 37 | 58 | if ref.id: |
@@ -51,6 +72,12 @@ class Resolver: | ||
| 51 | 72 | found = await self._by_alias(ref, org_id) |
| 52 | 73 | if found is None: |
| 53 | 74 | found = await self._by_slug(ref) |
| 75 | + if found is None and ref.entity_type == "model" and self.tier >= 2 and set(ref.identifiers) <= EVALUATOR_SCHEMES and not ref.attributes.get("hf_repo"): | |
| 76 | + # only evaluator-only references fold: a ref carrying official/hub identifiers (or a tier-1 source) names a real model | |
| 77 | + variant = await self.resolve_variant(ref, org_id=org_id) | |
| 78 | + if variant: | |
| 79 | + found = variant[0] | |
| 80 | + self.folded[key] = variant | |
| 54 | 81 | if found is None and not create: |
| 55 | 82 | return None |
| 56 | 83 | if found is None: |
@@ -64,12 +91,13 @@ class Resolver: | ||
| 64 | 91 | |
| 65 | 92 | # ---------------------------------------------------------------------------------------------- lookups |
| 66 | 93 | async def _by_identifiers(self, ref: EntityRef) -> str | None: |
| 94 | + types = compatible_types(ref.entity_type) | |
| 67 | 95 | for scheme, value in ref.identifiers.items(): |
| 68 | 96 | row = await fetch_one(self.conn, """select ei.entity_id, e.entity_type, e.merged_into from entity_identifiers ei |
| 69 | 97 | join entities e on e.id = ei.entity_id where ei.scheme = :s and ei.value = :v""", |
| 70 | 98 | s=scheme, v=str(value)) |
| 71 | 99 | if row: |
| 72 | − if row["entity_type"] != ref.entity_type: | |
| 100 | + if row["entity_type"] not in types: | |
| 73 | 101 | log.warning("identifier type mismatch", extra={"scheme": scheme, "value": value, "have": row["entity_type"], "want": ref.entity_type}) |
| 74 | 102 | continue |
| 75 | 103 | return row["merged_into"] or row["entity_id"] |
@@ -89,12 +117,27 @@ class Resolver: | ||
| 89 | 117 | norms = {normalize_alias(n) for n in names if n and normalize_alias(n)} |
| 90 | 118 | if not norms: |
| 91 | 119 | return None |
| 92 | − rows = await fetch_all(self.conn, """select distinct e.id, e.organization_id, e.canonical_name, e.merged_into from entity_aliases a | |
| 120 | + rows = await fetch_all(self.conn, """select distinct e.id, e.organization_id, e.canonical_name, e.merged_into, a.alias from entity_aliases a | |
| 93 | 121 | join entities e on e.id = a.entity_id |
| 94 | − where a.alias_norm = any(cast(:norms as text[])) and e.entity_type = :t""", norms=list(norms), t=ref.entity_type) | |
| 122 | + where a.alias_norm = any(cast(:norms as text[])) and e.entity_type = any(cast(:types as text[]))""", | |
| 123 | + norms=list(norms), types=list(compatible_types(ref.entity_type))) | |
| 95 | 124 | if not rows: |
| 96 | 125 | return None |
| 97 | 126 | rows = [{**r, "id": r["merged_into"] or r["id"]} for r in rows] |
| 127 | + # collision safety: when the normalised alias erases digit separators, the candidate must share the ref's variant key | |
| 128 | + # (or match one of the ref's names textually) — "Qwen 38B" must not land on "Qwen3-8B" | |
| 129 | + if ref.entity_type in ("model", "artifact") and any(_DIGIT_SEP_DIGIT.search(n) for n in names if n): | |
| 130 | + lower_names = {n.strip().lower() for n in names if n} | |
| 131 | + ref_vk = variant_key(ref.name) | |
| 132 | + safe = [] | |
| 133 | + for r in rows: | |
| 134 | + if (r["alias"] or "").strip().lower() in lower_names or variant_key(r["canonical_name"]) == ref_vk: | |
| 135 | + safe.append(r) | |
| 136 | + else: | |
| 137 | + log.info("alias collision refused", extra={"name": ref.name, "candidate": r["canonical_name"]}) | |
| 138 | + rows = safe | |
| 139 | + if not rows: | |
| 140 | + return None | |
| 98 | 141 | kept = [] |
| 99 | 142 | for r in rows: |
| 100 | 143 | if await self._conflicting_identifier(r["id"], ref): |
@@ -106,6 +149,13 @@ class Resolver: | ||
| 106 | 149 | if not rows: |
| 107 | 150 | return None |
| 108 | 151 | ids = {r["id"] for r in rows} |
| 152 | + # persisted decisions: a candidate an operator kept separate from another entity is only accepted with a matching organisation | |
| 153 | + separate = await self._kept_separate(sorted(ids)) | |
| 154 | + if separate: | |
| 155 | + rows = [r for r in rows if r["id"] not in separate or (org_id and r["organization_id"] == org_id)] | |
| 156 | + ids = {r["id"] for r in rows} | |
| 157 | + if not rows: | |
| 158 | + return None | |
| 109 | 159 | if len(ids) == 1: |
| 110 | 160 | return rows[0]["id"] |
| 111 | 161 | if org_id: |
@@ -117,22 +167,82 @@ class Resolver: | ||
| 117 | 167 | {"name": ref.name, "identifiers": ref.identifiers}) |
| 118 | 168 | return None |
| 119 | 169 | |
| 170 | + async def _kept_separate(self, ids: list[str]) -> set[str]: | |
| 171 | + """Ids among `ids` that carry a `keep_separate` decision with any entity.""" | |
| 172 | + if not ids: | |
| 173 | + return set() | |
| 174 | + rows = await fetch_all(self.conn, """select a_id, b_id from resolution_decisions where decision = 'keep_separate' | |
| 175 | + and (a_id = any(cast(:ids as text[])) or b_id = any(cast(:ids as text[])))""", ids=ids) | |
| 176 | + out: set[str] = set() | |
| 177 | + for r in rows: | |
| 178 | + out.update(x for x in (r["a_id"], r["b_id"]) if x in ids) | |
| 179 | + return out | |
| 180 | + | |
| 181 | + async def kept_separate(self, a: str, b: str) -> bool: | |
| 182 | + row = await fetch_one(self.conn, """select 1 from resolution_decisions where decision = 'keep_separate' | |
| 183 | + and ((a_id = :a and b_id = :b) or (a_id = :b and b_id = :a)) limit 1""", a=a, b=b) | |
| 184 | + return row is not None | |
| 185 | + | |
| 120 | 186 | async def _by_slug(self, ref: EntityRef) -> str | None: |
| 121 | 187 | slug = ref.slug_hint or slugify(ref.name) |
| 122 | 188 | row = await fetch_one(self.conn, "select id, entity_type, organization_id, merged_into from entities where slug = :s", s=slug) |
| 123 | − if row and row["entity_type"] == ref.entity_type and not await self._conflicting_identifier(row["merged_into"] or row["id"], ref): | |
| 189 | + if row and row["entity_type"] in compatible_types(ref.entity_type) and not await self._conflicting_identifier(row["merged_into"] or row["id"], ref): | |
| 124 | 190 | return row["merged_into"] or row["id"] |
| 125 | 191 | return None |
| 126 | 192 | |
| 193 | + async def resolve_variant(self, ref: EntityRef, *, org_id: str | None = None) -> tuple[str, dict[str, str]] | None: | |
| 194 | + """Evaluation-effort variant ("claude-opus-5-medium", "qwen3-6-27b-non-reasoning") → (canonical model id, effort config) when a | |
| 195 | + canonical model exists: same `artificial_analysis` identifier or alias as the base name, or the same `variant_key` | |
| 196 | + (through the optional precomputed index). Returns None when the base cannot be resolved or is ambiguous.""" | |
| 197 | + a = analyze_model_name(ref.name) | |
| 198 | + if not a.is_effort_variant: | |
| 199 | + return None | |
| 200 | + base = base_name(ref.name) | |
| 201 | + base_norms = {normalize_alias(base), normalize_alias(a.base_key)} - {""} | |
| 202 | + rows = await fetch_all(self.conn, """select distinct e.id, e.canonical_name, e.organization_id, e.merged_into from entities e | |
| 203 | + left join entity_identifiers ei on ei.entity_id = e.id | |
| 204 | + left join entity_aliases al on al.entity_id = e.id | |
| 205 | + where e.entity_type = 'model' and e.merged_into is null | |
| 206 | + and ((ei.scheme = 'artificial_analysis' and ei.value = any(cast(:bases as text[]))) | |
| 207 | + or al.alias_norm = any(cast(:norms as text[])) or e.slug = any(cast(:bases as text[])))""", | |
| 208 | + bases=sorted({base.lower(), a.base_key}), norms=sorted(base_norms)) | |
| 209 | + candidates = {r["id"]: r for r in rows if r["id"] != ref.id} | |
| 210 | + if self.variant_index: | |
| 211 | + vk = variant_key(base) | |
| 212 | + cid = self.variant_index.get(vk) | |
| 213 | + if cid and cid != ref.id and cid not in candidates: | |
| 214 | + row = await fetch_one(self.conn, "select id, canonical_name, organization_id, merged_into from entities where id = :id and merged_into is null", id=cid) | |
| 215 | + if row: | |
| 216 | + candidates[cid] = row | |
| 217 | + # never fold onto another effort variant | |
| 218 | + candidates = {k: v for k, v in candidates.items() if not analyze_model_name(v["canonical_name"]).is_effort_variant} | |
| 219 | + if not candidates: | |
| 220 | + return None | |
| 221 | + if len(candidates) > 1 and org_id: | |
| 222 | + same = {k: v for k, v in candidates.items() if v["organization_id"] == org_id} | |
| 223 | + if same: | |
| 224 | + candidates = same | |
| 225 | + if len(candidates) != 1: | |
| 226 | + await self._review("merge_candidate", sorted(candidates), f"effort variant '{ref.name}' matches {len(candidates)} base models", | |
| 227 | + {"name": ref.name, "base": base}) | |
| 228 | + return None | |
| 229 | + cid = next(iter(candidates)) | |
| 230 | + return cid, dict(a.effort) | |
| 231 | + | |
| 127 | 232 | # ---------------------------------------------------------------------------------------------- writes |
| 128 | 233 | async def _create(self, ref: EntityRef, org_id: str | None) -> str: |
| 129 | 234 | eid = new_id(ref.entity_type) |
| 130 | 235 | slug = await self._unique_slug(ref, org_id) |
| 131 | − first_seen = datetime.now(UTC) | |
| 132 | − await execute(self.conn, """insert into entities (id, entity_type, canonical_name, slug, description, status, organization_id, first_seen_at, last_seen_at) | |
| 133 | − values (:id, :t, :n, :slug, :d, :status, :org, :fs, now())""", | |
| 236 | + now = datetime.now(UTC) | |
| 237 | + first_seen = now | |
| 238 | + if ref.first_seen_hint is not None: | |
| 239 | + hint = ref.first_seen_hint if ref.first_seen_hint.tzinfo else ref.first_seen_hint.replace(tzinfo=UTC) | |
| 240 | + first_seen = min(now, hint) | |
| 241 | + await execute(self.conn, """insert into entities (id, entity_type, canonical_name, slug, description, status, organization_id, first_seen_at, last_seen_at, | |
| 242 | + identity_confidence, artifact_kind) | |
| 243 | + values (:id, :t, :n, :slug, :d, :status, :org, :fs, now(), :ic, :ak)""", | |
| 134 | 244 | id=eid, t=ref.entity_type, n=ref.name.strip()[:300], slug=slug, d=(ref.description or None), status=ref.status or "active", |
| 135 | − org=org_id, fs=first_seen) | |
| 245 | + org=org_id, fs=first_seen, ic=ref.identity_confidence or "high", ak=ref.artifact_kind if ref.entity_type == "artifact" else None) | |
| 136 | 246 | self.created.append(eid) |
| 137 | 247 | return eid |
| 138 | 248 | |
@@ -145,6 +255,9 @@ class Resolver: | ||
| 145 | 255 | if ref.description: |
| 146 | 256 | sets.append("description = case when description is null or length(description) < 40 then :d else description end") |
| 147 | 257 | params["d"] = ref.description |
| 258 | + if ref.first_seen_hint is not None: | |
| 259 | + sets.append("first_seen_at = least(first_seen_at, :fs)") | |
| 260 | + params["fs"] = ref.first_seen_hint if ref.first_seen_hint.tzinfo else ref.first_seen_hint.replace(tzinfo=UTC) | |
| 148 | 261 | await execute(self.conn, f"update entities set {', '.join(sets)} where id = :id", **params) |
| 149 | 262 | self.updated.add(eid) |
| 150 | 263 | |
@@ -185,4 +298,4 @@ class Resolver: | ||
| 185 | 298 | on conflict (dedupe_key) do nothing""", id=new_id("review"), k=kind, ids=entity_ids, p=jsonb(payload), r=reason, d=dedupe) |
| 186 | 299 | |
| 187 | 300 | |
| 188 | −__all__ = ["Resolver"] | |
| 301 | +__all__ = ["EVALUATOR_SCHEMES", "Resolver", "compatible_types"] | |
modified
src/aiatlas/sdk/writer.py
+226 −69
@@ -3,9 +3,14 @@ | ||
| 3 | 3 | Temporal rules (per entity × property): |
| 4 | 4 | * no current claim → insert current claim, set attribute, NEW_* event for new entities |
| 5 | 5 | * same value → confirm (observed_at bumped), nothing else |
| 6 | + * same assertion, older encoding (taxonomy) → re-encode the current claim in place (canonical value + value_raw), never an event | |
| 6 | 7 | * different, source ≥ tier → supersede (valid_to = observed), insert new current, update attribute, CHANGE event for material props |
| 8 | + * different, same source AND same extractor → supersede (a source correcting itself); an LLM claim never supersedes a deterministic | |
| 9 | + claim from the same URL — it is stored as conflicting | |
| 7 | 10 | * different, worse source → store as `conflicting`, flag confidence, review-queue item — never overwrite |
| 8 | −Prices and benchmark results have their own append-only tables with the same close/open semantics. | |
| 11 | +Prices and benchmark results have their own append-only tables with the same close/open semantics; results additionally keep ONE | |
| 12 | +current row per (model, benchmark, metric, config_key): a newer run group closes the older rows so leaderboards show the latest run. | |
| 13 | +Every fact carries `run_id` (batch rollback) and every event carries `recorded_at`, `is_backfill` and `group_key`. | |
| 9 | 14 | """ |
| 10 | 15 | from __future__ import annotations |
| 11 | 16 | |
@@ -19,6 +24,10 @@ from sqlalchemy.ext.asyncio import AsyncConnection | ||
| 19 | 24 | |
| 20 | 25 | from aiatlas.db import execute, fetch_one, jsonb |
| 21 | 26 | from aiatlas.ids import new_id, normalize_alias |
| 27 | +from aiatlas.ontology import benchmarks as bench_ontology | |
| 28 | +from aiatlas.ontology.anomalies import check_price_movement, check_result | |
| 29 | +from aiatlas.ontology.models import effort_config | |
| 30 | +from aiatlas.ontology.taxonomy import TAXONOMY_PROPERTIES, normalize_property | |
| 22 | 31 | from aiatlas.sdk.facts import ( |
| 23 | 32 | EVENT_CATEGORY_BY_TYPE, |
| 24 | 33 | MATERIAL_PROPERTIES, |
@@ -29,13 +38,18 @@ from aiatlas.sdk.facts import ( | ||
| 29 | 38 | ResultObs, |
| 30 | 39 | ) |
| 31 | 40 | from aiatlas.sdk.resolution import Resolver |
| 41 | +from aiatlas.services.events import classify_backfill, group_key_for, importance_for | |
| 32 | 42 | |
| 33 | 43 | log = logging.getLogger(__name__) |
| 34 | 44 | |
| 35 | 45 | TIER_CONFIDENCE = {1: "high", 2: "medium", 3: "low", 4: "low"} |
| 36 | 46 | SOFT_PROPERTIES = {"description", "summary", "tagline", "availability_note", "abstract", "notes", "training_data_notes", "safety_notes", "hardware_requirements"} |
| 47 | +# derived / secondary encodings of another property: never their own change event | |
| 48 | +NO_EVENT_PROPERTIES = {"license_key"} | |
| 37 | 49 | NEW_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} | |
| 50 | + "tool": 1, "repository": 0, "release": 2, "regulation": 2, "incident": 2, "organization": 1, "researcher": 0, | |
| 51 | + "artifact": 0, "model_family": 1, "license": 0} | |
| 52 | +SILENT_NEW_TYPES = {"researcher", "country", "license", "quantization"} | |
| 39 | 53 | |
| 40 | 54 | |
| 41 | 55 | def _norm_value(v: Any) -> Any: |
@@ -55,6 +69,14 @@ def _same(a: Any, b: Any) -> bool: | ||
| 55 | 69 | return json.dumps(_norm_value(a), sort_keys=True, default=str) == json.dumps(_norm_value(b), sort_keys=True, default=str) |
| 56 | 70 | |
| 57 | 71 | |
| 72 | +def _cfg_hash(config: dict[str, Any] | None) -> str: | |
| 73 | + return hashlib.sha1(json.dumps(config or {}, sort_keys=True, default=str).encode()).hexdigest()[:12] | |
| 74 | + | |
| 75 | + | |
| 76 | +def result_dedupe_key(model_id: str, benchmark_id: str, config: dict[str, Any] | None, metric: str | None) -> str: | |
| 77 | + return f"{model_id}:{benchmark_id}:{_cfg_hash(config)}:{metric or ''}" | |
| 78 | + | |
| 79 | + | |
| 58 | 80 | class WriteStats: |
| 59 | 81 | def __init__(self) -> None: |
| 60 | 82 | self.entities_created = 0 |
@@ -65,6 +87,7 @@ class WriteStats: | ||
| 65 | 87 | self.prices = 0 |
| 66 | 88 | self.results = 0 |
| 67 | 89 | self.conflicts = 0 |
| 90 | + self.folded_variants = 0 | |
| 68 | 91 | |
| 69 | 92 | def as_dict(self) -> dict[str, int]: |
| 70 | 93 | return dict(self.__dict__) |
@@ -73,7 +96,8 @@ class WriteStats: | ||
| 73 | 96 | class FactWriter: |
| 74 | 97 | def __init__(self, conn: AsyncConnection, *, source_id: str | None, snapshot_id: str | None, source_url: str | None, |
| 75 | 98 | tier: int = 2, connector_name: str | None = None, extractor: str = "deterministic", extractor_version: str = "1", |
| 76 | − observed_at: datetime | None = None): | |
| 99 | + observed_at: datetime | None = None, run_id: str | None = None, source_key: str | None = None, is_first_run: bool = False, | |
| 100 | + variant_index: dict[str, str] | None = None): | |
| 77 | 101 | self.conn = conn |
| 78 | 102 | self.source_id = source_id |
| 79 | 103 | self.snapshot_id = snapshot_id |
@@ -83,14 +107,29 @@ class FactWriter: | ||
| 83 | 107 | self.extractor = extractor |
| 84 | 108 | self.extractor_version = extractor_version |
| 85 | 109 | self.observed_at = observed_at or datetime.now(UTC) |
| 86 | − self.resolver = Resolver(conn, snapshot_id=snapshot_id, source_tier=tier) | |
| 110 | + self.run_id = run_id | |
| 111 | + self._source_key = source_key | |
| 112 | + self.is_first_run = is_first_run # connector's first successful run → every event is backfill (initial corpus) | |
| 113 | + self.resolver = Resolver(conn, snapshot_id=snapshot_id, source_tier=tier, variant_index=variant_index) | |
| 87 | 114 | self.stats = WriteStats() |
| 88 | 115 | |
| 116 | + @property | |
| 117 | + def derived(self) -> bool: | |
| 118 | + """Derived writers (canonicalization) store disagreements as conflicting claims but never open review items.""" | |
| 119 | + return self.extractor == "derived" | |
| 120 | + | |
| 121 | + async def source_key(self) -> str | None: | |
| 122 | + if self._source_key is None and self.source_id: | |
| 123 | + row = await fetch_one(self.conn, "select key from sources where id = :id", id=self.source_id) | |
| 124 | + self._source_key = row["key"] if row else "" | |
| 125 | + return self._source_key or None | |
| 126 | + | |
| 89 | 127 | # ---------------------------------------------------------------------------------------------- entry point |
| 90 | 128 | async def write(self, facts: Facts) -> WriteStats: |
| 91 | 129 | new_before = len(self.resolver.created) |
| 92 | 130 | for ref in facts.entities: |
| 93 | 131 | await self.resolver.resolve(ref) |
| 132 | + await self._apply_hierarchy(ref) | |
| 94 | 133 | for prop, value in ref.attributes.items(): |
| 95 | 134 | await self.write_claim(ref, prop, value) |
| 96 | 135 | for c in facts.claims: |
@@ -111,8 +150,31 @@ class FactWriter: | ||
| 111 | 150 | await self._new_entity_event(eid) |
| 112 | 151 | self.stats.entities_created = len(self.resolver.created) |
| 113 | 152 | self.stats.entities_updated = len(self.resolver.updated - set(self.resolver.created)) |
| 153 | + self.stats.folded_variants = len(self.resolver.folded) | |
| 114 | 154 | return self.stats |
| 115 | 155 | |
| 156 | + # ---------------------------------------------------------------------------------------------- hierarchy hints | |
| 157 | + async def _apply_hierarchy(self, ref: EntityRef) -> None: | |
| 158 | + """Materialise `family` / `canonical` hints: entities.family_id + member_of_family, entities.canonical_id + artifact_of, identity_confidence.""" | |
| 159 | + if not ref.id: | |
| 160 | + return | |
| 161 | + if ref.family is not None: | |
| 162 | + if ref.family.entity_type != "model_family": | |
| 163 | + ref.family.entity_type = "model_family" | |
| 164 | + fid = await self.resolver.resolve(ref.family) | |
| 165 | + if fid and fid != ref.id: | |
| 166 | + await execute(self.conn, "update entities set family_id = :f where id = :e and family_id is distinct from :f", f=fid, e=ref.id) | |
| 167 | + await self.write_relation(ref, "member_of_family", ref.family) | |
| 168 | + if ref.canonical is not None: | |
| 169 | + cid = await self.resolver.resolve(ref.canonical) | |
| 170 | + if cid and cid != ref.id: | |
| 171 | + await execute(self.conn, """update entities set canonical_id = :c, artifact_kind = coalesce(cast(:k as text), artifact_kind) | |
| 172 | + where id = :e and (canonical_id is distinct from :c or (cast(:k as text) is not null and artifact_kind is null))""", | |
| 173 | + c=cid, k=ref.artifact_kind, e=ref.id) | |
| 174 | + await self.write_relation(ref, "artifact_of", ref.canonical, {"artifact_kind": ref.artifact_kind} if ref.artifact_kind else None) | |
| 175 | + if ref.identity_confidence: | |
| 176 | + await execute(self.conn, "update entities set identity_confidence = :ic where id = :e and identity_confidence <> :ic", ic=ref.identity_confidence, e=ref.id) | |
| 177 | + | |
| 116 | 178 | # ---------------------------------------------------------------------------------------------- claims |
| 117 | 179 | async def write_claim(self, ref: EntityRef, prop: str, value: Any, *, unit: str | None = None, confidence: str | None = None, |
| 118 | 180 | observed_at: datetime | None = None, effective_at: datetime | None = None, source_url: str | None = None) -> None: |
@@ -121,67 +183,114 @@ class FactWriter: | ||
| 121 | 183 | eid = await self.resolver.resolve(ref) |
| 122 | 184 | assert eid |
| 123 | 185 | value = _norm_value(value) |
| 186 | + value_raw: str | None = None | |
| 187 | + if prop in TAXONOMY_PROPERTIES: | |
| 188 | + value, value_raw, mappings = normalize_property(ref.entity_type, prop, value) | |
| 189 | + for domain, raw, canon in mappings: | |
| 190 | + await self._record_mapping(domain, raw, canon) | |
| 191 | + value = _norm_value(value) | |
| 124 | 192 | observed = observed_at or self.observed_at |
| 125 | 193 | conf = confidence or TIER_CONFIDENCE.get(self.tier, "medium") |
| 126 | 194 | url = source_url or self.source_url |
| 127 | 195 | noisy = prop.startswith(NOISY_PREFIXES) |
| 128 | − current = await fetch_one(self.conn, """select id, value, tier, source_url, observed_at from claims | |
| 196 | + current = await fetch_one(self.conn, """select id, value, value_raw, tier, source_url, observed_at, extractor from claims | |
| 129 | 197 | where entity_id = :e and property = :p and status = 'current' order by valid_from desc limit 1""", |
| 130 | 198 | e=eid, p=prop) |
| 131 | 199 | is_new_entity = eid in self.resolver.created |
| 200 | + await self._apply_claim(ref, eid, prop, value, value_raw, current, unit=unit, conf=conf, observed=observed, effective_at=effective_at, url=url, | |
| 201 | + noisy=noisy, is_new_entity=is_new_entity) | |
| 202 | + if prop == "license" and isinstance(value, str): | |
| 203 | + from aiatlas.ontology.licenses import LICENSES | |
| 204 | + | |
| 205 | + if value in LICENSES: # canonical licence key as its own property (filterable), never its own event | |
| 206 | + await self.write_claim(ref, "license_key", value, confidence=confidence, observed_at=observed_at, effective_at=effective_at, source_url=source_url) | |
| 207 | + | |
| 208 | + async def _apply_claim(self, ref: EntityRef, eid: str, prop: str, value: Any, value_raw: str | None, current: dict[str, Any] | None, *, unit: str | None, | |
| 209 | + conf: str, observed: datetime, effective_at: datetime | None, url: str | None, noisy: bool, is_new_entity: bool) -> None: | |
| 132 | 210 | 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) | |
| 211 | + await self._insert_claim(eid, prop, value, unit, conf, "current", observed, effective_at, url, value_raw) | |
| 212 | + await self._set_attribute(eid, prop, value, unit, conf, url, observed, value_raw) | |
| 135 | 213 | return |
| 136 | 214 | if _same(current["value"], value): |
| 137 | 215 | await execute(self.conn, "update claims set observed_at = greatest(observed_at, :o) where id = :id", o=observed, id=current["id"]) |
| 138 | 216 | if noisy: |
| 139 | − await self._set_attribute(eid, prop, value, unit, conf, url, observed) | |
| 217 | + await self._set_attribute(eid, prop, value, unit, conf, url, observed, value_raw) | |
| 218 | + return | |
| 219 | + if prop in TAXONOMY_PROPERTIES and _same(normalize_property(ref.entity_type, prop, current["value"])[0], value): | |
| 220 | + # same assertion in an older encoding ("apache-2.0" → "Apache-2.0"): re-encode in place, keep the source label, no event | |
| 221 | + raw_keep = current["value_raw"] or (current["value"] if isinstance(current["value"], str) else json.dumps(current["value"], ensure_ascii=False)) | |
| 222 | + raw_keep = raw_keep if raw_keep != value else None | |
| 223 | + await execute(self.conn, """update claims set value = cast(:v as jsonb), value_text = :vt, value_raw = :raw, observed_at = greatest(observed_at, :o) where id = :id""", | |
| 224 | + v=jsonb(value), vt=value[:2000] if isinstance(value, str) else None, raw=raw_keep, o=observed, id=current["id"]) | |
| 225 | + await self._set_attribute(eid, prop, value, unit, conf, url, observed, raw_keep, keep_provenance=True) | |
| 140 | 226 | return |
| 141 | 227 | same_source = bool(url and current["source_url"] == url) |
| 228 | + same_extractor = (current["extractor"] or "deterministic") == self.extractor | |
| 142 | 229 | if prop in SOFT_PROPERTIES and not same_source: |
| 143 | 230 | # soft text (descriptions, notes): first statement wins until *its own* source changes; never an event, never a conflict |
| 144 | 231 | return |
| 145 | − if self.tier <= (current["tier"] or 2) or same_source: | |
| 232 | + if self.tier <= (current["tier"] or 2) or (same_source and same_extractor): | |
| 146 | 233 | # supersede |
| 147 | 234 | 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: | |
| 235 | + await self._insert_claim(eid, prop, value, unit, conf, "current", observed, effective_at, url, value_raw) | |
| 236 | + await self._set_attribute(eid, prop, value, unit, conf, url, observed, value_raw) | |
| 237 | + if not noisy and not is_new_entity and prop not in SOFT_PROPERTIES and prop not in NO_EVENT_PROPERTIES: | |
| 151 | 238 | 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 += 1 | |
| 239 | + return | |
| 240 | + # the same disagreement from the same source/extractor is recorded once (re-observed, not re-inserted) | |
| 241 | + dup = await fetch_one(self.conn, """select id from claims where entity_id = :e and property = :p and status = 'conflicting' and value = cast(:v as jsonb) | |
| 242 | + and extractor = :ex and coalesce(source_url, '') = :url limit 1""", e=eid, p=prop, v=jsonb(value), ex=self.extractor, url=url or "") | |
| 243 | + if dup: | |
| 244 | + await execute(self.conn, "update claims set observed_at = greatest(observed_at, :o) where id = :id", o=observed, id=dup["id"]) | |
| 245 | + return | |
| 246 | + await self._insert_claim(eid, prop, value, unit, conf, "conflicting", observed, effective_at, url, value_raw) | |
| 247 | + await execute(self.conn, "update claims set confidence = 'conflicted' where id = :id", id=current["id"]) | |
| 248 | + await execute(self.conn, """update entities set quality = quality || jsonb_build_object('conflicts', coalesce((quality->>'conflicts')::int, 0) + 1) | |
| 249 | + where id = :id""", id=eid) | |
| 250 | + self.stats.conflicts += 1 | |
| 251 | + if not self.derived: | |
| 158 | 252 | dedupe = f"conflict:{eid}:{prop}:{hashlib.sha1(json.dumps(value, sort_keys=True, default=str).encode()).hexdigest()[:10]}" |
| 159 | 253 | await execute(self.conn, """insert into review_queue (id, kind, entity_ids, payload, reason, dedupe_key) |
| 160 | 254 | values (:id, 'conflict', :ids, cast(:p as jsonb), :r, :d) on conflict (dedupe_key) do nothing""", |
| 161 | 255 | 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}), | |
| 256 | + "claimed": value, "claimed_source": url, "tier": self.tier, "extractor": self.extractor}), | |
| 163 | 257 | r=f"source disagrees on {prop}", d=dedupe) |
| 164 | 258 | |
| 259 | + async def _record_mapping(self, domain: str, raw: str, canonical: str | None) -> None: | |
| 260 | + if canonical is not None and raw == canonical: | |
| 261 | + return # identity mappings carry no information | |
| 262 | + await execute(self.conn, """insert into taxonomy_mappings (domain, raw, canonical) values (:d, :r, :c) | |
| 263 | + on conflict (domain, raw) do update set count = taxonomy_mappings.count + 1, last_seen_at = now(), | |
| 264 | + canonical = coalesce(excluded.canonical, taxonomy_mappings.canonical)""", d=domain, r=raw[:300], c=canonical) | |
| 265 | + | |
| 165 | 266 | 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: | |
| 267 | + effective_at: datetime | None, url: str | None, value_raw: str | None = None) -> None: | |
| 167 | 268 | text_val = value if isinstance(value, str) else None |
| 168 | 269 | num_val = float(value) if isinstance(value, (int, float)) and not isinstance(value, bool) else None |
| 169 | 270 | 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)""", | |
| 271 | + confidence, status, extractor, extractor_version, observed_at, effective_at, valid_from, run_id, value_raw) | |
| 272 | + values (:id, :e, :p, cast(:v as jsonb), :vt, :vn, :u, :src, :snap, :url, :tier, :conf, :status, :ex, :exv, :o, :eff, :vf, :run, :raw)""", | |
| 172 | 273 | 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 | 274 | 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) | |
| 275 | + o=observed, eff=effective_at, vf=effective_at or observed, run=self.run_id, raw=value_raw[:500] if value_raw else None) | |
| 175 | 276 | self.stats.claims += 1 |
| 176 | 277 | |
| 177 | − async def _set_attribute(self, eid: str, prop: str, value: Any, unit: str | None, conf: str, url: str | None, observed: datetime) -> None: | |
| 278 | + async def _set_attribute(self, eid: str, prop: str, value: Any, unit: str | None, conf: str, url: str | None, observed: datetime, | |
| 279 | + value_raw: str | None = None, *, keep_provenance: bool = False) -> None: | |
| 178 | 280 | prov = {"source_id": self.source_id, "snapshot_id": self.snapshot_id, "url": url, "observed_at": observed.isoformat(timespec="seconds"), |
| 179 | 281 | "tier": self.tier, "confidence": conf, "extractor": self.extractor} |
| 180 | 282 | if unit: |
| 181 | 283 | prov["unit"] = unit |
| 182 | − 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) | |
| 284 | + attrs = {prop: value} | |
| 285 | + if value_raw is not None: | |
| 286 | + attrs[f"{prop}_raw"] = value_raw | |
| 287 | + if keep_provenance: | |
| 288 | + await execute(self.conn, """update entities set attributes = (attributes - cast(:praw as text)) || cast(:a as jsonb), last_seen_at = greatest(last_seen_at, :o) where id = :id""", | |
| 289 | + praw=f"{prop}_raw", a=jsonb(attrs), o=observed, id=eid) | |
| 290 | + else: | |
| 291 | + await execute(self.conn, """update entities set attributes = (attributes - cast(:praw as text)) || cast(:a as jsonb), | |
| 292 | + provenance = provenance || jsonb_build_object(cast(:p as text), cast(:prov as jsonb)), last_seen_at = greatest(last_seen_at, :o) | |
| 293 | + where id = :id""", praw=f"{prop}_raw", a=jsonb(attrs), p=prop, prov=jsonb(prov), o=observed, id=eid) | |
| 185 | 294 | if prop == "status" and isinstance(value, str): |
| 186 | 295 | await execute(self.conn, "update entities set status = :s where id = :id", s=value[:40], id=eid) |
| 187 | 296 | if prop == "description" and isinstance(value, str): |
@@ -192,9 +301,9 @@ class FactWriter: | ||
| 192 | 301 | event_type, importance = MATERIAL_PROPERTIES.get(prop, ("PROPERTY_CHANGED", 0)) |
| 193 | 302 | row = await fetch_one(self.conn, "select canonical_name, entity_type from entities where id = :id", id=eid) |
| 194 | 303 | name = row["canonical_name"] if row else eid |
| 195 | − 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 = 3 | |
| 304 | + etype = row["entity_type"] if row else "" | |
| 305 | + category = EVENT_CATEGORY_BY_TYPE.get(etype, "update") | |
| 306 | + importance = importance_for(event_type, etype, old, new, tier=self.tier, default=importance) | |
| 198 | 307 | summary = f"{name}: {prop.replace('_', ' ')} changed from {_short(old)} to {_short(new)}" |
| 199 | 308 | dedupe = f"{event_type}:{eid}:{prop}:{hashlib.sha1(json.dumps([old, new], sort_keys=True, default=str).encode()).hexdigest()[:12]}" |
| 200 | 309 | await self.emit_event(event_type, category, summary, entity_id=eid, old_value=old, new_value=new, importance=importance, |
@@ -217,10 +326,10 @@ class FactWriter: | ||
| 217 | 326 | else: |
| 218 | 327 | await execute(self.conn, "update relations set observed_at = :o where id = :id", o=self.observed_at, id=existing["id"]) |
| 219 | 328 | return |
| 220 | − 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)""", | |
| 329 | + 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, run_id) | |
| 330 | + values (:id, :s, :p, :o, cast(:a as jsonb), :src, :snap, :url, :tier, :conf, :obs, :obs, :run)""", | |
| 222 | 331 | 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) | |
| 332 | + url=source_url or self.source_url, tier=self.tier, conf=conf, obs=self.observed_at, run=self.run_id) | |
| 224 | 333 | self.stats.relations += 1 |
| 225 | 334 | |
| 226 | 335 | # ---------------------------------------------------------------------------------------------- prices |
@@ -242,16 +351,17 @@ class FactWriter: | ||
| 242 | 351 | await execute(self.conn, "update prices set observed_at = :o where id = :id", o=self.observed_at, id=current["id"]) |
| 243 | 352 | return |
| 244 | 353 | await execute(self.conn, "update prices set valid_to = :o where id = :id", o=self.observed_at, id=current["id"]) |
| 354 | + price_id = new_id("price") | |
| 245 | 355 | 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 | 356 | 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, | |
| 357 | + observed_at, valid_from, source_id, snapshot_id, source_url, tier, meta, run_id) | |
| 358 | + 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), :run)""", | |
| 359 | + id=price_id, 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 | 360 | 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 | 361 | 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)) | |
| 362 | + url=url, tier=self.tier, meta=jsonb(p.meta), run=self.run_id) | |
| 253 | 363 | self.stats.prices += 1 |
| 254 | − model = await fetch_one(self.conn, "select canonical_name from entities where id = :id", id=mid) | |
| 364 | + model = await fetch_one(self.conn, "select canonical_name, entity_type from entities where id = :id", id=mid) | |
| 255 | 365 | provider = await fetch_one(self.conn, "select canonical_name from entities where id = :id", id=pid) |
| 256 | 366 | mname = model["canonical_name"] if model else mid |
| 257 | 367 | pname = provider["canonical_name"] if provider else pid |
@@ -260,8 +370,14 @@ class FactWriter: | ||
| 260 | 370 | new = {"input_per_mtok": p.input_per_mtok, "output_per_mtok": p.output_per_mtok} |
| 261 | 371 | summary = f"{pname} changed pricing for {mname}: {_fmt_price(old)} → {_fmt_price(new)}" |
| 262 | 372 | 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, | |
| 373 | + importance = importance_for("PRICE_CHANGED", model["entity_type"] if model else None, old, new, tier=self.tier) | |
| 374 | + await self.emit_event("PRICE_CHANGED", "price", summary, entity_id=mid, old_value=old, new_value=new, importance=importance, dedupe_key=dedupe, | |
| 264 | 375 | source_url=url, meta={"provider_id": pid, "provider": pname}) |
| 376 | + for anomaly in check_price_movement({**old, "model_id": mid, "provider_id": pid}, {**new, "model_id": mid, "provider_id": pid}): | |
| 377 | + from aiatlas.services.anomalies import record | |
| 378 | + | |
| 379 | + anomaly.detail["price_id"] = price_id | |
| 380 | + await record(self.conn, anomaly) | |
| 265 | 381 | else: |
| 266 | 382 | new = {"input_per_mtok": p.input_per_mtok, "output_per_mtok": p.output_per_mtok} |
| 267 | 383 | summary = f"{pname} lists {mname} at {_fmt_price(new)}" |
@@ -277,51 +393,87 @@ class FactWriter: | ||
| 277 | 393 | if not mid or not bid: |
| 278 | 394 | return |
| 279 | 395 | url = r.source_url or self.source_url |
| 280 | − 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) | |
| 396 | + # evaluation-effort variants ("gpt-5-4-mini-medium") are a *configuration* of the canonical model: the effort dict lands in the | |
| 397 | + # config so folded results stay distinguishable and comparable (reasoning_effort is a condition key, not a task key) | |
| 398 | + config = effort_config(r.model.name, r.config) | |
| 399 | + metric = r.metric | |
| 400 | + dedupe = result_dedupe_key(mid, bid, config, metric) | |
| 401 | + config_key = bench_ontology.config_key(config, metric) | |
| 402 | + trust = r.trust_level or bench_ontology.trust_level(await self.source_key(), config, extractor=self.extractor) | |
| 403 | + variant = r.variant or bench_ontology.variant_from_config(config) | |
| 404 | + run_group = r.run_group or bench_ontology.run_group_from_config(config) | |
| 283 | 405 | conf = r.confidence or TIER_CONFIDENCE.get(self.tier, "medium") |
| 406 | + lo, hi = bench_ontology.metric_bounds(metric, r.unit) | |
| 407 | + out_of_range = (hi is not None and r.score > hi + 1e-9) or (lo is not None and r.score < lo - 1e-9) | |
| 408 | + if out_of_range: | |
| 409 | + conf = "low" | |
| 410 | + existing = await fetch_one(self.conn, "select id, score from benchmark_results where dedupe_key = :d", d=dedupe) | |
| 411 | + names: dict[str, Any] | None = None | |
| 284 | 412 | if existing: |
| 285 | 413 | 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", | |
| 414 | + await execute(self.conn, "update benchmark_results set valid_to = :o, is_current = false, dedupe_key = dedupe_key || ':' || :suffix where id = :id", | |
| 287 | 415 | 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}", | |
| 416 | + names = await self._names(mid, bid) | |
| 417 | + await self.emit_event("BENCHMARK_UPDATED", "benchmark", f"{names['model']} on {names['bench']}: {existing['score']:g} → {r.score:g}", | |
| 292 | 418 | entity_id=mid, old_value=existing["score"], new_value=r.score, importance=1, |
| 293 | 419 | dedupe_key=f"BENCHMARK_UPDATED:{dedupe}:{r.score:g}", source_url=url, meta={"benchmark_id": bid}) |
| 294 | 420 | else: |
| 295 | 421 | await execute(self.conn, "update benchmark_results set observed_at = :o where id = :id", o=self.observed_at, id=existing["id"]) |
| 296 | 422 | return |
| 423 | + rid = new_id("result") | |
| 297 | 424 | 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) | |
| 425 | + source_id, snapshot_id, source_url, tier, confidence, dedupe_key, config_key, trust_level, variant, run_group, is_current, extractor, run_id) | |
| 426 | + values (:id, :m, :b, :s, :metric, :unit, :hib, cast(:cfg as jsonb), :ev, :obs, :src, :snap, :url, :tier, :conf, :d, :ck, :trust, :variant, :rg, true, :ex, :run)""", | |
| 427 | + id=rid, m=mid, b=bid, s=r.score, metric=metric, unit=r.unit, hib=r.higher_is_better, cfg=jsonb(config), ev=r.evaluated_at, obs=self.observed_at, | |
| 428 | + src=self.source_id, snap=self.snapshot_id, url=url, tier=self.tier, conf=conf, d=dedupe, ck=config_key, trust=trust, variant=variant, rg=run_group, | |
| 429 | + ex=self.extractor, run=self.run_id) | |
| 302 | 430 | self.stats.results += 1 |
| 431 | + # one current row per (model, benchmark, metric, config_key): a different (older) run group is closed so the leaderboard shows the latest run | |
| 432 | + await execute(self.conn, """update benchmark_results set is_current = false, valid_to = coalesce(valid_to, :obs) | |
| 433 | + where model_id = :m and benchmark_id = :b and coalesce(metric, '') = :metric and config_key = :ck and id <> :id and is_current | |
| 434 | + and coalesce(run_group, '') <> :rg and observed_at <= :obs""", | |
| 435 | + obs=self.observed_at, m=mid, b=bid, metric=metric or "", ck=config_key, id=rid, rg=run_group or "") | |
| 436 | + if out_of_range: | |
| 437 | + from aiatlas.services.anomalies import record | |
| 438 | + | |
| 439 | + names = names or await self._names(mid, bid) | |
| 440 | + for anomaly in check_result({"id": rid, "model_id": mid, "benchmark_id": bid, "score": r.score, "metric": metric, "unit": r.unit, | |
| 441 | + "model_name": names["model"], "benchmark_name": names["bench"]}): | |
| 442 | + await record(self.conn, anomaly) | |
| 303 | 443 | 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}", | |
| 444 | + names = names or await self._names(mid, bid) | |
| 445 | + await self.emit_event("BENCHMARK_RESULT", "benchmark", f"{names['model']} scores {r.score:g}{r.unit or ''} on {names['bench']}", | |
| 308 | 446 | 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}) | |
| 447 | + meta={"benchmark_id": bid, "metric": metric, "config_key": config_key}) | |
| 310 | 448 | await self.write_relation(r.model, "evaluated_on", r.benchmark, source_url=url) |
| 311 | 449 | |
| 450 | + async def _names(self, mid: str, bid: str) -> dict[str, Any]: | |
| 451 | + model = await fetch_one(self.conn, "select canonical_name from entities where id = :id", id=mid) | |
| 452 | + bench = await fetch_one(self.conn, "select canonical_name from entities where id = :id", id=bid) | |
| 453 | + return {"model": model["canonical_name"] if model else mid, "bench": bench["canonical_name"] if bench else bid} | |
| 454 | + | |
| 312 | 455 | # ---------------------------------------------------------------------------------------------- events |
| 313 | 456 | async def emit_event(self, event_type: str, category: str, summary: str, *, entity_id: str | None = None, old_value: Any = None, |
| 314 | 457 | 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: | |
| 458 | + source_url: str | None = None, meta: dict[str, Any] | None = None, observed_at: datetime | None = None, | |
| 459 | + entity_first_seen: datetime | None = None, is_backfill: bool | None = None) -> None: | |
| 316 | 460 | dedupe = dedupe_key or f"{event_type}:{entity_id or ''}:{normalize_alias(summary)[:120]}" |
| 461 | + observed = observed_at or self.observed_at | |
| 462 | + if self.derived: | |
| 463 | + # a derived writer re-encodes what we already know: its events are bookkeeping, never news | |
| 464 | + is_backfill, importance = True, min(importance, 1) | |
| 465 | + backfill = is_backfill if is_backfill is not None else classify_backfill(event_type, effective_at, observed, is_first_run=self.is_first_run, | |
| 466 | + entity_first_seen=entity_first_seen) | |
| 467 | + group_key = group_key_for(event_type, entity_id, effective_at, observed) | |
| 317 | 468 | 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)) | |
| 469 | + effective_at, source_id, snapshot_id, source_url, connector_name, dedupe_key, meta, run_id, recorded_at, is_backfill, group_key) | |
| 470 | + 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), | |
| 471 | + :run, now(), :bf, :gk) | |
| 320 | 472 | on conflict (dedupe_key) do nothing""", |
| 321 | 473 | 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, | |
| 474 | + n=jsonb(new_value) if new_value is not None else None, s=summary[:500], imp=max(0, min(3, importance)), obs=observed, | |
| 323 | 475 | 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 {})) | |
| 476 | + meta=jsonb(meta or {}), run=self.run_id, bf=backfill, gk=group_key) | |
| 325 | 477 | self.stats.events += 1 |
| 326 | 478 | |
| 327 | 479 | async def _new_entity_event(self, eid: str) -> None: |
@@ -329,25 +481,30 @@ class FactWriter: | ||
| 329 | 481 | if not row: |
| 330 | 482 | return |
| 331 | 483 | etype = row["entity_type"] |
| 332 | − if etype in ("researcher", "country", "license", "quantization"): | |
| 484 | + if etype in SILENT_NEW_TYPES: | |
| 333 | 485 | return |
| 334 | − importance = NEW_IMPORTANCE.get(etype, 1) | |
| 335 | − if self.tier > 2: | |
| 336 | − importance = max(0, importance - 1) | |
| 337 | 486 | org = None |
| 487 | + org_models = 0 | |
| 338 | 488 | if row["organization_id"]: |
| 339 | − o = await fetch_one(self.conn, "select canonical_name from entities where id = :id", id=row["organization_id"]) | |
| 489 | + o = await fetch_one(self.conn, """select canonical_name, (select count(*) from entities m where m.organization_id = e.id and m.entity_type = 'model' | |
| 490 | + and m.merged_into is null) as models from entities e where e.id = :id""", id=row["organization_id"]) | |
| 340 | 491 | org = o["canonical_name"] if o else None |
| 492 | + org_models = int(o["models"] or 0) if o else 0 | |
| 493 | + attrs = row["attributes"] or {} | |
| 494 | + if etype == "model": | |
| 495 | + importance = importance_for("NEW_MODEL", etype, tier=self.tier, org_model_count=org_models, openness=attrs.get("openness")) | |
| 496 | + else: | |
| 497 | + importance = importance_for(f"NEW_{etype.upper()}", etype, tier=self.tier, default=NEW_IMPORTANCE.get(etype, 1)) | |
| 341 | 498 | label = etype.replace("_", " ") |
| 342 | 499 | summary = f"New {label}: {row['canonical_name']}" + (f" ({org})" if org else "") |
| 343 | 500 | effective = None |
| 344 | − rel = row["attributes"].get("release_date") or row["attributes"].get("published_at") | |
| 501 | + rel = attrs.get("release_date") or attrs.get("published_at") | |
| 345 | 502 | if isinstance(rel, str): |
| 346 | 503 | from aiatlas.sdk.extract.dates import parse_datetime |
| 347 | 504 | |
| 348 | 505 | effective = parse_datetime(rel) |
| 349 | 506 | 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) | |
| 507 | + dedupe_key=f"NEW_{etype.upper()}:{eid}", effective_at=effective, entity_first_seen=effective) | |
| 351 | 508 | |
| 352 | 509 | |
| 353 | 510 | def _short(v: Any) -> str: |
@@ -365,4 +522,4 @@ def _fmt_price(p: dict[str, Any]) -> str: | ||
| 365 | 522 | return " / ".join(parts) + " per 1M tokens" if parts else "n/a" |
| 366 | 523 | |
| 367 | 524 | |
| 368 | −__all__ = ["FactWriter", "WriteStats"] | |
| 525 | +__all__ = ["FactWriter", "WriteStats", "result_dedupe_key"] | |
added
src/aiatlas/services/anomalies.py
+82 −0
@@ -0,0 +1,82 @@ | ||
| 1 | +"""Anomaly persistence: the ontology's deterministic checks (`aiatlas.ontology.anomalies`) produce flags; this module upserts them in the | |
| 2 | +`anomalies` table (one row per dedupe key, reopened when a check fires again, auto-resolved when it stops firing). Never deletes, | |
| 3 | +never edits the flagged value.""" | |
| 4 | +from __future__ import annotations | |
| 5 | + | |
| 6 | +import logging | |
| 7 | +from typing import Any | |
| 8 | + | |
| 9 | +from sqlalchemy.ext.asyncio import AsyncConnection | |
| 10 | +from ulid import ULID | |
| 11 | + | |
| 12 | +from aiatlas.db import execute, fetch_all, fetch_one, jsonb | |
| 13 | +from aiatlas.ontology.anomalies import Anomaly, check_hardware, check_model, check_price, check_result | |
| 14 | + | |
| 15 | +log = logging.getLogger(__name__) | |
| 16 | + | |
| 17 | + | |
| 18 | +async def record(conn: AsyncConnection, anomaly: Anomaly) -> str: | |
| 19 | + """Upsert one anomaly. Rows an operator marked `ignored`/`fixed` keep their status; `resolved` rows reopen.""" | |
| 20 | + row = await fetch_one(conn, """insert into anomalies (id, entity_id, check_name, severity, message, value, detail, status, dedupe_key) | |
| 21 | + values (:id, :e, :c, :sev, :m, cast(:v as jsonb), cast(:d as jsonb), 'open', :k) | |
| 22 | + on conflict (dedupe_key) do update set last_seen_at = now(), message = excluded.message, value = excluded.value, | |
| 23 | + detail = excluded.detail, severity = excluded.severity, | |
| 24 | + status = case when anomalies.status = 'resolved' then 'open' else anomalies.status end, | |
| 25 | + resolved_at = case when anomalies.status = 'resolved' then null else anomalies.resolved_at end | |
| 26 | + returning id""", | |
| 27 | + id=f"anom_{ULID()}", e=anomaly.entity_id, c=anomaly.check, sev=anomaly.severity, m=anomaly.message[:1000], | |
| 28 | + v=jsonb(anomaly.value) if anomaly.value is not None else None, d=jsonb(anomaly.detail or {}), k=anomaly.dedupe_key[:400]) | |
| 29 | + return row["id"] if row else "" | |
| 30 | + | |
| 31 | + | |
| 32 | +async def run_checks(conn: AsyncConnection, *, resolve_stale: bool = True) -> dict[str, Any]: | |
| 33 | + """Run every ontology check over live models/artifacts, hardware, live prices and live results; upsert flags; auto-resolve open flags whose | |
| 34 | + check no longer fires. Returns counts by severity and the number resolved.""" | |
| 35 | + found: list[Anomaly] = [] | |
| 36 | + rows = await fetch_all(conn, "select id, canonical_name, attributes from entities where entity_type in ('model','artifact') and merged_into is null") | |
| 37 | + for r in rows: | |
| 38 | + found.extend(check_model(r["id"], r["canonical_name"], r["attributes"] or {})) | |
| 39 | + rows = await fetch_all(conn, "select id, canonical_name, attributes from entities where entity_type = 'hardware' and merged_into is null") | |
| 40 | + for r in rows: | |
| 41 | + found.extend(check_hardware(r["id"], r["canonical_name"], r["attributes"] or {})) | |
| 42 | + rows = await fetch_all(conn, """select p.*, m.canonical_name as model_name, v.canonical_name as provider_name from prices p | |
| 43 | + join entities m on m.id = p.model_id join entities v on v.id = p.provider_id where p.valid_to is null""") | |
| 44 | + for r in rows: | |
| 45 | + found.extend(check_price(r)) | |
| 46 | + rows = await fetch_all(conn, """select r.id, r.model_id, r.benchmark_id, r.score, r.metric, r.unit, r.evaluated_at, m.canonical_name as model_name, | |
| 47 | + b.canonical_name as benchmark_name, m.attributes->>'release_date' as model_release_date | |
| 48 | + from benchmark_results r join entities m on m.id = r.model_id join entities b on b.id = r.benchmark_id | |
| 49 | + where r.valid_to is null and r.is_current""") | |
| 50 | + for r in rows: | |
| 51 | + found.extend(check_result(r)) | |
| 52 | + keys: set[str] = set() | |
| 53 | + by_sev: dict[str, int] = {} | |
| 54 | + for a in found: | |
| 55 | + key = a.dedupe_key[:400] | |
| 56 | + if key in keys: | |
| 57 | + continue | |
| 58 | + keys.add(key) | |
| 59 | + await record(conn, a) | |
| 60 | + by_sev[a.severity] = by_sev.get(a.severity, 0) + 1 | |
| 61 | + resolved = 0 | |
| 62 | + if resolve_stale: | |
| 63 | + row = await fetch_one(conn, """with s as (update anomalies set status = 'resolved', resolved_at = now(), resolution = 'check no longer fires' | |
| 64 | + where status = 'open' and not (dedupe_key = any(cast(:keys as text[]))) returning 1) select count(*) as n from s""", | |
| 65 | + keys=sorted(keys)) | |
| 66 | + resolved = int(row["n"]) if row else 0 | |
| 67 | + return {"flagged": len(keys), "by_severity": by_sev, "resolved": resolved} | |
| 68 | + | |
| 69 | + | |
| 70 | +async def list_anomalies(conn: AsyncConnection, *, severity: str | None = None, status: str = "open", limit: int = 200) -> list[dict[str, Any]]: | |
| 71 | + return await fetch_all(conn, """select a.*, e.canonical_name, e.slug, e.entity_type from anomalies a left join entities e on e.id = a.entity_id | |
| 72 | + where (cast(:st as text) = '' or a.status = :st) and (cast(:sev as text) = '' or a.severity = :sev) | |
| 73 | + order by case a.severity when 'critical' then 0 when 'warning' then 1 else 2 end, a.last_seen_at desc limit :n""", | |
| 74 | + st=status or "", sev=severity or "", n=limit) | |
| 75 | + | |
| 76 | + | |
| 77 | +async def set_status(conn: AsyncConnection, anomaly_id: str, status: str, *, resolution: str | None = None) -> None: | |
| 78 | + await execute(conn, "update anomalies set status = :s, resolution = coalesce(:r, resolution), resolved_at = case when :s in ('resolved','fixed','ignored') then now() else null end where id = :id", | |
| 79 | + s=status, r=resolution, id=anomaly_id) | |
| 80 | + | |
| 81 | + | |
| 82 | +__all__ = ["list_anomalies", "record", "run_checks", "set_status"] | |
added
src/aiatlas/services/events.py
+128 −0
@@ -0,0 +1,128 @@ | ||
| 1 | +"""Event semantics — deterministic, database-free rules shared by the writer and `aia canonicalize events`. | |
| 2 | + | |
| 3 | +Three clocks on every change event: | |
| 4 | + occurred_at = coalesce(effective_at, observed_at) when the thing happened (release date, price effective date) | |
| 5 | + observed_at when a connector saw it | |
| 6 | + recorded_at when the row was written | |
| 7 | + | |
| 8 | +`is_backfill` separates *history being loaded* from *news*: the "+N in 24 h" counters, /changes and the pulse feed only count live | |
| 9 | +(non-backfill) events; timelines and /asof use every event. | |
| 10 | +""" | |
| 11 | +from __future__ import annotations | |
| 12 | + | |
| 13 | +from datetime import datetime, timedelta | |
| 14 | +from typing import Any | |
| 15 | + | |
| 16 | +BACKFILL_LAG_DAYS = 3 | |
| 17 | +RELEASE_LIKE = {"RELEASE", "ANNOUNCEMENT", "NEW_MODEL", "VERSION_RELEASED"} | |
| 18 | + | |
| 19 | + | |
| 20 | +def classify_backfill(event_type: str, effective_at: datetime | None, observed_at: datetime, *, is_first_run: bool = False, | |
| 21 | + entity_first_seen: datetime | None = None) -> bool: | |
| 22 | + """True when the event describes something that happened well before we observed it. | |
| 23 | + | |
| 24 | + * `effective_at` more than BACKFILL_LAG_DAYS before `observed_at` → backfill (an old release date being loaded); | |
| 25 | + * the connector run is the connector's first successful run → backfill (initial corpus, not news); | |
| 26 | + * NEW_* events for an entity whose first-seen hint (release/publication date) predates observation by more than the lag → backfill. | |
| 27 | + """ | |
| 28 | + if is_first_run: | |
| 29 | + return True | |
| 30 | + lag = timedelta(days=BACKFILL_LAG_DAYS) | |
| 31 | + if effective_at is not None and observed_at is not None and _aware(effective_at) < _aware(observed_at) - lag: | |
| 32 | + return True | |
| 33 | + if event_type.startswith("NEW_") and entity_first_seen is not None and observed_at is not None and _aware(entity_first_seen) < _aware(observed_at) - lag: | |
| 34 | + return True | |
| 35 | + return False | |
| 36 | + | |
| 37 | + | |
| 38 | +def group_key_for(event_type: str, entity_id: str | None, effective_at: datetime | None, observed_at: datetime | None) -> str | None: | |
| 39 | + """One release seen in several documents (blog post, docs page, hub listing, provider listing) groups under one key per entity × month.""" | |
| 40 | + if event_type not in RELEASE_LIKE or not entity_id: | |
| 41 | + return None | |
| 42 | + when = effective_at or observed_at | |
| 43 | + if when is None: | |
| 44 | + return None | |
| 45 | + return f"release:{entity_id}:{when.strftime('%Y-%m')}" | |
| 46 | + | |
| 47 | + | |
| 48 | +def _aware(dt: datetime) -> datetime: | |
| 49 | + from datetime import UTC | |
| 50 | + | |
| 51 | + return dt if dt.tzinfo is not None else dt.replace(tzinfo=UTC) | |
| 52 | + | |
| 53 | + | |
| 54 | +# ---------------------------------------------------------------------------------------------- importance (0 minor … 3 major) | |
| 55 | +def _num(v: Any) -> float | None: | |
| 56 | + if isinstance(v, bool): | |
| 57 | + return None | |
| 58 | + if isinstance(v, (int, float)): | |
| 59 | + return float(v) | |
| 60 | + if isinstance(v, str): | |
| 61 | + try: | |
| 62 | + return float(v.replace(",", "")) | |
| 63 | + except ValueError: | |
| 64 | + return None | |
| 65 | + return None | |
| 66 | + | |
| 67 | + | |
| 68 | +def _price_change_ratio(old: Any, new: Any) -> float: | |
| 69 | + """Largest relative move among input/output per-million prices (0 when unknown).""" | |
| 70 | + if not isinstance(old, dict) or not isinstance(new, dict): | |
| 71 | + a, b = _num(old), _num(new) | |
| 72 | + return abs(b - a) / a if a and b is not None and a > 0 else 0.0 | |
| 73 | + best = 0.0 | |
| 74 | + for k in ("input_per_mtok", "output_per_mtok", "cached_input_per_mtok"): | |
| 75 | + a, b = _num(old.get(k)), _num(new.get(k)) | |
| 76 | + if a is not None and b is not None and a > 0: | |
| 77 | + best = max(best, abs(b - a) / a) | |
| 78 | + elif a in (None, 0) and b: | |
| 79 | + best = max(best, 1.0) # newly priced / free → paid counts as a full move | |
| 80 | + return best | |
| 81 | + | |
| 82 | + | |
| 83 | +def importance_for(event_type: str, entity_type: str | None, old: Any = None, new: Any = None, *, tier: int = 2, org_model_count: int = 0, | |
| 84 | + openness: str | None = None, leader_change: bool = False, default: int = 2) -> int: | |
| 85 | + """Deterministic importance: | |
| 86 | + artifact events 0 · model_family events 1 · metadata corrections (PROPERTY_CHANGED) 0 | |
| 87 | + frontier-model release (organisation with ≥ 3 models, source tier ≤ 2) 3 · open-weight release 3 · other NEW_MODEL 2 | |
| 88 | + price change ≥ 50 % → 3, ≥ 20 % → 2, else 1 · context ≥ 5× → 3 (shrink → 1) · deprecation / retirement 3 · benchmark leader change 2 | |
| 89 | + Tier > 2 sources lose one point (never below 0).""" | |
| 90 | + et = event_type.upper() | |
| 91 | + if entity_type == "artifact": | |
| 92 | + return 0 | |
| 93 | + if entity_type == "model_family": | |
| 94 | + return min(1, default) | |
| 95 | + if et == "PROPERTY_CHANGED": | |
| 96 | + return 0 | |
| 97 | + if et == "PRICE_CHANGED": | |
| 98 | + ratio = _price_change_ratio(old, new) | |
| 99 | + imp = 3 if ratio >= 0.5 else 2 if ratio >= 0.2 else 1 | |
| 100 | + elif et == "CONTEXT_CHANGED": | |
| 101 | + a, b = _num(old), _num(new) | |
| 102 | + if a and b and a > 0: | |
| 103 | + imp = 3 if b / a >= 5 else 1 if b < a else 2 | |
| 104 | + else: | |
| 105 | + imp = 2 | |
| 106 | + elif et in ("DEPRECATION_ANNOUNCED", "RETIREMENT_ANNOUNCED"): | |
| 107 | + imp = 3 | |
| 108 | + elif et == "STATUS_CHANGED": | |
| 109 | + imp = 3 if str(new).lower() in ("deprecated", "retired", "discontinued") else 2 | |
| 110 | + elif et == "OPENNESS_CHANGED": | |
| 111 | + imp = 3 if str(new).startswith("open") else 2 | |
| 112 | + elif et == "NEW_MODEL": | |
| 113 | + if openness and str(openness).startswith("open"): | |
| 114 | + imp = 3 | |
| 115 | + elif org_model_count >= 3 and tier <= 2: | |
| 116 | + imp = 3 | |
| 117 | + else: | |
| 118 | + imp = 2 | |
| 119 | + elif et in ("BENCHMARK_LEADER_CHANGED",) or (et.startswith("BENCHMARK") and leader_change): | |
| 120 | + imp = 2 | |
| 121 | + else: | |
| 122 | + imp = default | |
| 123 | + if tier > 2: | |
| 124 | + imp -= 1 | |
| 125 | + return max(0, min(3, imp)) | |
| 126 | + | |
| 127 | + | |
| 128 | +__all__ = ["BACKFILL_LAG_DAYS", "RELEASE_LIKE", "classify_backfill", "group_key_for", "importance_for"] | |
modified
src/aiatlas/services/handlers.py
+2 −1
@@ -55,7 +55,8 @@ async def llm_extract(payload: dict[str, Any], job: dict[str, Any]) -> dict[str, | ||
| 55 | 55 | tier = await fetch_one(conn, "select tier from sources where id = :id", id=snap["source_id"]) if snap["source_id"] else None |
| 56 | 56 | # LLM output never outranks a deterministic statement from the same source: one tier lower (disagreement → conflicting + review) |
| 57 | 57 | writer = FactWriter(conn, source_id=snap["source_id"], snapshot_id=snapshot_id, source_url=snap["url"], tier=min(4, (tier["tier"] if tier else 2) + 1), |
| 58 | − connector_name=snap["connector_name"], extractor="llm", extractor_version=res.model, observed_at=snap["observed_at"].astimezone(UTC)) | |
| 58 | + connector_name=snap["connector_name"], extractor="llm", extractor_version=res.model, observed_at=snap["observed_at"].astimezone(UTC), | |
| 59 | + run_id=snap["run_id"] or job["id"]) | |
| 59 | 60 | ws = await writer.write(facts) |
| 60 | 61 | await execute(conn, "update snapshots set processing_status = 'llm_done' where id = :id", id=snapshot_id) |
| 61 | 62 | return {"task": task, "model": res.model, **ws.as_dict()} |
modified
src/aiatlas/services/merge.py
+182 −11
@@ -1,27 +1,97 @@ | ||
| 1 | 1 | """Entity merging (curation): fold a duplicate `source` into `target`. Nothing is deleted — the source row stays with |
| 2 | −`status='merged'` and `merged_into=target` so old slugs and ids keep resolving; every dependent row is re-pointed.""" | |
| 2 | +`status='merged'` and `merged_into=target` so old slugs and ids keep resolving; every dependent row is re-pointed. | |
| 3 | + | |
| 4 | +Modes | |
| 5 | + merge (default) full merge; records a `merge` decision | |
| 6 | + alias same as merge; records an `alias` decision (the source was another name of the target) | |
| 7 | + variant the source is an *artifact* (quantisation / conversion / packaging) of the target: it keeps its own row, slug and | |
| 8 | + facts, gets `entity_type='artifact'`, `canonical_id=target` and an `artifact_of` relation | |
| 9 | + family_member the target is a `model_family`: the source gets `family_id=target` and a `member_of_family` relation | |
| 10 | + | |
| 11 | +Every call writes a `resolution_decisions` row (applied=true) and an `admin_audit_log` row; `keep_separate` decisions block merges. | |
| 12 | +""" | |
| 3 | 13 | from __future__ import annotations |
| 4 | 14 | |
| 15 | +import json | |
| 5 | 16 | from typing import Any |
| 6 | 17 | |
| 7 | 18 | from sqlalchemy.ext.asyncio import AsyncConnection |
| 8 | 19 | |
| 9 | −from aiatlas.db import execute, fetch_one, fetch_val | |
| 20 | +from aiatlas.db import execute, fetch_all, fetch_one, fetch_val, jsonb | |
| 10 | 21 | from aiatlas.ids import new_id, normalize_alias |
| 22 | +from aiatlas.ontology import benchmarks as bench_ontology | |
| 23 | + | |
| 24 | +ORG_TYPES = {"company", "organization", "lab", "university"} | |
| 25 | +MODES = ("merge", "alias", "variant", "family_member") | |
| 26 | +DERIVED_SOURCE_KEY = "ai-atlas.registry" | |
| 27 | + | |
| 28 | + | |
| 29 | +async def registry_source_id(conn: AsyncConnection) -> str | None: | |
| 30 | + row = await fetch_one(conn, "select id from sources where key = :k", k=DERIVED_SOURCE_KEY) | |
| 31 | + return row["id"] if row else None | |
| 32 | + | |
| 33 | + | |
| 34 | +async def record_decision(conn: AsyncConnection, a_id: str, b_id: str, decision: str, *, actor: str = "curation", note: str | None = None, | |
| 35 | + payload: dict[str, Any] | None = None, applied: bool = True) -> None: | |
| 36 | + await execute(conn, """insert into resolution_decisions (id, a_id, b_id, decision, actor, note, payload, applied) | |
| 37 | + values (:id, :a, :b, :d, :actor, :note, cast(:p as jsonb), :applied) | |
| 38 | + on conflict (a_id, b_id, decision) do update set applied = resolution_decisions.applied or excluded.applied, | |
| 39 | + note = coalesce(excluded.note, resolution_decisions.note), payload = resolution_decisions.payload || excluded.payload""", | |
| 40 | + id=f"rd_{new_id('review').split('_', 1)[1]}", a=a_id, b=b_id, d=decision, actor=actor, note=note, p=jsonb(payload or {}), applied=applied) | |
| 41 | + | |
| 42 | + | |
| 43 | +async def audit(conn: AsyncConnection, action: str, target: str | None, payload: dict[str, Any] | None = None, *, actor: str = "curation") -> None: | |
| 44 | + await execute(conn, "insert into admin_audit_log (actor, action, target, payload) values (:actor, :action, :target, cast(:p as jsonb))", | |
| 45 | + actor=actor, action=action, target=target, p=jsonb(payload or {})) | |
| 46 | + | |
| 47 | + | |
| 48 | +async def kept_separate(conn: AsyncConnection, a: str, b: str) -> bool: | |
| 49 | + row = await fetch_one(conn, """select 1 from resolution_decisions where decision = 'keep_separate' | |
| 50 | + and ((a_id = :a and b_id = :b) or (a_id = :b and b_id = :a)) limit 1""", a=a, b=b) | |
| 51 | + return row is not None | |
| 52 | + | |
| 11 | 53 | |
| 54 | +async def upsert_relation(conn: AsyncConnection, subject_id: str, predicate: str, object_id: str, attributes: dict[str, Any] | None = None, *, | |
| 55 | + source_id: str | None = None, tier: int = 2, confidence: str = "high") -> bool: | |
| 56 | + """Insert a live relation unless an identical live edge exists. Returns True when a row was inserted.""" | |
| 57 | + if subject_id == object_id: | |
| 58 | + return False | |
| 59 | + existing = await fetch_one(conn, "select id from relations where subject_id = :s and predicate = :p and object_id = :o and valid_to is null", | |
| 60 | + s=subject_id, p=predicate, o=object_id) | |
| 61 | + if existing: | |
| 62 | + if attributes: | |
| 63 | + await execute(conn, "update relations set attributes = attributes || cast(:a as jsonb) where id = :id", a=jsonb(attributes), id=existing["id"]) | |
| 64 | + return False | |
| 65 | + await execute(conn, """insert into relations (id, subject_id, predicate, object_id, attributes, source_id, tier, confidence, observed_at, valid_from) | |
| 66 | + values (:id, :s, :p, :o, cast(:a as jsonb), :src, :tier, :conf, now(), now())""", | |
| 67 | + id=new_id("relation"), s=subject_id, p=predicate, o=object_id, a=jsonb(attributes or {}), src=source_id, tier=tier, conf=confidence) | |
| 68 | + return True | |
| 12 | 69 | |
| 13 | −async def merge_entities(conn: AsyncConnection, source_id: str, target_id: str) -> dict[str, Any]: | |
| 70 | + | |
| 71 | +async def merge_entities(conn: AsyncConnection, source_id: str, target_id: str, *, mode: str = "merge", actor: str = "curation", | |
| 72 | + note: str | None = None, payload: dict[str, Any] | None = None) -> dict[str, Any]: | |
| 73 | + if mode not in MODES: | |
| 74 | + raise ValueError(f"unknown merge mode {mode!r}") | |
| 14 | 75 | if source_id == target_id: |
| 15 | 76 | raise ValueError("source and target are the same entity") |
| 16 | 77 | src = await fetch_one(conn, "select id, entity_type, canonical_name, slug, merged_into, attributes, provenance from entities where id = :id", id=source_id) |
| 17 | − dst = await fetch_one(conn, "select id, entity_type, merged_into from entities where id = :id", id=target_id) | |
| 78 | + dst = await fetch_one(conn, "select id, entity_type, canonical_name, merged_into from entities where id = :id", id=target_id) | |
| 18 | 79 | if not src or not dst: |
| 19 | 80 | raise LookupError("source or target entity not found") |
| 20 | 81 | if dst["merged_into"]: |
| 21 | 82 | raise ValueError("target is itself merged; merge into its survivor instead") |
| 22 | 83 | if src["merged_into"]: |
| 23 | 84 | raise ValueError("source is already merged") |
| 24 | − if src["entity_type"] != dst["entity_type"]: | |
| 85 | + if mode in ("merge", "alias") and await kept_separate(conn, source_id, target_id): | |
| 86 | + raise ValueError("a keep_separate decision exists for this pair; refusing to merge") | |
| 87 | + if mode == "variant": | |
| 88 | + return await _mark_artifact(conn, src, dst, actor=actor, note=note, payload=payload) | |
| 89 | + if mode == "family_member": | |
| 90 | + return await _mark_family_member(conn, src, dst, actor=actor, note=note, payload=payload) | |
| 91 | + same_type = src["entity_type"] == dst["entity_type"] | |
| 92 | + org_pair = src["entity_type"] in ORG_TYPES and dst["entity_type"] in ORG_TYPES | |
| 93 | + model_pair = {src["entity_type"], dst["entity_type"]} <= {"model", "artifact"} | |
| 94 | + if not (same_type or org_pair or model_pair): | |
| 25 | 95 | raise ValueError(f"cannot merge a {src['entity_type']} into a {dst['entity_type']}") |
| 26 | 96 | |
| 27 | 97 | moved: dict[str, int] = {} |
@@ -56,22 +126,123 @@ async def merge_entities(conn: AsyncConnection, source_id: str, target_id: str) | ||
| 56 | 126 | "provider_id = case when provider_id = :s then :t else provider_id end where model_id = :s or provider_id = :s") |
| 57 | 127 | await count_update("results", "update benchmark_results set model_id = case when model_id = :s then :t else model_id end, " |
| 58 | 128 | "benchmark_id = case when benchmark_id = :s then :t else benchmark_id end where model_id = :s or benchmark_id = :s") |
| 129 | + moved["result_dedupe_collisions"] = await recompute_result_keys(conn, target_id, source_id=source_id) | |
| 130 | + moved["results_closed"] = await enforce_current_results(conn, model_id=target_id) | |
| 59 | 131 | await count_update("documents", "update documents set entity_id = :t where entity_id = :s") |
| 60 | 132 | await count_update("children", "update entities set organization_id = :t where organization_id = :s") |
| 133 | + await count_update("family_members", "update entities set family_id = :t where family_id = :s") | |
| 134 | + await count_update("artifacts", "update entities set canonical_id = :t where canonical_id = :s") | |
| 61 | 135 | await count_update("llm_jobs", "update llm_jobs set entity_id = :t where entity_id = :s") |
| 136 | + await count_update("sources", "update sources set organization_id = :t where organization_id = :s") | |
| 137 | + await count_update("domains", "update domains set organization_id = :t where organization_id = :s") | |
| 138 | + await count_update("review_items", "update review_queue set entity_ids = array_replace(entity_ids, :s, :t) where :s = any(entity_ids)") | |
| 139 | + # embeddings (table only exists where pgvector is installed): the target keeps its own vector; a source-only vector moves over | |
| 140 | + if await fetch_val(conn, "select to_regclass('entity_embeddings') is not null"): | |
| 141 | + has_target_vec = await fetch_one(conn, "select 1 from entity_embeddings where entity_id = :t", t=target_id) | |
| 142 | + if has_target_vec: | |
| 143 | + await count_update("embeddings", "delete from entity_embeddings where entity_id = :s") | |
| 144 | + else: | |
| 145 | + await count_update("embeddings", "update entity_embeddings set entity_id = :t where entity_id = :s") | |
| 62 | 146 | # attributes the target lacks are inherited (with their provenance); target values always win |
| 63 | 147 | await execute(conn, """update entities t set attributes = coalesce(s.attributes, '{}'::jsonb) || t.attributes, |
| 64 | 148 | provenance = coalesce(s.provenance, '{}'::jsonb) || t.provenance, last_seen_at = greatest(t.last_seen_at, s.last_seen_at), |
| 65 | 149 | first_seen_at = least(t.first_seen_at, s.first_seen_at), updated_at = now() |
| 66 | 150 | from entities s where t.id = :t and s.id = :s""", s=source_id, t=target_id) |
| 67 | − await execute(conn, "update entities set merged_into = :t, status = 'merged', updated_at = now() where id = :s", s=source_id, t=target_id) | |
| 151 | + await execute(conn, "update entities set merged_into = :t, status = 'merged', canonical_id = null, family_id = null, updated_at = now() where id = :s", s=source_id, t=target_id) | |
| 68 | 152 | await execute(conn, "update entities set merged_into = :t where merged_into = :s", s=source_id, t=target_id) |
| 69 | − await execute(conn, """insert into change_events (id, entity_id, event_type, category, summary, importance, connector_name, dedupe_key, meta) | |
| 70 | − values (:id, :t, 'ENTITY_MERGED', 'source', :sum, 1, 'curation', :dk, cast(:m as jsonb)) | |
| 153 | + await execute(conn, """insert into change_events (id, entity_id, event_type, category, summary, importance, connector_name, dedupe_key, meta, is_backfill) | |
| 154 | + values (:id, :t, 'ENTITY_MERGED', 'source', :sum, 0, 'curation', :dk, cast(:m as jsonb), true) | |
| 71 | 155 | on conflict (dedupe_key) do nothing""", |
| 72 | 156 | id=new_id("change_event"), t=target_id, sum=f"Merged duplicate '{src['canonical_name']}' ({src['slug']})", dk=f"merge:{source_id}:{target_id}", |
| 73 | − m=f'{{"source_id": "{source_id}", "source_slug": "{src["slug"]}"}}') | |
| 74 | − return {"source_id": source_id, "target_id": target_id, "moved": moved} | |
| 157 | + m=json.dumps({"source_id": source_id, "source_slug": src["slug"], "mode": mode})) | |
| 158 | + decision = "alias" if mode == "alias" else "merge" | |
| 159 | + await record_decision(conn, source_id, target_id, decision, actor=actor, note=note, payload={"moved": moved, **(payload or {})}) | |
| 160 | + await audit(conn, f"entity.{decision}", target_id, {"source_id": source_id, "source_slug": src["slug"], "target_id": target_id, "moved": moved, "note": note}, actor=actor) | |
| 161 | + return {"source_id": source_id, "target_id": target_id, "mode": mode, "moved": moved} | |
| 162 | + | |
| 163 | + | |
| 164 | +async def _mark_artifact(conn: AsyncConnection, src: dict[str, Any], dst: dict[str, Any], *, actor: str, note: str | None, payload: dict[str, Any] | None) -> dict[str, Any]: | |
| 165 | + if dst["entity_type"] not in ("model",): | |
| 166 | + raise ValueError("the canonical entity of an artifact must be a model") | |
| 167 | + kind = (payload or {}).get("artifact_kind") or (src["attributes"] or {}).get("artifact_kind") or "conversion" | |
| 168 | + await execute(conn, """update entities set entity_type = 'artifact', canonical_id = :t, artifact_kind = coalesce(artifact_kind, :k), | |
| 169 | + identity_confidence = coalesce(cast(:ic as text), identity_confidence), updated_at = now() where id = :s""", | |
| 170 | + t=dst["id"], k=kind, ic=(payload or {}).get("identity_confidence"), s=src["id"]) | |
| 171 | + inserted = await upsert_relation(conn, src["id"], "artifact_of", dst["id"], {"artifact_kind": kind}, source_id=await registry_source_id(conn)) | |
| 172 | + await record_decision(conn, src["id"], dst["id"], "variant_of", actor=actor, note=note, payload={"artifact_kind": kind, **(payload or {})}) | |
| 173 | + await audit(conn, "entity.artifact_of", src["id"], {"canonical_id": dst["id"], "artifact_kind": kind, "note": note}, actor=actor) | |
| 174 | + return {"source_id": src["id"], "target_id": dst["id"], "mode": "variant", "moved": {"relations": int(inserted)}} | |
| 175 | + | |
| 176 | + | |
| 177 | +async def _mark_family_member(conn: AsyncConnection, src: dict[str, Any], dst: dict[str, Any], *, actor: str, note: str | None, payload: dict[str, Any] | None) -> dict[str, Any]: | |
| 178 | + if dst["entity_type"] != "model_family": | |
| 179 | + raise ValueError("family_member requires a model_family target") | |
| 180 | + await execute(conn, "update entities set family_id = :t, updated_at = now() where id = :s and family_id is distinct from :t", t=dst["id"], s=src["id"]) | |
| 181 | + inserted = await upsert_relation(conn, src["id"], "member_of_family", dst["id"], source_id=await registry_source_id(conn)) | |
| 182 | + await record_decision(conn, src["id"], dst["id"], "family_member", actor=actor, note=note, payload=payload) | |
| 183 | + await audit(conn, "entity.family_member", src["id"], {"family_id": dst["id"], "note": note}, actor=actor) | |
| 184 | + return {"source_id": src["id"], "target_id": dst["id"], "mode": "family_member", "moved": {"relations": int(inserted)}} | |
| 185 | + | |
| 186 | + | |
| 187 | +# ---------------------------------------------------------------------------------------------- benchmark result bookkeeping | |
| 188 | +def _cfg_hash(config: dict[str, Any] | None) -> str: | |
| 189 | + import hashlib | |
| 190 | + | |
| 191 | + return hashlib.sha1(json.dumps(config or {}, sort_keys=True, default=str).encode()).hexdigest()[:12] | |
| 192 | + | |
| 193 | + | |
| 194 | +async def recompute_result_keys(conn: AsyncConnection, model_id: str, *, source_id: str | None = None) -> int: | |
| 195 | + """Recompute `dedupe_key` (= model:benchmark:cfg_hash:metric) and `config_key` for the rows now attached to `model_id`. | |
| 196 | + Collisions (the target already had the same result from the same source) close the older row. Returns the number of collisions.""" | |
| 197 | + rows = await fetch_all(conn, """select id, benchmark_id, metric, config, dedupe_key, config_key, run_group, variant, observed_at, valid_to, is_current from benchmark_results | |
| 198 | + where model_id = :m order by observed_at asc, id asc""", m=model_id) | |
| 199 | + collisions = 0 | |
| 200 | + seen: dict[str, dict[str, Any]] = {} | |
| 201 | + for r in rows: | |
| 202 | + base = f"{model_id}:{r['benchmark_id']}:{_cfg_hash(r['config'])}:{r['metric'] or ''}" | |
| 203 | + ck = bench_ontology.config_key(r["config"], r["metric"]) | |
| 204 | + rg = r["run_group"] or bench_ontology.run_group_from_config(r["config"]) | |
| 205 | + variant = r["variant"] or bench_ontology.variant_from_config(r["config"]) | |
| 206 | + if (rg, variant) != (r["run_group"], r["variant"]): | |
| 207 | + await execute(conn, "update benchmark_results set run_group = :rg, variant = :v where id = :id", rg=rg, v=variant, id=r["id"]) | |
| 208 | + if r["valid_to"] is not None: | |
| 209 | + # historical rows keep a unique suffixed key | |
| 210 | + key = base if r["dedupe_key"] == base and base not in seen else f"{base}:{r['id'][-10:]}" | |
| 211 | + await execute(conn, "update benchmark_results set dedupe_key = :d, config_key = :ck where id = :id and (dedupe_key <> :d or config_key is distinct from :ck)", | |
| 212 | + d=key, ck=ck, id=r["id"]) | |
| 213 | + continue | |
| 214 | + prev = seen.get(base) | |
| 215 | + if prev is not None: | |
| 216 | + # two live rows for the same key: keep the newer (rows are sorted by observed_at asc), close the older | |
| 217 | + older, newer = prev, r | |
| 218 | + await execute(conn, "update benchmark_results set valid_to = :o, is_current = false, dedupe_key = :d, config_key = :ck where id = :id", | |
| 219 | + o=newer["observed_at"], d=f"{base}:{older['id'][-10:]}", ck=ck, id=older["id"]) | |
| 220 | + collisions += 1 | |
| 221 | + seen[base] = r | |
| 222 | + await execute(conn, "update benchmark_results set dedupe_key = :d, config_key = :ck where id = :id and (dedupe_key <> :d or config_key is distinct from :ck)", | |
| 223 | + d=base, ck=ck, id=r["id"]) | |
| 224 | + return collisions | |
| 225 | + | |
| 226 | + | |
| 227 | +async def enforce_current_results(conn: AsyncConnection, *, model_id: str | None = None, dry_run: bool = False) -> int: | |
| 228 | + """One current row per (model, benchmark, metric, config_key): rows from an older run group than the latest observed one are closed | |
| 229 | + (`is_current=false`, `valid_to` = the newer observation). Returns the number of rows (that would be) closed.""" | |
| 230 | + scope = "and r.model_id = :m" if model_id else "" | |
| 231 | + sql = f"""with latest as ( | |
| 232 | + select distinct on (model_id, benchmark_id, coalesce(metric, ''), config_key) id, model_id, benchmark_id, coalesce(metric, '') as metric, config_key, | |
| 233 | + coalesce(run_group, '') as run_group, observed_at | |
| 234 | + from benchmark_results r where valid_to is null and is_current and config_key is not null {scope} | |
| 235 | + order by model_id, benchmark_id, coalesce(metric, ''), config_key, observed_at desc, coalesce(run_group, '') desc, id desc) | |
| 236 | + select r.id, l.observed_at as close_at from benchmark_results r join latest l | |
| 237 | + on l.model_id = r.model_id and l.benchmark_id = r.benchmark_id and l.metric = coalesce(r.metric, '') and l.config_key = r.config_key | |
| 238 | + where r.valid_to is null and r.is_current and r.id <> l.id and coalesce(r.run_group, '') <> l.run_group {scope}""" | |
| 239 | + rows = await fetch_all(conn, sql, m=model_id) if model_id else await fetch_all(conn, sql) | |
| 240 | + if dry_run or not rows: | |
| 241 | + return len(rows) | |
| 242 | + for r in rows: | |
| 243 | + await execute(conn, "update benchmark_results set is_current = false, valid_to = :o where id = :id", o=r["close_at"], id=r["id"]) | |
| 244 | + return len(rows) | |
| 75 | 245 | |
| 76 | 246 | |
| 77 | −__all__ = ["merge_entities"] | |
| 247 | +__all__ = ["MODES", "ORG_TYPES", "audit", "enforce_current_results", "kept_separate", "merge_entities", "recompute_result_keys", "record_decision", | |
| 248 | + "registry_source_id", "upsert_relation"] | |
added
tests/test_events.py
+47 −0
@@ -0,0 +1,47 @@ | ||
| 1 | +"""Event semantics: backfill classification, release grouping, deterministic importance.""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +from datetime import UTC, datetime, timedelta | |
| 5 | + | |
| 6 | +from aiatlas.services.events import classify_backfill, group_key_for, importance_for | |
| 7 | + | |
| 8 | +NOW = datetime(2026, 9, 12, 12, 0, tzinfo=UTC) | |
| 9 | + | |
| 10 | + | |
| 11 | +def test_backfill_rules() -> None: | |
| 12 | + assert classify_backfill("NEW_MODEL", None, NOW) is False | |
| 13 | + assert classify_backfill("NEW_MODEL", NOW - timedelta(days=1), NOW) is False # 1 day lag: news | |
| 14 | + assert classify_backfill("NEW_MODEL", NOW - timedelta(days=4), NOW) is True # > 3 days: history being loaded | |
| 15 | + assert classify_backfill("PRICE_CHANGED", None, NOW, is_first_run=True) is True # connector's first run = initial corpus | |
| 16 | + assert classify_backfill("NEW_MODEL", None, NOW, entity_first_seen=NOW - timedelta(days=30)) is True | |
| 17 | + assert classify_backfill("PRICE_CHANGED", None, NOW, entity_first_seen=NOW - timedelta(days=30)) is False # hint only applies to NEW_* | |
| 18 | + assert classify_backfill("NEW_PAPER", datetime(2026, 9, 11), NOW) is False # naive datetimes are treated as UTC | |
| 19 | + | |
| 20 | + | |
| 21 | +def test_group_key() -> None: | |
| 22 | + assert group_key_for("NEW_MODEL", "model_x", datetime(2026, 5, 3, tzinfo=UTC), NOW) == "release:model_x:2026-05" | |
| 23 | + assert group_key_for("RELEASE", "model_x", None, NOW) == "release:model_x:2026-09" | |
| 24 | + assert group_key_for("PRICE_CHANGED", "model_x", None, NOW) is None | |
| 25 | + assert group_key_for("RELEASE", None, None, NOW) is None | |
| 26 | + | |
| 27 | + | |
| 28 | +def test_importance_matrix() -> None: | |
| 29 | + assert importance_for("NEW_MODEL", "artifact") == 0 | |
| 30 | + assert importance_for("NEW_MODEL_FAMILY", "model_family") == 1 | |
| 31 | + assert importance_for("PROPERTY_CHANGED", "model") == 0 | |
| 32 | + assert importance_for("NEW_MODEL", "model", org_model_count=5, tier=1) == 3 # frontier lab | |
| 33 | + assert importance_for("NEW_MODEL", "model", org_model_count=1, tier=1) == 2 | |
| 34 | + assert importance_for("NEW_MODEL", "model", org_model_count=1, openness="open-weights") == 3 | |
| 35 | + assert importance_for("NEW_MODEL", "model", org_model_count=5, tier=3) == 1 # not frontier at tier 3, and a community source loses a point | |
| 36 | + old, new = {"input_per_mtok": 10, "output_per_mtok": 30}, {"input_per_mtok": 5, "output_per_mtok": 30} | |
| 37 | + assert importance_for("PRICE_CHANGED", "model", old, new) == 3 # -50 % | |
| 38 | + assert importance_for("PRICE_CHANGED", "model", old, {"input_per_mtok": 7.5, "output_per_mtok": 30}) == 2 # -25 % | |
| 39 | + assert importance_for("PRICE_CHANGED", "model", old, {"input_per_mtok": 9.5, "output_per_mtok": 30}) == 1 # -5 % | |
| 40 | + assert importance_for("CONTEXT_CHANGED", "model", 200_000, 1_000_000) == 3 | |
| 41 | + assert importance_for("CONTEXT_CHANGED", "model", 200_000, 400_000) == 2 | |
| 42 | + assert importance_for("CONTEXT_CHANGED", "model", 200_000, 100_000) == 1 | |
| 43 | + assert importance_for("DEPRECATION_ANNOUNCED", "model") == 3 and importance_for("RETIREMENT_ANNOUNCED", "model") == 3 | |
| 44 | + assert importance_for("STATUS_CHANGED", "model", "active", "deprecated") == 3 and importance_for("STATUS_CHANGED", "model", "preview", "active") == 2 | |
| 45 | + assert importance_for("OPENNESS_CHANGED", "model", "proprietary", "open-weights") == 3 | |
| 46 | + assert importance_for("BENCHMARK_RESULT", "model", leader_change=True) == 2 and importance_for("BENCHMARK_RESULT", "model", default=1) == 1 | |
| 47 | + assert importance_for("PRICE_CHANGED", "artifact", old, new) == 0 | |
added
tests/test_fetch_ssrf.py
+102 −0
@@ -0,0 +1,102 @@ | ||
| 1 | +"""SSRF guard — pure validator tests (no network: DNS answers are injected).""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +import pytest | |
| 5 | + | |
| 6 | +from aiatlas.sdk.fetch import MAX_REDIRECTS, BlockedDestination, validate_destination | |
| 7 | + | |
| 8 | + | |
| 9 | +@pytest.mark.parametrize("url", [ | |
| 10 | + "ftp://example.com/x", "file:///etc/passwd", "gopher://example.com", "javascript:alert(1)", "data:text/html,hi", | |
| 11 | + "http://localhost/", "http://LOCALHOST:8080/", "https://foo.local/", "http://db.internal/", "http://intranet/", "http://metadata.google.internal/computeMetadata/v1/", | |
| 12 | + "http://127.0.0.1/", "http://127.1.2.3/", "http://10.0.0.5/", "http://172.16.4.4/", "http://172.31.255.1/", "http://192.168.2.10/", | |
| 13 | + "http://169.254.169.254/latest/meta-data/", "http://169.254.170.2/", "http://100.64.1.1/", "http://100.127.255.255/", "http://0.0.0.0/", | |
| 14 | + "http://[::1]/", "http://[fe80::1]/", "http://[fc00::1]/", "http://[fd12:3456::1]/", "http://[::]/", "http://[::ffff:127.0.0.1]/", "http://[::ffff:10.0.0.1]/", | |
| 15 | + "http://user:pass@example.com/", | |
| 16 | +]) | |
| 17 | +def test_blocked_literals(url: str) -> None: | |
| 18 | + with pytest.raises(BlockedDestination): | |
| 19 | + validate_destination(url, resolved_ips=[]) | |
| 20 | + | |
| 21 | + | |
| 22 | +@pytest.mark.parametrize("url", ["https://huggingface.co/Qwen/Qwen3-8B", "http://arxiv.org/abs/2505.09388", "https://8.8.8.8/", "https://[2606:4700::6810:84e5]/"]) | |
| 23 | +def test_public_allowed(url: str) -> None: | |
| 24 | + validate_destination(url, resolved_ips=["104.16.0.1"]) | |
| 25 | + | |
| 26 | + | |
| 27 | +def test_dns_rebinding_to_private_is_blocked() -> None: | |
| 28 | + with pytest.raises(BlockedDestination): | |
| 29 | + validate_destination("https://evil.example.com/", resolved_ips=["10.0.0.9"]) | |
| 30 | + with pytest.raises(BlockedDestination): | |
| 31 | + validate_destination("https://evil.example.com/", resolved_ips=["104.16.0.1", "169.254.169.254"]) # any private answer blocks | |
| 32 | + with pytest.raises(BlockedDestination): | |
| 33 | + validate_destination("https://evil.example.com/", resolved_ips=["fd00::1"]) | |
| 34 | + | |
| 35 | + | |
| 36 | +def test_unresolvable_host_is_not_ssrf() -> None: | |
| 37 | + # resolution failures surface later as transport errors, not as blocked destinations | |
| 38 | + validate_destination("https://does-not-exist.example.com/", resolved_ips=[]) | |
| 39 | + | |
| 40 | + | |
| 41 | +def test_redirect_cap_constant() -> None: | |
| 42 | + assert MAX_REDIRECTS == 5 | |
| 43 | + | |
| 44 | + | |
| 45 | +async def test_fetcher_get_blocks_before_connecting() -> None: | |
| 46 | + from aiatlas.sdk.fetch import Fetcher, FetchError | |
| 47 | + | |
| 48 | + async with Fetcher(robots=False) as f: | |
| 49 | + with pytest.raises(FetchError) as exc: | |
| 50 | + await f.get("http://169.254.169.254/latest/meta-data/") | |
| 51 | + assert "blocked destination" in str(exc.value) | |
| 52 | + with pytest.raises(FetchError): | |
| 53 | + await f.get("file:///etc/hosts") | |
| 54 | + | |
| 55 | + | |
| 56 | +async def test_fetcher_blocks_redirect_to_private(monkeypatch: pytest.MonkeyPatch) -> None: | |
| 57 | + """A public host redirecting to an internal address is refused on the redirect hop (uses respx-free httpx MockTransport).""" | |
| 58 | + import httpx | |
| 59 | + | |
| 60 | + from aiatlas.sdk import fetch as fetch_mod | |
| 61 | + from aiatlas.sdk.fetch import Fetcher, FetchError | |
| 62 | + | |
| 63 | + async def fake_validate(url: str) -> None: | |
| 64 | + # public → ok, private literal → raise, exactly like the real validator without DNS | |
| 65 | + fetch_mod.validate_destination(url, resolved_ips=["104.16.0.1"]) | |
| 66 | + | |
| 67 | + monkeypatch.setattr(fetch_mod, "validate_destination_async", fake_validate) | |
| 68 | + | |
| 69 | + def handler(request: httpx.Request) -> httpx.Response: | |
| 70 | + if request.url.host == "public.example.com": | |
| 71 | + return httpx.Response(302, headers={"location": "http://169.254.169.254/latest/"}) | |
| 72 | + return httpx.Response(200, content=b"x" * 100) | |
| 73 | + | |
| 74 | + async with Fetcher(robots=False) as f: | |
| 75 | + f._client = httpx.AsyncClient(transport=httpx.MockTransport(handler), follow_redirects=False) | |
| 76 | + with pytest.raises(FetchError) as exc: | |
| 77 | + await f.get("https://public.example.com/start", min_bytes=1) | |
| 78 | + assert "blocked destination" in str(exc.value) | |
| 79 | + | |
| 80 | + | |
| 81 | +async def test_fetcher_caps_redirects(monkeypatch: pytest.MonkeyPatch) -> None: | |
| 82 | + import httpx | |
| 83 | + | |
| 84 | + from aiatlas.sdk import fetch as fetch_mod | |
| 85 | + from aiatlas.sdk.fetch import Fetcher, FetchError | |
| 86 | + | |
| 87 | + async def fake_validate(url: str) -> None: | |
| 88 | + fetch_mod.validate_destination(url, resolved_ips=["104.16.0.1"]) | |
| 89 | + | |
| 90 | + monkeypatch.setattr(fetch_mod, "validate_destination_async", fake_validate) | |
| 91 | + n = {"hops": 0} | |
| 92 | + | |
| 93 | + def handler(request: httpx.Request) -> httpx.Response: | |
| 94 | + n["hops"] += 1 | |
| 95 | + return httpx.Response(301, headers={"location": f"https://public.example.com/{n['hops']}"}) | |
| 96 | + | |
| 97 | + async with Fetcher(robots=False) as f: | |
| 98 | + f._client = httpx.AsyncClient(transport=httpx.MockTransport(handler), follow_redirects=False) | |
| 99 | + with pytest.raises(FetchError) as exc: | |
| 100 | + await f.get("https://public.example.com/start", min_bytes=1) | |
| 101 | + assert "too many redirects" in str(exc.value) | |
| 102 | + assert n["hops"] == MAX_REDIRECTS + 1 | |
added
tests/test_resolver.py
+123 −0
@@ -0,0 +1,123 @@ | ||
| 1 | +"""Resolver guards against the local dev database (rolled back).""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +import uuid | |
| 5 | +from collections.abc import AsyncIterator | |
| 6 | + | |
| 7 | +import pytest | |
| 8 | +from sqlalchemy.ext.asyncio import AsyncConnection | |
| 9 | + | |
| 10 | +from aiatlas import db | |
| 11 | +from aiatlas.db import execute, fetch_one | |
| 12 | +from aiatlas.sdk.facts import EntityRef, Facts | |
| 13 | +from aiatlas.sdk.resolution import Resolver, compatible_types | |
| 14 | +from aiatlas.sdk.writer import FactWriter | |
| 15 | +from aiatlas.services.merge import merge_entities, record_decision | |
| 16 | + | |
| 17 | + | |
| 18 | +@pytest.fixture | |
| 19 | +async def conn() -> AsyncIterator[AsyncConnection]: | |
| 20 | + async with db.engine().connect() as c: | |
| 21 | + trans = await c.begin() | |
| 22 | + try: | |
| 23 | + yield c | |
| 24 | + finally: | |
| 25 | + await trans.rollback() | |
| 26 | + await db.dispose() | |
| 27 | + | |
| 28 | + | |
| 29 | +def _tag() -> str: | |
| 30 | + return uuid.uuid4().hex[:8] | |
| 31 | + | |
| 32 | + | |
| 33 | +async 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 w | |
| 37 | + | |
| 38 | + | |
| 39 | +def test_compatible_types() -> None: | |
| 40 | + assert set(compatible_types("model")) == {"model", "artifact"} and compatible_types("paper") == ("paper",) | |
| 41 | + | |
| 42 | + | |
| 43 | +async 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 model | |
| 50 | + other = await r.resolve(EntityRef("model", f"Qwen{t} 38B"), create=False) | |
| 51 | + assert other is None | |
| 52 | + same = await r.resolve(EntityRef("model", f"Qwen{t} 8B"), create=False) | |
| 53 | + assert same == a.id | |
| 54 | + | |
| 55 | + | |
| 56 | +async 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 automatically | |
| 67 | + assert await r.resolve(EntityRef("model", f"Nimbus {t}"), create=False) is None | |
| 68 | + assert await r.resolve(EntityRef("model", f"Nimbus {t}", organization=EntityRef("company", f"Org One {t}", id=org1.id)), create=False) == a.id | |
| 69 | + assert await r.resolve(EntityRef("model", f"Nimbus {t}", organization=EntityRef("company", f"Org Two {t}", id=org2.id)), create=False) == b.id | |
| 70 | + with pytest.raises(ValueError): | |
| 71 | + await merge_entities(conn, a.id, b.id) | |
| 72 | + | |
| 73 | + | |
| 74 | +async 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.id | |
| 81 | + | |
| 82 | + | |
| 83 | +async 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 variant | |
| 93 | + assert await r.resolve_variant(EntityRef("model", f"omega-{t}-high")) is None # base unknown | |
| 94 | + # a ref with its own hub identifier is a real model even if its name ends in "-thinking": created, never folded | |
| 95 | + real = await r.resolve(EntityRef("model", f"zeta-{t}-thinking", identifiers={"hf_repo": f"zorg/Zeta-{t}-Thinking"})) | |
| 96 | + assert real != base.id | |
| 97 | + # a tier-1 resolver never folds either | |
| 98 | + r1 = Resolver(conn, source_tier=1) | |
| 99 | + official = await r1.resolve(EntityRef("model", f"zeta-{t}-high"), create=False) | |
| 100 | + assert official is None | |
| 101 | + # an evaluator-only ref folds and records the fold | |
| 102 | + 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"}) | |
| 105 | + | |
| 106 | + | |
| 107 | +async def test_first_seen_hint_and_touch(conn: AsyncConnection) -> None: | |
| 108 | + from datetime import UTC, datetime | |
| 109 | + | |
| 110 | + 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 forward | |
| 120 | + 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) | |
added
tests/test_writer.py
+266 −0
@@ -0,0 +1,266 @@ | ||
| 1 | +"""FactWriter rules against the local dev database, inside a transaction that is always rolled back (nothing persists).""" | |
| 2 | +from __future__ import annotations | |
| 3 | + | |
| 4 | +import uuid | |
| 5 | +from collections.abc import AsyncIterator | |
| 6 | +from datetime import UTC, datetime, timedelta | |
| 7 | + | |
| 8 | +import pytest | |
| 9 | +from sqlalchemy.ext.asyncio import AsyncConnection | |
| 10 | + | |
| 11 | +from aiatlas import db | |
| 12 | +from aiatlas.db import fetch_all, fetch_one | |
| 13 | +from aiatlas.sdk.facts import Claim, EntityRef, Event, Facts, PriceObs, ResultObs, Target, facts_from_json, facts_to_json | |
| 14 | +from aiatlas.sdk.writer import FactWriter | |
| 15 | + | |
| 16 | +pytestmark = pytest.mark.usefixtures("conn") | |
| 17 | + | |
| 18 | + | |
| 19 | +@pytest.fixture | |
| 20 | +async def conn() -> AsyncIterator[AsyncConnection]: | |
| 21 | + async with db.engine().connect() as c: | |
| 22 | + trans = await c.begin() | |
| 23 | + try: | |
| 24 | + yield c | |
| 25 | + finally: | |
| 26 | + await trans.rollback() | |
| 27 | + await db.dispose() | |
| 28 | + | |
| 29 | + | |
| 30 | +def _tag() -> str: | |
| 31 | + return uuid.uuid4().hex[:8] | |
| 32 | + | |
| 33 | + | |
| 34 | +def _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) | |
| 38 | + | |
| 39 | + | |
| 40 | +# ---------------------------------------------------------------------------------------------- pure: Facts JSON round-trip | |
| 41 | +def 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 = m | |
| 53 | + 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) == data | |
| 64 | + | |
| 65 | + | |
| 66 | +# ---------------------------------------------------------------------------------------------- taxonomy at write time | |
| 67 | +async 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 event | |
| 84 | + 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.id | |
| 88 | + 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"] == 1 | |
| 90 | + 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"} | |
| 92 | + | |
| 93 | + | |
| 94 | +async 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 None | |
| 101 | + 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 None | |
| 103 | + | |
| 104 | + | |
| 105 | +# ---------------------------------------------------------------------------------------------- same-source loophole & supersede | |
| 106 | +async 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 lower | |
| 113 | + 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 == 1 | |
| 120 | + 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 importance | |
| 122 | + 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.id | |
| 127 | + 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 None | |
| 129 | + | |
| 130 | + | |
| 131 | +async 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 == 1 | |
| 141 | + 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) | |
| 143 | + | |
| 144 | + | |
| 145 | +# ---------------------------------------------------------------------------------------------- results: comparability, one current row per key, bounds | |
| 146 | +async 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 True | |
| 155 | + 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_key | |
| 165 | + 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] | |
| 172 | + | |
| 173 | + | |
| 174 | +async 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" | |
| 185 | + | |
| 186 | + | |
| 187 | +async 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 == 1 | |
| 200 | + 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") | |
| 205 | + | |
| 206 | + | |
| 207 | +# ---------------------------------------------------------------------------------------------- events: backfill, group key, NEW_* importance | |
| 208 | +async 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"] == 2 | |
| 215 | + # first run of a connector → everything is backfill even without dates | |
| 216 | + 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 True | |
| 221 | + # artifacts and families never make importance-3 news | |
| 222 | + 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) | |
| 231 | + | |
| 232 | + | |
| 233 | +async 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" | |
| 243 | + | |
| 244 | + | |
| 245 | +async 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)] | |
| 261 | + | |
| 262 | + | |
| 263 | +async 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 None | |
| 265 | + 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 None | |
| 267 | ||