From eb13426aae7b92fd8f1a3532be73422cf82a939b Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 12 Aug 2025 14:14:18 +0200 Subject: [PATCH] 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