lorebase-store 0.2.3__py3-none-any.whl

lorebase_store/extraction/items.py

@@ -0,0 +1,187 @@
+"""Extraction items as data: element identity, reviewer edits and the Proposal summary.
+
+Shared by Scribe, which writes the items of a new extraction Proposal, and Keeper, which
+records reviewer decisions; neither needs a model for this.
+
+Reviewer edits, by item type (Attributes named by id or api name, values in API form):
+
+- `record`: `{"values": {attribute: value}}`, merged over the proposed values;
+- `value`: `{"value": value, "key": key}`, the value, or the Record it belongs to (the
+  ambiguous match picker);
+- `link`: `{"from": key, "to": key, "values": {attribute: value}}`.
+
+`normalize_edits` validates them against the schema and returns them with Attribute ids and
+canonical values, the form stored in `proposal_items.edits` and applied by
+`effective_subject`.
+"""
+
+from __future__ import annotations
+
+from collections import Counter
+from collections.abc import Iterable, Mapping
+from typing import Any
+
+from lorebase_core.model import Attribute, Schema
+from lorebase_core.tree import kind_partition, link_key, relation_partition
+from lorebase_core.types import EncodingError, encode_value, normalize_key
+from lorebase_store.models.enums import ItemDecision
+from lorebase_store.repositories.proposal_items import (
+    LINK_ITEM,
+    RECORD_ITEM,
+    VALUE_ITEM,
+    ItemRow,
+)
+
+
+class ItemEditError(ValueError):
+    """Reviewer edits that do not fit the item or the schema."""
+
+    def __init__(self, item_id: str, message: str) -> None:
+        super().__init__(f"item {item_id}: {message}")
+        self.item_id = item_id
+        self.detail = message
+
+
+def record_element(kind_id: str, key: str) -> tuple[str, str]:
+    """`(partition, element key)` of a Record."""
+    return kind_partition(kind_id), key
+
+
+def link_element(relation_id: str, from_key: str, to_key: str) -> tuple[str, str]:
+    """`(partition, element key)` of a Link."""
+    return relation_partition(relation_id), link_key(from_key, to_key)
+
+
+def _attributes(schema: Schema, item: ItemRow) -> Mapping[str, Attribute]:
+    subject = item.subject
+    if item.item_type == LINK_ITEM:
+        relation = schema.relations.get(str(subject.get("relation")))
+        if relation is None:
+            raise ItemEditError(item.item_id, "the Relation no longer exists")
+        return relation.attributes
+    kind = schema.kinds.get(str(subject.get("kind")))
+    if kind is None:
+        raise ItemEditError(item.item_id, "the Kind no longer exists")
+    return kind.attributes
+
+
+def _encode(item_id: str, attribute: Attribute, raw: Any) -> Any:
+    if attribute.is_derived:
+        raise ItemEditError(item_id, f"{attribute.api_name} is derived")
+    try:
+        return encode_value(attribute.type, raw)
+    except EncodingError as error:
+        raise ItemEditError(item_id, f"{attribute.api_name}: {error}") from None
+
+
+def _values(item_id: str, attributes: Mapping[str, Attribute], raw: Any) -> dict[str, Any]:
+    if not isinstance(raw, Mapping):
+        raise ItemEditError(item_id, "values must be an object")
+    by_name = {a.api_name: a for a in attributes.values()}
+    encoded: dict[str, Any] = {}
+    for name, value in raw.items():
+        attribute = attributes.get(name) or by_name.get(name)
+        if attribute is None:
+            raise ItemEditError(item_id, f"unknown Attribute {name!r}")
+        encoded[attribute.id] = _encode(item_id, attribute, value)
+    return encoded
+
+
+def _key(item_id: str, raw: Any) -> str:
+    if not isinstance(raw, str) or not raw.strip():
+        raise ItemEditError(item_id, "keys must be non-empty strings")
+    return normalize_key(raw.strip())
+
+
+def normalize_edits(
+    item: ItemRow, edits: Mapping[str, Any] | None, schema: Schema
+) -> dict[str, Any]:
+    """Edits validated against the schema, with Attribute ids and canonical values."""
+    if not edits:
+        return {}
+    allowed = {
+        RECORD_ITEM: {"values"},
+        VALUE_ITEM: {"value", "key"},
+        LINK_ITEM: {"from", "to", "values"},
+    }[item.item_type]
+    unknown = sorted(set(edits) - allowed)
+    if unknown:
+        raise ItemEditError(
+            item.item_id, f"{item.item_type} items accept edits of {sorted(allowed)}, not {unknown}"
+        )
+    attributes = _attributes(schema, item)
+    normalized: dict[str, Any] = {}
+    if "values" in edits:
+        normalized["values"] = _values(item.item_id, attributes, edits["values"])
+    if "value" in edits:
+        attribute = attributes.get(str(item.subject.get("attribute")))
+        if attribute is None:
+            raise ItemEditError(item.item_id, "the Attribute no longer exists")
+        normalized["value"] = _encode(item.item_id, attribute, edits["value"])
+    for name in ("key", "from", "to"):
+        if name in edits:
+            normalized[name] = _key(item.item_id, edits[name])
+    return normalized
+
+
+def effective_subject(item: ItemRow) -> dict[str, Any]:
+    """The subject with the stored reviewer edits applied."""
+    subject = dict(item.subject)
+    edits = item.edits or {}
+    if "values" in edits:
+        subject["values"] = {**(subject.get("values") or {}), **edits["values"]}
+    for name in ("value", "key", "from", "to"):
+        if name in edits:
+            subject[name] = edits[name]
+    return subject
+
+
+def summarize(items: Iterable[ItemRow], schema: Schema | None = None) -> dict[str, Any]:
+    """Counts of a Proposal's items: per decision, per item type and per Kind or Relation (api
+    names when `schema` is given), and the mean confidence."""
+    rows = list(items)
+    decisions: Counter[str] = Counter({d.value: 0 for d in ItemDecision})
+    types: Counter[str] = Counter()
+    targets: Counter[str] = Counter()
+    confidences = [row.confidence for row in rows if row.confidence is not None]
+    for row in rows:
+        decisions[row.decision.value] += 1
+        types[row.item_type] += 1
+        targets[_target_name(row, schema)] += 1
+    return {
+        "items": len(rows),
+        "decisions": dict(decisions),
+        "types": dict(sorted(types.items())),
+        "targets": dict(sorted(targets.items())),
+        "mean_confidence": (round(sum(confidences) / len(confidences), 4) if confidences else None),
+    }
+
+
+def _target_name(row: ItemRow, schema: Schema | None) -> str:
+    target = row.target
+    if schema is None:
+        return target
+    kind = schema.kinds.get(target)
+    if kind is not None:
+        return kind.api_name
+    relation = schema.relations.get(target)
+    return relation.api_name if relation is not None else target
+
+
+def proposal_text(summary: Mapping[str, Any], documents: Iterable[str]) -> tuple[str, str]:
+    """Title and description of an extraction Proposal from its summary."""
+    decisions = summary.get("decisions", {})
+    names = sorted(documents)
+    shown = ", ".join(names[:3]) + (f" and {len(names) - 3} more" if len(names) > 3 else "")
+    title = (
+        f"Extract {shown}: {summary.get('items', 0)} item(s), "
+        f"{decisions.get('accepted', 0)} accepted, {decisions.get('pending', 0)} pending"
+    )
+    targets = ", ".join(f"{name} {count}" for name, count in summary.get("targets", {}).items())
+    description = (
+        f"Proposed from {len(names)} document(s). Items by type: "
+        + ", ".join(f"{name} {count}" for name, count in summary.get("types", {}).items())
+        + (f"; by Kind or Relation: {targets}" if targets else "")
+        + f". Rejected: {decisions.get('rejected', 0)}. Review the items, then merge."
+    )
+    return title[:300], description

