feat(preprocess): publish incremental Evidence corpus
This commit is contained in:
@@ -0,0 +1,175 @@
|
||||
"""Incremental Evidence preprocessing with generation-isolated vector writes."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import uuid
|
||||
from dataclasses import asdict, dataclass
|
||||
|
||||
from tht.corpus.chunk import ChunkPolicy, chunk
|
||||
from tht.corpus.models import CanonicalChunk, CanonicalDocument, CorpusManifest
|
||||
from tht.corpus.normalize import normalize
|
||||
from tht.corpus.store import CorpusStore
|
||||
from tht.ports.evidence import EvidenceSource, SourceObject
|
||||
from tht.ports.vector import VectorStore, VectorWriteRecord
|
||||
from tht.vectorstore.records import VectorRecord
|
||||
|
||||
|
||||
class PipelineError(RuntimeError):
|
||||
"""Credential-free failure at the preprocessing boundary."""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PipelineResult:
|
||||
status: str
|
||||
generation: str | None
|
||||
published: bool
|
||||
changed: tuple[str, ...]
|
||||
unchanged: tuple[str, ...]
|
||||
removed: tuple[str, ...]
|
||||
manifest: CorpusManifest
|
||||
|
||||
def model_dump(self, mode=None):
|
||||
value = asdict(self)
|
||||
value["manifest"] = self.manifest.model_dump(mode="json")
|
||||
return value
|
||||
|
||||
|
||||
def _fingerprint(value) -> str:
|
||||
payload = json.dumps(value, sort_keys=True, separators=(",", ":"), default=str)
|
||||
return "sha256:" + hashlib.sha256(payload.encode()).hexdigest()
|
||||
|
||||
|
||||
class CorpusPipeline:
|
||||
def __init__(
|
||||
self, *, store: CorpusStore, sources: list[EvidenceSource], embedder,
|
||||
vector_store: VectorStore, embedding_model: str, embedding_dimensions: int,
|
||||
chunk_policy: ChunkPolicy, pipeline_version: str,
|
||||
) -> None:
|
||||
self.store = store
|
||||
self.sources = sources
|
||||
self.embedder = embedder
|
||||
self.vector_store = vector_store
|
||||
self.embedding_model = embedding_model
|
||||
self.embedding_dimensions = embedding_dimensions
|
||||
self.chunk_policy = chunk_policy
|
||||
self.pipeline_version = pipeline_version
|
||||
|
||||
def _discover(self) -> list[tuple[EvidenceSource, SourceObject]]:
|
||||
discovered = []
|
||||
seen = set()
|
||||
for source in self.sources:
|
||||
for item in source.discover():
|
||||
if item.source_id in seen:
|
||||
raise PipelineError("duplicate Evidence source identity")
|
||||
seen.add(item.source_id)
|
||||
discovered.append((source, item))
|
||||
return sorted(discovered, key=lambda pair: pair[1].source_id)
|
||||
|
||||
def run(self, *, dry_run: bool = False, resume: str | None = None) -> PipelineResult:
|
||||
with self.store.writer_lock():
|
||||
return self._run(dry_run=dry_run, resume=resume)
|
||||
|
||||
def _run(self, *, dry_run: bool = False, resume: str | None = None) -> PipelineResult:
|
||||
previous = self.store.active_manifest()
|
||||
try:
|
||||
discovered = self._discover()
|
||||
except Exception as error:
|
||||
raise PipelineError("Evidence discovery failed") from error
|
||||
prior_documents = {doc.source_id: doc for doc in previous.documents} if previous else {}
|
||||
fingerprints = {item.source_id: item.fingerprint for _, item in discovered}
|
||||
compatibility = _fingerprint({
|
||||
"pipeline": self.pipeline_version, "model": self.embedding_model,
|
||||
"dimensions": self.embedding_dimensions, "chunk_policy": asdict(self.chunk_policy),
|
||||
})
|
||||
previous_compatibility = previous.metadata.get("compatibility_fingerprint") if previous else None
|
||||
rebuild = previous is not None and compatibility != previous_compatibility
|
||||
changed = tuple(item.source_id for _, item in discovered if rebuild or prior_documents.get(item.source_id) is None or prior_documents[item.source_id].source_fingerprint != item.fingerprint)
|
||||
unchanged = tuple(item.source_id for _, item in discovered if item.source_id not in changed)
|
||||
removed = tuple(sorted(set(prior_documents) - set(fingerprints)))
|
||||
if dry_run:
|
||||
manifest = previous or CorpusManifest(pipeline_version=self.pipeline_version)
|
||||
return PipelineResult("succeeded", None, False, changed, unchanged, removed, manifest)
|
||||
if previous is not None and not changed and not removed:
|
||||
return PipelineResult(
|
||||
"succeeded", previous.manifest_id, False, changed, unchanged, removed, previous
|
||||
)
|
||||
|
||||
documents: list[CanonicalDocument] = [prior_documents[source_id] for source_id in unchanged]
|
||||
changed_set = set(changed)
|
||||
try:
|
||||
for source, item in discovered:
|
||||
if item.source_id in changed_set:
|
||||
documents.append(normalize(source.acquire(item), self.pipeline_version))
|
||||
documents.sort(key=lambda document: document.source_id)
|
||||
chunks: list[CanonicalChunk] = []
|
||||
for document in documents:
|
||||
chunks.extend(chunk(document, self.chunk_policy))
|
||||
generation = resume or f"gen:{uuid.uuid4().hex}"
|
||||
previous_generations = dict(previous.metadata.get("document_generations", {})) if previous else {}
|
||||
document_generations = {
|
||||
document.document_id: (
|
||||
generation if document.source_id in changed_set
|
||||
else previous_generations.get(document.document_id, previous.vector_generation)
|
||||
)
|
||||
for document in documents
|
||||
}
|
||||
manifest = CorpusManifest(
|
||||
pipeline_version=self.pipeline_version,
|
||||
embedding_model=self.embedding_model,
|
||||
embedding_dimensions=self.embedding_dimensions,
|
||||
vector_generation=generation,
|
||||
documents=tuple(documents), chunks=tuple(chunks),
|
||||
metadata={
|
||||
"compatibility_fingerprint": compatibility,
|
||||
"fingerprints": fingerprints,
|
||||
"removed": list(removed),
|
||||
"document_generations": document_generations,
|
||||
},
|
||||
)
|
||||
changed_documents = {document.document_id for document in documents if document.source_id in changed_set}
|
||||
changed_chunks = [part for part in chunks if part.document_id in changed_documents]
|
||||
embeddings = self.embedder.embed_documents([part.content for part in changed_chunks])
|
||||
if len(embeddings) != len(changed_chunks):
|
||||
raise PipelineError("embedding count mismatch")
|
||||
if any(len(vector) != self.embedding_dimensions for vector in embeddings):
|
||||
raise PipelineError("embedding dimension mismatch")
|
||||
records = [self._vector_record(part, vector, generation) for part, vector in zip(changed_chunks, embeddings, strict=True)]
|
||||
if records:
|
||||
written = self.vector_store.upsert("evidence", records)
|
||||
if written != len(records):
|
||||
raise PipelineError("vector write count mismatch")
|
||||
generation_path = self.store.generation_path(generation)
|
||||
if resume is not None and generation_path.exists():
|
||||
staged_manifest = self.store.manifest(generation)
|
||||
expected = manifest.model_dump(mode="json", exclude={"created_at", "manifest_id"})
|
||||
actual = staged_manifest.model_dump(mode="json", exclude={"created_at", "manifest_id"})
|
||||
actual["metadata"].pop("files", None)
|
||||
if actual != expected:
|
||||
raise PipelineError("resume generation is incompatible")
|
||||
staged = generation
|
||||
else:
|
||||
staged = self.store.stage(
|
||||
manifest, {document.document_id: document.content for document in documents},
|
||||
generation=generation,
|
||||
)
|
||||
self.store.publish(staged)
|
||||
except PipelineError:
|
||||
raise
|
||||
except Exception as error:
|
||||
raise PipelineError("Evidence preprocessing failed") from error
|
||||
return PipelineResult("succeeded", generation, True, changed, unchanged, removed, self.store.manifest(generation))
|
||||
|
||||
@staticmethod
|
||||
def _vector_record(chunk: CanonicalChunk, embedding: list[float], generation: str):
|
||||
record = VectorRecord(
|
||||
id=f"{generation}:{chunk.chunk_id}", kind="evidence", ref=chunk.document_id,
|
||||
title=str(chunk.metadata.get("title", "")), content=chunk.content,
|
||||
metadata={
|
||||
**dict(chunk.metadata), "document_id": chunk.document_id,
|
||||
"source_uri": chunk.source_uri, "ordinal": chunk.ordinal,
|
||||
"vector_generation": generation,
|
||||
},
|
||||
)
|
||||
return VectorWriteRecord(record=record, embedding=embedding, content_hash=chunk.content_hash)
|
||||
Reference in New Issue
Block a user