expand_context() by using the adjacent chunks

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.
 

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"
 

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