From c221a355479352e2cab7ef9be2859c7b5d48f833 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 18 Sep 2025 18:00:46 +0300 Subject: [PATCH] Upgrade fts index for phrase queries --- src/haiku/rag/store/engine.py | 12 ++++++--- src/haiku/rag/store/repositories/chunk.py | 10 +++++--- src/haiku/rag/store/upgrades/__init__.py | 30 +++++++++++++++++------ src/haiku/rag/store/upgrades/v0_9_3.py | 28 ++++++++++++++------- 4 files changed, 56 insertions(+), 24 deletions(-) diff --git a/src/haiku/rag/store/engine.py b/src/haiku/rag/store/engine.py index 7485e3b3..aabb7360 100644 --- a/src/haiku/rag/store/engine.py +++ b/src/haiku/rag/store/engine.py @@ -118,8 +118,10 @@ 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 - self.chunks_table.create_fts_index("content", replace=True) + # Create FTS index on the new table with phrase query support + self.chunks_table.create_fts_index( + "content", replace=True, with_position=True, remove_stop_words=False + ) # Create or get settings table if "settings" in existing_tables: @@ -222,8 +224,10 @@ 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 - self.chunks_table.create_fts_index("content", replace=True) + # Create FTS index on the new table with phrase query support + self.chunks_table.create_fts_index( + "content", replace=True, with_position=True, remove_stop_words=False + ) def close(self): """Close the database connection.""" diff --git a/src/haiku/rag/store/repositories/chunk.py b/src/haiku/rag/store/repositories/chunk.py index 1526b485..567a5a34 100644 --- a/src/haiku/rag/store/repositories/chunk.py +++ b/src/haiku/rag/store/repositories/chunk.py @@ -28,7 +28,9 @@ class ChunkRepository: def _ensure_fts_index(self) -> None: """Ensure FTS index exists on the content column.""" try: - self.store.chunks_table.create_fts_index("content", replace=True) + self.store.chunks_table.create_fts_index( + "content", 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}") @@ -236,8 +238,10 @@ class ChunkRepository: self.store.chunks_table = self.store.db.create_table( "chunks", schema=self.store.ChunkRecord ) - # Create FTS index on the new table - self.store.chunks_table.create_fts_index("content", replace=True) + # Create FTS index on the new table with phrase query support + self.store.chunks_table.create_fts_index( + "content", replace=True, with_position=True, remove_stop_words=False + ) async def delete_by_document_id(self, document_id: str) -> bool: """Delete all chunks for a document.""" diff --git a/src/haiku/rag/store/upgrades/__init__.py b/src/haiku/rag/store/upgrades/__init__.py index a39283ea..5888cdba 100644 --- a/src/haiku/rag/store/upgrades/__init__.py +++ b/src/haiku/rag/store/upgrades/__init__.py @@ -1,5 +1,4 @@ -from __future__ import annotations - +import logging from collections.abc import Callable from dataclasses import dataclass @@ -7,6 +6,8 @@ from packaging.version import Version, parse from haiku.rag.store.engine import Store +logger = logging.getLogger(__name__) + @dataclass class Upgrade: @@ -33,14 +34,27 @@ def run_pending_upgrades(store: Store, from_version: str, to_version: str) -> No if highest_step_version > v_to: v_to = highest_step_version - # Ensure upgrades are applied in ascending version order - for step in sorted(upgrades, key=lambda u: parse(u.version)): - v_step = parse(step.version) - if v_from < v_step <= v_to: - step.apply(store) + # Determine applicable steps + sorted_steps = sorted(upgrades, key=lambda u: parse(u.version)) + applicable = [s for s in sorted_steps if v_from < parse(s.version) <= v_to] + if applicable: + logger.info("%d upgrade step(s) pending", len(applicable)) + + # Apply in ascending order + for idx, step in enumerate(applicable, start=1): + logger.info( + "Applying upgrade %s: %s (%d/%d)", + step.version, + step.description or "", + idx, + len(applicable), + ) + step.apply(store) + logger.info("Completed upgrade %s", step.version) -# Import concrete upgrade modules (module names cannot start with a digit) +from .v0_9_3 import upgrade_fts_phrase as upgrade_0_9_3_fts # noqa: E402 from .v0_9_3 import upgrade_order as upgrade_0_9_3_order # noqa: E402 upgrades.append(upgrade_0_9_3_order) +upgrades.append(upgrade_0_9_3_fts) diff --git a/src/haiku/rag/store/upgrades/v0_9_3.py b/src/haiku/rag/store/upgrades/v0_9_3.py index 7090a38e..fc3dccd5 100644 --- a/src/haiku/rag/store/upgrades/v0_9_3.py +++ b/src/haiku/rag/store/upgrades/v0_9_3.py @@ -1,15 +1,10 @@ -from __future__ import annotations - import json -from typing import TYPE_CHECKING from lancedb.pydantic import LanceModel, Vector from pydantic import Field -from . import Upgrade - -if TYPE_CHECKING: # pragma: no cover - for type hints only - from haiku.rag.store.engine import Store +from haiku.rag.store.engine import Store +from haiku.rag.store.upgrades import Upgrade def _infer_vector_dim(store: Store) -> int: @@ -32,7 +27,6 @@ def _apply_chunk_order(store: Store) -> None: vector_dim = _infer_vector_dim(store) - # ============== Chunks: add 'order' column and backfill ============== class ChunkRecordV2(LanceModel): id: str document_id: str @@ -88,7 +82,6 @@ def _apply_chunk_order(store: Store) -> None: pass store.chunks_table = store.db.create_table("chunks", schema=ChunkRecordV2) - # Recreate FTS index on content store.chunks_table.create_fts_index("content", replace=True) if new_chunk_records: @@ -100,3 +93,20 @@ upgrade_order = Upgrade( apply=_apply_chunk_order, description="Add 'order' column to chunks and backfill from metadata", ) + + +def _apply_fts_phrase_support(store: Store) -> None: + """Recreate FTS index with phrase query support and no stop-word removal.""" + try: + store.chunks_table.create_fts_index( + "content", replace=True, with_position=True, remove_stop_words=False + ) + except Exception: + pass + + +upgrade_fts_phrase = Upgrade( + version="0.9.3", + apply=_apply_fts_phrase_support, + description="Enable FTS phrase queries (with positions) and keep stop-words", +)