danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Use FTS on contextualized text

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-12-27 16:17:04 +02:00
parent 02f5dc4fb9
commit c9cb005d07
No known key found for this signature in database
12 changed files with 217 additions and 20 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
evaluations/pyproject.toml
View file

@ -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"

2
haiku_rag_slim/haiku/rag/chunkers/docling_local.py
View file

@ -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)

3
haiku_rag_slim/haiku/rag/client.py
View file

@ -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,

5
haiku_rag_slim/haiku/rag/embeddings/__init__.py
View file

@ -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:

9
haiku_rag_slim/haiku/rag/store/engine.py
View file

@ -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):

19
haiku_rag_slim/haiku/rag/store/repositories/chunk.py
View file

@ -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,
)

4
haiku_rag_slim/haiku/rag/store/upgrades/__init__.py
View file

@ -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)

100
haiku_rag_slim/haiku/rag/store/upgrades/v0_23_1.py Normal file
View file

@ -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",
)

2
haiku_rag_slim/pyproject.toml
View file

@ -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" }

2
pyproject.toml
View file

@ -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" }

83
tests/test_chunk.py
View file

@ -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."

6
uv.lock
View file

@ -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" },