chatbot_v3/tests/unit/application/test_embedding.py

"""Batching, concurrency bounding, and failure translation for `embed_chunks`
(ADR-0001, ADR-0017, plan 001 Phase 4).
"""

import asyncio
import uuid
from collections.abc import Sequence
from dataclasses import dataclass, field

import pytest
from anyio import CapacityLimiter

from src.application.ingestion.embedding import embed_chunks
from src.application.ingestion.errors import EmbedderError
from src.application.ingestion.models import Chunk, ContentType, SparseVector
from src.config import IngestionSettings

pytestmark = [pytest.mark.unit, pytest.mark.asyncio]


def _chunk(index: int, content: str = "hello world") -> Chunk:
    return Chunk(
        chunk_id=uuid.uuid4(),
        chunk_index=index,
        order_id=float(index + 1),
        content=content,
        content_type=ContentType.PARAGRAPH,
        token_count=2,
        character_count=len(content),
    )


@dataclass
class _TrackingDenseEmbedder:
    """Records each batch's texts and the peak number of batches in flight
    at once, to prove concurrency is bounded, not serial.
    """

    name: str
    model_version: str = "stub-v1"
    dimensions: int = 3
    batches: list[list[str]] = field(default_factory=list)
    in_flight: int = 0
    peak_in_flight: int = 0

    async def embed_batch(self, texts: Sequence[str]) -> list[list[float]]:
        self.in_flight += 1
        self.peak_in_flight = max(self.peak_in_flight, self.in_flight)
        self.batches.append(list(texts))
        await asyncio.sleep(0)  # yield so overlapping calls can interleave
        self.in_flight -= 1
        return [[0.0] * self.dimensions for _ in texts]


@dataclass
class _FailingDenseEmbedder:
    name: str
    model_version: str = "stub-v1"

    async def embed_batch(self, texts: Sequence[str]) -> list[list[float]]:
        raise RuntimeError("boom")


@dataclass
class _StubSparseEmbedder:
    name: str = "sparse"
    model_version: str = "stub-sparse-v1"
    calls: list[list[str]] = field(default_factory=list)

    def embed_batch(self, texts: Sequence[str], *, query: bool = False) -> list[SparseVector]:
        self.calls.append(list(texts))
        return [SparseVector(indices=[i], values=[1.0]) for i in range(len(texts))]


@dataclass
class _FailingSparseEmbedder:
    name: str = "sparse"
    model_version: str = "stub-sparse-v1"

    def embed_batch(self, texts: Sequence[str], *, query: bool = False) -> list[SparseVector]:
        raise RuntimeError("boom")


@pytest.fixture
def settings() -> IngestionSettings:
    return IngestionSettings(embed_batch_size=2, embed_concurrency=2)


@pytest.fixture
def thread_limiter() -> CapacityLimiter:
    return CapacityLimiter(4)


async def test_embed_chunks_empty_input_returns_empty(
    settings: IngestionSettings, thread_limiter: CapacityLimiter
) -> None:
    result = await embed_chunks(
        [],
        dense_embedders=[_TrackingDenseEmbedder(name="dense_nomic")],
        sparse_embedder=_StubSparseEmbedder(),
        settings=settings,
        thread_limiter=thread_limiter,
    )
    assert result == []


async def test_embed_chunks_batches_before_parallelizing(
    settings: IngestionSettings, thread_limiter: CapacityLimiter
) -> None:
    chunks = [_chunk(i) for i in range(5)]
    embedder = _TrackingDenseEmbedder(name="dense_nomic")

    await embed_chunks(
        chunks,
        dense_embedders=[embedder],
        sparse_embedder=_StubSparseEmbedder(),
        settings=settings,
        thread_limiter=thread_limiter,
    )

    # embed_batch_size=2 over 5 chunks -> 3 batches (2, 2, 1), never one call
    # per chunk.
    assert len(embedder.batches) == 3
    assert [len(batch) for batch in embedder.batches] == [2, 2, 1]


async def test_embed_chunks_bounds_concurrency_by_embed_concurrency(
    thread_limiter: CapacityLimiter,
) -> None:
    settings = IngestionSettings(embed_batch_size=1, embed_concurrency=2)
    chunks = [_chunk(i) for i in range(6)]
    embedder = _TrackingDenseEmbedder(name="dense_nomic")

    await embed_chunks(
        chunks,
        dense_embedders=[embedder],
        sparse_embedder=_StubSparseEmbedder(),
        settings=settings,
        thread_limiter=thread_limiter,
    )

    assert embedder.peak_in_flight <= settings.embed_concurrency
    assert embedder.peak_in_flight > 1  # proves it isn't serial either


async def test_embed_chunks_passes_chunk_text_through_unmodified(
    settings: IngestionSettings, thread_limiter: CapacityLimiter
) -> None:
    """Text shaping (task prefixes, `keep_alive`) is the adapter's job, not
    this module's -- see `src/infrastructure/embedding/openai_compatible.py`.
    Orchestration here must stay provider-agnostic.
    """
    chunks = [_chunk(0, content="salam")]
    nomic = _TrackingDenseEmbedder(name="dense_nomic")
    openai = _TrackingDenseEmbedder(name="dense_openai")

    await embed_chunks(
        chunks,
        dense_embedders=[nomic, openai],
        sparse_embedder=_StubSparseEmbedder(),
        settings=settings,
        thread_limiter=thread_limiter,
    )

    assert nomic.batches[0] == ["salam"]
    assert openai.batches[0] == ["salam"]


async def test_embed_chunks_assembles_dense_and_sparse_per_chunk_in_order(
    settings: IngestionSettings, thread_limiter: CapacityLimiter
) -> None:
    chunks = [_chunk(0, "a"), _chunk(1, "b")]
    nomic = _TrackingDenseEmbedder(name="dense_nomic", dimensions=3)

    result = await embed_chunks(
        chunks,
        dense_embedders=[nomic],
        sparse_embedder=_StubSparseEmbedder(),
        settings=settings,
        thread_limiter=thread_limiter,
    )

    assert [item.chunk.chunk_id for item in result] == [c.chunk_id for c in chunks]
    assert all(len(item.dense["dense_nomic"]) == 3 for item in result)
    assert all("sparse" not in item.dense for item in result)


async def test_embed_chunks_dense_failure_raises_embedder_error(
    settings: IngestionSettings, thread_limiter: CapacityLimiter
) -> None:
    chunks = [_chunk(0)]

    with pytest.raises(EmbedderError):
        await embed_chunks(
            chunks,
            dense_embedders=[_FailingDenseEmbedder(name="dense_nomic")],
            sparse_embedder=_StubSparseEmbedder(),
            settings=settings,
            thread_limiter=thread_limiter,
        )


async def test_embed_chunks_sparse_failure_raises_embedder_error(
    settings: IngestionSettings, thread_limiter: CapacityLimiter
) -> None:
    chunks = [_chunk(0)]

    with pytest.raises(EmbedderError):
        await embed_chunks(
            chunks,
            dense_embedders=[_TrackingDenseEmbedder(name="dense_nomic")],
            sparse_embedder=_FailingSparseEmbedder(),
            settings=settings,
            thread_limiter=thread_limiter,
        )