Merge pull request #212 from ggozad/feat/contextualize-fts

Add additional context (headers/sections etc) to the FTS index

CHANGELOG.md

@@ -8,6 +8,11 @@
   - LLM tests (QA, embeddings, research graph) replay from cassettes without real API calls
   - docling-serve tests run without Docker container in CI
   - Uses pytest-recording with custom JSON body serializer
+- **Contextualized FTS Search**: Full-text search now includes section headings
+  - New `content_fts` column stores contextualized content (headings + body text)
+  - FTS index now searches `content_fts` for better keyword matching on section context
+  - Original `content` column preserved for display and context expansion
+  - Migration automatically populates `content_fts` for existing databases
 
 ### Removed
 

evaluations/pyproject.toml

@@ -2,7 +2,7 @@
 
 name = "haiku.rag-evals"
 description = "Benchmarking and evaluation scripts for haiku.rag"
-version = "0.23.0"
+version = "0.23.1"
 authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }]
 license = { text = "MIT" }
 requires-python = ">=3.12"

haiku_rag_slim/haiku/rag/chunkers/docling_local.py

@@ -117,8 +117,6 @@ class DoclingLocalChunker(DocumentChunker):
         result: list[Chunk] = []
 
         for chunk in raw_chunks:
-            # Use raw chunk text - headings are stored separately in metadata
-            # and prepended at embedding time for better semantic search
             text = chunk.text
 
             # Extract metadata from DocChunk.meta (cast to DocMeta for type safety)

haiku_rag_slim/haiku/rag/client.py

@@ -1410,12 +1410,13 @@ class HaikuRAG:
             embeddings = await self.chunk_repository.embedder.embed_documents(texts)
 
             # Build updated records
