Merge pull request #61 from ggozad/feat/chunk-order

Support phrase search. Move "order" from metadata to its own column.

pyproject.toml

@@ -1,7 +1,8 @@
 [project]
+
 name = "haiku.rag"
-version = "0.9.2"
 description = "Agentic Retrieval Augmented Generation (RAG) with LanceDB"
+version = "0.9.2"
 authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }]
 license = { text = "MIT" }
 readme = { file = "README.md", content-type = "text/markdown" }

src/haiku/rag/client.py

@@ -388,7 +388,7 @@ class HaikuRAG:
                 all_chunks = adjacent_chunks + [chunk]
 
                 # Get the range of orders for this expanded chunk
-                orders = [c.metadata.get("order", 0) for c in all_chunks]
+                orders = [c.order for c in all_chunks]
                 min_order = min(orders)
                 max_order = max(orders)
 
@@ -398,9 +398,7 @@ class HaikuRAG:
                         "score": score,
                         "min_order": min_order,
                         "max_order": max_order,
-                        "all_chunks": sorted(
+                        "all_chunks": sorted(all_chunks, key=lambda c: c.order),
-                            all_chunks, key=lambda c: c.metadata.get("order", 0)
-                        ),
                     }
                 )
 
@@ -459,7 +457,7 @@ class HaikuRAG:
                 # Merge all_chunks and deduplicate by order
                 all_chunks_dict = {}
                 for chunk in current["all_chunks"] + range_info["all_chunks"]:
-                    order = chunk.metadata.get("order", 0)
+                    order = chunk.order
                     all_chunks_dict[order] = chunk
                 current["all_chunks"] = [
                     all_chunks_dict[order] for order in sorted(all_chunks_dict.keys())

src/haiku/rag/store/engine.py

@@ -35,6 +35,7 @@ def create_chunk_model(vector_dim: int):
         document_id: str
         content: str
         metadata: str = Field(default="{}")
+        order: int = Field(default=0)
         vector: Vector(vector_dim) = Field(default_factory=lambda: [0.0] * vector_dim)  # type: ignore
 
     return ChunkRecord
@@ -117,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
+            # Create FTS index on the new table with phrase query support
-            self.chunks_table.create_fts_index("content", replace=True)
+            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:
@@ -133,21 +136,41 @@ class Store:
                 [SettingsRecord(id="settings", settings=json.dumps(settings_data))]
             )
 
-        # Set current version in settings
+        # Run pending upgrades based on stored version and package version
-        current_version = metadata.version("haiku.rag")
-        self.set_haiku_version(current_version)
-
-        # Check if we need to perform upgrades
         try:
-            existing_settings = list(
+            from haiku.rag.store.upgrades import run_pending_upgrades
-                self.settings_table.search().limit(1).to_pydantic(SettingsRecord)
+
+            current_version = metadata.version("haiku.rag")
+            db_version = self.get_haiku_version()
+
+            run_pending_upgrades(self, db_version, current_version)
+
+            # After upgrades complete (or if none), set stored version
+            # to the greater of the installed package version and the
+            # highest available upgrade step version in code.
+            try:
+                from packaging.version import parse as _v
+
+                from haiku.rag.store.upgrades import upgrades as _steps
+
+                highest_step = max((_v(u.version) for u in _steps), default=None)
+                effective_version = (
+                    str(max(_v(current_version), highest_step))
+                    if highest_step is not None
+                    else current_version
+                )
+            except Exception:
+                effective_version = current_version
+
+            self.set_haiku_version(effective_version)
+        except Exception as e:
+            # Avoid hard failure on initial connection; log and continue so CLI remains usable.
+            logger.warning(
+                "Skipping upgrade due to error (db=%s -> pkg=%s): %s",
+                self.get_haiku_version(),
+                metadata.version("haiku.rag") if hasattr(metadata, "version") else "",
+                e,
             )
-            if existing_settings:
-                db_version = self.get_haiku_version()  # noqa: F841
-                # TODO: Add upgrade logic here similar to SQLite version when needed
-        except Exception:
-            # Settings table might not exist yet in fresh databases
-            pass
 
     def get_haiku_version(self) -> str:
         """Returns the user version stored in settings."""
@@ -201,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
+        # Create FTS index on the new table with phrase query support
-        self.chunks_table.create_fts_index("content", replace=True)
+        self.chunks_table.create_fts_index(
+            "content", replace=True, with_position=True, remove_stop_words=False
+        )
 
     def close(self):
         """Close the database connection."""

src/haiku/rag/store/models/chunk.py

@@ -10,6 +10,7 @@ class Chunk(BaseModel):
     document_id: str | None = None
     content: str
     metadata: dict = {}
+    order: int = 0
     document_uri: str | None = None
     document_meta: dict = {}
     embedding: list[float] | None = None

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}")
@@ -59,11 +61,16 @@ class ChunkRepository:
             embedding = entity.embedding
         else:
             embedding = await self.embedder.embed(entity.content)
