Merge pull request #9 from ggozad/feat/document-meta-in-search

Include document uri and meta in search results

src/haiku/rag/store/models/chunk.py

@@ -3,10 +3,12 @@ from pydantic import BaseModel
 
 class Chunk(BaseModel):
     """
-    Represents a document with an ID, content, and metadata.
+    Represents a chunk with content, metadata, and optional document information.
     """
 
     id: int | None = None
     document_id: int
     content: str
     metadata: dict = {}
+    document_uri: str | None = None
+    document_meta: dict = {}

src/haiku/rag/store/repositories/chunk.py

@@ -240,9 +240,10 @@ class ChunkRepository(BaseRepository[Chunk]):
         # Search for similar chunks using sqlite-vec
         cursor.execute(
             """
-            SELECT c.id, c.document_id, c.content, c.metadata, distance
+            SELECT c.id, c.document_id, c.content, c.metadata, distance, d.uri, d.metadata as document_metadata
             FROM chunk_embeddings
             JOIN chunks c ON c.id = chunk_embeddings.chunk_id
+            JOIN documents d ON c.document_id = d.id
             WHERE embedding MATCH :embedding AND k = :k
             ORDER BY distance
             """,
@@ -257,10 +258,14 @@ class ChunkRepository(BaseRepository[Chunk]):
                     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 {},
                 ),
                 1.0 / (1.0 + distance),
             )
-            for chunk_id, document_id, content, metadata_json, distance in results
+            for chunk_id, document_id, content, metadata_json, distance, document_uri, document_metadata_json in results
         ]
 
     async def search_chunks_fts(
@@ -281,9 +286,10 @@ class ChunkRepository(BaseRepository[Chunk]):
         # Search using FTS5
         cursor.execute(
             """
-            SELECT c.id, c.document_id, c.content, c.metadata, rank
+            SELECT c.id, c.document_id, c.content, c.metadata, rank, d.uri, d.metadata as document_metadata
             FROM chunks_fts
             JOIN chunks c ON c.id = chunks_fts.rowid
+            JOIN documents d ON c.document_id = d.id
             WHERE chunks_fts MATCH :query
             ORDER BY rank
             LIMIT :limit
@@ -300,10 +306,14 @@ class ChunkRepository(BaseRepository[Chunk]):
                     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 {},
                 ),
                 -rank,
             )
-            for chunk_id, document_id, content, metadata_json, rank in results
+            for chunk_id, document_id, content, metadata_json, rank, document_uri, document_metadata_json in results
             # FTS5 rank is negative BM25 score
         ]
 
@@ -368,9 +378,10 @@ class ChunkRepository(BaseRepository[Chunk]):
                 LEFT JOIN vector_search v ON a.id = v.id
                 LEFT JOIN fts_search f ON a.id = f.id
             )
-            SELECT id, document_id, content, metadata, rrf_score
+            SELECT r.id, r.document_id, r.content, r.metadata, r.rrf_score, d.uri, d.metadata as document_metadata
-            FROM rrf_scores
+            FROM rrf_scores r
-            ORDER BY rrf_score DESC
+            JOIN documents d ON r.document_id = d.id
+            ORDER BY r.rrf_score DESC
             LIMIT :limit
             """,
             {
@@ -390,10 +401,14 @@ class ChunkRepository(BaseRepository[Chunk]):
                     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 {},
                 ),
                 rrf_score,
             )
-            for chunk_id, document_id, content, metadata_json, rrf_score in results
+            for chunk_id, document_id, content, metadata_json, rrf_score, document_uri, document_metadata_json in results
         ]
 
     async def get_by_document_id(self, document_id: int) -> list[Chunk]:
@@ -404,9 +419,11 @@ class ChunkRepository(BaseRepository[Chunk]):
         cursor = self.store._connection.cursor()
         cursor.execute(
             """
-            SELECT id, document_id, content, metadata
+            SELECT c.id, c.document_id, c.content, c.metadata, d.uri, d.metadata as document_metadata
-            FROM chunks WHERE document_id = :document_id
+            FROM chunks c
-            ORDER BY JSON_EXTRACT(metadata, '$.order')
+            JOIN documents d ON c.document_id = d.id
+            WHERE c.document_id = :document_id
+            ORDER BY JSON_EXTRACT(c.metadata, '$.order')
             """,
             {"document_id": document_id},
         )
@@ -418,6 +435,10 @@ class ChunkRepository(BaseRepository[Chunk]):
                 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 in rows
+            for chunk_id, document_id, content, metadata_json, document_uri, document_metadata_json in rows
         ]

tests/test_search.py

@@ -56,3 +56,33 @@ async def test_search_qa_corpus(qa_corpus: Dataset):
         assert target_document.id in target_document_ids
 
     store.close()
+
+
+@pytest.mark.asyncio
+async def test_chunks_include_document_info():
+    """Test that search results include document URI and metadata."""
+    store = Store(":memory:")
+    doc_repo = DocumentRepository(store)
+    chunk_repo = ChunkRepository(store)
+
+    # Create a document with URI and metadata
+    document = Document(
+        content="This is a test document with some content for searching.",
+        uri="https://example.com/test.html",
+        metadata={"title": "Test Document", "author": "Test Author"},
+    )
+
+    created_document = await doc_repo.create(document)
+
+    # Search for chunks
+    results = await chunk_repo.search_chunks_hybrid("test document", limit=1)
+
+    assert len(results) > 0
+    chunk, score = results[0]
+
+    # Verify the chunk includes document information
+    assert chunk.document_uri == "https://example.com/test.html"
+    assert chunk.document_meta == {"title": "Test Document", "author": "Test Author"}
+    assert chunk.document_id == created_document.id
+
+    store.close()