-            for chunk, embedding in zip(chunks, embeddings):
+            for chunk, content_fts, embedding in zip(chunks, texts, embeddings):
                 pending_records.append(
                     self.store.ChunkRecord(
                         id=chunk.id,  # type: ignore[arg-type]
                         document_id=chunk.document_id,  # type: ignore[arg-type]
                         content=chunk.content,
+                        content_fts=content_fts,
                         metadata=json.dumps(chunk.metadata),
                         order=chunk.order,
                         vector=embedding,

haiku_rag_slim/haiku/rag/embeddings/__init__.py

@@ -32,16 +32,15 @@ class EmbedderWrapper:
 
 
 def contextualize(chunks: list["Chunk"]) -> list[str]:
-    """Prepare chunk content for embedding by adding context.
+    """Prepare chunk content for embedding/FTS by adding context.
 
     Prepends section headings to chunk content for better semantic search.
-    The embeddings will capture section context while stored content stays raw.
 
     Args:
         chunks: List of chunks to contextualize.
 
     Returns:
-        List of contextualized text strings for embedding.
+        List of contextualized text strings.
     """
     texts = []
     for chunk in chunks:

haiku_rag_slim/haiku/rag/store/engine.py

@@ -40,6 +40,7 @@ def create_chunk_model(vector_dim: int):
         id: str = Field(default_factory=lambda: str(uuid4()))
         document_id: str
         content: str
+        content_fts: str = Field(default="")
         metadata: str = Field(default="{}")
         order: int = Field(default=0)
         vector: Vector(vector_dim) = Field(default_factory=lambda: [0.0] * vector_dim)  # type: ignore
@@ -288,9 +289,9 @@ class Store:
             self.chunks_table = self.db.open_table("chunks")
         else:
             self.chunks_table = self.db.create_table("chunks", schema=self.ChunkRecord)
-            # Create FTS index on the new table with phrase query support
+            # Create FTS index on content_fts (contextualized content) for better search
             self.chunks_table.create_fts_index(
-                "content", replace=True, with_position=True, remove_stop_words=False
+                "content_fts", replace=True, with_position=True, remove_stop_words=False
             )
 
         # Create or get settings table
@@ -392,9 +393,9 @@ class Store:
         self.ChunkRecord = create_chunk_model(self.embedder._vector_dim)
         self.chunks_table = self.db.create_table("chunks", schema=self.ChunkRecord)
 
-        # Create FTS index on the new table with phrase query support
+        # Create FTS index on content_fts (contextualized content) for better search
         self.chunks_table.create_fts_index(
-            "content", replace=True, with_position=True, remove_stop_words=False
+            "content_fts", replace=True, with_position=True, remove_stop_words=False
         )
 
     def close(self):

haiku_rag_slim/haiku/rag/store/repositories/chunk.py

@@ -27,15 +27,22 @@ class ChunkRepository:
         self.embedder = store.embedder
 
     def _ensure_fts_index(self) -> None:
-        """Ensure FTS index exists on the content column."""
+        """Ensure FTS index exists on the content_fts column."""
         try:
             self.store.chunks_table.create_fts_index(
-                "content", replace=True, with_position=True, remove_stop_words=False
+                "content_fts", replace=True, with_position=True, remove_stop_words=False
             )
         except Exception as e:
             # Log the error but don't fail - FTS might already exist
             logger.debug(f"FTS index creation skipped: {e}")
 
+    def _contextualize_content(self, chunk: Chunk) -> str:
+        """Generate contextualized content for FTS by prepending headings."""
+        meta = chunk.get_chunk_metadata()
+        if meta.headings:
+            return "\n".join(meta.headings) + "\n" + chunk.content
+        return chunk.content
+
     async def create(self, entity: Chunk | list[Chunk]) -> Chunk | list[Chunk]:
         """Create one or more chunks in the database.
 
@@ -54,6 +61,7 @@ class ChunkRepository:
                 id=chunk_id,
                 document_id=entity.document_id,
                 content=entity.content,
+                content_fts=self._contextualize_content(entity),
                 metadata=json.dumps(
                     {k: v for k, v in entity.metadata.items() if k != "order"}
                 ),
@@ -86,6 +94,7 @@ class ChunkRepository:
                 id=chunk_id,
                 document_id=chunk.document_id,
                 content=chunk.content,
+                content_fts=self._contextualize_content(chunk),
                 metadata=json.dumps(
                     {k: v for k, v in chunk.metadata.items() if k != "order"}
                 ),
@@ -136,6 +145,7 @@ class ChunkRepository:
             values={
                 "document_id": entity.document_id,
                 "content": entity.content,
+                "content_fts": self._contextualize_content(entity),
                 "metadata": json.dumps(
                     {k: v for k, v in entity.metadata.items() if k != "order"}
                 ),
@@ -190,9 +200,9 @@ class ChunkRepository:
         self.store.chunks_table = self.store.db.create_table(
             "chunks", schema=self.store.ChunkRecord
         )
-        # Create FTS index on the new table with phrase query support
+        # Create FTS index on content_fts (contextualized content) for better search
         self.store.chunks_table.create_fts_index(
-            "content", replace=True, with_position=True, remove_stop_words=False
+            "content_fts", replace=True, with_position=True, remove_stop_words=False
         )
 
     async def delete_by_document_id(self, document_id: str) -> bool:
@@ -409,6 +419,7 @@ class ChunkRepository:
                 id=str(row["id"]),
                 document_id=str(row["document_id"]),
                 content=str(row["content"]),
+                content_fts=str(row.get("content_fts", "")),
                 metadata=str(row["metadata"]),
                 order=int(row["order"]) if "order" in row else 0,
             )

haiku_rag_slim/haiku/rag/store/upgrades/__init__.py

@@ -66,9 +66,13 @@ from haiku.rag.store.upgrades.v0_19_6 import (
 from haiku.rag.store.upgrades.v0_20_0 import (
     upgrade_add_docling_document as upgrade_0_20_0_docling,
 )
+from haiku.rag.store.upgrades.v0_23_1 import (
+    upgrade_contextualize_chunks as upgrade_0_23_1_contextualize,
+)
 
 upgrades.append(upgrade_0_9_3_order)
 upgrades.append(upgrade_0_9_3_fts)
 upgrades.append(upgrade_0_10_1_add_title)
 upgrades.append(upgrade_0_19_6_embeddings)
 upgrades.append(upgrade_0_20_0_docling)
+upgrades.append(upgrade_0_23_1_contextualize)

haiku_rag_slim/haiku/rag/store/upgrades/v0_23_1.py

@@ -0,0 +1,100 @@
+import json
+
+from lancedb.pydantic import LanceModel, Vector
+from pydantic import Field
+
+from haiku.rag.store.engine import Store
+from haiku.rag.store.upgrades import Upgrade
+
+
+def _apply_add_content_fts(store: Store) -> None:  # pragma: no cover
+    """Add content_fts column with contextualized content for better FTS."""
+    # Read existing chunks
+    try:
+        chunks_arrow = store.chunks_table.search().to_arrow()
+        rows = chunks_arrow.to_pylist()
+    except Exception:
+        return
+
+    if not rows:
+        return
+
+    # Infer vector dimensions from first row
+    vec = rows[0].get("vector")
+    if not isinstance(vec, list) or not vec:
+        return
+    vector_dim = len(vec)
+
+    class ChunkRecord(LanceModel):
+        id: str
+        document_id: str
+        content: str
+        content_fts: str = Field(default="")
+        metadata: str = Field(default="{}")
+        order: int = Field(default=0)
+        vector: Vector(vector_dim) = Field(  # type: ignore
+            default_factory=lambda: [0.0] * vector_dim
+        )
+
+    # Drop and recreate table with new schema
+    try:
+        store.db.drop_table("chunks")
+    except Exception:
+        pass
+
+    store.chunks_table = store.db.create_table("chunks", schema=ChunkRecord)
+
+    # Populate content_fts with contextualized content
+    new_records: list[ChunkRecord] = []
+    for row in rows:
+        metadata_raw = row.get("metadata") or "{}"
+        try:
+            metadata = (
+                json.loads(metadata_raw)
+                if isinstance(metadata_raw, str)
+                else metadata_raw
+            )
+        except Exception:
+            metadata = {}
+
+        headings = metadata.get("headings") if isinstance(metadata, dict) else None
+        content = row.get("content", "")
+
+        # Build contextualized content for FTS
+        if headings:
+            content_fts = "\n".join(headings) + "\n" + content
+        else:
+            content_fts = content
+
+        new_records.append(
+            ChunkRecord(
+                id=row.get("id"),
+                document_id=row.get("document_id"),
+                content=content,
+                content_fts=content_fts,
+                metadata=metadata_raw,
+                order=row.get("order", 0),
+                vector=row.get("vector") or [0.0] * vector_dim,
+            )
+        )
+
+    if new_records:
+        store.chunks_table.add(new_records)
+
+    # Drop old FTS index on content column if it exists
+    try:
+        store.chunks_table.drop_index("content_idx")
+    except Exception:
+        pass
+
+    # Create FTS index on content_fts
+    store.chunks_table.create_fts_index(
+        "content_fts", replace=True, with_position=True, remove_stop_words=False
+    )
+
+
+upgrade_contextualize_chunks = Upgrade(
+    version="0.23.1",
+    apply=_apply_add_content_fts,
+    description="Add content_fts column for contextualized FTS search",
+)

haiku_rag_slim/pyproject.toml

@@ -2,7 +2,7 @@
 
 name = "haiku.rag-slim"
 description = "Opinionated agentic RAG powered by LanceDB, Pydantic AI, and Docling - Minimal dependencies"
-version = "0.23.0"
+version = "0.23.1"
 authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }]
 license = { text = "MIT" }
 readme = { file = "README.md", content-type = "text/markdown" }

pyproject.toml

@@ -2,7 +2,7 @@
 
 name = "haiku.rag"
 description = "Opinionated agentic RAG powered by LanceDB, Pydantic AI, and Docling"
-version = "0.23.0"
+version = "0.23.1"
 authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }]
 license = { text = "MIT" }
 readme = { file = "README.md", content-type = "text/markdown" }