+        order_val = int(entity.order)
+
         chunk_record = self.store.ChunkRecord(
             id=chunk_id,
             document_id=entity.document_id,
             content=entity.content,
-            metadata=json.dumps(entity.metadata),
+            metadata=json.dumps(
+                {k: v for k, v in entity.metadata.items() if k != "order"}
+            ),
+            order=order_val,
             vector=embedding,
         )
 
@@ -90,11 +97,13 @@ class ChunkRepository:
             return None
 
         chunk_record = results[0]
+        md = json.loads(chunk_record.metadata)
         return Chunk(
             id=chunk_record.id,
             document_id=chunk_record.document_id,
             content=chunk_record.content,
-            metadata=json.loads(chunk_record.metadata) if chunk_record.metadata else {},
+            metadata=md,
+            order=chunk_record.order,
         )
 
     async def update(self, entity: Chunk) -> Chunk:
@@ -102,13 +111,17 @@ class ChunkRepository:
         assert entity.id, "Chunk ID is required for update"
 
         embedding = await self.embedder.embed(entity.content)
+        order_val = int(entity.order)
 
         self.store.chunks_table.update(
             where=f"id = '{entity.id}'",
             values={
                 "document_id": entity.document_id,
                 "content": entity.content,
-                "metadata": json.dumps(entity.metadata),
+                "metadata": json.dumps(
+                    {k: v for k, v in entity.metadata.items() if k != "order"}
+                ),
+                "order": order_val,
                 "vector": embedding,
             },
         )
@@ -140,15 +153,19 @@ class ChunkRepository:
 
         results = list(query.to_pydantic(self.store.ChunkRecord))
 
-        return [
+        chunks: list[Chunk] = []
-            Chunk(
+        for rec in results:
-                id=chunk.id,
+            md = json.loads(rec.metadata)
-                document_id=chunk.document_id,
+            chunks.append(
-                content=chunk.content,
+                Chunk(
-                metadata=json.loads(chunk.metadata) if chunk.metadata else {},
+                    id=rec.id,
+                    document_id=rec.document_id,
+                    content=rec.content,
+                    metadata=md,
+                    order=rec.order,
+                )
             )
-            for chunk in results
+        return chunks
-        ]
 
     async def create_chunks_for_document(
         self, document_id: str, document: DoclingDocument
@@ -191,7 +208,8 @@ class ChunkRepository:
                 id=chunk_id,
                 document_id=document_id,
                 content=chunk_text,
-                metadata=json.dumps({"order": order}),
+                metadata=json.dumps({}),
+                order=order,
                 vector=embedding,
             )
             chunk_records.append(chunk_record)
@@ -200,7 +218,8 @@ class ChunkRepository:
                 id=chunk_id,
                 document_id=document_id,
                 content=chunk_text,
-                metadata={"order": order},
+                metadata={},
+                order=order,
             )
             created_chunks.append(chunk)
 
@@ -219,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
+        # Create FTS index on the new table with phrase query support
-        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
+        )
 
     async def delete_by_document_id(self, document_id: str) -> bool:
         """Delete all chunks for a document."""
@@ -298,37 +319,36 @@ class ChunkRepository:
         doc_uri = doc_results[0].uri if doc_results else None
         doc_meta = doc_results[0].metadata if doc_results else "{}"
 
