rag-wright 0.1.0__py3-none-any.whl

rag_wright/store/chunk_text.py

@@ -0,0 +1,66 @@
+"""The chunk-text sidecar (T40, FR-I.3): the per-chunk full text the retrieval index omits.
+
+The index is dense-over-summary by design: `ChunkRecord` holds the summary and the vectors, never the
+raw chunk text (the text is used once to compute the `chunk_id` content hash, then discarded). But
+synthesis (FR-Q.5) extracts over full chunk text, so the text must be persisted somewhere keyed by
+`chunk_id` and rehydrated at query time by `chunk_read` (T38). This is that store: one JSON file per
+source document, mapping the canonical `chunk_id` string to its full text.
+
+It is deliberately NOT part of the swappable `Store` seam (ArcadeDB / LanceDB): it holds no vectors and
+answers no query, it only round-trips text by id, so it does not belong behind the retrieval seam. Its
+lifecycle is coupled to the chunk lifecycle: written at chunk write under the same content-hash gate as
+the index upsert (so the two never diverge), and deleted per `source_doc_id` when a document is
+re-chunked (`delete_document`, the T34 seam), so stale text is not orphaned relative to the index.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from pathlib import Path
+from typing import Optional
+
+from rag_wright.contracts.identifiers import ChunkId
+
+
+class ChunkTextStore:
+    """A per-source-document JSON sidecar mapping `chunk_id` -> full chunk text (T40, FR-I.3)."""
+
+    def __init__(self, root: Path) -> None:
+        self._root = Path(root)
+        self._root.mkdir(parents=True, exist_ok=True)
+
+    def put(self, chunk_id: ChunkId, text: str) -> None:
+        """Persist `text` for `chunk_id`, upserting by id under the document's sidecar file.
+
+        Integrity check: the `chunk_id` already carries the content hash of its text (the same hash
+        `ChunkId.of` computed at chunking, over these exact bytes), so `text` MUST hash to it. This is
+        the sidecar's core promise made provable — a mismatch is a silent-wrong-text bug (a loop index
+        error, a mismatched map) that would otherwise surface only as synthesis citing a valid-looking
+        but wrong `chunk_id`; it is rejected here at the boundary, before it can be persisted.
+        """
+        digest = hashlib.sha256(text.encode("utf-8")).hexdigest()
+        if digest != chunk_id.content_hash:
+            raise ValueError(
+                f"chunk-text integrity: text for {chunk_id.value!r} hashes to {digest!r}, "
+                f"not the chunk_id's content_hash {chunk_id.content_hash!r}"
+            )
+        path = self._path(chunk_id.source_doc_id)
+        mapping = json.loads(path.read_text(encoding="utf-8")) if path.exists() else {}
+        mapping[chunk_id.value] = text
+        path.write_text(json.dumps(mapping, ensure_ascii=False), encoding="utf-8")
+
+    def get(self, chunk_id: str) -> Optional[str]:
+        """The persisted text for `chunk_id` (canonical string form), or None if absent."""
+        source_doc_id = chunk_id.rsplit(":", 2)[0]  # <source_doc_id>:<chunk_index>:<content_hash>
+        path = self._path(source_doc_id)
+        if not path.exists():
+            return None
+        return json.loads(path.read_text(encoding="utf-8")).get(chunk_id)
+
+    def delete_document(self, source_doc_id: str) -> None:
+        """Remove all persisted text for a document (the T34 delete-and-re-chunk seam). Idempotent."""
+        self._path(source_doc_id).unlink(missing_ok=True)
+
+    def _path(self, source_doc_id: str) -> Path:
+        return self._root / f"{source_doc_id}.json"

rag_wright/store/seam.py

@@ -0,0 +1,213 @@
+"""The store seam: the swappable interface every store implementation binds (T13, FR-S.5).
+
+The store is reached only through this interface so it can be swapped without touching capability
+code: ArcadeDB is the default (one multi-model store for both the hybrid index and the graph,
+FR-S.1), and a separate hybrid vector store (LanceDB) is the eval-gated Phase 1 fallback for the
+retrieval leg (GATE-2). The seam is intentionally semantic, not SQL: it exposes schema readiness and
+introspection, never a query string, so a second implementation can bind it without inheriting
+ArcadeDB's SQL dialect. Write-side and query-side methods are added by the tasks that need them
+(T20, T21, T26); T13 defines only the schema-management surface the foundation needs.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional, Protocol, runtime_checkable
+
+from rag_wright.contracts.chunk import ChunkRecord, MetadataValue
+
+
+class _NotNull:
+    """Sentinel for a `kg_edges` where-value meaning `<field> IS NOT NULL` (vs an equality/membership match)."""
+
+    def __repr__(self) -> str:  # pragma: no cover - debug aid
+        return "NOT_NULL"
+
+
+NOT_NULL = _NotNull()  # kg_edges where-value: presence test, e.g. edge_where={"dimension": NOT_NULL}
+
+
+@dataclass(frozen=True)
+class GraphNode:
+    """A graph entity node to write (T25). `node_key` is the vertex identity (the resolver's canonical id when
+    linked, an `UNLINKED:<key>` surrogate otherwise); `entity_id` is the canonical id, or empty when unlinked."""
+
+    node_key: str
+    entity_id: str
+    name: str
+    entity_type: str
+    confidence: str
+    chunk_id: str
+
+
+@dataclass(frozen=True)
+class GraphEdge:
+    """A relationship edge between two entity nodes (by `node_key`), carrying provenance + confidence."""
+
+    source_key: str
+    target_key: str
+    relationship_type: str
+    confidence: str
+    chunk_id: str
+
+
+@dataclass(frozen=True)
+class KgNode:
+    """A typed KG node to upsert (DD-1b, ADR-0117): `type` is the vertex type, `key_field` the identity field to
+    upsert on, `props` the fields (including `key_field`) as DOMAIN-NATIVE values. The store encodes each prop by its
+    pack-declared storage type -- the caller never serializes to the store's wire format."""
+
+    type: str
+    key_field: str
+    props: dict[str, object]
+
+
+@dataclass(frozen=True)
+class KgEdge:
+    """A typed KG edge to create between two nodes identified by (type, key_field, key). `props` are native values
+    (edge properties are type-driven: edges declare no storage schema)."""
+
+    type: str
+    from_type: str
+    from_key_field: str
+    from_key: object
+    to_type: str
+    to_key_field: str
+    to_key: object
+    props: dict[str, object]
+
+
+@runtime_checkable
+class Store(Protocol):
+    """A swappable store. The ArcadeDB implementation is the default; a stub proves swappability."""
+
+    def ensure_schema(self) -> None:
+        """Create the chunk-record and graph-node types and their indexes, idempotently."""
+
+    def type_names(self) -> set[str]:
+        """The names of the types (tables/classes) present in the store."""
+
+    def property_names(self, type_name: str) -> set[str]:
+        """The property names declared on `type_name` (empty if the type is absent)."""
+
+    def index_names(self) -> set[str]:
+        """The names of the indexes present in the store."""
+
+    def ping(self) -> bool:
+        """True if the store is reachable."""
+
+    def close(self) -> None:
+        """Release any resources held by the implementation."""
+
+    # --- write-side (T20): the chunk-record write leg. Semantic, not SQL: the seam takes the T3
+    # ChunkRecord and each store maps it to its own representation (ArcadeDB decomposes the sparse
+    # vector into two parallel arrays; a LanceDB fallback would store it its own way).
+
+    def upsert_chunk(self, record: ChunkRecord) -> None:
+        """Write a chunk record, upserting by `chunk_id` (re-write of the same id updates in place)."""
+
+    def get_chunk(self, chunk_id: str) -> Optional[dict]:
+        """The stored row for `chunk_id` (store-native fields), or None if absent."""
+
+    def chunk_count(self) -> int:
+        """The number of chunk records in the store."""
+
+    # --- query-side (T21): the hybrid retrieval leg. Semantic, not SQL: the seam takes the two
+    # query vectors and each store fuses them its own way (ArcadeDB by server-side RRF over its
+    # dense/sparse indexes; a LanceDB fallback by its own hybrid query), so FR-C.3 is swappable.
+
+    def hybrid_search(
+        self,
+        dense_query: list[float],
+        sparse_query: dict[int, float],
+        *,
+        k: int,
+        filters: Optional[dict[str, MetadataValue]] = None,
+    ) -> list[dict]:
+        """Fuse the dense and sparse legs into one ranked candidate list (Reciprocal Rank Fusion),
+        honoring equality metadata filters, returning up to `k` rows (each with at least `chunk_id`
+        and `source_doc_id`) in ranked order, best first."""
+
+    # --- graph-write (T25): the knowledge-graph leg. Semantic, not SQL: the seam takes entity nodes
+    # and relationship edges and each store writes them its own way (ArcadeDB as vertices/edges in
+    # one transaction; a fallback store however it models a graph). Nodes/edges carry chunk_id (FR-I.4).
+
+    def write_graph(self, nodes: list[GraphNode], edges: list[GraphEdge]) -> None:
+        """Write entity nodes (upsert by `node_key`) and relationship edges between them in ONE
+        transaction (FR-S.1: a chunk and its entities land together), connecting each entity to its
+        source chunk. Nodes and edges carry `chunk_id` and confidence (FR-I.4)."""
+
+    def graph_counts(self) -> dict[str, int]:
+        """Counts for introspection/tests: `{'entities': n, 'relationships': m}`."""
+
+    def entities_by_name(self, name: str) -> list[dict]:
+        """Resolve a party NAME to its graph entities (issue 0030): the first step before
+        `graph_neighbors`/`graph_query`, which take a `start_entity_id` (an exact node key) and cannot be
+        reached from a name otherwise. Returns `[{entity_id, name, entity_type}]` for every stored entity
+        whose name normalizes to the same clustering key as `name`, via the SAME `normalize_entity_name`
+        the ingestion side uses to merge 'Acme Corp' / 'Acme Corporation' / 'ACME, Inc.' into one entity.
+        Normalization is the engine's rule and is applied HERE, so a caller never re-implements it (a raw
+        or an already-normalized name both work; the normalization is idempotent). `entity_id` is exactly
+        the node key `graph_neighbors`/`graph_query` take as `start_entity_id`. Empty list on no match."""
+
+    # --- graph-query (T26): relationship traversal. Semantic, not SQL: returns store-agnostic path
+    # rows (target + the entity_ids/chunk_ids/confidences along the path) so the capability can shape
+    # the cited evidence. Confidence is SURFACED on every path, not filtered on (FR-C.5/FR-Q.3).
+
+    def graph_neighbors(
+        self, entity_id: str, *, relationship_type: str, max_hops: int, documents: Optional[list[str]] = None
+    ) -> list[dict]:
+        """Traverse `relationship_type` edges from the start entity up to `max_hops`, returning one row
+        per reached entity+path: `target_id`, `target_name`, `path_entity_ids`, `path_chunk_ids`,
+        `path_confidences`, `hops`. Every edge is surfaced regardless of confidence (T26 does not gate).
+        `documents` (issue 0031): scope the traversal to those source documents -- EVERY edge on a path must
+        belong to one of them; `None` = the whole graph, `[]` = scope-to-nothing (no rows)."""
+
+    # --- generic typed-node read (DD-1a, ADR-0117): the backend-agnostic primitive a domain store extension
+    # delegates to, so a domain pack never emits store-native SQL. Semantic, not SQL.
+
+    def kg_read(
+        self,
+        node_type: str,
+        *,
+        where: Optional[dict[str, object]] = None,
+        fields: Optional[list[str]] = None,
+        distinct: Optional[str] = None,
+        order_by: Optional[str] = None,
+        limit: Optional[int] = None,
+    ) -> list[dict]:
+        """Read typed nodes of `node_type`. `where` maps a field to a scalar (equality) or a list (membership);
+        a list value that is EMPTY means scope-to-nothing and returns `[]` without a query. `distinct` returns the
+        distinct values of one field; `fields=None` returns all fields. Equality/membership clauses are AND-ed."""
+
+    def kg_edges(
+        self,
+        from_type: Optional[str] = None,
+        *,
+        where: Optional[dict[str, object]] = None,
+        key_range: Optional[tuple[str, object, object]] = None,
+        direction: str = "out",
+        edge_type: Optional[str] = None,
+        edge_where: Optional[dict[str, object]] = None,
+        target_where: Optional[dict[str, object]] = None,
+        select: dict[str, str],
+    ) -> list[dict]:
+        """Generic, backend-agnostic edge TRAVERSAL -- the primitive a domain store extension delegates to so it
+        never emits store-native traversal SQL (EP-REF-1a, ADR-0117). Three idioms behind one surface:
+
+        - **node-start traversal** (give a start selector): from the `from_type` nodes matched by `where`
+          (scalar=equality, list=membership, `NOT_NULL`=presence; a field in `key_range=(field, lo, hi)` adds
+          `field >= lo AND field < hi`, the contract-scope range), follow `direction="out"` (outgoing edges to
+          the TARGET vertex) or `"in"` (incoming edges to the SOURCE vertex). `edge_type=None` means every edge
+          in that direction. `edge_where` filters the edge, `target_where` the reached vertex.
+        - **edge scan** (give NEITHER `where` NOR `key_range`): scan the `edge_type` table directly, filtered by
+          `edge_where` -- for edge properties that are not reachable from a node key (e.g. a span id on the edge).
+
+        `select` maps each output alias to an expression: `c.<f>` (start node), `e.<f>` / `e.@type` (edge),
+        `v.<f>` (far vertex) for a traversal; or a bare edge field / `inV().<f>` / `outV().<f>` for an edge scan.
+        A membership value that is an EMPTY list means scope-to-nothing -> `[]` with no query. Returns the rows."""
+
+    def kg_write(self, nodes: list["KgNode"], edges: "Iterable[KgEdge]" = ()) -> None:
+        """Upsert typed `nodes` (by each node's `key_field`) then create typed `edges` (FROM/TO by node key), ALL in
+        ONE transaction, nodes first so endpoints exist. The caller passes DOMAIN-NATIVE values; the store owns all
+        wire encoding, driven by each node type's pack-declared property storage type. Empty input is a no-op."""

