biblicus 0.15.0__py3-none-any.whl → 0.16.0__py3-none-any.whl

biblicus/__init__.py

@@ -27,4 +27,4 @@ __all__ = [
     "RetrievalRun",
 ]
 
-__version__ = "0.15.0"
+__version__ = "0.16.0"

biblicus/analysis/markov.py

@@ -686,6 +686,15 @@ def _build_observations(
         llm = config.llm_observations
         assert llm.client is not None and llm.prompt_template is not None
         for index, observation in enumerate(observations):
+            if observation.segment_text in {"START", "END"}:
+                observations[index] = observation.model_copy(
+                    update={
+                        "llm_label": observation.segment_text,
+                        "llm_label_confidence": 1.0,
+                        "llm_summary": observation.segment_text,
+                    }
+                )
+                continue
             prompt = llm.prompt_template.format(segment=observation.segment_text)
             response_text = generate_completion(
                 client=llm.client,
@@ -707,8 +716,12 @@ def _build_observations(
     if config.embeddings.enabled:
         embedding_config = config.embeddings
         assert embedding_config.client is not None
+        embed_indices: List[int] = []
         embed_texts: List[str] = []
-        for observation in observations:
+        for index, observation in enumerate(observations):
+            if observation.segment_text in {"START", "END"}:
+                continue
+            embed_indices.append(index)
             if embedding_config.text_source == "segment_text":
                 embed_texts.append(observation.segment_text)
             else:
@@ -717,10 +730,29 @@ def _build_observations(
                         "embeddings.text_source is 'llm_summary' but llm_summary is missing"
                     )
                 embed_texts.append(observation.llm_summary)
+
+        if not embed_indices:
+            raise ValueError("Embeddings require at least one non-boundary segment")
+
         vectors = generate_embeddings_batch(client=embedding_config.client, texts=embed_texts)
+        if len(vectors) != len(embed_indices):
+            raise ValueError(
+                "Embedding provider returned unexpected vector count: "
+                f"expected {len(embed_indices)} but got {len(vectors)}"
+            )
+
+        vector_by_observation_index: Dict[int, List[float]] = {}
+        for observation_index, vector in zip(embed_indices, vectors):
+            vector_by_observation_index[observation_index] = list(vector)
+
+        embedding_dimension = len(next(iter(vector_by_observation_index.values())))
+        boundary_embedding = [0.0 for _ in range(embedding_dimension)]
         updated: List[MarkovAnalysisObservation] = []
-        for observation, vector in zip(observations, vectors):
-            updated.append(observation.model_copy(update={"embedding": vector}))
+        for index, observation in enumerate(observations):
+            vector = vector_by_observation_index.get(index)
+            updated.append(
+                observation.model_copy(update={"embedding": vector or boundary_embedding})
+            )
         observations = updated
 
     return observations

biblicus/backends/__init__.py

@@ -7,10 +7,12 @@ from __future__ import annotations
 from typing import Dict, Type
 
 from .base import RetrievalBackend
+from .embedding_index_file import EmbeddingIndexFileBackend
+from .embedding_index_inmemory import EmbeddingIndexInMemoryBackend
 from .hybrid import HybridBackend
 from .scan import ScanBackend
 from .sqlite_full_text_search import SqliteFullTextSearchBackend
-from .vector import VectorBackend
+from .tf_vector import TfVectorBackend
 
 
 def available_backends() -> Dict[str, Type[RetrievalBackend]]:
@@ -21,10 +23,12 @@ def available_backends() -> Dict[str, Type[RetrievalBackend]]:
     :rtype: dict[str, Type[RetrievalBackend]]
     """
     return {
+        EmbeddingIndexFileBackend.backend_id: EmbeddingIndexFileBackend,
+        EmbeddingIndexInMemoryBackend.backend_id: EmbeddingIndexInMemoryBackend,
         HybridBackend.backend_id: HybridBackend,
         ScanBackend.backend_id: ScanBackend,
         SqliteFullTextSearchBackend.backend_id: SqliteFullTextSearchBackend,
-        VectorBackend.backend_id: VectorBackend,
+        TfVectorBackend.backend_id: TfVectorBackend,
     }
 
 

biblicus/backends/embedding_index_common.py

@@ -0,0 +1,301 @@
+"""
+Shared primitives for embedding-index retrieval backends.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Dict, Iterable, Iterator, List, Optional, Tuple
+
+import numpy as np
+from pydantic import BaseModel, ConfigDict, Field, model_validator
+
+from ..chunking import ChunkerConfig, TextChunk, TokenizerConfig
+from ..corpus import CORPUS_DIR_NAME, RUNS_DIR_NAME, Corpus
+from ..embedding_providers import EmbeddingProviderConfig, _l2_normalize_rows
+from ..frontmatter import parse_front_matter
+from ..models import ExtractionRunReference, parse_extraction_run_reference
+
+
+class ChunkRecord(BaseModel):
+    """
+    Minimal persisted representation of a chunk.
+
+    :ivar item_id: Item identifier that produced the chunk.
+    :vartype item_id: str
+    :ivar span_start: Inclusive start character offset.
+    :vartype span_start: int
+    :ivar span_end: Exclusive end character offset.
+    :vartype span_end: int
+    """
+
+    model_config = ConfigDict(extra="forbid")
+
+    item_id: str = Field(min_length=1)
+    span_start: int = Field(ge=0)
+    span_end: int = Field(ge=0)
+
+    @model_validator(mode="after")
+    def _validate_span(self) -> "ChunkRecord":
+        if self.span_end <= self.span_start:
+            raise ValueError("chunk span_end must be greater than span_start")
+        return self
+
+
+class EmbeddingIndexRecipeConfig(BaseModel):
+    """
+    Configuration for embedding-index retrieval backends.
+
+    :ivar snippet_characters: Maximum characters to include in evidence snippets.
+    :vartype snippet_characters: int
+    :ivar extraction_run: Optional extraction run reference in the form extractor_id:run_id.
+    :vartype extraction_run: str or None
+    :ivar chunker: Chunker configuration.
+    :vartype chunker: biblicus.chunking.ChunkerConfig
+    :ivar tokenizer: Optional tokenizer configuration.
+    :vartype tokenizer: biblicus.chunking.TokenizerConfig or None
+    :ivar embedding_provider: Embedding provider configuration.
+    :vartype embedding_provider: biblicus.embedding_providers.EmbeddingProviderConfig
+    """
+
+    model_config = ConfigDict(extra="forbid")
+
+    snippet_characters: int = Field(default=400, ge=1)
+    extraction_run: Optional[str] = None
+    chunker: ChunkerConfig = Field(default_factory=lambda: ChunkerConfig(chunker_id="paragraph"))
+    tokenizer: Optional[TokenizerConfig] = None
+    embedding_provider: EmbeddingProviderConfig
+
+
+def resolve_extraction_reference(
+    corpus: Corpus, recipe_config: EmbeddingIndexRecipeConfig
+) -> Optional[ExtractionRunReference]:
+    """
+    Resolve an extraction run reference from an embedding-index recipe config.
+
+    :param corpus: Corpus associated with the recipe.
+    :type corpus: Corpus
+    :param recipe_config: Parsed embedding-index recipe configuration.
+    :type recipe_config: EmbeddingIndexRecipeConfig
+    :return: Parsed extraction reference or None.
+    :rtype: ExtractionRunReference or None
+    :raises FileNotFoundError: If an extraction run is referenced but not present.
+    """
+    if not recipe_config.extraction_run:
+        return None
+    extraction_reference = parse_extraction_run_reference(recipe_config.extraction_run)
+    run_dir = corpus.extraction_run_dir(
+        extractor_id=extraction_reference.extractor_id,
+        run_id=extraction_reference.run_id,
+    )
+    if not run_dir.is_dir():
+        raise FileNotFoundError(f"Missing extraction run: {extraction_reference.as_string()}")
+    return extraction_reference
+
+
+def _load_text_from_item(
+    corpus: Corpus,
+    *,
+    item_id: str,
+    relpath: str,
+    media_type: str,
+    extraction_reference: Optional[ExtractionRunReference],
+) -> Optional[str]:
+    if extraction_reference:
+        extracted_text = corpus.read_extracted_text(
+            extractor_id=extraction_reference.extractor_id,
+            run_id=extraction_reference.run_id,
+            item_id=item_id,
+        )
+        if isinstance(extracted_text, str):
+            return extracted_text
+
+    if media_type == "text/markdown":
+        raw = (corpus.root / relpath).read_text(encoding="utf-8")
+        return parse_front_matter(raw).body
+    if media_type.startswith("text/"):
+        return (corpus.root / relpath).read_text(encoding="utf-8")
+    return None
+
+
+def iter_text_payloads(
+    corpus: Corpus, *, extraction_reference: Optional[ExtractionRunReference]
+) -> Iterator[Tuple[object, str]]:
+    """
+    Yield catalog items and their text payloads.
+
+    :param corpus: Corpus containing the items.
+    :type corpus: Corpus
+    :param extraction_reference: Optional extraction reference.
+    :type extraction_reference: ExtractionRunReference or None
+    :yield: (catalog_item, text) pairs.
+    :rtype: Iterator[tuple[object, str]]
+    """
+    catalog = corpus.load_catalog()
+    for catalog_item in catalog.items.values():
+        item_id = str(getattr(catalog_item, "id", ""))
+        relpath = str(getattr(catalog_item, "relpath", ""))
+        media_type = str(getattr(catalog_item, "media_type", ""))
+        if not item_id or not relpath or not media_type:
+            continue
+        text = _load_text_from_item(
+            corpus,
+            item_id=item_id,
+            relpath=relpath,
+            media_type=media_type,
+            extraction_reference=extraction_reference,
+        )
+        if not isinstance(text, str) or not text.strip():
+            continue
+        yield catalog_item, text
+
+
+def collect_chunks(
+    corpus: Corpus, *, recipe_config: EmbeddingIndexRecipeConfig
+) -> Tuple[List[TextChunk], int]:
+    """
+    Collect chunks from text payloads in a corpus.
+
+    :param corpus: Corpus to chunk.
+    :type corpus: Corpus
+    :param recipe_config: Parsed embedding-index recipe configuration.
+    :type recipe_config: EmbeddingIndexRecipeConfig
+    :return: (chunks, text_item_count)
+    :rtype: tuple[list[TextChunk], int]
+    """
+    tokenizer = recipe_config.tokenizer.build_tokenizer() if recipe_config.tokenizer else None
+    chunker = recipe_config.chunker.build_chunker(tokenizer=tokenizer)
+    extraction_reference = resolve_extraction_reference(corpus, recipe_config)
+
+    chunks: List[TextChunk] = []
+    next_chunk_id = 0
+    text_items = 0
+    for catalog_item, text in iter_text_payloads(corpus, extraction_reference=extraction_reference):
+        text_items += 1
+        item_id = str(getattr(catalog_item, "id"))
+        item_chunks = chunker.chunk_text(
+            item_id=item_id, text=text, starting_chunk_id=next_chunk_id
+        )
+        if item_chunks:
+            next_chunk_id = item_chunks[-1].chunk_id + 1
+            chunks.extend(item_chunks)
+    return chunks, text_items
+
+
+def chunks_to_records(chunks: Iterable[TextChunk]) -> List[ChunkRecord]:
+    """
+    Convert chunk objects to persisted chunk records.
+
+    :param chunks: Chunk list.
+    :type chunks: Iterable[TextChunk]
+    :return: Chunk record list.
+    :rtype: list[ChunkRecord]
+    """
+    records: List[ChunkRecord] = []
+    for chunk in chunks:
+        records.append(
+            ChunkRecord(
+                item_id=chunk.item_id,
+                span_start=chunk.span_start,
+                span_end=chunk.span_end,
+            )
+        )
+    return records
+
+
+def write_chunks_jsonl(path: Path, records: Iterable[ChunkRecord]) -> None:
+    """
+    Write chunk records as newline-delimited JSON.
+
+    :param path: Destination path.
+    :type path: pathlib.Path
+    :param records: Chunk records.
+    :type records: Iterable[ChunkRecord]
+    :return: None.
+    :rtype: None
+    """
+    with path.open("w", encoding="utf-8") as handle:
+        for record in records:
+            handle.write(json.dumps(record.model_dump(), separators=(",", ":")) + "\n")
+
+
+def read_chunks_jsonl(path: Path) -> List[ChunkRecord]:
+    """
+    Read chunk records from a JSON Lines file.
+
+    :param path: Source path.
+    :type path: pathlib.Path
+    :return: Chunk record list.
+    :rtype: list[ChunkRecord]
+    """
+    records: List[ChunkRecord] = []
+    for line in path.read_text(encoding="utf-8").splitlines():
+        if not line.strip():
+            continue
+        records.append(ChunkRecord.model_validate(json.loads(line)))
+    return records
+
+
+def write_embeddings(path: Path, embeddings: np.ndarray) -> None:
+    """
+    Write embeddings to disk.
+
+    :param path: Destination path.
+    :type path: pathlib.Path
+    :param embeddings: Embedding matrix.
+    :type embeddings: numpy.ndarray
+    :return: None.
+    :rtype: None
+    """
+    np.save(path, embeddings.astype(np.float32))
+
+
+def read_embeddings(path: Path, *, mmap: bool) -> np.ndarray:
+    """
+    Read embeddings from disk.
+
+    :param path: Source path.
+    :type path: pathlib.Path
+    :param mmap: Whether to memory-map the file.
+    :type mmap: bool
+    :return: Embedding matrix.
+    :rtype: numpy.ndarray
+    """
+    mode = "r" if mmap else None
+    return np.load(path, mmap_mode=mode)
+
+
+def cosine_similarity_scores(embeddings: np.ndarray, query_vector: np.ndarray) -> np.ndarray:
+    """
+    Compute cosine similarity scores for a query vector.
+
+    The embedding matrix must already be L2-normalized.
+
+    :param embeddings: Embedding matrix of shape (n, d).
+    :type embeddings: numpy.ndarray
+    :param query_vector: Query vector of shape (d,).
+    :type query_vector: numpy.ndarray
+    :return: Score vector of shape (n,).
+    :rtype: numpy.ndarray
+    """
+    query_vector = query_vector.astype(np.float32).reshape(-1)
+    query_vector = _l2_normalize_rows(query_vector.reshape(1, -1)).reshape(-1)
+    return embeddings @ query_vector
+
+
+def artifact_paths_for_run(*, run_id: str, backend_id: str) -> Dict[str, str]:
+    """
+    Build deterministic artifact relative paths for an embedding index run.
+
+    :param run_id: Run identifier.
+    :type run_id: str
+    :param backend_id: Backend identifier.
+    :type backend_id: str
+    :return: Mapping with keys embeddings and chunks.
+    :rtype: dict[str, str]
+    """
+    prefix = f"{run_id}.{backend_id}"
+    embeddings_relpath = str(Path(CORPUS_DIR_NAME) / RUNS_DIR_NAME / f"{prefix}.embeddings.npy")
+    chunks_relpath = str(Path(CORPUS_DIR_NAME) / RUNS_DIR_NAME / f"{prefix}.chunks.jsonl")
+    return {"embeddings": embeddings_relpath, "chunks": chunks_relpath}

biblicus/backends/embedding_index_file.py

@@ -0,0 +1,266 @@
+"""
+Embedding-index retrieval backend that reads the embedding matrix via memory mapping.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Dict, List, Optional
+
+import numpy as np
+
+from ..corpus import Corpus
+from ..models import Evidence, ExtractionRunReference, QueryBudget, RetrievalResult, RetrievalRun
+from ..retrieval import apply_budget, create_recipe_manifest, create_run_manifest, hash_text
+from ..time import utc_now_iso
+from .embedding_index_common import (
+    ChunkRecord,
+    EmbeddingIndexRecipeConfig,
+    artifact_paths_for_run,
+    chunks_to_records,
+    collect_chunks,
+    cosine_similarity_scores,
+    read_chunks_jsonl,
+    read_embeddings,
+    resolve_extraction_reference,
+    write_chunks_jsonl,
+    write_embeddings,
+)
+from .scan import _build_snippet
+
+
+class EmbeddingIndexFileBackend:
+    """
+    Embedding retrieval backend using memory-mapped similarity scanning.
+    """
+
+    backend_id = "embedding-index-file"
+
+    def build_run(
+        self, corpus: Corpus, *, recipe_name: str, config: Dict[str, object]
+    ) -> RetrievalRun:
+        """
+        Build an embedding index run by chunking text payloads and materializing embeddings.
+
+        :param corpus: Corpus to build against.
+        :type corpus: Corpus
+        :param recipe_name: Human-readable recipe name.
+        :type recipe_name: str
+        :param config: Backend-specific configuration values.
+        :type config: dict[str, object]
+        :return: Run manifest describing the build.
+        :rtype: biblicus.models.RetrievalRun
+        """
+        recipe_config = EmbeddingIndexRecipeConfig.model_validate(config)
+        chunks, text_items = collect_chunks(corpus, recipe_config=recipe_config)
+
+        provider = recipe_config.embedding_provider.build_provider()
+        embeddings = provider.embed_texts([chunk.text for chunk in chunks]).astype(np.float32)
+
+        recipe = create_recipe_manifest(
+            backend_id=self.backend_id,
+            name=recipe_name,
+            config=recipe_config.model_dump(),
+        )
+        run = create_run_manifest(corpus, recipe=recipe, stats={}, artifact_paths=[])
+
+        paths = artifact_paths_for_run(run_id=run.run_id, backend_id=self.backend_id)
+        embeddings_path = corpus.root / paths["embeddings"]
+        chunks_path = corpus.root / paths["chunks"]
+        corpus.runs_dir.mkdir(parents=True, exist_ok=True)
+
+        write_embeddings(embeddings_path, embeddings)
+        write_chunks_jsonl(chunks_path, chunks_to_records(chunks))
+
+        stats = {
+            "items": len(corpus.load_catalog().items),
+            "text_items": text_items,
+            "chunks": len(chunks),
+            "dimensions": (
+                int(embeddings.shape[1])
+                if embeddings.size
+                else recipe_config.embedding_provider.dimensions
+            ),
+        }
+        run = run.model_copy(
+            update={"artifact_paths": [paths["embeddings"], paths["chunks"]], "stats": stats}
+        )
+        corpus.write_run(run)
+        return run
+
+    def query(
+        self,
+        corpus: Corpus,
+        *,
+        run: RetrievalRun,
+        query_text: str,
+        budget: QueryBudget,
+    ) -> RetrievalResult:
+        """
+        Query an embedding index run and return ranked evidence.
+
+        :param corpus: Corpus associated with the run.
+        :type corpus: Corpus
+        :param run: Run manifest to use for querying.
+        :type run: biblicus.models.RetrievalRun
+        :param query_text: Query text to embed.
+        :type query_text: str
+        :param budget: Evidence selection budget.
+        :type budget: biblicus.models.QueryBudget
+        :return: Retrieval results containing evidence.
+        :rtype: biblicus.models.RetrievalResult
+        """
+        recipe_config = EmbeddingIndexRecipeConfig.model_validate(run.recipe.config)
+        extraction_reference = resolve_extraction_reference(corpus, recipe_config)
+
+        paths = artifact_paths_for_run(run_id=run.run_id, backend_id=self.backend_id)
+        embeddings_path = corpus.root / paths["embeddings"]
+        chunks_path = corpus.root / paths["chunks"]
+        if not embeddings_path.is_file() or not chunks_path.is_file():
+            raise FileNotFoundError("Embedding index artifacts are missing for this run")
+
+        embeddings = read_embeddings(embeddings_path, mmap=True).astype(np.float32)
+        chunk_records = read_chunks_jsonl(chunks_path)
+        if embeddings.shape[0] != len(chunk_records):
+            raise ValueError(
+                "Embedding index artifacts are inconsistent: "
+                "embeddings row count does not match chunk record count"
+            )
+
+        provider = recipe_config.embedding_provider.build_provider()
+        query_embedding = provider.embed_texts([query_text]).astype(np.float32)
+        if query_embedding.shape[0] != 1:
+            raise ValueError("Embedding provider returned an invalid query embedding shape")
+
+        candidates = _top_indices_batched(
+            embeddings=embeddings,
+            query_vector=query_embedding[0],
+            limit=_candidate_limit(budget.max_total_items + budget.offset),
+        )
+        evidence_items = _build_evidence(
+            corpus,
+            run=run,
+            recipe_config=recipe_config,
+            candidates=candidates,
+            embeddings=embeddings,
+            query_vector=query_embedding[0],
+            chunk_records=chunk_records,
+            extraction_reference=extraction_reference,
+        )
+        ranked = [
+            item.model_copy(
+                update={"rank": index, "recipe_id": run.recipe.recipe_id, "run_id": run.run_id}
+            )
+            for index, item in enumerate(evidence_items, start=1)
+        ]
+        evidence = apply_budget(ranked, budget)
+        return RetrievalResult(
+            query_text=query_text,
+            budget=budget,
+            run_id=run.run_id,
+            recipe_id=run.recipe.recipe_id,
+            backend_id=self.backend_id,
+            generated_at=utc_now_iso(),
+            evidence=evidence,
+            stats={"candidates": len(evidence_items), "returned": len(evidence)},
+        )
+
+
+def _candidate_limit(max_total_items: int, *, multiplier: int = 10) -> int:
+    return max(1, int(max_total_items) * int(multiplier))
+
+
+@dataclass(frozen=True)
+class _ScoredIndex:
+    score: float
+    index: int
+
+
+def _top_indices_batched(
+    *, embeddings: np.ndarray, query_vector: np.ndarray, limit: int, batch_rows: int = 4096
+) -> List[int]:
+    if embeddings.size == 0:
+        return []
+    limit = min(int(limit), int(embeddings.shape[0]))
+
+    best: List[_ScoredIndex] = []
+    for start in range(0, embeddings.shape[0], int(batch_rows)):
+        end = min(start + int(batch_rows), embeddings.shape[0])
+        scores = cosine_similarity_scores(embeddings[start:end], query_vector)
+        batch_limit = min(limit, int(scores.size))
+        if batch_limit <= 0:
+            continue
+        indices = np.argpartition(-scores, batch_limit - 1)[:batch_limit]
+        for local_index in indices:
+            global_index = int(start + int(local_index))
+            best.append(_ScoredIndex(score=float(scores[int(local_index)]), index=global_index))
+
+    best.sort(key=lambda item: (-item.score, item.index))
+    return [int(item.index) for item in best[:limit]]
+
+
+def _build_evidence(
+    corpus: Corpus,
+    *,
+    run: RetrievalRun,
+    recipe_config: EmbeddingIndexRecipeConfig,
+    candidates: List[int],
+    embeddings: np.ndarray,
+    query_vector: np.ndarray,
+    chunk_records: List[ChunkRecord],
+    extraction_reference: Optional[ExtractionRunReference],
+) -> List[Evidence]:
+    catalog = corpus.load_catalog()
+    evidence_items: List[Evidence] = []
+    for idx in candidates:
+        record = chunk_records[idx]
+        catalog_item = catalog.items[record.item_id]
+        text = _load_text_for_evidence(
+            corpus,
+            item_id=record.item_id,
+            relpath=str(getattr(catalog_item, "relpath")),
+            media_type=str(getattr(catalog_item, "media_type")),
+            extraction_reference=extraction_reference,
+        )
+        snippet = _build_snippet(
+            text, (record.span_start, record.span_end), max_chars=recipe_config.snippet_characters
+        )
+        score = float(cosine_similarity_scores(embeddings[idx : idx + 1], query_vector)[0])
+        evidence_items.append(
+            Evidence(
+                item_id=record.item_id,
+                source_uri=getattr(catalog_item, "source_uri", None),
+                media_type=str(getattr(catalog_item, "media_type")),
+                score=score,
+                rank=1,
+                text=snippet,
+                content_ref=None,
+                span_start=record.span_start,
+                span_end=record.span_end,
+                stage=EmbeddingIndexFileBackend.backend_id,
+                stage_scores=None,
+                recipe_id=run.recipe.recipe_id,
+                run_id=run.run_id,
+                hash=hash_text(snippet),
+            )
+        )
+    return evidence_items
+
+
+def _load_text_for_evidence(
+    corpus: Corpus,
+    *,
+    item_id: str,
+    relpath: str,
+    media_type: str,
+    extraction_reference: Optional[ExtractionRunReference],
+) -> Optional[str]:
+    from .embedding_index_common import _load_text_from_item
+
+    return _load_text_from_item(
+        corpus,
+        item_id=item_id,
+        relpath=relpath,
+        media_type=media_type,
+        extraction_reference=extraction_reference,
+    )