tests/test_chunk.py

@@ -342,3 +342,86 @@ def test_search_result_get_primary_label():
     # Empty labels
     result = SearchResult(content="x", score=0.5, labels=[])
     assert result._get_primary_label() is None
+
+
+@pytest.mark.vcr()
+async def test_chunk_content_fts_populated(temp_db_path):
+    """Test that content_fts column is populated with contextualized content."""
+    from haiku.rag.embeddings import get_embedder
+
+    async with HaikuRAG(db_path=temp_db_path, config=Config, create=True) as client:
+        # Create a chunk with headings
+        chunk = Chunk(
+            document_id="test-doc",
+            content="This is the raw chunk content.",
+            metadata={"headings": ["Chapter 1", "Section 1.1"]},
+            order=0,
+        )
+
+        # Generate embedding
+        embedder = get_embedder(Config)
+        embedding = (await embedder.embed_documents([chunk.content]))[0]
+        chunk.embedding = embedding
+
+        # Store the chunk
+        await client.chunk_repository.create(chunk)
+
+        # Read the raw record from the database
+        records = list(
+            client.store.chunks_table.search()
+            .where(f"id = '{chunk.id}'")
+            .limit(1)
+            .to_arrow()
+            .to_pylist()
+        )
+
+        assert len(records) == 1
+        record = records[0]
+
+        # Verify content is raw (no headings)
+        assert record["content"] == "This is the raw chunk content."
+
+        # Verify content_fts is contextualized (headings + content)
+        assert (
+            record["content_fts"]
+            == "Chapter 1\nSection 1.1\nThis is the raw chunk content."
+        )
+
+
+@pytest.mark.vcr()
+async def test_chunk_content_fts_without_headings(temp_db_path):
+    """Test that content_fts equals content when no headings present."""
+    from haiku.rag.embeddings import get_embedder
+
+    async with HaikuRAG(db_path=temp_db_path, config=Config, create=True) as client:
+        # Create a chunk without headings
+        chunk = Chunk(
+            document_id="test-doc",
+            content="Plain content without headings.",
+            metadata={},
+            order=0,
+        )
+
+        # Generate embedding
+        embedder = get_embedder(Config)
+        embedding = (await embedder.embed_documents([chunk.content]))[0]
+        chunk.embedding = embedding
+
+        # Store the chunk
+        await client.chunk_repository.create(chunk)
+
+        # Read the raw record from the database
+        records = list(
+            client.store.chunks_table.search()
+            .where(f"id = '{chunk.id}'")
+            .limit(1)
+            .to_arrow()
+            .to_pylist()
+        )
+
+        assert len(records) == 1
+        record = records[0]
+
+        # Both should be the same when no headings
+        assert record["content"] == "Plain content without headings."
+        assert record["content_fts"] == "Plain content without headings."

uv.lock

@@ -1248,7 +1248,7 @@ wheels = [
 
 [[package]]
 name = "haiku-rag"
-version = "0.23.0"
+version = "0.23.1"
 source = { editable = "." }
 dependencies = [
     { name = "haiku-rag-slim", extra = ["cohere", "docling", "inspector", "mxbai", "voyageai", "zeroentropy"] },
@@ -1303,7 +1303,7 @@ dev = [
 
 [[package]]
 name = "haiku-rag-evals"
-version = "0.23.0"
+version = "0.23.1"
 source = { editable = "evaluations" }
 dependencies = [
     { name = "datasets" },
@@ -1324,7 +1324,7 @@ requires-dist = [
 
 [[package]]
 name = "haiku-rag-slim"
-version = "0.23.0"
+version = "0.23.1"
 source = { editable = "haiku_rag_slim" }
 dependencies = [
     { name = "docling-core" },