diff --git a/evaluations/pyproject.toml b/evaluations/pyproject.toml index 630974db..c1b0e202 100644 --- a/evaluations/pyproject.toml +++ b/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" diff --git a/haiku_rag_slim/haiku/rag/chunkers/docling_local.py b/haiku_rag_slim/haiku/rag/chunkers/docling_local.py index 1bebbe07..2ca7a2f4 100644 --- a/haiku_rag_slim/haiku/rag/chunkers/docling_local.py +++ b/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) diff --git a/haiku_rag_slim/haiku/rag/client.py b/haiku_rag_slim/haiku/rag/client.py index 338182f9..f7a7b32f 100644 --- a/haiku_rag_slim/haiku/rag/client.py +++ b/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, diff --git a/haiku_rag_slim/haiku/rag/embeddings/__init__.py b/haiku_rag_slim/haiku/rag/embeddings/__init__.py index 7894ea69..f17416e1 100644 --- a/haiku_rag_slim/haiku/rag/embeddings/__init__.py +++ b/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: diff --git a/haiku_rag_slim/haiku/rag/store/engine.py b/haiku_rag_slim/haiku/rag/store/engine.py index ed1b3dbb..6e2b9ace 100644 --- a/haiku_rag_slim/haiku/rag/store/engine.py +++ b/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): diff --git a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py index b5e7799e..ed47af5f 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py +++ b/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, ) diff --git a/haiku_rag_slim/haiku/rag/store/upgrades/__init__.py b/haiku_rag_slim/haiku/rag/store/upgrades/__init__.py index ef8661fb..43a93dc4 100644 --- a/haiku_rag_slim/haiku/rag/store/upgrades/__init__.py +++ b/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) diff --git a/haiku_rag_slim/haiku/rag/store/upgrades/v0_23_1.py b/haiku_rag_slim/haiku/rag/store/upgrades/v0_23_1.py new file mode 100644 index 00000000..5797ee7e --- /dev/null +++ b/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", +) diff --git a/haiku_rag_slim/pyproject.toml b/haiku_rag_slim/pyproject.toml index e7e65abf..e1b922c1 100644 --- a/haiku_rag_slim/pyproject.toml +++ b/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" } diff --git a/pyproject.toml b/pyproject.toml index 770556a1..d5db3079 100644 --- a/pyproject.toml +++ b/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" } diff --git a/tests/test_chunk.py b/tests/test_chunk.py index 2738b383..b4e0ac8e 100644 --- a/tests/test_chunk.py +++ b/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.asyncio +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.asyncio +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." diff --git a/uv.lock b/uv.lock index 2889f0cd..534821a0 100644 --- a/uv.lock +++ b/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" },