-        # Sort by order in metadata
+        chunks: list[Chunk] = []
-        chunks = [
+        for rec in results:
-            Chunk(
+            md = json.loads(rec.metadata)
-                id=chunk.id,
+            chunks.append(
-                document_id=chunk.document_id,
+                Chunk(
-                content=chunk.content,
+                    id=rec.id,
-                metadata=json.loads(chunk.metadata) if chunk.metadata else {},
+                    document_id=rec.document_id,
-                document_uri=doc_uri,
+                    content=rec.content,
-                document_meta=json.loads(doc_meta) if doc_meta else {},
+                    metadata=md,
+                    order=rec.order,
+                    document_uri=doc_uri,
+                    document_meta=json.loads(doc_meta),
+                )
             )
-            for chunk in results
-        ]
 
-        chunks.sort(key=lambda c: c.metadata.get("order", 0))
+        chunks.sort(key=lambda c: c.order)
         return chunks
 
     async def get_adjacent_chunks(self, chunk: Chunk, num_adjacent: int) -> list[Chunk]:
         """Get adjacent chunks before and after the given chunk within the same document."""
         assert chunk.document_id, "Document id is required for adjacent chunk finding"
 
-        chunk_order = chunk.metadata.get("order")
+        chunk_order = chunk.order
-        if chunk_order is None:
-            return []
 
-        # Get all chunks for the document
+        # Fetch chunks for the same document and filter by order proximity
         all_chunks = await self.get_by_document_id(chunk.document_id)
 
-        # Filter to adjacent chunks
+        adjacent_chunks: list[Chunk] = []
-        adjacent_chunks = []
         for c in all_chunks:
-            c_order = c.metadata.get("order", 0)
+            c_order = c.order
             if c.id != chunk.id and abs(c_order - chunk_order) <= num_adjacent:
                 adjacent_chunks.append(c)
 
@@ -380,15 +400,16 @@ class ChunkRepository:
             doc_uri = doc.uri if doc else None
             doc_meta = doc.metadata if doc else "{}"
 
+            md = json.loads(chunk_record.metadata)
+
             chunk = Chunk(
                 id=chunk_record.id,
                 document_id=chunk_record.document_id,
                 content=chunk_record.content,
-                metadata=json.loads(chunk_record.metadata)
+                metadata=md,
-                if chunk_record.metadata
+                order=chunk_record.order,
-                else {},
                 document_uri=doc_uri,
-                document_meta=json.loads(doc_meta) if doc_meta else {},
+                document_meta=json.loads(doc_meta),
             )
 
             # Get score from arrow result

src/haiku/rag/store/repositories/document.py

@@ -34,7 +34,7 @@ class DocumentRepository:
             id=record.id,
             content=record.content,
             uri=record.uri,
-            metadata=json.loads(record.metadata) if record.metadata else {},
+            metadata=json.loads(record.metadata),
             created_at=datetime.fromisoformat(record.created_at)
             if record.created_at
             else datetime.now(),
@@ -194,7 +194,7 @@ class DocumentRepository:
                 )
                 for order, chunk in enumerate(chunks):
                     chunk.document_id = created_doc.id
-                    chunk.metadata["order"] = order
+                    chunk.order = order
                     await self.chunk_repository.create(chunk)
 
             return created_doc

src/haiku/rag/store/repositories/settings.py

@@ -84,11 +84,18 @@ class SettingsRepository:
         )
 
         if existing:
