ai-pipeline-core 0.3.4__py3-none-any.whl → 0.4.0__py3-none-any.whl

ai_pipeline_core/document_store/local.py

@@ -0,0 +1,312 @@
+"""Local filesystem document store for CLI/debug mode.
+
+Layout:
+    {base_path}/{canonical_name}/{filename}           <- raw content
+    {base_path}/{canonical_name}/{filename}.meta.json  <- metadata
+    {base_path}/{canonical_name}/{filename}.att/        <- attachments directory
+"""
+
+import asyncio
+import json
+from pathlib import Path
+from typing import Any
+
+from ai_pipeline_core.document_store._summary import SummaryGenerator
+from ai_pipeline_core.document_store._summary_worker import SummaryWorker
+from ai_pipeline_core.documents._context_vars import suppress_registration
+from ai_pipeline_core.documents._hashing import compute_content_sha256, compute_document_sha256
+from ai_pipeline_core.documents.attachment import Attachment
+from ai_pipeline_core.documents.document import Document
+from ai_pipeline_core.logging import get_pipeline_logger
+
+logger = get_pipeline_logger(__name__)
+
+
+class LocalDocumentStore:
+    """Filesystem-backed document store for local development and debugging.
+
+    Documents are stored as browsable files organized by canonical type name.
+    Write order (content before meta) ensures crash safety — load() ignores
+    content files without a valid .meta.json.
+    """
+
+    def __init__(
+        self,
+        base_path: Path | None = None,
+        *,
+        summary_generator: SummaryGenerator | None = None,
+    ) -> None:
+        self._base_path = base_path or Path.cwd()
+        self._meta_path_cache: dict[str, Path] = {}  # "{run_scope}:{sha256}" -> meta file path
+        self._summary_worker: SummaryWorker | None = None
+        if summary_generator:
+            self._summary_worker = SummaryWorker(
+                generator=summary_generator,
+                update_fn=self.update_summary,
+            )
+            self._summary_worker.start()
+
+    @property
+    def base_path(self) -> Path:
+        """Root directory for all stored documents."""
+        return self._base_path
+
+    async def save(self, document: Document, run_scope: str) -> None:
+        """Save a document to disk. Idempotent — same SHA256 is a no-op."""
+        written = await asyncio.to_thread(self._save_sync, document, run_scope)
+        if written and self._summary_worker:
+            self._summary_worker.schedule(run_scope, document)
+
+    async def save_batch(self, documents: list[Document], run_scope: str) -> None:
+        """Save multiple documents sequentially."""
+        for doc in documents:
+            await self.save(doc, run_scope)
+
+    async def load(self, run_scope: str, document_types: list[type[Document]]) -> list[Document]:
+        """Load documents by type from the run scope directory."""
+        return await asyncio.to_thread(self._load_sync, run_scope, document_types)
+
+    async def has_documents(self, run_scope: str, document_type: type[Document]) -> bool:
+        """Check for meta files in the type's directory without loading content."""
+        return await asyncio.to_thread(self._has_documents_sync, run_scope, document_type)
+
+    async def check_existing(self, sha256s: list[str]) -> set[str]:
+        """Scan all meta files to find matching document_sha256 values."""
+        return await asyncio.to_thread(self._check_existing_sync, sha256s)
+
+    async def update_summary(self, run_scope: str, document_sha256: str, summary: str) -> None:
+        """Update summary in the document's .meta.json file."""
+        await asyncio.to_thread(self._update_summary_sync, run_scope, document_sha256, summary)
+
+    async def load_summaries(self, run_scope: str, document_sha256s: list[str]) -> dict[str, str]:
+        """Load summaries from .meta.json files."""
+        return await asyncio.to_thread(self._load_summaries_sync, run_scope, document_sha256s)
+
+    def flush(self) -> None:
+        """Block until all pending document summaries are processed."""
+        if self._summary_worker:
+            self._summary_worker.flush()
+
+    def shutdown(self) -> None:
+        """Flush pending summaries and stop the summary worker."""
+        if self._summary_worker:
+            self._summary_worker.shutdown()
+
+    # --- Sync implementation (called via asyncio.to_thread) ---
+
+    def _scope_path(self, run_scope: str) -> Path:
+        return self._base_path / run_scope
+
+    def _save_sync(self, document: Document, run_scope: str) -> bool:
+        canonical = document.canonical_name()
+        doc_dir = self._scope_path(run_scope) / canonical
+        doc_dir.mkdir(parents=True, exist_ok=True)
+
+        content_path = doc_dir / document.name
+        if not content_path.resolve().is_relative_to(doc_dir.resolve()):
+            raise ValueError(f"Path traversal detected: document name '{document.name}' escapes store directory")
+        meta_path = doc_dir / f"{document.name}.meta.json"
+
+        doc_sha256 = compute_document_sha256(document)
+        content_sha256 = compute_content_sha256(document.content)
+
+        # Check for concurrent access: if meta exists with different SHA256, log warning
+        if meta_path.exists():
+            existing_meta = self._read_meta(meta_path)
+            if existing_meta and existing_meta.get("document_sha256") == doc_sha256:
+                # Populate cache even for idempotent saves
+                self._meta_path_cache[f"{run_scope}:{doc_sha256}"] = meta_path
+                return False  # Idempotent — same document already saved
+            if existing_meta:
+                logger.warning(
+                    f"Overwriting document '{document.name}' in '{canonical}': "
+                    f"existing SHA256 {existing_meta.get('document_sha256', '?')[:12]}... "
+                    f"differs from new {doc_sha256[:12]}..."
+                )
+
+        # Write content before meta (crash safety)
+        content_path.write_bytes(document.content)
+
+        # Write attachments
+        att_meta_list: list[dict[str, Any]] = []
+        if document.attachments:
+            att_dir = doc_dir / f"{document.name}.att"
+            att_dir.mkdir(exist_ok=True)
+            for att in document.attachments:
+                att_path = att_dir / att.name
+                if not att_path.resolve().is_relative_to(att_dir.resolve()):
+                    raise ValueError(f"Path traversal detected: attachment name '{att.name}' escapes store directory")
+                att_path.write_bytes(att.content)
+                att_meta_list.append({
+                    "name": att.name,
+                    "description": att.description,
+                    "sha256": compute_content_sha256(att.content),
+                })
+
+        # Write meta last (crash safety — content is already on disk)
+        meta = {
+            "document_sha256": doc_sha256,
+            "content_sha256": content_sha256,
+            "class_name": document.__class__.__name__,
+            "description": document.description,
+            "sources": list(document.sources),
+            "origins": list(document.origins),
+            "mime_type": document.mime_type,
+            "attachments": att_meta_list,
+        }
+        meta_path.write_text(json.dumps(meta, indent=2), encoding="utf-8")
+
+        # Cache meta path for summary updates
+        self._meta_path_cache[f"{run_scope}:{doc_sha256}"] = meta_path
+        return True
+
+    def _load_sync(self, run_scope: str, document_types: list[type[Document]]) -> list[Document]:
+        scope_path = self._scope_path(run_scope)
+        if not scope_path.exists():
+            return []
+
+        # Build reverse map: canonical_name -> document type
+        type_by_canonical: dict[str, type[Document]] = {}
+        for doc_type in document_types:
+            cn = doc_type.canonical_name()
+            type_by_canonical[cn] = doc_type
+
+        documents: list[Document] = []
+
+        with suppress_registration():
+            for canonical, doc_type in type_by_canonical.items():
+                type_dir = scope_path / canonical
+                if not type_dir.is_dir():
+                    continue
+                self._load_type_dir(type_dir, doc_type, documents)
+
+        return documents
+
+    def _load_type_dir(self, type_dir: Path, doc_type: type[Document], out: list[Document]) -> None:
+        """Load all documents of a single type from its directory."""
+        for meta_path in type_dir.glob("*.meta.json"):
+            meta = self._read_meta(meta_path)
+            if meta is None:
+                continue
+
+            content_name = meta_path.name.removesuffix(".meta.json")
+            content_path = type_dir / content_name
+
+            if not content_path.exists():
+                logger.warning(f"Meta file {meta_path} has no corresponding content file, skipping")
+                continue
+
+            content = content_path.read_bytes()
+
+            # Load attachments
+            attachments: tuple[Attachment, ...] = ()
+            att_meta_list = meta.get("attachments", [])
+            if att_meta_list:
+                att_dir = type_dir / f"{content_name}.att"
+                att_list: list[Attachment] = []
+                for att_meta in att_meta_list:
+                    att_path = att_dir / att_meta["name"]
+                    if not att_path.exists():
+                        logger.warning(f"Attachment file {att_path} missing, skipping")
+                        continue
+                    att_list.append(
+                        Attachment(
+                            name=att_meta["name"],
+                            content=att_path.read_bytes(),
+                            description=att_meta.get("description"),
+                        )
+                    )
+                attachments = tuple(att_list)
+
+            doc = doc_type(
+                name=content_name,
+                content=content,
+                description=meta.get("description"),
+                sources=tuple(meta.get("sources", ())),
+                origins=tuple(meta.get("origins", ())),
+                attachments=attachments or None,
+            )
+            out.append(doc)
+
+    def _has_documents_sync(self, run_scope: str, document_type: type[Document]) -> bool:
+        """Check for meta files in the type's directory without loading content."""
+        scope_path = self._scope_path(run_scope)
+        canonical = document_type.canonical_name()
+        type_dir = scope_path / canonical
+        if not type_dir.is_dir():
+            return False
+        return any(type_dir.glob("*.meta.json"))
+
+    def _check_existing_sync(self, sha256s: list[str]) -> set[str]:
+        """Scan all meta files to find matching document_sha256 values."""
+        target = set(sha256s)
+        found: set[str] = set()
+        if not self._base_path.exists():
+            return found
+
+        for meta_path in self._base_path.rglob("*.meta.json"):
+            meta = self._read_meta(meta_path)
+            if meta and meta.get("document_sha256") in target:
+                found.add(meta["document_sha256"])
+                if found == target:
+                    break
+        return found
+
+    def _update_summary_sync(self, run_scope: str, document_sha256: str, summary: str) -> None:
+        """Update summary in the document's .meta.json file."""
+        cache_key = f"{run_scope}:{document_sha256}"
+        meta_path = self._meta_path_cache.get(cache_key)
+
+        # Fallback: scan for the meta file if not cached
+        if meta_path is None or not meta_path.exists():
+            meta_path = self._find_meta_by_sha256(run_scope, document_sha256)
+            if meta_path is None:
+                return
+
+        meta = self._read_meta(meta_path)
+        if meta is None:
+            return
+
+        meta["summary"] = summary
+        meta_path.write_text(json.dumps(meta, indent=2), encoding="utf-8")
+        self._meta_path_cache[cache_key] = meta_path
+
+    def _load_summaries_sync(self, run_scope: str, document_sha256s: list[str]) -> dict[str, str]:
+        """Scan meta files for matching sha256s and return their summaries."""
+        target = set(document_sha256s)
+        result: dict[str, str] = {}
+        scope_path = self._scope_path(run_scope)
+        if not scope_path.exists():
+            return result
+
+        for meta_path in scope_path.rglob("*.meta.json"):
+            meta = self._read_meta(meta_path)
+            if meta is None:
+                continue
+            sha = meta.get("document_sha256")
+            summary = meta.get("summary")
+            if sha in target and summary:
+                result[sha] = summary
+                if len(result) == len(target):
+                    break
+        return result
+
+    def _find_meta_by_sha256(self, run_scope: str, document_sha256: str) -> Path | None:
+        """Scan meta files in run_scope to find one matching the given sha256."""
+        scope_path = self._scope_path(run_scope)
+        if not scope_path.exists():
+            return None
+        for meta_path in scope_path.rglob("*.meta.json"):
+            meta = self._read_meta(meta_path)
+            if meta and meta.get("document_sha256") == document_sha256:
+                return meta_path
+        return None
+
+    @staticmethod
+    def _read_meta(meta_path: Path) -> dict[str, Any] | None:
+        """Read and parse a meta.json file, returning None on any error."""
+        try:
+            return json.loads(meta_path.read_text(encoding="utf-8"))
+        except (json.JSONDecodeError, OSError) as e:
+            logger.warning(f"Failed to read meta file {meta_path}: {e}")
+            return None

