From e299aed788d290a3d0c06ac98fcbc6e9ef957f17 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 5 Dec 2025 11:58:50 +0200 Subject: [PATCH] contextualize() and embed_chunks() utilities in embeddings --- CHANGELOG.md | 3 + .../haiku/rag/embeddings/__init__.py | 68 ++++++++++ tests/test_embedder.py | 127 ++++++++++++++++++ 3 files changed, 198 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1c49456b..282d25c3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -45,6 +45,9 @@ - Supports `file://` URIs - **New `chunk()` Method**: Chunk a DoclingDocument into Chunk objects - `client.chunk(docling_doc)` - returns `list[Chunk]` without embeddings +- **New `contextualize()` and `embed_chunks()` Utilities**: Standalone embedding utilities in `haiku.rag.embeddings` + - `contextualize(chunks)` - prepend section headings to chunk content for better semantic search + - `embed_chunks(chunks)` - generate embeddings for chunks, returns new Chunk objects with embeddings set ### Changed diff --git a/haiku_rag_slim/haiku/rag/embeddings/__init__.py b/haiku_rag_slim/haiku/rag/embeddings/__init__.py index 4a62e4e5..510cf01b 100644 --- a/haiku_rag_slim/haiku/rag/embeddings/__init__.py +++ b/haiku_rag_slim/haiku/rag/embeddings/__init__.py @@ -1,7 +1,75 @@ +from typing import TYPE_CHECKING + from haiku.rag.config import AppConfig, Config from haiku.rag.embeddings.base import EmbedderBase from haiku.rag.embeddings.ollama import Embedder as OllamaEmbedder +if TYPE_CHECKING: + from haiku.rag.store.models.chunk import Chunk + + +def contextualize(chunks: list["Chunk"]) -> list[str]: + """Prepare chunk content for embedding by adding context. + + Prepends section headings to chunk content for better semantic search. + The embeddings will capture section context while stored content stays raw. + + Args: + chunks: List of chunks to contextualize. + + Returns: + List of contextualized text strings for embedding. + """ + texts = [] + for chunk in chunks: + meta = chunk.get_chunk_metadata() + if meta.headings: + text = "\n".join(meta.headings) + "\n" + chunk.content + else: + text = chunk.content + texts.append(text) + return texts + + +async def embed_chunks( + chunks: list["Chunk"], config: AppConfig = Config +) -> list["Chunk"]: + """Generate embeddings for chunks. + + Contextualizes chunks (prepends headings) before embedding for better + semantic search. Returns new Chunk objects with embeddings set. + + Args: + chunks: List of chunks to embed. + config: Configuration for embedder selection. + + Returns: + New list of Chunk objects with embedding field populated. + """ + if not chunks: + return [] + + from haiku.rag.store.models.chunk import Chunk + + embedder = get_embedder(config) + texts = contextualize(chunks) + embeddings = await embedder.embed(texts) + + return [ + Chunk( + id=chunk.id, + document_id=chunk.document_id, + content=chunk.content, + metadata=chunk.metadata, + order=chunk.order, + document_uri=chunk.document_uri, + document_title=chunk.document_title, + document_meta=chunk.document_meta, + embedding=embedding, + ) + for chunk, embedding in zip(chunks, embeddings) + ] + def get_embedder(config: AppConfig = Config) -> EmbedderBase: """ diff --git a/tests/test_embedder.py b/tests/test_embedder.py index 4ef9e7bc..4d5d0bfa 100644 --- a/tests/test_embedder.py +++ b/tests/test_embedder.py @@ -4,9 +4,11 @@ import numpy as np import pytest from haiku.rag.config import Config +from haiku.rag.embeddings import contextualize, embed_chunks from haiku.rag.embeddings.ollama import Embedder as OllamaEmbedder from haiku.rag.embeddings.openai import Embedder as OpenAIEmbedder from haiku.rag.embeddings.vllm import Embedder as VLLMEmbedder +from haiku.rag.store.models.chunk import Chunk OPENAI_AVAILABLE = bool(os.getenv("OPENAI_API_KEY")) VOYAGEAI_AVAILABLE = bool(os.getenv("VOYAGE_API_KEY")) @@ -171,3 +173,128 @@ async def test_vllm_embedder(): sims = similarities(embeddings, test_embedding) assert max(sims) == sims[1] + + +def test_contextualize_with_headings(): + """Test that contextualize prepends headings to chunk content.""" + chunks = [ + Chunk( + content="This is the content.", + metadata={"headings": ["Chapter 1", "Section 1.1"]}, + ), + Chunk( + content="More content here.", + metadata={"headings": ["Chapter 2"]}, + ), + ] + + texts = contextualize(chunks) + + assert len(texts) == 2 + assert texts[0] == "Chapter 1\nSection 1.1\nThis is the content." + assert texts[1] == "Chapter 2\nMore content here." + + +def test_contextualize_without_headings(): + """Test that contextualize returns raw content when no headings.""" + chunks = [ + Chunk(content="Plain content."), + Chunk(content="Another chunk.", metadata={}), + Chunk(content="With empty headings.", metadata={"headings": None}), + ] + + texts = contextualize(chunks) + + assert len(texts) == 3 + assert texts[0] == "Plain content." + assert texts[1] == "Another chunk." + assert texts[2] == "With empty headings." + + +def test_contextualize_empty_list(): + """Test that contextualize handles empty list.""" + texts = contextualize([]) + assert texts == [] + + +@pytest.mark.asyncio +async def test_embed_chunks_basic(): + """Test that embed_chunks generates embeddings for chunks.""" + chunks = [ + Chunk( + id="chunk1", + document_id="doc1", + content="I enjoy eating great food.", + metadata={"headings": ["Food"]}, + order=0, + ), + Chunk( + id="chunk2", + document_id="doc1", + content="Python is my favorite programming language.", + metadata={"headings": ["Programming"]}, + order=1, + ), + ] + + embedded_chunks = await embed_chunks(chunks) + + assert len(embedded_chunks) == 2 + # Check that all original fields are preserved + assert embedded_chunks[0].id == "chunk1" + assert embedded_chunks[0].document_id == "doc1" + assert embedded_chunks[0].content == "I enjoy eating great food." + assert embedded_chunks[0].metadata == {"headings": ["Food"]} + assert embedded_chunks[0].order == 0 + # Check that embeddings are generated + assert embedded_chunks[0].embedding is not None + assert len(embedded_chunks[0].embedding) > 0 + assert embedded_chunks[1].embedding is not None + + +@pytest.mark.asyncio +async def test_embed_chunks_returns_new_objects(): + """Test that embed_chunks returns new Chunk objects (immutable pattern).""" + original = Chunk(id="orig", content="Test content.") + embedded = await embed_chunks([original]) + + # Original should be unchanged + assert original.embedding is None + # New chunk should have embedding + assert embedded[0].embedding is not None + # They should be different objects + assert embedded[0] is not original + + +@pytest.mark.asyncio +async def test_embed_chunks_empty_list(): + """Test that embed_chunks handles empty list.""" + result = await embed_chunks([]) + assert result == [] + + +@pytest.mark.asyncio +async def test_embed_chunks_preserves_all_fields(): + """Test that embed_chunks preserves all chunk fields.""" + chunk = Chunk( + id="test-id", + document_id="doc-id", + content="Test content", + metadata={"key": "value", "headings": ["Heading"]}, + order=5, + document_uri="https://example.com/doc", + document_title="Test Document", + document_meta={"author": "Test"}, + ) + + embedded = await embed_chunks([chunk]) + + assert embedded[0].id == "test-id" + assert embedded[0].document_id == "doc-id" + assert embedded[0].content == "Test content" + assert embedded[0].metadata == {"key": "value", "headings": ["Heading"]} + assert embedded[0].order == 5 + assert embedded[0].document_uri == "https://example.com/doc" + assert embedded[0].document_title == "Test Document" + assert embedded[0].document_meta == {"author": "Test"} + assert embedded[0].embedding is not None