-            # Only update when configuration actually changed to avoid needless new versions
+            # Preserve existing version if present to avoid interfering with upgrade flow
-            existing_payload = (
+            try:
-                json.loads(existing[0].settings) if existing[0].settings else {}
+                existing_settings = (
-            )
+                    json.loads(existing[0].settings) if existing[0].settings else {}
-            if existing_payload != current_config:
+                )
+            except Exception:
+                existing_settings = {}
+            if "version" in existing_settings:
+                current_config["version"] = existing_settings["version"]
+
+            # Update existing settings
+            if existing_settings != current_config:
                 self.store.settings_table.update(
                     where="id = 'settings'",
                     values={"settings": json.dumps(current_config)},

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

@@ -1 +1,60 @@
-upgrades = []
+import logging
+from collections.abc import Callable
+from dataclasses import dataclass
+
+from packaging.version import Version, parse
+
+from haiku.rag.store.engine import Store
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class Upgrade:
+    """Represents a database upgrade step."""
+
+    version: str
+    apply: Callable[[Store], None]
+    description: str = ""
+
+
+# Registry of upgrade steps (ordered by version)
+upgrades: list[Upgrade] = []
+
+
+def run_pending_upgrades(store: Store, from_version: str, to_version: str) -> None:
+    """Run upgrades where from_version < step.version <= to_version."""
+    v_from: Version = parse(from_version)
+    v_to: Version = parse(to_version)
+
+    # Ensure that tests/development run available code upgrades even if the
+    # installed package version hasn't been bumped to include them yet.
+    if upgrades:
+        highest_step_version: Version = max(parse(u.version) for u in upgrades)
+        if highest_step_version > v_to:
+            v_to = highest_step_version
+
+    # 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)
+
+
+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)

src/haiku/rag/store/upgrades/v0_9_3.py

@@ -0,0 +1,112 @@
+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 _infer_vector_dim(store: Store) -> int:
+    """Infer vector dimension from existing data; fallback to embedder config."""
+    try:
+        arrow = store.chunks_table.search().limit(1).to_arrow()
+        rows = arrow.to_pylist()
+        if rows:
+            vec = rows[0].get("vector")
+            if isinstance(vec, list) and vec:
+                return len(vec)
+    except Exception:
+        pass
+    # Fallback to configured embedder vector dim
+    return getattr(store.embedder, "_vector_dim", 1024)
+
+
+def _apply_chunk_order(store: Store) -> None:
+    """Add integer 'order' column to chunks and backfill from metadata."""
+
+    vector_dim = _infer_vector_dim(store)
+
+    class ChunkRecordV2(LanceModel):
+        id: str
+        document_id: str
+        content: str
+        metadata: str = Field(default="{}")
+        order: int = Field(default=0)
+        vector: Vector(vector_dim) = Field(  # type: ignore
+            default_factory=lambda: [0.0] * vector_dim
+        )
+
+    # Read existing chunks
+    try:
+        chunks_arrow = store.chunks_table.search().to_arrow()
+        rows = chunks_arrow.to_pylist()
+    except Exception:
+        rows = []
+
+    new_chunk_records: list[ChunkRecordV2] = []
+    for row in rows:
+        md_raw = row.get("metadata") or "{}"
+        try:
+            md = json.loads(md_raw) if isinstance(md_raw, str) else md_raw
+        except Exception:
+            md = {}
+        # Extract and normalize order
+        order_val = 0
+        try:
+            if isinstance(md, dict) and "order" in md:
+                order_val = int(md["order"])  # type: ignore[arg-type]
+        except Exception:
+            order_val = 0
+
+        if isinstance(md, dict) and "order" in md:
+            md = {k: v for k, v in md.items() if k != "order"}
+
+        vec = row.get("vector") or [0.0] * vector_dim
+
+        new_chunk_records.append(
+            ChunkRecordV2(
+                id=row.get("id"),
+                document_id=row.get("document_id"),
+                content=row.get("content", ""),
+                metadata=json.dumps(md),
+                order=order_val,
+                vector=vec,
+            )
+        )
+
+    # Recreate chunks table with new schema
+    try:
+        store.db.drop_table("chunks")
+    except Exception:
+        pass
+
+    store.chunks_table = store.db.create_table("chunks", schema=ChunkRecordV2)
+    store.chunks_table.create_fts_index("content", replace=True)
+
+    if new_chunk_records:
+        store.chunks_table.add(new_chunk_records)
+
+
+upgrade_order = Upgrade(
+    version="0.9.3",
+    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",
+)

tests/test_chunk.py

@@ -80,9 +80,9 @@ async def test_create_chunks_for_document(qa_corpus: Dataset, temp_db_path):
     assert all(chunk.document_id == document_id for chunk in chunks)
     assert all(chunk.id is not None for chunk in chunks)
 