ai_pipeline_core/document_store/memory.py

@@ -0,0 +1,85 @@
+"""In-memory document store for testing.
+
+Simple dict-based storage implementing the full DocumentStore protocol.
+Not for production use — all data is lost when the process exits.
+"""
+
+from ai_pipeline_core.document_store._summary import SummaryGenerator
+from ai_pipeline_core.document_store._summary_worker import SummaryWorker
+from ai_pipeline_core.documents.document import Document
+
+
+class MemoryDocumentStore:
+    """Dict-based document store for unit tests.
+
+    Storage layout: dict[run_scope, dict[document_sha256, Document]].
+    """
+
+    def __init__(
+        self,
+        *,
+        summary_generator: SummaryGenerator | None = None,
+    ) -> None:
+        self._data: dict[str, dict[str, Document]] = {}
+        self._summaries: dict[str, dict[str, str]] = {}  # run_scope -> sha256 -> summary
+        self._summary_worker: SummaryWorker | None = None
+        if summary_generator:
+            self._summary_worker = SummaryWorker(
+                generator=summary_generator,
+                update_fn=self.update_summary,
+            )
+            self._summary_worker.start()
+
+    async def save(self, document: Document, run_scope: str) -> None:
+        """Store document in memory, keyed by SHA256."""
+        scope = self._data.setdefault(run_scope, {})
+        if document.sha256 in scope:
+            return  # Idempotent — same document already saved
+        scope[document.sha256] = document
+        if self._summary_worker:
+            self._summary_worker.schedule(run_scope, document)
+
+    async def save_batch(self, documents: list[Document], run_scope: str) -> None:
+        """Save multiple documents sequentially."""
+        for doc in documents:
+            await self.save(doc, run_scope)
+
+    async def load(self, run_scope: str, document_types: list[type[Document]]) -> list[Document]:
+        """Return all documents matching the given types from a run scope."""
+        scope = self._data.get(run_scope, {})
+        type_tuple = tuple(document_types)
+        return [doc for doc in scope.values() if isinstance(doc, type_tuple)]
+
+    async def has_documents(self, run_scope: str, document_type: type[Document]) -> bool:
+        """Check if any documents of this type exist in the run scope."""
+        scope = self._data.get(run_scope, {})
+        return any(isinstance(doc, document_type) for doc in scope.values())
+
+    async def check_existing(self, sha256s: list[str]) -> set[str]:
+        """Return the subset of sha256s that exist across all scopes."""
+        all_hashes: set[str] = set()
+        for scope in self._data.values():
+            all_hashes.update(scope.keys())
+        return all_hashes & set(sha256s)
+
+    async def update_summary(self, run_scope: str, document_sha256: str, summary: str) -> None:
+        """Update summary for a stored document. No-op if document doesn't exist."""
+        scope = self._data.get(run_scope, {})
+        if document_sha256 not in scope:
+            return
+        self._summaries.setdefault(run_scope, {})[document_sha256] = summary
+
+    async def load_summaries(self, run_scope: str, document_sha256s: list[str]) -> dict[str, str]:
+        """Load summaries by SHA256."""
+        scope_summaries = self._summaries.get(run_scope, {})
+        return {sha: scope_summaries[sha] for sha in document_sha256s if sha in scope_summaries}
+
+    def flush(self) -> None:
+        """Block until all pending document summaries are processed."""
+        if self._summary_worker:
+            self._summary_worker.flush()
+
+    def shutdown(self) -> None:
+        """Flush pending summaries and stop the summary worker."""
+        if self._summary_worker:
+            self._summary_worker.shutdown()

