gatewaysdk 0.3.2__py3-none-any.whl

gatewaysdk/environment/schema/extract.py

@@ -0,0 +1,263 @@
+"""``data extract``: rows out of what an agent already has, shaped for the world.
+
+Three inputs an agent holds, one output every import accepts (tagged JSONL)::
+
+    tool-results     one {..., "raw_output": {...}} per line; rows are the objects of
+                     raw_output.results[] (or raw_output.items[]) that carry the
+                     entity's primary key
+    vendor-envelope  the vendor's own response, one per line or one document; rows sit
+                     under the world's [api].envelope key ("" = the bare array)
+    session-export   calls.jsonl from `worlds session export`; rows are
+                     result.<envelope>[] of successful `call` records
+
+Every adapter then drops keys whose value is null or "", strips keys the entity's
+schema does not declare (counted, never silent), dedupes by the entity's primary
+key (``fill``: first row wins and later ones only add missing fields; ``upsert``:
+the last row wins) and applies the world's ``[redact]`` policy from
+``connector.toml`` (``drop`` removed, ``hash`` spelled ``sha256:<hex>`` over
+``<slug>:<value>``, ``round`` rounded) unless told not to. The summary goes to
+stderr::
+
+    {"read": N, "rows": M, "unknownFields": {name: n}, "redacted": {"dropped": n, "hashed": n}}
+
+Input is streamed. The deduplicated rows are held keyed until the input ends,
+because a later row may still change an earlier one.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import IO, Any, Iterable, Iterator
+
+from .batch import _Reader, _stream_array, encode_row
+from .compiler import Contract, build
+from .connector import SPEC, Spec, parse_spec, redact
+from .validation import ContractError, fail
+
+SHAPES = ("tool-results", "vendor-envelope", "session-export")
+MERGES = ("fill", "upsert")
+
+
+@dataclass
+class Summary:
+    read: int = 0
+    rows: int = 0
+    unknown: dict[str, int] = field(default_factory=dict)
+    dropped: int = 0
+    hashed: int = 0
+
+    def to_json(self) -> dict:
+        return {
+            "read": self.read,
+            "rows": self.rows,
+            "unknownFields": dict(sorted(self.unknown.items())),
+            "redacted": {"dropped": self.dropped, "hashed": self.hashed},
+        }
+
+
+# ---------------------------------------------------------------- the world
+
+
+def envelope_key(world: Path) -> str:
+    """``[api].envelope`` of the world's connector.toml; "" when unset or absent."""
+    try:
+        import tomllib
+    except ModuleNotFoundError:  # pragma: no cover - python 3.10
+        import tomli as tomllib
+
+    spec = world / SPEC
+    if not spec.is_file():
+        return ""
+    api = tomllib.loads(spec.read_text(encoding="utf-8")).get("api", {})
+    value = api.get("envelope", "") if isinstance(api, dict) else ""
+    return value if isinstance(value, str) else ""
+
+
+def redaction_spec(world: Path) -> Spec | None:
+    spec = world / SPEC
+    return parse_spec(spec.read_text(encoding="utf-8")) if spec.is_file() else None
+
+
+# ---------------------------------------------------------------- adapters
+
+
+def _lines(stream: IO[str], where: str) -> Iterator[Any]:
+    for number, line in enumerate(stream, start=1):
+        if not line.strip():
+            continue
+        try:
+            yield json.loads(line)
+        except ValueError:
+            fail("extract", f"{where}:{number} is not JSON", where)
+
+
+def _objects(value: Any) -> Iterator[dict]:
+    if isinstance(value, list):
+        for item in value:
+            if isinstance(item, dict):
+                yield item
+    elif isinstance(value, dict):
+        yield value
+
+
+def _dig(value: Any, path: str) -> Any:
+    for part in [p for p in path.split(".") if p]:
+        if not isinstance(value, dict) or part not in value:
+            return None
+        value = value[part]
+    return value
+
+
+def tool_results(stream: IO[str], where: str, *, keys: list[str], envelope: str = "") -> Iterator[dict]:
+    """Every object of ``raw_output.results[]`` (or ``.items[]``) that carries the key."""
+    for record in _lines(stream, where):
+        raw = record.get("raw_output") if isinstance(record, dict) else None
+        if not isinstance(raw, dict):
+            continue
+        listed = raw.get("results")
+        if listed is None:
+            listed = raw.get("items")
+        for item in _objects(listed):
+            if all(name in item for name in keys):
+                yield item
+
+
+def vendor_envelope(stream: IO[str], where: str, *, keys: list[str], envelope: str) -> Iterator[dict]:
+    """Rows under the envelope key of each response line, or of the one document."""
+    reader = _Reader(stream)
+    head = reader.peek()
+    if head is None:
+        return
+    if head == "[":
+        if envelope:
+            fail("extract", f"{where} is a bare array but the world's envelope is {envelope!r}", where)
+        yield from _stream_array(reader, where)
+        return
+    if head != "{":
+        fail("extract", f"{where} is neither a JSON object nor an array", where)
+    if not envelope:
+        fail("extract", f"{where} is an object but the world's envelope is empty: expected a bare array", where)
+    start = reader.pos
+    first = reader.line(cap=1024 * 1024)
+    one_per_line = False
+    if first is not None and len(first) <= 1024 * 1024:
+        try:
+            json.loads(first)
+            one_per_line = True
+        except ValueError:
+            one_per_line = False
+    reader.pos = start
+    if one_per_line:
+        line = reader.line()
+        number = 0
+        while line is not None:
+            number += 1
+            if line.strip():
+                try:
+                    response = json.loads(line)
+                except ValueError:
+                    fail("extract", f"{where}:{number} is not JSON", where)
+                yield from _objects(_dig(response, envelope))
+            reader.compact()
+            line = reader.line()
+        return
+    # One pretty-printed response: a single page, read whole.
+    yield from _objects(_dig(reader.value(where), envelope))
+
+
+def session_export(stream: IO[str], where: str, *, keys: list[str], envelope: str) -> Iterator[dict]:
+    """``result.<envelope>[]`` of successful ``call`` records."""
+    for record in _lines(stream, where):
+        if not isinstance(record, dict) or record.get("kind") != "call" or record.get("error") is not None:
+            continue
+        result = record.get("result")
+        yield from _objects(_dig(result, envelope) if envelope else result)
+
+
+ADAPTERS = {"tool-results": tool_results, "vendor-envelope": vendor_envelope, "session-export": session_export}
+
+
+# ---------------------------------------------------------------- the pipeline
+
+
+def _clean(row: dict, declared: set[str], summary: Summary) -> dict:
+    out: dict = {}
+    for name, value in row.items():
+        if value is None or value == "":
+            continue
+        if name not in declared:
+            summary.unknown[name] = summary.unknown.get(name, 0) + 1
+            continue
+        out[name] = value
+    return out
+
+
+def _count_redaction(before: dict, after: dict, spec: Spec, summary: Summary) -> None:
+    for name in before:
+        if name not in after:
+            summary.dropped += 1
+        elif name in spec.hash_fields and before[name] is not None:
+            summary.hashed += 1
+
+
+def extract(
+    world: str | Path,
+    sources: Iterable[str | Path],
+    *,
+    shape: str,
+    entity: str,
+    merge: str = "fill",
+    redaction: bool = True,
+    out: IO[str],
+) -> Summary:
+    """Run one adapter over the sources, write tagged JSONL to ``out``, return the summary."""
+    if shape not in SHAPES:
+        fail("extract", f"shape must be one of {', '.join(SHAPES)}", "/shape")
+    if merge not in MERGES:
+        fail("extract", f"merge must be one of {', '.join(MERGES)}", "/merge")
+    root = Path(world)
+    contract: Contract = build(root)
+    if entity not in contract.entities:
+        fail("entity", f"Unknown entity {entity}; the world has {', '.join(sorted(contract.entities))}", f"/{entity}")
+    spec_entity = contract.entities[entity]
+    keys = list(spec_entity["primary_key"])
+    declared = set(spec_entity["schema"].get("properties", {}))
+    envelope = envelope_key(root)
+    policy = redaction_spec(root) if redaction else None
+    adapter = ADAPTERS[shape]
+    summary = Summary()
+    held: dict[str, dict] = {}
+    from .batch import _open_text
+
+    for source in sources:
+        where = str(source)
+        with _open_text(source) as stream:
+            for raw in adapter(stream, where, keys=keys, envelope=envelope):
+                summary.read += 1
+                row = _clean(raw, declared, summary)
+                if not all(name in row for name in keys):
+                    continue
+                key = json.dumps([row[name] for name in keys], sort_keys=True, separators=(",", ":"))
+                if key in held:
+                    if merge == "upsert":
+                        held[key] = row
+                    else:
+                        first = held[key]
+                        for name, value in row.items():
+                            if name not in first:
+                                first[name] = value
+                else:
+                    held[key] = row
+    for row in held.values():
+        if policy is not None:
+            redacted = redact(row, policy)
+            _count_redaction(row, redacted, policy, summary)
+            row = redacted
+        out.write(encode_row({"entity": entity, "row": row}) + "\n")
+        summary.rows += 1
+    return summary
+
+
+__all__ = ["ADAPTERS", "MERGES", "SHAPES", "ContractError", "Summary", "extract"]