danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

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

Browse source 
Include document uri and meta in search results
This commit is contained in:
Yiorgis Gozadinos 2025-06-28 09:41:37 +03:00 committed by GitHub
commit 7e540583ce
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 66 additions and 13 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
src/haiku/rag/store/models/chunk.py
View file

@ -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 = {}

45
src/haiku/rag/store/repositories/chunk.py
View file

@ -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
]

30
tests/test_search.py
View file

@ -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()