ai_pipeline_core/document_store/protocol.py

@@ -0,0 +1,68 @@
+"""Document store protocol and singleton management.
+
+Defines the DocumentStore protocol that all storage backends must implement,
+along with get/set helpers for the process-global singleton.
+"""
+
+from typing import Protocol, runtime_checkable
+
+from ai_pipeline_core.documents.document import Document
+
+
+@runtime_checkable
+class DocumentStore(Protocol):
+    """Protocol for document storage backends.
+
+    Implementations: ClickHouseDocumentStore (production), LocalDocumentStore (CLI/debug),
+    MemoryDocumentStore (testing).
+    """
+
+    async def save(self, document: Document, run_scope: str) -> None:
+        """Save a single document to the store. Idempotent — same SHA256 is a no-op."""
+        ...
+
+    async def save_batch(self, documents: list[Document], run_scope: str) -> None:
+        """Save multiple documents. Dependencies must be sorted (caller's responsibility)."""
+        ...
+
+    async def load(self, run_scope: str, document_types: list[type[Document]]) -> list[Document]:
+        """Load all documents of the given types from a run scope."""
+        ...
+
+    async def has_documents(self, run_scope: str, document_type: type[Document]) -> bool:
+        """Check if any documents of this type exist in the run scope."""
+        ...
+
+    async def check_existing(self, sha256s: list[str]) -> set[str]:
+        """Return the subset of sha256s that already exist in the store."""
+        ...
+
+    async def update_summary(self, run_scope: str, document_sha256: str, summary: str) -> None:
+        """Update summary for a stored document. No-op if document doesn't exist."""
+        ...
+
+    async def load_summaries(self, run_scope: str, document_sha256s: list[str]) -> dict[str, str]:
+        """Load summaries by SHA256. Returns {sha256: summary} for docs that have summaries."""
+        ...
+
+    def flush(self) -> None:
+        """Block until all pending background work (summaries) is processed."""
+        ...
+
+    def shutdown(self) -> None:
+        """Flush pending work and stop background workers."""
+        ...
+
+
+_document_store: DocumentStore | None = None
+
+
+def get_document_store() -> DocumentStore | None:
+    """Get the process-global document store singleton."""
+    return _document_store
+
+
+def set_document_store(store: DocumentStore | None) -> None:
+    """Set the process-global document store singleton."""
+    global _document_store
+    _document_store = store

