From 2f42b1a3ac42f03512eb8b92e8a1176b93309a8f Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 12 Aug 2025 12:15:23 +0200 Subject: [PATCH 1/4] Get adjacent chunks method --- src/haiku/rag/store/repositories/chunk.py | 46 ++++++++++++++++++++ tests/test_chunk.py | 53 +++++++++++++++++++++++ tests/test_search.py | 2 +- 3 files changed, 100 insertions(+), 1 deletion(-) diff --git a/src/haiku/rag/store/repositories/chunk.py b/src/haiku/rag/store/repositories/chunk.py index bb9e1c2b..49c32aff 100644 --- a/src/haiku/rag/store/repositories/chunk.py +++ b/src/haiku/rag/store/repositories/chunk.py @@ -468,3 +468,49 @@ class ChunkRepository(BaseRepository[Chunk]): ) for chunk_id, document_id, content, metadata_json, document_uri, document_metadata_json in rows ] + + 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.""" + if self.store._connection is None: + raise ValueError("Store connection is not available") + if chunk.document_id is None: + return [] + + cursor = self.store._connection.cursor() + chunk_order = chunk.metadata.get("order") + if chunk_order is None: + return [] + + # Get adjacent chunks within the same document + cursor.execute( + """ + SELECT c.id, c.document_id, c.content, c.metadata, d.uri, d.metadata as document_metadata + FROM chunks c + JOIN documents d ON c.document_id = d.id + WHERE c.document_id = :document_id + AND JSON_EXTRACT(c.metadata, '$.order') BETWEEN :start_order AND :end_order + AND c.id != :chunk_id + ORDER BY JSON_EXTRACT(c.metadata, '$.order') + """, + { + "document_id": chunk.document_id, + "start_order": max(0, chunk_order - num_adjacent), + "end_order": chunk_order + num_adjacent, + "chunk_id": chunk.id, + }, + ) + + rows = cursor.fetchall() + return [ + Chunk( + id=chunk_id, + document_id=document_id, + content=content, + metadata=json.loads(metadata_json) if metadata_json else {}, + document_uri=document_uri, + document_meta=json.loads(document_metadata_json) + if document_metadata_json + else {}, + ) + for chunk_id, document_id, content, metadata_json, document_uri, document_metadata_json in rows + ] diff --git a/tests/test_chunk.py b/tests/test_chunk.py index 902dbdf3..044449de 100644 --- a/tests/test_chunk.py +++ b/tests/test_chunk.py @@ -159,3 +159,56 @@ async def test_chunk_repository_crud(): assert retrieved_chunk is None store.close() + + +@pytest.mark.asyncio +async def test_adjacent_chunks(): + """Test the get_adjacent_chunks repository method.""" + store = Store(":memory:") + doc_repo = DocumentRepository(store) + chunk_repo = ChunkRepository(store) + + # Create a simple document first + document_content = "Test document for chunking" + document = Document(content=document_content) + created_document = await doc_repo.create(document) + + # Manually create multiple chunks with order metadata + chunks_data = [ + ("First chunk content", 0), + ("Second chunk content", 1), + ("Third chunk content", 2), + ("Fourth chunk content", 3), + ("Fifth chunk content", 4), + ] + + created_chunks = [] + for content, order in chunks_data: + chunk = Chunk( + document_id=created_document.id, content=content, metadata={"order": order} + ) + created_chunk = await chunk_repo.create(chunk) + created_chunks.append(created_chunk) + + # Test with the middle chunk (index 2, order 2) + middle_chunk = created_chunks[2] + + # Get adjacent chunks (1 before and after) + adjacent_chunks = await chunk_repo.get_adjacent_chunks(middle_chunk, 1) + + # Should have 2 chunks (one before, one after) + assert len(adjacent_chunks) == 2 + + # Should not include the original chunk + 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] + assert 1 in orders + assert 3 in orders + + # All adjacent chunks should be from the same document + for chunk in adjacent_chunks: + assert chunk.document_id == created_document.id + + store.close() diff --git a/tests/test_search.py b/tests/test_search.py index bcc418bd..c2967730 100644 --- a/tests/test_search.py +++ b/tests/test_search.py @@ -82,7 +82,7 @@ async def test_chunks_include_document_info(): results = await chunk_repo.search_chunks_hybrid("test document", limit=1) assert len(results) > 0 - chunk, score = results[0] + chunk, _ = results[0] # Verify the chunk includes document information assert chunk.document_uri == "https://example.com/test.html" From eb13426aae7b92fd8f1a3532be73422cf82a939b Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 12 Aug 2025 14:14:18 +0200 Subject: [PATCH 2/4] expand_context() by using the adjacent chunks --- src/haiku/rag/client.py | 49 ++++++++++++++++ src/haiku/rag/config.py | 1 + tests/test_client.py | 120 ++++++++++++++++++++++++++++++++++++++++ 3 files changed, 170 insertions(+) diff --git a/src/haiku/rag/client.py b/src/haiku/rag/client.py index c5a6eb96..d128cfca 100644 --- a/src/haiku/rag/client.py +++ b/src/haiku/rag/client.py @@ -348,6 +348,55 @@ class HaikuRAG: # Return reranked results with scores from reranker return reranked_results + async def expand_context( + self, search_results: list[tuple[Chunk, float]] + ) -> list[tuple[Chunk, float]]: + """Expand search results with adjacent chunks based on CONTEXT_CHUNK_RADIUS. + + Args: + search_results: List of (chunk, score) tuples from search. + + Returns: + List of (chunk, score) tuples with expanded context chunks. + """ + if Config.CONTEXT_CHUNK_RADIUS == 0: + return search_results + + results = [] + + for chunk, score in search_results: + adjacent_chunks = await self.chunk_repository.get_adjacent_chunks( + chunk, Config.CONTEXT_CHUNK_RADIUS + ) + + chunk_order = chunk.metadata.get("order", 0) + before_chunks = [ + c for c in adjacent_chunks if c.metadata.get("order", 0) < chunk_order + ] + after_chunks = [ + c for c in adjacent_chunks if c.metadata.get("order", 0) > chunk_order + ] + + combined_content_parts = ( + [c.content for c in before_chunks] + + [chunk.content] + + [c.content for c in after_chunks] + ) + + # Create expanded chunk with combined content + expanded_chunk = Chunk( + id=chunk.id, + document_id=chunk.document_id, + content="".join(combined_content_parts), + metadata=chunk.metadata, + document_uri=chunk.document_uri, + document_meta=chunk.document_meta, + ) + + results.append((expanded_chunk, score)) + + return results + async def ask(self, question: str, cite: bool = False) -> str: """Ask a question using the configured QA agent. diff --git a/src/haiku/rag/config.py b/src/haiku/rag/config.py index e4732f48..78328fe6 100644 --- a/src/haiku/rag/config.py +++ b/src/haiku/rag/config.py @@ -26,6 +26,7 @@ class AppConfig(BaseModel): QA_MODEL: str = "qwen3" CHUNK_SIZE: int = 256 + CONTEXT_CHUNK_RADIUS: int = 0 OLLAMA_BASE_URL: str = "http://localhost:11434" diff --git a/tests/test_client.py b/tests/test_client.py index 7c20c517..e8fcd0e8 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -519,3 +519,123 @@ async def test_client_ask_with_cite(): assert answer == "Test answer with citations [1]" mock_qa_agent.answer.assert_called_once_with("What is Python?") + + +@pytest.mark.asyncio +async def test_client_expand_context(): + """Test expanding search results with adjacent chunks.""" + # Mock Config to have CONTEXT_CHUNK_RADIUS = 2 + with patch("haiku.rag.client.Config.CONTEXT_CHUNK_RADIUS", 2): + async with HaikuRAG(":memory:") 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}), + ] + + doc = await client.create_document( + content="Full document content", + uri="test_doc.txt", + chunks=manual_chunks, + ) + + # Get all chunks for the document + assert doc.id is not None + chunks = await client.chunk_repository.get_by_document_id(doc.id) + assert len(chunks) == 5 + + # Find the middle chunk (order=2) + middle_chunk = next(c for c in chunks if c.metadata.get("order") == 2) + search_results = [(middle_chunk, 0.8)] + + # Test expand_context + expanded_results = await client.expand_context(search_results) + + assert len(expanded_results) == 1 + expanded_chunk, score = expanded_results[0] + + # Check that the expanded chunk has combined content + assert expanded_chunk.id == middle_chunk.id + assert score == 0.8 + assert "Chunk 2 content" in expanded_chunk.content + + # Should include all chunks (radius=2 from chunk 2 = chunks 0,1,2,3,4) + assert "Chunk 0 content" in expanded_chunk.content + assert "Chunk 1 content" in expanded_chunk.content + assert "Chunk 2 content" in expanded_chunk.content + assert "Chunk 3 content" in expanded_chunk.content + assert "Chunk 4 content" in expanded_chunk.content + + +@pytest.mark.asyncio +async def test_client_expand_context_radius_zero(): + """Test expand_context with radius 0 returns original results.""" + with patch("haiku.rag.client.Config.CONTEXT_CHUNK_RADIUS", 0): + async with HaikuRAG(":memory:") as client: + # Create a simple document + doc = await client.create_document(content="Simple test content") + assert doc.id is not None + chunks = await client.chunk_repository.get_by_document_id(doc.id) + + search_results = [(chunks[0], 0.9)] + expanded_results = await client.expand_context(search_results) + + # Should return exactly the same results + assert expanded_results == search_results + + +@pytest.mark.asyncio +async def test_client_expand_context_multiple_chunks(): + """Test expand_context with multiple search results.""" + with patch("haiku.rag.client.Config.CONTEXT_CHUNK_RADIUS", 1): + async with HaikuRAG(":memory:") 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 + ) + + # 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 + ) + + assert doc1.id is not None + assert doc2.id is not None + chunks1 = await client.chunk_repository.get_by_document_id(doc1.id) + 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) + + search_results = [(chunk1, 0.8), (chunk2, 0.7)] + expanded_results = await client.expand_context(search_results) + + assert len(expanded_results) == 2 + + # Check first expanded result (should include chunks 0,1,2 from doc1) + expanded1, score1 = expanded_results[0] + assert expanded1.id == chunk1.id + assert score1 == 0.8 + assert "Doc1 Part A" in expanded1.content + assert "Doc1 Part B" in expanded1.content + assert "Doc1 Part C" in expanded1.content + + # Check second expanded result (should include chunks 0,1 from doc2) + expanded2, score2 = expanded_results[1] + assert expanded2.id == chunk2.id + assert score2 == 0.7 + assert "Doc2 Section X" in expanded2.content + assert "Doc2 Section Y" in expanded2.content From e7c2a105ed99247a768d44b072263f716ff19839 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 12 Aug 2025 14:50:29 +0200 Subject: [PATCH 3/4] Use the expand_context when searching in the qa agent --- docs/configuration.md | 4 ++++ docs/python.md | 18 ++++++++++++++++++ src/haiku/rag/qa/anthropic.py | 4 +--- src/haiku/rag/qa/base.py | 6 ++++++ src/haiku/rag/qa/ollama.py | 4 +--- src/haiku/rag/qa/openai.py | 6 +----- 6 files changed, 31 insertions(+), 11 deletions(-) diff --git a/docs/configuration.md b/docs/configuration.md index c1507e49..1302ade6 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -172,4 +172,8 @@ DEFAULT_DATA_DIR="/path/to/data" ```bash # Chunk size for document processing CHUNK_SIZE=256 + +# Number of adjacent chunks to include before/after retrieved chunks for context +# 0 = no expansion (default), 1 = include 1 chunk before and after, etc. +CONTEXT_CHUNK_RADIUS=0 ``` diff --git a/docs/python.md b/docs/python.md index a498c22a..fb46f0ad 100644 --- a/docs/python.md +++ b/docs/python.md @@ -130,6 +130,24 @@ for chunk, relevance_score in results: print(f"Document metadata: {chunk.document_meta}") ``` +### Expanding Search Context + +Expand search results with adjacent chunks for more complete context: + +```python +# Get initial search results +search_results = await client.search("machine learning", limit=3) + +# Expand with adjacent chunks based on CONTEXT_CHUNK_RADIUS setting +expanded_results = await client.expand_context(search_results) + +# The expanded results contain chunks with combined content from adjacent chunks +for chunk, score in expanded_results: + print(f"Expanded content: {chunk.content}") # Now includes before/after chunks +``` + +This is automatically used by the QA system when `CONTEXT_CHUNK_RADIUS > 0` to provide better answers with more complete context. + ## Question Answering Ask questions about your documents: diff --git a/src/haiku/rag/qa/anthropic.py b/src/haiku/rag/qa/anthropic.py index 6c516aff..c3138e1f 100644 --- a/src/haiku/rag/qa/anthropic.py +++ b/src/haiku/rag/qa/anthropic.py @@ -79,12 +79,10 @@ try: else 3 ) - search_results = await self._client.search( + context = await self._search_and_expand( query, limit=limit ) - context = self._format_search_results(search_results) - tool_results.append( { "type": "tool_result", diff --git a/src/haiku/rag/qa/base.py b/src/haiku/rag/qa/base.py index eea670df..92a53060 100644 --- a/src/haiku/rag/qa/base.py +++ b/src/haiku/rag/qa/base.py @@ -20,6 +20,12 @@ class QuestionAnswerAgentBase: "QABase is an abstract class. Please implement the answer method in a subclass." ) + async def _search_and_expand(self, query: str, limit: int = 3) -> str: + """Search for documents and expand context, then format as JSON""" + search_results = await self._client.search(query, limit=limit) + expanded_results = await self._client.expand_context(search_results) + return self._format_search_results(expanded_results) + def _format_search_results(self, search_results) -> str: """Format search results as JSON list of {content, score, document_uri}""" formatted_results = [] diff --git a/src/haiku/rag/qa/ollama.py b/src/haiku/rag/qa/ollama.py index c993b510..6b363542 100644 --- a/src/haiku/rag/qa/ollama.py +++ b/src/haiku/rag/qa/ollama.py @@ -44,9 +44,7 @@ class QuestionAnswerOllamaAgent(QuestionAnswerAgentBase): query = args.get("query", question) limit = int(args.get("limit", 3)) - search_results = await self._client.search(query, limit=limit) - - context = self._format_search_results(search_results) + context = await self._search_and_expand(query, limit=limit) messages.append( { "role": "tool", diff --git a/src/haiku/rag/qa/openai.py b/src/haiku/rag/qa/openai.py index 27396d7a..0ab45793 100644 --- a/src/haiku/rag/qa/openai.py +++ b/src/haiku/rag/qa/openai.py @@ -77,11 +77,7 @@ try: query = args.get("query", question) limit = int(args.get("limit", 3)) - search_results = await self._client.search( - query, limit=limit - ) - - context = self._format_search_results(search_results) + context = await self._search_and_expand(query, limit=limit) messages.append( ChatCompletionToolMessageParam( From 99f7401817069180a2fb8a4168e281a07f7454f8 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 12 Aug 2025 15:02:07 +0200 Subject: [PATCH 4/4] Merge expanded context if possible --- docs/configuration.md | 2 + docs/python.md | 2 + src/haiku/rag/client.py | 133 +++++++++++++++++++++++++++++++--------- tests/test_client.py | 111 +++++++++++++++++++++++++++++++++ 4 files changed, 220 insertions(+), 28 deletions(-) diff --git a/docs/configuration.md b/docs/configuration.md index 1302ade6..50fdd36d 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -175,5 +175,7 @@ CHUNK_SIZE=256 # Number of adjacent chunks to include before/after retrieved chunks for context # 0 = no expansion (default), 1 = include 1 chunk before and after, etc. +# When expanded chunks overlap or are adjacent, they are automatically merged +# into single chunks with continuous content to eliminate duplication CONTEXT_CHUNK_RADIUS=0 ``` diff --git a/docs/python.md b/docs/python.md index fb46f0ad..9721cfbd 100644 --- a/docs/python.md +++ b/docs/python.md @@ -146,6 +146,8 @@ for chunk, score in expanded_results: print(f"Expanded content: {chunk.content}") # Now includes before/after chunks ``` +**Smart Merging**: When expanded chunks overlap or are adjacent within the same document, they are automatically merged into single chunks with continuous content. This eliminates duplication and provides coherent text blocks. The merged chunk uses the highest relevance score from the original chunks. + This is automatically used by the QA system when `CONTEXT_CHUNK_RADIUS > 0` to provide better answers with more complete context. ## Question Answering diff --git a/src/haiku/rag/client.py b/src/haiku/rag/client.py index d128cfca..46c6ab0d 100644 --- a/src/haiku/rag/client.py +++ b/src/haiku/rag/client.py @@ -351,52 +351,129 @@ class HaikuRAG: async def expand_context( self, search_results: list[tuple[Chunk, float]] ) -> list[tuple[Chunk, float]]: - """Expand search results with adjacent chunks based on CONTEXT_CHUNK_RADIUS. + """Expand search results with adjacent chunks, merging overlapping chunks. Args: search_results: List of (chunk, score) tuples from search. Returns: - List of (chunk, score) tuples with expanded context chunks. + List of (chunk, score) tuples with expanded and merged context chunks. """ if Config.CONTEXT_CHUNK_RADIUS == 0: return search_results + # Group chunks by document_id to handle merging within documents + document_groups = {} + for chunk, score in search_results: + doc_id = chunk.document_id + if doc_id not in document_groups: + document_groups[doc_id] = [] + document_groups[doc_id].append((chunk, score)) + results = [] - for chunk, score in search_results: - adjacent_chunks = await self.chunk_repository.get_adjacent_chunks( - chunk, Config.CONTEXT_CHUNK_RADIUS - ) + for doc_id, doc_chunks in document_groups.items(): + # Get all expanded ranges for this document + expanded_ranges = [] + for chunk, score in doc_chunks: + adjacent_chunks = await self.chunk_repository.get_adjacent_chunks( + chunk, Config.CONTEXT_CHUNK_RADIUS + ) - chunk_order = chunk.metadata.get("order", 0) - before_chunks = [ - c for c in adjacent_chunks if c.metadata.get("order", 0) < chunk_order - ] - after_chunks = [ - c for c in adjacent_chunks if c.metadata.get("order", 0) > chunk_order - ] + all_chunks = adjacent_chunks + [chunk] - combined_content_parts = ( - [c.content for c in before_chunks] - + [chunk.content] - + [c.content for c in after_chunks] - ) + # Get the range of orders for this expanded chunk + orders = [c.metadata.get("order", 0) for c in all_chunks] + min_order = min(orders) + max_order = max(orders) - # Create expanded chunk with combined content - expanded_chunk = Chunk( - id=chunk.id, - document_id=chunk.document_id, - content="".join(combined_content_parts), - metadata=chunk.metadata, - document_uri=chunk.document_uri, - document_meta=chunk.document_meta, - ) + expanded_ranges.append( + { + "original_chunk": chunk, + "score": score, + "min_order": min_order, + "max_order": max_order, + "all_chunks": sorted( + all_chunks, key=lambda c: c.metadata.get("order", 0) + ), + } + ) - results.append((expanded_chunk, score)) + # Merge overlapping/adjacent ranges + merged_ranges = self._merge_overlapping_ranges(expanded_ranges) + + # Create merged chunks + for merged_range in merged_ranges: + combined_content_parts = [c.content for c in merged_range["all_chunks"]] + + # Use the first original chunk for metadata + original_chunk = merged_range["original_chunks"][0] + + merged_chunk = Chunk( + id=original_chunk.id, + document_id=original_chunk.document_id, + content="".join(combined_content_parts), + metadata=original_chunk.metadata, + document_uri=original_chunk.document_uri, + document_meta=original_chunk.document_meta, + ) + + # Use the highest score from merged chunks + best_score = max(merged_range["scores"]) + results.append((merged_chunk, best_score)) return results + def _merge_overlapping_ranges(self, expanded_ranges): + """Merge overlapping or adjacent expanded ranges.""" + if not expanded_ranges: + return [] + + # Sort by min_order + sorted_ranges = sorted(expanded_ranges, key=lambda x: x["min_order"]) + merged = [] + + current = { + "min_order": sorted_ranges[0]["min_order"], + "max_order": sorted_ranges[0]["max_order"], + "original_chunks": [sorted_ranges[0]["original_chunk"]], + "scores": [sorted_ranges[0]["score"]], + "all_chunks": sorted_ranges[0]["all_chunks"], + } + + for range_info in sorted_ranges[1:]: + # Check if ranges overlap or are adjacent (max_order + 1 >= min_order) + if current["max_order"] >= range_info["min_order"] - 1: + # Merge ranges + current["max_order"] = max( + current["max_order"], range_info["max_order"] + ) + current["original_chunks"].append(range_info["original_chunk"]) + current["scores"].append(range_info["score"]) + + # 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) + all_chunks_dict[order] = chunk + current["all_chunks"] = [ + all_chunks_dict[order] for order in sorted(all_chunks_dict.keys()) + ] + else: + # No overlap, add current to merged and start new + merged.append(current) + current = { + "min_order": range_info["min_order"], + "max_order": range_info["max_order"], + "original_chunks": [range_info["original_chunk"]], + "scores": [range_info["score"]], + "all_chunks": range_info["all_chunks"], + } + + # Add the last range + merged.append(current) + return merged + async def ask(self, question: str, cite: bool = False) -> str: """Ask a question using the configured QA agent. diff --git a/tests/test_client.py b/tests/test_client.py index e8fcd0e8..3a39c2ce 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -639,3 +639,114 @@ async def test_client_expand_context_multiple_chunks(): assert score2 == 0.7 assert "Doc2 Section X" in expanded2.content assert "Doc2 Section Y" in expanded2.content + + +@pytest.mark.asyncio +async def test_client_expand_context_merges_overlapping_chunks(): + """Test that overlapping expanded chunks are merged into one.""" + with patch("haiku.rag.client.Config.CONTEXT_CHUNK_RADIUS", 1): + async with HaikuRAG(":memory:") 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}), + ] + + doc = await client.create_document( + content="Full document content", chunks=manual_chunks + ) + + assert doc.id is not None + 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) + + # With radius=1: + # chunk1 expanded would be [0,1,2] + # chunk2 expanded would be [1,2,3] + # These should merge into one chunk containing [0,1,2,3] + search_results = [(chunk1, 0.8), (chunk2, 0.7)] + expanded_results = await client.expand_context(search_results) + + # Should have only 1 merged result instead of 2 overlapping ones + assert len(expanded_results) == 1 + + merged_chunk, score = expanded_results[0] + + # Should contain all chunks from 0 to 3 + assert "Chunk 0" in merged_chunk.content + assert "Chunk 1" in merged_chunk.content + assert "Chunk 2" in merged_chunk.content + assert "Chunk 3" in merged_chunk.content + assert "Chunk 4" not in merged_chunk.content # Should not include chunk 4 + + # Should use the higher score (0.8) + assert score == 0.8 + + +@pytest.mark.asyncio +async def test_client_expand_context_keeps_separate_non_overlapping(): + """Test that non-overlapping expanded chunks remain separate.""" + with patch("haiku.rag.client.Config.CONTEXT_CHUNK_RADIUS", 1): + async with HaikuRAG(":memory:") 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}), + ] + + doc = await client.create_document( + content="Full document content", chunks=manual_chunks + ) + + assert doc.id is not None + chunks = await client.chunk_repository.get_by_document_id(doc.id) + + # 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" + chunk5 = next( + c for c in chunks if c.metadata.get("order") == 5 + ) # Content: "Chunk 7" + + # 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) + + # Should have 2 separate results + assert len(expanded_results) == 2 + + # Sort by score to ensure predictable order + expanded_results.sort(key=lambda x: x[1], reverse=True) + + chunk0_expanded, score1 = expanded_results[0] + chunk5_expanded, score2 = expanded_results[1] + + # First chunk (order=0) expanded should contain orders [0,1] + # Content should be "Chunk 0" + "Chunk 1" + assert "Chunk 0" in chunk0_expanded.content + assert "Chunk 1" in chunk0_expanded.content + assert ( + "Chunk 7" 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" (orders 4 and 5) + assert "Chunk 6" in chunk5_expanded.content # Order 4 content + assert "Chunk 7" in chunk5_expanded.content # Order 5 content + assert "Chunk 0" not in chunk5_expanded.content + assert score2 == 0.7