From d1e2470d9468e8711d33a98f0dd5e523d119df0c Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 9 Sep 2025 16:48:21 +0300 Subject: [PATCH] Upgrades are back. Make chunk order a table column rather than a metadata property --- pyproject.toml | 3 +- src/haiku/rag/client.py | 8 +- src/haiku/rag/store/engine.py | 47 +++++-- src/haiku/rag/store/models/chunk.py | 1 + src/haiku/rag/store/repositories/chunk.py | 89 ++++++++------ src/haiku/rag/store/repositories/document.py | 4 +- src/haiku/rag/store/repositories/settings.py | 17 ++- src/haiku/rag/store/upgrades/__init__.py | 47 ++++++- src/haiku/rag/store/upgrades/v0_9_3.py | 86 +++++++++++++ tests/test_chunk.py | 10 +- tests/test_client.py | 121 ++++++++++--------- tests/test_document.py | 5 +- 12 files changed, 309 insertions(+), 129 deletions(-) create mode 100644 src/haiku/rag/store/upgrades/v0_9_3.py diff --git a/pyproject.toml b/pyproject.toml index 584d2ea6..af89e7ff 100644 --- a/pyproject.toml +++ b/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" } diff --git a/src/haiku/rag/client.py b/src/haiku/rag/client.py index 7a360835..20e5b089 100644 --- a/src/haiku/rag/client.py +++ b/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, key=lambda c: c.metadata.get("order", 0) - ), + "all_chunks": sorted(all_chunks, key=lambda c: c.order), } ) @@ -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()) diff --git a/src/haiku/rag/store/engine.py b/src/haiku/rag/store/engine.py index 737d90e7..7485e3b3 100644 --- a/src/haiku/rag/store/engine.py +++ b/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 @@ -133,21 +134,41 @@ class Store: [SettingsRecord(id="settings", settings=json.dumps(settings_data))] ) - # Set current version in settings - current_version = metadata.version("haiku.rag") - self.set_haiku_version(current_version) - - # Check if we need to perform upgrades + # Run pending upgrades based on stored version and package version try: - existing_settings = list( - self.settings_table.search().limit(1).to_pydantic(SettingsRecord) + from haiku.rag.store.upgrades import run_pending_upgrades + + 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.""" diff --git a/src/haiku/rag/store/models/chunk.py b/src/haiku/rag/store/models/chunk.py index 106e0cef..383f178e 100644 --- a/src/haiku/rag/store/models/chunk.py +++ b/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 diff --git a/src/haiku/rag/store/repositories/chunk.py b/src/haiku/rag/store/repositories/chunk.py index 842c06f4..1526b485 100644 --- a/src/haiku/rag/store/repositories/chunk.py +++ b/src/haiku/rag/store/repositories/chunk.py @@ -59,11 +59,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 +95,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 +109,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 +151,19 @@ class ChunkRepository: results = list(query.to_pydantic(self.store.ChunkRecord)) - return [ - Chunk( - id=chunk.id, - document_id=chunk.document_id, - content=chunk.content, - metadata=json.loads(chunk.metadata) if chunk.metadata else {}, + chunks: list[Chunk] = [] + for rec in results: + md = json.loads(rec.metadata) + chunks.append( + Chunk( + 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 +206,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 +216,8 @@ class ChunkRepository: id=chunk_id, document_id=document_id, content=chunk_text, - metadata={"order": order}, + metadata={}, + order=order, ) created_chunks.append(chunk) @@ -298,37 +315,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 = [ - Chunk( - id=chunk.id, - document_id=chunk.document_id, - content=chunk.content, - metadata=json.loads(chunk.metadata) if chunk.metadata else {}, - document_uri=doc_uri, - document_meta=json.loads(doc_meta) if doc_meta else {}, + chunks: list[Chunk] = [] + for rec in results: + md = json.loads(rec.metadata) + chunks.append( + Chunk( + id=rec.id, + document_id=rec.document_id, + content=rec.content, + 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") - if chunk_order is None: - return [] + chunk_order = chunk.order - # 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 = [] + adjacent_chunks: list[Chunk] = [] 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 +396,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) - if chunk_record.metadata - else {}, + metadata=md, + order=chunk_record.order, 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 diff --git a/src/haiku/rag/store/repositories/document.py b/src/haiku/rag/store/repositories/document.py index 02d0be6c..bf8160e8 100644 --- a/src/haiku/rag/store/repositories/document.py +++ b/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 diff --git a/src/haiku/rag/store/repositories/settings.py b/src/haiku/rag/store/repositories/settings.py index af752fda..49b0649b 100644 --- a/src/haiku/rag/store/repositories/settings.py +++ b/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 - existing_payload = ( - json.loads(existing[0].settings) if existing[0].settings else {} - ) - if existing_payload != current_config: + # Preserve existing version if present to avoid interfering with upgrade flow + try: + existing_settings = ( + json.loads(existing[0].settings) if existing[0].settings else {} + ) + 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)}, diff --git a/src/haiku/rag/store/upgrades/__init__.py b/src/haiku/rag/store/upgrades/__init__.py index 25115643..a39283ea 100644 --- a/src/haiku/rag/store/upgrades/__init__.py +++ b/src/haiku/rag/store/upgrades/__init__.py @@ -1 +1,46 @@ -upgrades = [] +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass + +from packaging.version import Version, parse + +from haiku.rag.store.engine import Store + + +@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 + + # 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) + + +# Import concrete upgrade modules (module names cannot start with a digit) +from .v0_9_3 import upgrade_order as upgrade_0_9_3_order # noqa: E402 + +upgrades.append(upgrade_0_9_3_order) diff --git a/src/haiku/rag/store/upgrades/v0_9_3.py b/src/haiku/rag/store/upgrades/v0_9_3.py new file mode 100644 index 00000000..241ac54b --- /dev/null +++ b/src/haiku/rag/store/upgrades/v0_9_3.py @@ -0,0 +1,86 @@ +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 + + +def _apply_chunk_order(store: Store) -> None: + """Add integer 'order' column to chunks and backfill from metadata.""" + vector_dim = store.embedder._vector_dim + + # ============== Chunks: add 'order' column and backfill ============== + 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) + # Recreate FTS index on content + 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", +) diff --git a/tests/test_chunk.py b/tests/test_chunk.py index 6f4e7bbc..0a44d5ca 100644 --- a/tests/test_chunk.py +++ b/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( - document_id=created_document.id, content=content, metadata={"order": order} - ) + chunk = Chunk(document_id=created_document.id, content=content, 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 diff --git a/tests/test_client.py b/tests/test_client.py index b9e03cf4..60038067 100644 --- a/tests/test_client.py +++ b/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 ( - chunk.metadata["order"] == i - ) # Order should be set from list position + assert chunk.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: - # Create chunks manually - manual_chunks = [ - Chunk(content="Chunk 0 content", metadata={"order": 0}), - Chunk(content="Chunk 1 content", metadata={"order": 1}), - Chunk(content="Chunk 2 content", metadata={"order": 2}), - Chunk(content="Chunk 3 content", metadata={"order": 3}), - Chunk(content="Chunk 4 content", metadata={"order": 4}), - ] + # Mock Config to have CONTEXT_CHUNK_RADIUS = 2 + with patch("haiku.rag.client.Config.CONTEXT_CHUNK_RADIUS", 2): + async with HaikuRAG(temp_db_path) as client: + # Create chunks manually + manual_chunks = [ + Chunk(content="Chunk 0 content", order=0), + Chunk(content="Chunk 1 content", order=1), + 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: - # Create first document with manual chunks - doc1_chunks = [ - Chunk(content="Doc1 Part A", metadata={"order": 0}), - Chunk(content="Doc1 Part B", metadata={"order": 1}), - Chunk(content="Doc1 Part C", metadata={"order": 2}), - ] - doc1 = await client.create_document( - content="Doc1 content", uri="doc1.txt", chunks=doc1_chunks - ) + with patch("haiku.rag.client.Config.CONTEXT_CHUNK_RADIUS", 1): + async with HaikuRAG(temp_db_path) as client: + # Create first document with manual chunks + doc1_chunks = [ + Chunk(content="Doc1 Part A", order=0), + 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 + ) - # Create second document with manual chunks - doc2_chunks = [ - Chunk(content="Doc2 Section X", metadata={"order": 0}), - Chunk(content="Doc2 Section Y", metadata={"order": 1}), - ] - doc2 = await client.create_document( - content="Doc2 content", uri="doc2.txt", chunks=doc2_chunks - ) + # Create second document with manual chunks + doc2_chunks = [ + Chunk(content="Doc2 Section X", order=0), + Chunk(content="Doc2 Section Y", order=1), + ] + doc2 = await client.create_document( + 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) - chunk2 = next(c for c in chunks2 if c.metadata.get("order") == 0) + chunk1 = next(c for c in chunks1 if c.order == 1) + 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 1", metadata={"order": 1}), - Chunk(content="Chunk 2", metadata={"order": 2}), - Chunk(content="Chunk 3", metadata={"order": 3}), - Chunk(content="Chunk 4", metadata={"order": 4}), + Chunk(content="Chunk 0", order=0), + Chunk(content="Chunk 1", order=1), + Chunk(content="Chunk 2", order=2), + Chunk(content="Chunk 3", order=3), + 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) - chunk2 = next(c for c in chunks if c.metadata.get("order") == 2) + chunk1 = next(c for c in chunks if c.order == 1) + 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 1", metadata={"order": 1}), - Chunk(content="Chunk 2", metadata={"order": 2}), - Chunk(content="Chunk 5", metadata={"order": 5}), # Gap here - Chunk(content="Chunk 6", metadata={"order": 6}), - Chunk(content="Chunk 7", metadata={"order": 7}), + Chunk(content="Chunk 0", order=0), + Chunk(content="Chunk 1", order=1), + Chunk(content="Chunk 2", order=2), + Chunk(content="Chunk 5", order=5), # Gap here + Chunk(content="Chunk 6", order=6), + 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( - c for c in chunks if c.metadata.get("order") == 0 - ) # Content: "Chunk 0" + chunk0 = next(c for c in chunks if c.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) - assert "Chunk 6" in chunk5_expanded.content # Order 4 content - assert "Chunk 7" in chunk5_expanded.content # Order 5 content + # Content should be "Chunk 6" (order 4) + "Chunk 7" (order 5) + assert "Chunk 6" in chunk5_expanded.content + assert "Chunk 7" in chunk5_expanded.content assert "Chunk 0" not in chunk5_expanded.content assert score2 == 0.7 diff --git a/tests/test_document.py b/tests/test_document.py index 748fa26d..61051505 100644 --- a/tests/test_document.py +++ b/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.metadata["order"] == i + assert chunk.order == i store.close()