From 2f42b1a3ac42f03512eb8b92e8a1176b93309a8f Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Tue, 12 Aug 2025 12:15:23 +0200
Subject: [PATCH] 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"