Merge pull request #37 from ggozad/feat/adjacent_chunks

Allow expanding the search context with adjacent chunks through CONTEXT_CHUNK_RADIUS

docs/configuration.md

@@ -172,4 +172,10 @@ 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.
+# When expanded chunks overlap or are adjacent, they are automatically merged
+# into single chunks with continuous content to eliminate duplication
+CONTEXT_CHUNK_RADIUS=0
 ```

docs/python.md

@@ -130,6 +130,26 @@ 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
+```
+
+**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
 
 Ask questions about your documents:

src/haiku/rag/client.py

@@ -348,6 +348,132 @@ 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, merging overlapping chunks.
+
+        Args:
+            search_results: List of (chunk, score) tuples from search.
+
+        Returns:
+            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 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
+                )
+
+                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]
+                min_order = min(orders)
+                max_order = max(orders)
+
+                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)
+                        ),
+                    }
+                )
+
+            # 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.
 

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"
 

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

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 = []

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

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(

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
+        ]

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

tests/test_client.py

@@ -519,3 +519,234 @@ 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
+
+
+@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

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"