-    # Verify chunk order metadata
+    # Verify chunk order
     for i, chunk in enumerate(chunks):
-        assert chunk.metadata.get("order") == i
+        assert chunk.order == i
 
     # Verify chunks exist in database
     db_chunks = await chunk_repo.get_by_document_id(document_id)
@@ -167,9 +167,7 @@ async def test_adjacent_chunks(temp_db_path):
 
     created_chunks = []
     for content, order in chunks_data:
-        chunk = Chunk(
+        chunk = Chunk(document_id=created_document.id, content=content, order=order)
-            document_id=created_document.id, content=content, metadata={"order": order}
-        )
         created_chunk = await chunk_repo.create(chunk)
         created_chunks.append(created_chunk)
 
@@ -186,7 +184,7 @@ async def test_adjacent_chunks(temp_db_path):
     assert middle_chunk.id not in [chunk.id for chunk in adjacent_chunks]
 
     # Should include chunks with order 1 and 3
-    orders = [chunk.metadata.get("order") for chunk in adjacent_chunks]
+    orders = [chunk.order for chunk in adjacent_chunks]
     assert 1 in orders
     assert 3 in orders
 

tests/test_client.py

@@ -460,13 +460,22 @@ async def test_client_create_document_with_custom_chunks(temp_db_path):
     async with HaikuRAG(temp_db_path) as client:
         # Create some custom chunks with and without embeddings
         chunks = [
-            Chunk(content="This is the first chunk", metadata={"custom": "metadata1"}),
+            Chunk(
+                content="This is the first chunk",
+                metadata={"custom": "metadata1"},
+                order=0,
+            ),
             Chunk(
                 content="This is the second chunk",
                 metadata={"custom": "metadata2"},
                 embedding=[0.1] * 1024,
+                order=1,
             ),  # With embedding
-            Chunk(content="This is the third chunk", metadata={"custom": "metadata3"}),
+            Chunk(
+                content="This is the third chunk",
+                metadata={"custom": "metadata3"},
+                order=2,
+            ),
         ]
 
         # Create document with custom chunks
@@ -485,9 +494,7 @@ async def test_client_create_document_with_custom_chunks(temp_db_path):
         for i, chunk in enumerate(doc_chunks):
             assert chunk.document_id == document.id
             assert chunk.content == chunks[i].content
-            assert (
+            assert chunk.order == i  # Order should be set from list position
-                chunk.metadata["order"] == i
-            )  # Order should be set from list position
             assert (
                 chunk.metadata["custom"] == f"metadata{i + 1}"
             )  # Original metadata preserved
@@ -526,15 +533,17 @@ async def test_client_ask_with_cite(temp_db_path):
 @pytest.mark.asyncio
 async def test_client_expand_context(temp_db_path):
     """Test expanding search results with adjacent chunks."""
