Make radius a parameter we can pass to expand_context()

docs/python.md

@@ -138,9 +138,12 @@ Expand search results with adjacent chunks for more complete context:
 # Get initial search results
 search_results = await client.search("machine learning", limit=3)
 
-# Expand with adjacent chunks based on CONTEXT_CHUNK_RADIUS setting
+# Expand with adjacent chunks using config setting
 expanded_results = await client.expand_context(search_results)
 
+# Or specify a custom radius
+expanded_results = await client.expand_context(search_results, radius=2)
+
 # 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

src/haiku/rag/client.py

@@ -349,17 +349,21 @@ class HaikuRAG:
         return reranked_results
 
     async def expand_context(
-        self, search_results: list[tuple[Chunk, float]]
+        self,
+        search_results: list[tuple[Chunk, float]],
+        radius: int = Config.CONTEXT_CHUNK_RADIUS,
     ) -> list[tuple[Chunk, float]]:
         """Expand search results with adjacent chunks, merging overlapping chunks.
 
         Args:
             search_results: List of (chunk, score) tuples from search.
+            radius: Number of adjacent chunks to include before/after each chunk.
+                   Defaults to CONTEXT_CHUNK_RADIUS config setting.
 
         Returns:
             List of (chunk, score) tuples with expanded and merged context chunks.
         """
-        if Config.CONTEXT_CHUNK_RADIUS == 0:
+        if radius == 0:
             return search_results
 
         # Group chunks by document_id to handle merging within documents
@@ -377,7 +381,7 @@ class HaikuRAG:
             expanded_ranges = []
             for chunk, score in doc_chunks:
                 adjacent_chunks = await self.chunk_repository.get_adjacent_chunks(
-                    chunk, Config.CONTEXT_CHUNK_RADIUS
+                    chunk, radius
                 )
 
                 all_chunks = adjacent_chunks + [chunk]

tests/test_client.py

@@ -644,7 +644,6 @@ async def test_client_expand_context_multiple_chunks():
 @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 = [
@@ -671,7 +670,7 @@ async def test_client_expand_context_merges_overlapping_chunks():
         # 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)
+        expanded_results = await client.expand_context(search_results, radius=1)
 
         # Should have only 1 merged result instead of 2 overlapping ones
         assert len(expanded_results) == 1
@@ -692,7 +691,6 @@ async def test_client_expand_context_merges_overlapping_chunks():
 @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 = [
@@ -724,7 +722,7 @@ async def test_client_expand_context_keeps_separate_non_overlapping():
         # 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)
+        expanded_results = await client.expand_context(search_results, radius=1)
 
         # Should have 2 separate results
         assert len(expanded_results) == 2