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

Move embedding logic from ChunkRepository to client._ensure_chunks_embedded()

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-12-05 14:19:47 +02:00
parent e97d235c98
commit 2572125804
No known key found for this signature in database
3 changed files with 55 additions and 135 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

37
haiku_rag_slim/haiku/rag/client.py
View file

@ -193,6 +193,37 @@ class HaikuRAG:
return chunks
async def _ensure_chunks_embedded(self, chunks: list[Chunk]) -> list[Chunk]:
"""Ensure all chunks have embeddings, embedding any that don't.
Args:
chunks: List of chunks, some may have embeddings already.
Returns:
List of chunks with all embeddings populated.
"""
from haiku.rag.embeddings import embed_chunks
# Find chunks that need embedding
chunks_to_embed = [c for c in chunks if c.embedding is None]
if not chunks_to_embed:
return chunks
# Embed chunks that don't have embeddings (returns new Chunk objects)
embedded = await embed_chunks(chunks_to_embed, self._config)
# Build result maintaining original order
embedded_map = {(c.content, c.order): c for c in embedded}
result = []
for chunk in chunks:
if chunk.embedding is not None:
result.append(chunk)
else:
result.append(embedded_map[(chunk.content, chunk.order)])
return result
async def _store_document_with_chunks(
self,
document: Document,
@ -211,6 +242,9 @@ class HaikuRAG:
"""
import asyncio
# Ensure all chunks have embeddings before storing
chunks = await self._ensure_chunks_embedded(chunks)
# Snapshot table versions for versioned rollback (if supported)
versions = self.store.current_table_versions()
@ -258,6 +292,9 @@ class HaikuRAG:
assert document.id is not None, "Document ID is required for update"
# Ensure all chunks have embeddings before storing
chunks = await self._ensure_chunks_embedded(chunks)
# Snapshot table versions for versioned rollback
versions = self.store.current_table_versions()

44
haiku_rag_slim/haiku/rag/store/repositories/chunk.py
View file

@ -37,20 +37,18 @@ class ChunkRepository:
logger.debug(f"FTS index creation skipped: {e}")
async def create(self, entity: Chunk | list[Chunk]) -> Chunk | list[Chunk]:
"""Create one or more chunks in the database."""
"""Create one or more chunks in the database.
Chunks must have embeddings set before calling this method.
Use client._ensure_chunks_embedded() to embed chunks if needed.
"""
# Handle single chunk
if isinstance(entity, Chunk):
assert entity.document_id, "Chunk must have a document_id to be created"
assert entity.embedding is not None, "Chunk must have an embedding"
chunk_id = str(uuid4())
# Generate embedding if not provided
if entity.embedding is not None:
embedding = entity.embedding
else:
embedding = await self.embedder.embed(entity.content)
order_val = int(entity.order)
chunk_record = self.store.ChunkRecord(
id=chunk_id,
document_id=entity.document_id,
@ -58,8 +56,8 @@ class ChunkRepository:
metadata=json.dumps(
{k: v for k, v in entity.metadata.items() if k != "order"}
),
order=order_val,
vector=embedding,
order=int(entity.order),
vector=entity.embedding,
)
self.store.chunks_table.add([chunk_record])
@ -72,22 +70,15 @@ class ChunkRepository:
if not chunks:
return []
# Validate all chunks have document_id
# Validate all chunks have document_id and embedding
for chunk in chunks:
assert chunk.document_id, "All chunks must have a document_id to be created"
# Batch generate embeddings for chunks that need them
texts_to_embed = [chunk.content for chunk in chunks if chunk.embedding is None]
embeddings = await self.embedder.embed(texts_to_embed) if texts_to_embed else []
embedding_iter = iter(embeddings)
assert chunk.embedding is not None, "All chunks must have embeddings"
# Prepare all chunk records
chunk_records = []
for chunk in chunks:
chunk_id = str(uuid4())
embedding = (
chunk.embedding if chunk.embedding is not None else next(embedding_iter)
)
assert chunk.document_id is not None
chunk_record = self.store.ChunkRecord(
@ -98,7 +89,7 @@ class ChunkRepository:
{k: v for k, v in chunk.metadata.items() if k != "order"}
),
order=int(chunk.order),
vector=embedding,
vector=chunk.embedding,
)
chunk_records.append(chunk_record)
chunk.id = chunk_id
@ -131,11 +122,12 @@ class ChunkRepository:
)
async def update(self, entity: Chunk) -> Chunk:
"""Update an existing chunk."""
assert entity.id, "Chunk ID is required for update"
"""Update an existing chunk.
embedding = await self.embedder.embed(entity.content)
order_val = int(entity.order)
Chunk must have embedding set before calling this method.
"""
assert entity.id, "Chunk ID is required for update"
assert entity.embedding is not None, "Chunk must have an embedding"
self.store.chunks_table.update(
where=f"id = '{entity.id}'",
@ -145,8 +137,8 @@ class ChunkRepository:
"metadata": json.dumps(
{k: v for k, v in entity.metadata.items() if k != "order"}
),
"order": order_val,
"vector": embedding,
"order": int(entity.order),
"vector": entity.embedding,
},
)
return entity

109
tests/test_chunk.py
View file

@ -3,11 +3,7 @@ from datasets import Dataset
from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config
from haiku.rag.store.engine import Store
from haiku.rag.store.models.chunk import Chunk, ChunkMetadata
from haiku.rag.store.models.document import Document
from haiku.rag.store.repositories.chunk import ChunkRepository
from haiku.rag.store.repositories.document import DocumentRepository
@pytest.mark.asyncio
@ -77,111 +73,6 @@ async def test_chunking_pipeline(qa_corpus: Dataset, temp_db_path):
assert chunk.order == i
@pytest.mark.asyncio
async def test_chunk_repository_crud(temp_db_path):
"""Test basic CRUD operations in ChunkRepository."""
# Create a store
store = Store(temp_db_path, create=True)
chunk_repo = ChunkRepository(store)
doc_repo = DocumentRepository(store)
# First create a document to reference
document = Document(content="Test document content", metadata={})
created_document = await doc_repo.create(document)
document_id = created_document.id
assert document_id is not None, "Document ID should not be None"
# Test create chunk manually
chunk = Chunk(
document_id=document_id,
content="Test chunk content",
metadata={"test": "value"},
)
created_chunk = await chunk_repo.create(chunk)
assert isinstance(created_chunk, Chunk)
assert created_chunk.id is not None
assert created_chunk.content == "Test chunk content"
# Test get by ID
retrieved_chunk = await chunk_repo.get_by_id(created_chunk.id)
assert retrieved_chunk is not None
assert retrieved_chunk.content == "Test chunk content"
assert retrieved_chunk.metadata["test"] == "value"
# Test update
retrieved_chunk.content = "Updated chunk content"
updated_chunk = await chunk_repo.update(retrieved_chunk)
assert updated_chunk.content == "Updated chunk content"
# Test list all
all_chunks = await chunk_repo.list_all()
assert len(all_chunks) >= 1
assert any(chunk.id == created_chunk.id for chunk in all_chunks)
# Test delete
deleted = await chunk_repo.delete(created_chunk.id)
assert deleted is True
# Verify chunk is gone
retrieved_chunk = await chunk_repo.get_by_id(created_chunk.id)
assert retrieved_chunk is None
store.close()
@pytest.mark.asyncio
async def test_adjacent_chunks(temp_db_path):
"""Test the get_adjacent_chunks repository method."""
store = Store(temp_db_path, create=True)
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, 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.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()
def test_chunk_metadata_parsing():
"""Test ChunkMetadata parsing from chunk metadata dict."""
metadata_dict = {