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

Simplify repos

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-06-19 08:55:34 +02:00
parent 8a5c505b86
commit 9f5ba912d7
No known key found for this signature in database
2 changed files with 60 additions and 85 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -178,20 +178,15 @@ class ChunkRepository(BaseRepository[Chunk]):
cursor.execute(query, params) cursor.execute(query, params)
rows = cursor.fetchall() rows = cursor.fetchall()
chunks = [] return [
for row in rows: Chunk(
chunk_id, document_id, content, metadata_json = row id=chunk_id,
metadata = json.loads(metadata_json) if metadata_json else {} document_id=document_id,
chunks.append( content=content,
Chunk( metadata=json.loads(metadata_json) if metadata_json else {},
id=chunk_id,
document_id=document_id,
content=content,
metadata=metadata,
)
) )
for chunk_id, document_id, content, metadata_json in rows
return chunks ]
async def create_chunks_for_document( async def create_chunks_for_document(
self, document_id: int, content: str, commit: bool = True self, document_id: int, content: str, commit: bool = True
@ -255,22 +250,18 @@ class ChunkRepository(BaseRepository[Chunk]):
) )
results = cursor.fetchall() results = cursor.fetchall()
chunks = [] return [
(
for row in results: Chunk(
chunk_id, document_id, content, metadata_json, distance = row id=chunk_id,
metadata = json.loads(metadata_json) if metadata_json else {} document_id=document_id,
chunk = Chunk( content=content,
id=chunk_id, metadata=json.loads(metadata_json) if metadata_json else {},
document_id=document_id, ),
content=content, 1.0 / (1.0 + distance),
metadata=metadata,
) )
for chunk_id, document_id, content, metadata_json, distance in results
similarity_score = 1.0 / (1.0 + distance) ]
chunks.append((chunk, similarity_score))
return chunks
async def search_chunks_fts( async def search_chunks_fts(
self, query: str, limit: int = 5 self, query: str, limit: int = 5
@ -301,23 +292,20 @@ class ChunkRepository(BaseRepository[Chunk]):
) )
results = cursor.fetchall() results = cursor.fetchall()
chunks = []
for row in results: return [
chunk_id, document_id, content, metadata_json, rank = row (
metadata = json.loads(metadata_json) if metadata_json else {} Chunk(
chunk = Chunk( id=chunk_id,
id=chunk_id, document_id=document_id,
document_id=document_id, content=content,
content=content, metadata=json.loads(metadata_json) if metadata_json else {},
metadata=metadata, ),
-rank,
) )
# Convert rank to a score - FTS5 rank is negative BM25 score for chunk_id, document_id, content, metadata_json, rank in results
# More negative = better match, so we negate it to get positive scores # FTS5 rank is negative BM25 score
fts_score = -rank ]
chunks.append((chunk, fts_score))
return chunks
async def search_chunks_hybrid( async def search_chunks_hybrid(
self, query: str, limit: int = 5, k: int = 60 self, query: str, limit: int = 5, k: int = 60
@ -396,20 +384,18 @@ class ChunkRepository(BaseRepository[Chunk]):
) )
results = cursor.fetchall() results = cursor.fetchall()
chunks = [] return [
(
for row in results: Chunk(
chunk_id, document_id, content, metadata_json, rrf_score = row id=chunk_id,
metadata = json.loads(metadata_json) if metadata_json else {} document_id=document_id,
chunk = Chunk( content=content,
id=chunk_id, metadata=json.loads(metadata_json) if metadata_json else {},
document_id=document_id, ),
content=content, rrf_score,
metadata=metadata,
) )
chunks.append((chunk, rrf_score)) for chunk_id, document_id, content, metadata_json, rrf_score in results
]
return chunks
async def get_by_document_id(self, document_id: int) -> list[Chunk]: async def get_by_document_id(self, document_id: int) -> list[Chunk]:
"""Get all chunks for a specific document.""" """Get all chunks for a specific document."""
@ -427,18 +413,12 @@ class ChunkRepository(BaseRepository[Chunk]):
) )
rows = cursor.fetchall() rows = cursor.fetchall()
chunks = [] return [
Chunk(
for row in rows: id=chunk_id,
chunk_id, document_id, content, metadata_json = row document_id=document_id,
metadata = json.loads(metadata_json) if metadata_json else {} content=content,
chunks.append( metadata=json.loads(metadata_json) if metadata_json else {},
Chunk(
id=chunk_id,
document_id=document_id,
content=content,
metadata=metadata,
)
) )
for chunk_id, document_id, content, metadata_json in rows
return chunks ]

25
src/haiku/rag/store/repositories/document.py
View file

@ -197,19 +197,14 @@ class DocumentRepository(BaseRepository[Document]):
cursor.execute(query, params) cursor.execute(query, params)
rows = cursor.fetchall() rows = cursor.fetchall()
documents = [] return [
for row in rows: Document(
document_id, content, uri, metadata_json, created_at, updated_at = row id=document_id,
metadata = json.loads(metadata_json) if metadata_json else {} content=content,
documents.append( uri=uri,
Document( metadata=json.loads(metadata_json) if metadata_json else {},
id=document_id, created_at=created_at,
content=content, updated_at=updated_at,
uri=uri,
metadata=metadata,
created_at=created_at,
updated_at=updated_at,
)
) )
for document_id, content, uri, metadata_json, created_at, updated_at in rows
return documents ]