lorebase_store/extraction/materialize.py

@@ -0,0 +1,364 @@
+"""Materialize the accepted items of an extraction Proposal on its branch (design 9.2 step 12).
+
+The extraction branch holds, for every element an item touches, the target branch's current
+content with the accepted items applied, and nothing else changes. Each call compares that
+desired state with the branch head and commits the difference as one Revision, so accepting
+an item adds its change and rejecting an accepted one reverts it. Elements no item touches keep
+whatever the branch holds. Calling it again without new decisions commits nothing.
+
+Rules:
+
+- a `record` item creates the Record (or sets its values on a Record that exists by now);
+- a `value` item sets one Attribute of an existing Record; it is skipped when the Record does
+  not exist;
+- a `link` item adds the Link when both endpoints exist (or are created by accepted items);
+  on a Relation that allows one Link per `from` (or `to`) Record, the Links it replaces are
+  removed; Links of foreign-key backed Relations are skipped (they derive from an Attribute);
+- every value set by an item is attributed to an `extraction` provenance entry with the job,
+  document, evidence offsets, model and prompt version.
+
+Validation is the commit pipeline's: an accepted Record that misses a required value fails the
+commit with `ValidationFailedError` and nothing is written.
+"""
+
+from __future__ import annotations
+
+from collections import defaultdict
+from collections.abc import Callable, Mapping, Sequence
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Any
+
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from lorebase_core.model import (
+    ExtractionOrigin,
+    IdPrefix,
+    Link,
+    ProvenanceEntry,
+    Record,
+    RelationBacking,
+    Schema,
+    new_id,
+)
+from lorebase_core.revision import Author, Origin
+from lorebase_core.types.encoding import format_timestamp
+from lorebase_core.validate import DataChange, LinkDelete, LinkUpsert, RecordDelete, RecordUpsert
+from lorebase_store.commit import Change, commit
+from lorebase_store.extraction.items import effective_subject
+from lorebase_store.models.enums import ItemDecision
+from lorebase_store.repositories.materialized import MaterializedRepository
+from lorebase_store.repositories.proposal_items import LINK_ITEM, RECORD_ITEM, VALUE_ITEM, ItemRow
+from lorebase_store.runtime import StoreRuntime
+
+EXTRACTION_ORIGIN = "extraction"
+RULES_MODEL = "rules"
+"""Model id recorded for items produced without a model (tier 0)."""
+
+RecordId = tuple[str, str]
+LinkId = tuple[str, str, str]
+Clock = Callable[[], datetime]
+
+
+@dataclass(frozen=True, slots=True)
+class MaterializeContext:
+    """Where and as whom to materialize."""
+
+    domain_id: str
+    branch: str
+    target: str
+    job_id: str
+    prompt_version: str
+    actor: str
+    runtime: StoreRuntime
+    clock: Clock
+
+
+@dataclass
+class Materialized:
+    """Outcome: the Revision committed (none when the branch already matched), the number of
+    changes, and the accepted items that could not be applied with the reason."""
+
+    revision: bytes | None = None
+    changes: int = 0
+    applied: list[str] = field(default_factory=list)
+    skipped: dict[str, str] = field(default_factory=dict)
+
+    def as_dict(self) -> dict[str, Any]:
+        return {
+            "revision": None if self.revision is None else self.revision.hex(),
+            "changes": self.changes,
+            "applied": len(self.applied),
+            "skipped": dict(self.skipped),
+        }
+
+
+@dataclass
+class _Plan:
+    records: dict[RecordId, dict[str, Any] | None] = field(default_factory=dict)
+    links: dict[LinkId, dict[str, Any] | None] = field(default_factory=dict)
+    record_sources: dict[RecordId, dict[str, str]] = field(default_factory=dict)
+    link_sources: dict[LinkId, str] = field(default_factory=dict)
+
+
+class _Provenance:
+    """One provenance entry per distinct evidence span of an item."""
+
+    def __init__(self, context: MaterializeContext) -> None:
+        self._context = context
+        self._entries: dict[tuple[str, int, int, str], ProvenanceEntry] = {}
+
+    def entry(self, item: ItemRow, evidence: Mapping[str, Any] | None) -> str:
+        found = evidence or item.evidence or {}
+        start = int(found.get("start") or 0)
+        end = max(int(found.get("end") or 0), start)
+        document = str(found.get("document_id") or "")
+        model = str(found.get("model") or item.model or RULES_MODEL)
+        key = (document, start, end, model)
+        if key not in self._entries:
+            self._entries[key] = ProvenanceEntry(
+                id=new_id(IdPrefix.PROVENANCE),
+                origin=ExtractionOrigin(
+                    job_id=item.job_id or self._context.job_id,
+                    document_id=document,
+                    span_start=start,
+                    span_end=end,
+                    model_id=model,
+                    prompt_version=self._context.prompt_version,
+                ),
+                actor=self._context.actor,
+                at=format_timestamp(self._context.clock()),
+            )
+        return self._entries[key].id
+
+    def used(self, ids: set[str]) -> list[ProvenanceEntry]:
+        return [entry for entry in self._entries.values() if entry.id in ids]
+
+
+async def materialize_items(
+    session: AsyncSession,
+    context: MaterializeContext,
+    items: Sequence[ItemRow],
+    schema: Schema,
+    *,
+    message: str,
+) -> Materialized:
+    """Make the extraction branch hold the accepted `items` (see the module docstring)."""
+    writes = context.runtime.settings.writes
+    target = MaterializedRepository(session, context.domain_id, context.target, writes=writes)
+    branch = MaterializedRepository(session, context.domain_id, context.branch, writes=writes)
+    ordered = sorted(items, key=lambda item: (item.seq, item.item_id))
+    subjects = {item.item_id: effective_subject(item) for item in ordered}
+
+    record_keys: dict[str, set[str]] = defaultdict(set)
+    link_pairs: dict[str, set[tuple[str, str]]] = defaultdict(set)
+    endpoints: dict[str, set[str]] = defaultdict(set)
+    for item in ordered:
+        for subject in (item.subject, subjects[item.item_id]):
+            if item.item_type == LINK_ITEM:
+                relation = schema.relations.get(str(subject["relation"]))
+                if relation is None:
+                    continue
+                pair = (str(subject["from"]), str(subject["to"]))
+                link_pairs[relation.id].add(pair)
+                endpoints[relation.from_kind].add(pair[0])
+                endpoints[relation.to_kind].add(pair[1])
+            else:
+                record_keys[str(subject["kind"])].add(str(subject["key"]))
+    await _replaced_links(target, schema, ordered, subjects, link_pairs)
+
+    plan = _Plan()
+    current_records: dict[RecordId, dict[str, Any]] = {}
+    for kind_id, keys in record_keys.items():
+        on_target = await target.records(kind_id, keys)
+        on_branch = await branch.records(kind_id, keys)
+        for key in keys:
+            row = on_target.get(key)
+            plan.records[(kind_id, key)] = None if row is None else dict(row.values)
+            if key in on_branch:
+                current_records[(kind_id, key)] = dict(on_branch[key].values)
+    current_links: dict[LinkId, dict[str, Any]] = {}
+    for relation_id, pairs in link_pairs.items():
+        on_target_links = await target.links(relation_id, pairs)
+        on_branch_links = await branch.links(relation_id, pairs)
+        for pair in pairs:
+            link = on_target_links.get(pair)
+            plan.links[(relation_id, *pair)] = None if link is None else dict(link.values)
+            if pair in on_branch_links:
+                current_links[(relation_id, *pair)] = dict(on_branch_links[pair].values)
+    existing: dict[str, set[str]] = {}
+    for kind_id, keys in endpoints.items():
+        existing[kind_id] = set(await branch.records(kind_id, keys))
+
+    provenance = _Provenance(context)
+    outcome = Materialized()
+    accepted = [item for item in ordered if item.decision is ItemDecision.ACCEPTED]
+    for item in accepted:
+        if item.item_type != LINK_ITEM:
+            _apply_record_item(item, subjects[item.item_id], plan, provenance, outcome)
+    for item in accepted:
+        if item.item_type == LINK_ITEM:
+            _apply_link_item(
+                item, subjects[item.item_id], schema, plan, existing, provenance, outcome
+            )
+
+    changes = _diff(plan, current_records, current_links)
+    outcome.changes = len(changes)
+    if changes:
+        used = {
+            value
+            for change in changes
+            for value in (
+                change.provenance.values()
+                if isinstance(change.provenance, Mapping)
+                else [change.provenance]
+            )
+            if value is not None
+        }
+        result = await commit(
+            session,
+            context.domain_id,
+            context.branch,
+            changes,
+            author=Author(principal=context.actor),
+            message=message,
+            origin=Origin(type=EXTRACTION_ORIGIN, ref=context.job_id),
+            provenance=provenance.used(used),
+            runtime=context.runtime,
+            clock=context.clock,
+        )
+        if not result.no_op:
+            outcome.revision = result.revision_id
+    return outcome
+
+
+async def _replaced_links(
+    target: MaterializedRepository,
+    schema: Schema,
+    items: Sequence[ItemRow],
+    subjects: Mapping[str, Mapping[str, Any]],
+    link_pairs: dict[str, set[tuple[str, str]]],
+) -> None:
+    """Add the target's Links that a Link item would replace on single-valued Relations."""
+    for item in items:
+        if item.item_type != LINK_ITEM:
+            continue
+        subject = subjects[item.item_id]
+        relation = schema.relations.get(str(subject["relation"]))
+        if relation is None:
+            continue
+        from_key, to_key = str(subject["from"]), str(subject["to"])
+        if relation.single_per_from:
+            for _near, far, _values in await target.linked_keys(
+                relation.id, near="from_key", keys=[from_key]
+            ):
+                link_pairs[relation.id].add((from_key, far))
+        if relation.single_per_to:
+            for _near, far, _values in await target.linked_keys(
+                relation.id, near="to_key", keys=[to_key]
+            ):
+                link_pairs[relation.id].add((far, to_key))
+
+
+def _apply_record_item(
+    item: ItemRow,
+    subject: Mapping[str, Any],
+    plan: _Plan,
+    provenance: _Provenance,
+    outcome: Materialized,
+) -> None:
+    element = (str(subject["kind"]), str(subject["key"]))
+    desired = plan.records.get(element)
+    sources = plan.record_sources.setdefault(element, {})
+    if item.item_type == RECORD_ITEM:
+        values = dict(subject.get("values") or {})
+        evidence = subject.get("evidence") or {}
+        desired = {**(desired or {}), **values}
+        for attribute_id in values:
+            sources[attribute_id] = provenance.entry(item, evidence.get(attribute_id))
+    elif item.item_type == VALUE_ITEM:
+        if desired is None:
+            outcome.skipped[item.item_id] = "the Record does not exist"
+            return
+        attribute_id = str(subject["attribute"])
+        desired = {**desired, attribute_id: subject.get("value")}
+        sources[attribute_id] = provenance.entry(item, None)
+    plan.records[element] = {k: v for k, v in (desired or {}).items() if v is not None}
+    outcome.applied.append(item.item_id)
+
+
+def _apply_link_item(
+    item: ItemRow,
+    subject: Mapping[str, Any],
+    schema: Schema,
+    plan: _Plan,
+    existing: Mapping[str, set[str]],
+    provenance: _Provenance,
+    outcome: Materialized,
+) -> None:
+    relation = schema.relations.get(str(subject["relation"]))
+    if relation is None:
+        outcome.skipped[item.item_id] = "the Relation no longer exists"
+        return
+    if relation.backing is RelationBacking.FOREIGN_KEY:
+        outcome.skipped[item.item_id] = "Links of this Relation derive from an Attribute"
+        return
+    from_key, to_key = str(subject["from"]), str(subject["to"])
+
+    def present(kind_id: str, key: str) -> bool:
+        if (kind_id, key) in plan.records:
+            return plan.records[(kind_id, key)] is not None
+        return key in existing.get(kind_id, set())
+
+    if not (present(relation.from_kind, from_key) and present(relation.to_kind, to_key)):
+        outcome.skipped[item.item_id] = "an endpoint Record does not exist"
+        return
+    element = (relation.id, from_key, to_key)
+    for other in list(plan.links):
+        if other[0] != relation.id or other == element:
+            continue
+        if (relation.single_per_from and other[1] == from_key) or (
+            relation.single_per_to and other[2] == to_key
+        ):
+            plan.links[other] = None
+    plan.links[element] = {**(plan.links.get(element) or {}), **(subject.get("values") or {})}
+    plan.link_sources[element] = provenance.entry(item, None)
+    outcome.applied.append(item.item_id)
+
+
+def _diff(
+    plan: _Plan,
+    current_records: Mapping[RecordId, dict[str, Any]],
+    current_links: Mapping[LinkId, dict[str, Any]],
+) -> list[Change]:
+    removals: list[Change] = []
+    upserts: list[Change] = []
+    for (relation_id, from_key, to_key), values in sorted(
+        plan.links.items(), key=lambda entry: entry[0]
+    ):
+        current = current_links.get((relation_id, from_key, to_key))
+        if values is None and current is not None:
+            removals.append(
+                Change(LinkDelete(relation=relation_id, from_key=from_key, to_key=to_key))
+            )
+    for (kind_id, key), values in sorted(plan.records.items(), key=lambda entry: entry[0]):
+        current = current_records.get((kind_id, key))
+        if values is None:
+            if current is not None:
+                removals.append(Change(RecordDelete(kind=kind_id, key=key)))
+        elif current != values:
+            sources = plan.record_sources.get((kind_id, key)) or {}
+            change: DataChange = RecordUpsert(record=Record(kind=kind_id, key=key, values=values))
+            upserts.append(Change(change, {a: p for a, p in sources.items() if a in values}))
+    for (relation_id, from_key, to_key), values in sorted(
+        plan.links.items(), key=lambda entry: entry[0]
+    ):
+        current = current_links.get((relation_id, from_key, to_key))
+        if values is not None and current != values:
+            link = Link(relation=relation_id, from_key=from_key, to_key=to_key, values=values)
+            upserts.append(
+                Change(
+                    LinkUpsert(link=link), plan.link_sources.get((relation_id, from_key, to_key))
+                )
+            )
+    return removals + upserts

