diff --git a/src/haiku/rag/store/models/chunk.py b/src/haiku/rag/store/models/chunk.py index 5b17bea4..bd26aaed 100644 --- a/src/haiku/rag/store/models/chunk.py +++ b/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 = {} diff --git a/src/haiku/rag/store/repositories/chunk.py b/src/haiku/rag/store/repositories/chunk.py index 1ffec47b..1f33ec68 100644 --- a/src/haiku/rag/store/repositories/chunk.py +++ b/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 - FROM rrf_scores - ORDER BY rrf_score DESC + SELECT r.id, r.document_id, r.content, r.metadata, r.rrf_score, d.uri, d.metadata as document_metadata + FROM rrf_scores r + 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 - FROM chunks WHERE document_id = :document_id - ORDER BY JSON_EXTRACT(metadata, '$.order') + 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 + 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 ] diff --git a/tests/test_search.py b/tests/test_search.py index 340a2525..57ffb74f 100644 --- a/tests/test_search.py +++ b/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()