rag_wright/subgraphs/async_ingestion.py

@@ -0,0 +1,204 @@
+"""PROD-3 (ADR-0050): async, job-based ingestion -- submit returns a job_id immediately, the corpus ingests in
+the background with bounded document parallelism, and progress is POLLED via a status read (never a blocking run
+with a spinner).
+
+This module holds the durable job model + store (2a). The async runner (2b) and the submit/status MCP tools (2c)
+build on it. The per-document pipeline stays the existing LangGraph subgraph (with its RetryPolicy +
+Increment-1 dead-letter/partial); this is the thin async envelope around `run_corpus_ingestion`'s work.
+
+The JobStore is file-based (one <job_id>.json per job): process-independent, so a status reader in another
+process (an MCP call) sees live progress; durable, so a crashed/restarted runner's job record survives; and it
+needs no KG schema change for the MVP. The full version can move the store to ArcadeDB/Postgres (ADR-0050).
+"""
+from __future__ import annotations
+
+import asyncio
+import os
+import threading
+from datetime import datetime, timezone
+from enum import Enum
+from pathlib import Path
+from typing import Any, Callable
+
+from pydantic import BaseModel, Field
+
+from rag_wright.subgraphs.contract_ingestion_pipeline import (  # ENG-1 shape + 0009 deferred parse
+    PendingDocument,
+    aparse_pending,
+    build_partial_entry,
+)
+
+
+def _now() -> str:
+    return datetime.now(timezone.utc).isoformat()
+
+
+class JobStatus(str, Enum):
+    QUEUED = "queued"      # submitted, not yet started
+    RUNNING = "running"    # the background runner is processing documents
+    SUCCEEDED = "succeeded"  # all documents accounted for (ingested / partial / dead-lettered), runner finished
+    FAILED = "failed"      # the runner itself errored irrecoverably (NOT a per-document failure -> that's dead_lettered)
+
+
+class IngestionJob(BaseModel):
+    """A durable ingestion job. `dead_lettered` / `partial` carry the PROD-3 lossless outcomes (a failure is
+    KNOWN here at completion, never grep-only). Progress = `documents_done` / `documents_total`."""
+
+    job_id: str
+    corpus_ref: dict                      # e.g. {"kind": "gcs", "bucket": ..., "prefix": ..., "include": [...]}
+    db: str                               # the target KG database
+    status: JobStatus = JobStatus.QUEUED
+    documents_total: int = 0
+    documents_done: int = 0               # ingested (incl. partial) + dead-lettered + resume-skipped
+    ingested: int = 0
+    dead_lettered: list[dict] = Field(default_factory=list)
+    partial: list[dict] = Field(default_factory=list)
+    party_links: int = 0
+    created_at: str = Field(default_factory=_now)
+    updated_at: str = Field(default_factory=_now)
+    error: str | None = None              # a runner-level (not per-document) failure
+
+    @property
+    def done(self) -> bool:
+        return self.status in (JobStatus.SUCCEEDED, JobStatus.FAILED)
+
+
+def _atomic_write(path: Path, text: str) -> None:
+    """Write `text` to `path` atomically: write a temp file, then `os.replace` (an atomic rename on POSIX and
+    Windows). A concurrent reader therefore sees either the old complete file or the new complete one -- never a
+    truncated/empty file mid-write (the JobStore read-mid-write race, ADR-0057 B2e/B4)."""
+    tmp = path.with_name(f"{path.name}.tmp.{os.getpid()}")
+    tmp.write_text(text, encoding="utf-8")
+    os.replace(tmp, path)
+
+
+class JobStore:
+    """File-based job store: one `<job_id>.json` per job under `jobs_dir`. Read-modify-write `update` is safe for
+    a SINGLE writer per job (the runner's orchestrator coroutine updates the record; parallel document workers
+    report back to it, they do not write the file) -- so no cross-writer race in the MVP."""
+
+    def __init__(self, jobs_dir: Any) -> None:
+        self._dir = Path(jobs_dir)
+        self._dir.mkdir(parents=True, exist_ok=True)
+
+    def _path(self, job_id: str) -> Path:
+        return self._dir / f"{job_id}.json"
+
+    def create(self, job: IngestionJob) -> IngestionJob:
+        _atomic_write(self._path(job.job_id), job.model_dump_json(indent=2))
+        return job
+
+    def get(self, job_id: str) -> IngestionJob | None:
+        path = self._path(job_id)
+        if not path.exists():
+            return None
+        text = path.read_text(encoding="utf-8")
+        # atomic writes (create) mean a reader never sees a partial file; tolerate an empty read defensively
+        # (e.g. a truncated legacy write) as "not ready yet" rather than raising.
+        return IngestionJob.model_validate_json(text) if text.strip() else None
+
+    def update(self, job_id: str, **fields: Any) -> IngestionJob:
+        """Read-modify-write the job's fields, stamping `updated_at`. Raises KeyError if the job is unknown."""
+        job = self.get(job_id)
+        if job is None:
+            raise KeyError(job_id)
+        updated = job.model_copy(update={**fields, "updated_at": _now()})
+        return self.create(updated)
+
+    def list_jobs(self) -> list[IngestionJob]:
+        return [IngestionJob.model_validate_json(p.read_text(encoding="utf-8")) for p in sorted(self._dir.glob("*.json"))]
+
+
+# --- 2b: the async runner (submit returns immediately; a background thread ingests in parallel) ----------------
+
+
+async def run_job(
+    job_id: str,
+    documents: list,
+    ingest_graph: Any,
+    store: JobStore,
+    *,
+    link_fn: Callable[[], int] = lambda: 0,
+    is_done: Callable[[Any], bool] = lambda _doc: False,
+    max_concurrency: int = 8,
+) -> IngestionJob:
+    """Ingest a materialized document list in the background with bounded parallelism, updating the job record
+    as each document completes (so status polling sees live progress). Per-document dead-letter / partial (the
+    Increment-1 lossless outcomes) are accumulated onto the job. A per-document failure NEVER fails the job -- it
+    is dead_lettered; only a runner-level error sets status=FAILED. Updates come from THIS single orchestrator
+    coroutine (asyncio is single-threaded; `store.update` has no await), so there is no cross-writer file race."""
+    try:
+        store.update(job_id, status=JobStatus.RUNNING, documents_total=len(documents))
+        sem = asyncio.Semaphore(max_concurrency)
+        dead_lettered: list[dict] = []
+        partial: list[dict] = []
+        done = 0
+        ingested = 0
+
+        async def _one(item: Any) -> tuple[str, Any, Any]:
+            async with sem:  # bound concurrent LLM/DB + PARSE work (rate limits)
+                if is_done(item):  # RESUME: a prior run already wrote this document (skip before parsing)
+                    return ("skip", item, None)
+                doc = item
+                try:
+                    # 0009-ASYNC-INGEST: parse a deferred doc HERE -- concurrently + deadline-bounded, off the loop
+                    # (its tiered OCR/VLM escalation is the slowest call) -- so it never blocks the others.
+                    if isinstance(item, PendingDocument):
+                        doc = await aparse_pending(item)
+                    # ASYNC-B2e (ADR-0057): ainvoke runs an async-node graph on the loop (true deadline) and a
+                    # sync-node graph in LangGraph's threadpool -- so it is safe on any compiled graph.
+                    out = await ingest_graph.ainvoke({"document": doc})  # per-doc LangGraph graph
+                except Exception as exc:  # noqa: BLE001 - a per-doc CRASH/parse-timeout must dead-letter THAT doc,
+                    out = {"dead_letter": {                                          # never fail the whole job
+                        "source_doc_id": item.source_doc_id, "stage": "invoke",
+                        "reason": "ingest_crashed", "error": str(exc)[:200]}}
+                return ("out", doc, out)
+
+        for coro in asyncio.as_completed([_one(doc) for doc in documents]):
+            kind, doc, out = await coro
+            done += 1
+            if kind == "out" and out.get("dead_letter"):
+                dead_lettered.append(out["dead_letter"])
+            else:
+                ingested += 1
+                ocr_failures = [{"page": pg, "reason": "unreadable scan (OCR + VLM failed)"}  # 0009-WIRE2
+                                for pg in (getattr(doc, "ocr_unreadable_pages", None) or [])]
+                entry = build_partial_entry(  # ENG-1: same shape as the blocking driver; also surfaces span/ocr losses
+                    doc.source_doc_id, (out or {}).get("clause_failures"), (out or {}).get("span_failures"),
+                    ocr_failures)
+                if entry is not None:
+                    partial.append(entry)
+            store.update(job_id, documents_done=done, ingested=ingested,
+                         dead_lettered=dead_lettered, partial=partial)
+
+        links = link_fn()  # KG-7 party linking ONCE, after all documents
+        return store.update(job_id, status=JobStatus.SUCCEEDED, party_links=links)
+    except Exception as exc:  # noqa: BLE001 - a RUNNER-level failure (not a per-document one) -> job FAILED, visible
+        return store.update(job_id, status=JobStatus.FAILED, error=str(exc)[:500])
+
+
+def submit_ingestion(
+    adapter: Any,
+    ingest_graph: Any,
+    store: JobStore,
+    *,
+    job_id: str,
+    db: str,
+    corpus_ref: dict,
+    link_fn: Callable[[], int] = lambda: 0,
+    is_done: Callable[[Any], bool] = lambda _doc: False,
+    max_concurrency: int = 8,
+) -> str:
+    """Create a QUEUED job, start the background runner, and return the `job_id` IMMEDIATELY (non-blocking). The
+    runner ingests in a daemon thread with its own event loop; poll `store.get(job_id)` (or the status MCP tool)
+    for progress. NOTE (MVP): the daemon thread lives with the submitting process -- a persistent service/worker
+    (or LangGraph Platform / Pub-Sub, the 'full' version) is what survives process exit; documented in ADR-0050."""
+    store.create(IngestionJob(job_id=job_id, corpus_ref=corpus_ref, db=db))
+
+    def _worker() -> None:
+        documents = list(adapter.documents())  # materialize inside the worker (may hit the network)
+        asyncio.run(run_job(job_id, documents, ingest_graph, store,
+                            link_fn=link_fn, is_done=is_done, max_concurrency=max_concurrency))
+
+    threading.Thread(target=_worker, daemon=True, name=f"ingest-{job_id}").start()
+    return job_id