ai_pipeline_core/documents/__init__.py

@@ -1,26 +1,24 @@
-"""Document abstraction system for AI pipeline flows.
+"""Document system for AI pipeline flows.
 
-@public
-
-The documents package provides immutable, type-safe data structures for handling
-various content types in AI pipelines, including text, images, PDFs, and other
-binary data with automatic MIME type detection.
+Provides the Document base class (immutable, content-addressed), Attachment for
+binary sub-documents, and RunContext/TaskDocumentContext for document lifecycle
+management within pipeline tasks.
 """
 
+from .attachment import Attachment
+from .context import RunContext, TaskDocumentContext, get_run_context, reset_run_context, set_run_context
 from .document import Document
-from .document_list import DocumentList
-from .flow_document import FlowDocument
-from .task_document import TaskDocument
-from .temporary_document import TemporaryDocument
 from .utils import canonical_name_key, is_document_sha256, sanitize_url
 
 __all__ = [
+    "Attachment",
     "Document",
-    "DocumentList",
-    "FlowDocument",
-    "TaskDocument",
-    "TemporaryDocument",
+    "RunContext",
+    "TaskDocumentContext",
     "canonical_name_key",
+    "get_run_context",
     "is_document_sha256",
+    "reset_run_context",
     "sanitize_url",
+    "set_run_context",
 ]

ai_pipeline_core/documents/_context_vars.py

@@ -0,0 +1,85 @@
+"""Low-level ContextVar declarations for document registration and task context.
+
+Extracted into a separate module to break the circular dependency between
+document.py (which needs suppression/task-context checks) and context.py
+(which defines the full TaskDocumentContext that depends on Document).
+"""
+
+from collections.abc import Iterator
+from contextlib import contextmanager
+from contextvars import ContextVar, Token
+from dataclasses import dataclass
+
+# --- Run context ---
+
+
+@dataclass(frozen=True, slots=True)
+class RunContext:
+    """Immutable context for a pipeline run, carried via ContextVar."""
+
+    run_scope: str
+
+
+_run_context: ContextVar[RunContext | None] = ContextVar("_run_context", default=None)
+
+
+def get_run_context() -> RunContext | None:
+    """Get the current run context, or None if not set."""
+    return _run_context.get()
+
+
+def set_run_context(ctx: RunContext) -> Token[RunContext | None]:
+    """Set the run context. Returns a token for restoring the previous value."""
+    return _run_context.set(ctx)
+
+
+def reset_run_context(token: Token[RunContext | None]) -> None:
+    """Reset the run context to its previous value using a token from set_run_context."""
+    _run_context.reset(token)
+
+
+# --- Suppression flag ---
+
+_suppression_flag: ContextVar[bool] = ContextVar("_document_registration_suppressed", default=False)
+
+
+def is_registration_suppressed() -> bool:
+    """Check if document registration is currently suppressed."""
+    return _suppression_flag.get()
+
+
+@contextmanager
+def suppress_registration() -> Iterator[None]:
+    """Context manager that suppresses Document registration with TaskDocumentContext.
+
+    Used during model_validate() and other internal Pydantic operations that
+    construct intermediate Document objects that should not be tracked.
+    """
+    token = _suppression_flag.set(True)
+    try:
+        yield
+    finally:
+        _suppression_flag.reset(token)
+
+
+# --- Task document context ContextVar ---
+
+# Forward reference: the actual TaskDocumentContext class lives in context.py.
+# Here we only manage the ContextVar holding it.
+
+_task_context: ContextVar[object | None] = ContextVar("_task_context", default=None)
+
+
+def get_task_context() -> object | None:
+    """Get the current task document context, or None if not inside a pipeline task."""
+    return _task_context.get()
+
+
+def set_task_context(ctx: object) -> Token[object | None]:
+    """Set the task document context. Returns a token for restoring the previous value."""
+    return _task_context.set(ctx)
+
+
+def reset_task_context(token: Token[object | None]) -> None:
+    """Reset the task document context to its previous value."""
+    _task_context.reset(token)

ai_pipeline_core/documents/_hashing.py

@@ -0,0 +1,52 @@
+"""Document hashing utilities for store implementations.
+
+Computes document_sha256 and content_sha256 as defined in the document store
+design: length-prefixed fields with null-byte separators, BASE32 encoded.
+"""
+
+import hashlib
+from base64 import b32encode
+from typing import Any, Protocol
+
+
+class _Hashable(Protocol):
+    """Protocol for objects whose identity hash can be computed."""
+
+    name: str
+    content: bytes
+    attachments: Any
+
+
+def compute_document_sha256(doc: _Hashable) -> str:
+    """Compute the document identity hash: hash(name + content + sorted_attachments).
+
+    Uses length-prefixed fields with null-byte separators for collision resistance.
+    Attachments are sorted by name. Result is BASE32 encoded (uppercase, no padding),
+    consistent with Document.sha256.
+
+    Excluded from hash: description, sources, origins, class_name.
+    """
+    h = hashlib.sha256()
+
+    name_bytes = doc.name.encode("utf-8")
+    _hash_field(h, name_bytes)
+    _hash_field(h, doc.content)
+
+    for att in sorted(doc.attachments, key=lambda a: a.name):
+        att_name_bytes = att.name.encode("utf-8")
+        _hash_field(h, att_name_bytes)
+        _hash_field(h, att.content)
+
+    return b32encode(h.digest()).decode("ascii").upper().rstrip("=")
+
+
+def compute_content_sha256(content: bytes) -> str:
+    """Compute SHA256 of raw content bytes, BASE32 encoded."""
+    return b32encode(hashlib.sha256(content).digest()).decode("ascii").upper().rstrip("=")
+
+
+def _hash_field(h: Any, data: bytes) -> None:
+    """Append a length-prefixed, null-separated field to the hash."""
+    h.update(str(len(data)).encode("ascii"))
+    h.update(b"\x00")
+    h.update(data)