contextualize() and embed_chunks() utilities in embeddings

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
 

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:
     """

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