-    async with HaikuRAG(temp_db_path) as client:
+    # Mock Config to have CONTEXT_CHUNK_RADIUS = 2
-        # Create chunks manually
+    with patch("haiku.rag.client.Config.CONTEXT_CHUNK_RADIUS", 2):
-        manual_chunks = [
+        async with HaikuRAG(temp_db_path) as client:
-            Chunk(content="Chunk 0 content", metadata={"order": 0}),
+            # Create chunks manually
-            Chunk(content="Chunk 1 content", metadata={"order": 1}),
+            manual_chunks = [
-            Chunk(content="Chunk 2 content", metadata={"order": 2}),
+                Chunk(content="Chunk 0 content", order=0),
-            Chunk(content="Chunk 3 content", metadata={"order": 3}),
+                Chunk(content="Chunk 1 content", order=1),
-            Chunk(content="Chunk 4 content", metadata={"order": 4}),
+                Chunk(content="Chunk 2 content", order=2),
-        ]
+                Chunk(content="Chunk 3 content", order=3),
+                Chunk(content="Chunk 4 content", order=4),
+            ]
 
         doc = await client.create_document(
             content="Full document content",
@@ -548,7 +557,7 @@ async def test_client_expand_context(temp_db_path):
         assert len(chunks) == 5
 
         # Find the middle chunk (order=2)
-        middle_chunk = next(c for c in chunks if c.metadata.get("order") == 2)
+        middle_chunk = next(c for c in chunks if c.order == 2)
         search_results = [(middle_chunk, 0.8)]
 
         # Test expand_context with radius=2
@@ -589,25 +598,26 @@ async def test_client_expand_context_radius_zero(temp_db_path):
 @pytest.mark.asyncio
 async def test_client_expand_context_multiple_chunks(temp_db_path):
     """Test expand_context with multiple search results."""
-    async with HaikuRAG(temp_db_path) as client:
+    with patch("haiku.rag.client.Config.CONTEXT_CHUNK_RADIUS", 1):
-        # Create first document with manual chunks
+        async with HaikuRAG(temp_db_path) as client:
-        doc1_chunks = [
+            # Create first document with manual chunks
-            Chunk(content="Doc1 Part A", metadata={"order": 0}),
+            doc1_chunks = [
-            Chunk(content="Doc1 Part B", metadata={"order": 1}),
+                Chunk(content="Doc1 Part A", order=0),
-            Chunk(content="Doc1 Part C", metadata={"order": 2}),
+                Chunk(content="Doc1 Part B", order=1),
-        ]
+                Chunk(content="Doc1 Part C", order=2),
-        doc1 = await client.create_document(
+            ]
-            content="Doc1 content", uri="doc1.txt", chunks=doc1_chunks
+            doc1 = await client.create_document(
-        )
+                content="Doc1 content", uri="doc1.txt", chunks=doc1_chunks
+            )
 
-        # Create second document with manual chunks
+            # Create second document with manual chunks
-        doc2_chunks = [
+            doc2_chunks = [
-            Chunk(content="Doc2 Section X", metadata={"order": 0}),
+                Chunk(content="Doc2 Section X", order=0),
-            Chunk(content="Doc2 Section Y", metadata={"order": 1}),
+                Chunk(content="Doc2 Section Y", order=1),
-        ]
+            ]
-        doc2 = await client.create_document(
+            doc2 = await client.create_document(
-            content="Doc2 content", uri="doc2.txt", chunks=doc2_chunks
+                content="Doc2 content", uri="doc2.txt", chunks=doc2_chunks
-        )
+            )
 
         assert doc1.id is not None
         assert doc2.id is not None
@@ -615,8 +625,8 @@ async def test_client_expand_context_multiple_chunks(temp_db_path):
         chunks2 = await client.chunk_repository.get_by_document_id(doc2.id)
 
         # Get middle chunk from doc1 (order=1) and first chunk from doc2 (order=0)
-        chunk1 = next(c for c in chunks1 if c.metadata.get("order") == 1)
+        chunk1 = next(c for c in chunks1 if c.order == 1)
-        chunk2 = next(c for c in chunks2 if c.metadata.get("order") == 0)
+        chunk2 = next(c for c in chunks2 if c.order == 0)
 
         search_results = [(chunk1, 0.8), (chunk2, 0.7)]
         expanded_results = await client.expand_context(search_results, radius=1)
@@ -645,11 +655,11 @@ async def test_client_expand_context_merges_overlapping_chunks(temp_db_path):
     async with HaikuRAG(temp_db_path) as client:
         # Create document with 5 chunks
         manual_chunks = [
-            Chunk(content="Chunk 0", metadata={"order": 0}),
+            Chunk(content="Chunk 0", order=0),
-            Chunk(content="Chunk 1", metadata={"order": 1}),
+            Chunk(content="Chunk 1", order=1),
-            Chunk(content="Chunk 2", metadata={"order": 2}),
+            Chunk(content="Chunk 2", order=2),
-            Chunk(content="Chunk 3", metadata={"order": 3}),
+            Chunk(content="Chunk 3", order=3),
-            Chunk(content="Chunk 4", metadata={"order": 4}),
+            Chunk(content="Chunk 4", order=4),
         ]
 
         doc = await client.create_document(
@@ -660,8 +670,8 @@ async def test_client_expand_context_merges_overlapping_chunks(temp_db_path):
         chunks = await client.chunk_repository.get_by_document_id(doc.id)
 
         # Get adjacent chunks (orders 1 and 2) - these will overlap when expanded
-        chunk1 = next(c for c in chunks if c.metadata.get("order") == 1)
+        chunk1 = next(c for c in chunks if c.order == 1)
-        chunk2 = next(c for c in chunks if c.metadata.get("order") == 2)
+        chunk2 = next(c for c in chunks if c.order == 2)
 
         # With radius=1:
         # chunk1 expanded would be [0,1,2]
@@ -692,12 +702,12 @@ async def test_client_expand_context_keeps_separate_non_overlapping(temp_db_path
     async with HaikuRAG(temp_db_path) as client:
         # Create document with chunks far apart
         manual_chunks = [
-            Chunk(content="Chunk 0", metadata={"order": 0}),
+            Chunk(content="Chunk 0", order=0),
-            Chunk(content="Chunk 1", metadata={"order": 1}),
+            Chunk(content="Chunk 1", order=1),
-            Chunk(content="Chunk 2", metadata={"order": 2}),
+            Chunk(content="Chunk 2", order=2),
-            Chunk(content="Chunk 5", metadata={"order": 5}),  # Gap here
+            Chunk(content="Chunk 5", order=5),  # Gap here
-            Chunk(content="Chunk 6", metadata={"order": 6}),
+            Chunk(content="Chunk 6", order=6),
-            Chunk(content="Chunk 7", metadata={"order": 7}),
+            Chunk(content="Chunk 7", order=7),
         ]
 
         doc = await client.create_document(
@@ -709,16 +719,13 @@ async def test_client_expand_context_keeps_separate_non_overlapping(temp_db_path
 
         # Get chunks by index - they will have sequential orders 0,1,2,3,4,5
         # So get chunk with order=0 and chunk with order=5 (far enough apart)
-        chunk0 = next(
+        chunk0 = next(c for c in chunks if c.order == 0)  # Content: "Chunk 0"
-            c for c in chunks if c.metadata.get("order") == 0
-        )  # Content: "Chunk 0"
         chunk5 = next(
-            c for c in chunks if c.metadata.get("order") == 5
+            c for c in chunks if c.order == 5
         )  # Content: "Chunk 7" but now at order 5
 
         # chunk0 expanded: [0,1] with radius=1 (orders 0,1)
         # chunk5 expanded: [4,5] with radius=1 (orders 4,5)
-        # These should remain separate (max_order 1 < min_order 4 - 1)
         search_results = [(chunk0, 0.8), (chunk5, 0.7)]
         expanded_results = await client.expand_context(search_results, radius=1)
 
@@ -736,13 +743,13 @@ async def test_client_expand_context_keeps_separate_non_overlapping(temp_db_path
         assert "Chunk 0" in chunk0_expanded.content
         assert "Chunk 1" in chunk0_expanded.content
         assert (
-            "Chunk 7" not in chunk0_expanded.content
+            "Chunk 5" not in chunk0_expanded.content
         )  # Should not have chunk 7 content
         assert score1 == 0.8
 
         # Second chunk (order=5) expanded should contain orders [4,5]
-        # Content should be "Chunk 6" + "Chunk 7" (but they are now at orders 4 and 5)
+        # Content should be "Chunk 6" (order 4) + "Chunk 7" (order 5)
-        assert "Chunk 6" in chunk5_expanded.content  # Order 4 content
+        assert "Chunk 6" in chunk5_expanded.content
-        assert "Chunk 7" in chunk5_expanded.content  # Order 5 content
+        assert "Chunk 7" in chunk5_expanded.content
         assert "Chunk 0" not in chunk5_expanded.content
         assert score2 == 0.7

tests/test_document.py

@@ -43,10 +43,9 @@ async def test_create_document_with_chunks(qa_corpus: Dataset, temp_db_path):
 
     assert len(chunks) > 0
 
-    # Verify chunk metadata contains order information
+    # Verify chunk order is set correctly
     for i, chunk in enumerate(chunks):
-        assert "order" in chunk.metadata
+        assert chunk.order == i
-        assert chunk.metadata["order"] == i
 
     store.close()