lorebase_store/extraction/proposal.py

@@ -0,0 +1,65 @@
+"""Keep an extraction Proposal in step with its items: summary, checks and merge preview."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Any
+
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from lorebase_core.merge import MergePolicy
+from lorebase_core.model import Schema
+from lorebase_store.extraction.items import proposal_text, summarize
+from lorebase_store.extraction.materialize import MaterializeContext, Materialized
+from lorebase_store.merging import AttemptResult, MergeRequest, MergeService
+from lorebase_store.models.enums import ProposalStatus
+from lorebase_store.repositories.proposal_items import ItemRow
+from lorebase_store.repositories.proposals import ProposalRepository, ProposalRow
+
+EXTRACTION_CHECK = "extraction"
+"""Key of the item summary in `proposals.checks`."""
+
+
+def extraction_check(
+    items: Sequence[ItemRow], schema: Schema, outcome: Materialized
+) -> dict[str, Any]:
+    """The `extraction` entry of a Proposal's checks."""
+    return {
+        **summarize(items, schema),
+        "applied": len(outcome.applied),
+        "skipped": dict(outcome.skipped),
+        "revision": None if outcome.revision is None else outcome.revision.hex(),
+    }
+
+
+async def refresh_proposal(
+    session: AsyncSession,
+    context: MaterializeContext,
+    proposal: ProposalRow,
+    items: Sequence[ItemRow],
+    schema: Schema,
+    outcome: Materialized,
+    *,
+    documents: Sequence[str],
+    policy: MergePolicy | None = None,
+) -> tuple[ProposalRow, AttemptResult | None]:
+    """Rewrite the Proposal's title, description and `extraction` check from its items, then
+    recompute the merge preview when the Proposal is open."""
+    summary = summarize(items, schema)
+    title, description = proposal_text(summary, documents)
+    checks = dict(proposal.checks or {})
+    checks[EXTRACTION_CHECK] = extraction_check(items, schema, outcome)
+    proposals = ProposalRepository(session, context.domain_id)
+    updated = await proposals.update(
+        proposal.id, at=context.clock(), title=title, description=description, checks=checks
+    )
+    assert updated is not None
+    if updated.status is not ProposalStatus.OPEN:
+        return updated, None
+    service = MergeService(session, context.domain_id, runtime=context.runtime, clock=context.clock)
+    result = await service.recompute(
+        MergeRequest(proposal=updated, actor=context.actor, policy=policy or MergePolicy())
+    )
+    final = await proposals.get(proposal.id)
+    assert final is not None
+    return final, result