RAG/backend/tests/unit/test_indexing_service.py

from __future__ import annotations

import hashlib
import math
import uuid
from collections.abc import Callable, Mapping, Sequence
from dataclasses import dataclass, field, replace

import pytest

from app.persistence.job_queue import JobLease, LeaseLostError
from app.persistence.retrieval import ActiveEmbeddingProfile
from app.ports.model_providers import (
    EmbeddingResult,
    ModelProviderError,
    ProviderErrorKind,
    ProviderUsage,
)
from app.services.indexing import (
    ApprovedIndexingPlan,
    AssignmentProgress,
    AssignmentStatus,
    CachedEmbedding,
    DocumentIndexingService,
    EmbeddingCacheLookup,
    EmbeddingWrite,
    IndexingItem,
    IndexingNotReadyError,
    InvalidEmbeddingResponseError,
    InvocationStatus,
    embedding_cache_key,
)

DOCUMENT_VERSION_ID = uuid.UUID("10000000-0000-0000-0000-000000000001")
KNOWLEDGE_BASE_ID = uuid.UUID("20000000-0000-0000-0000-000000000002")
JOB_ID = uuid.UUID("30000000-0000-0000-0000-000000000003")
LEASE_TOKEN = uuid.UUID("40000000-0000-0000-0000-000000000004")
TRACE_ID = uuid.UUID("50000000-0000-0000-0000-000000000005")
LEASE = JobLease(JOB_ID, "embedding-worker-a", LEASE_TOKEN)
PROFILE = ActiveEmbeddingProfile(
    profile_hash="a" * 64,
    model="text-embedding-v4",
    dimension=1024,
)


def _text_hash(text: str) -> str:
    return hashlib.sha256(text.encode()).hexdigest()


def _item(index: int, *, status: AssignmentStatus = "PENDING") -> IndexingItem:
    text = f"第 {index} 个已批准地质文本"
    return IndexingItem(
        chunk_id=uuid.UUID(int=1_000 + index),
        ordinal=index,
        embedding_text=text,
        embedding_text_sha256=_text_hash(text),
        assignment_status=status,
    )


def _plan(count: int, *, profile: ActiveEmbeddingProfile = PROFILE) -> ApprovedIndexingPlan:
    return ApprovedIndexingPlan(
        knowledge_base_id=KNOWLEDGE_BASE_ID,
        document_version_id=DOCUMENT_VERSION_ID,
        review_state="CLOUD_APPROVED",
        outbound_manifest_sha256="b" * 64,
        expected_count=count,
        profile=profile,
        items=tuple(_item(index) for index in range(count)),
    )


def _vector(index: int = 0) -> tuple[float, ...]:
    return (float(index + 1),) + (0.0,) * 1023


@dataclass
class SpyRepository:
    plan: ApprovedIndexingPlan
    events: list[str] = field(default_factory=list)
    ready_chunks: set[uuid.UUID] = field(default_factory=set)
    cache: dict[str, CachedEmbedding] = field(default_factory=dict)
    writes: list[tuple[EmbeddingWrite, ...]] = field(default_factory=list)
    begin_calls: list[dict[str, object]] = field(default_factory=list)
    finish_calls: list[dict[str, object]] = field(default_factory=list)
    write_leases: list[JobLease] = field(default_factory=list)
    in_repository: bool = False
    activation_allowed: bool = True
    report_incomplete: bool = False
    lose_lease_on: str | None = None
    persist_count: int = 0

    def _enter(self, event: str) -> None:
        assert self.in_repository is False
        self.in_repository = True
        self.events.append(event)

    def _leave(self) -> None:
        self.in_repository = False

    def _maybe_lose(self, operation: str) -> None:
        if self.lose_lease_on == operation:
            raise LeaseLostError("lease moved")

    def load_approved_plan(
        self,
        *,
        lease: JobLease,
        document_version_id: uuid.UUID,
    ) -> ApprovedIndexingPlan:
        self._enter("repo.load")
        try:
            assert lease == LEASE
            assert document_version_id == DOCUMENT_VERSION_ID
            items = tuple(
                replace(
                    item,
                    assignment_status=(
                        "READY" if item.chunk_id in self.ready_chunks else item.assignment_status
                    ),
                )
                for item in self.plan.items
            )
            return replace(self.plan, items=items)
        finally:
            self._leave()

    def lookup_cache(
        self,
        *,
        lease: JobLease,
        lookups: Sequence[EmbeddingCacheLookup],
    ) -> Mapping[str, CachedEmbedding]:
        self._enter(f"repo.cache:{len(lookups)}")
        try:
            assert lease == LEASE
            cache_keys = tuple(lookup.cache_key for lookup in lookups)
            return {key: self.cache[key] for key in cache_keys if key in self.cache}
        finally:
            self._leave()

    def begin_model_invocation(
        self,
        *,
        lease: JobLease,
        trace_id: uuid.UUID,
        profile_hash: str,
        model: str,
        item_count: int,
    ) -> uuid.UUID:
        self._enter(f"repo.begin:{item_count}")
        try:
            self.write_leases.append(lease)
            self._maybe_lose("begin")
            call = {
                "lease": lease,
                "trace_id": trace_id,
                "profile_hash": profile_hash,
                "model": model,
                "item_count": item_count,
            }
            self.begin_calls.append(call)
            return uuid.UUID(int=9_000 + len(self.begin_calls))
        finally:
            self._leave()

    def finish_model_invocation(
        self,
        *,
        lease: JobLease,
        invocation_id: uuid.UUID,
        status: InvocationStatus,
        provider_request_id: str | None,
        usage: ProviderUsage,
        elapsed_ms: float,
        error_code: str | None,
    ) -> None:
        self._enter(f"repo.finish:{status}")
        try:
            self.write_leases.append(lease)
            self._maybe_lose("finish")
            self.finish_calls.append(
                {
                    "lease": lease,
                    "invocation_id": invocation_id,
                    "status": status,
                    "provider_request_id": provider_request_id,
                    "usage": usage,
                    "elapsed_ms": elapsed_ms,
                    "error_code": error_code,
                }
            )
        finally:
            self._leave()

    def fenced_persist_batch(
        self,
        *,
        lease: JobLease,
        document_version_id: uuid.UUID,
        profile_hash: str,
        writes: Sequence[EmbeddingWrite],
    ) -> AssignmentProgress:
        self._enter(f"repo.persist:{len(writes)}")
        try:
            self.write_leases.append(lease)
            self._maybe_lose("persist")
            assert document_version_id == DOCUMENT_VERSION_ID
            assert profile_hash == self.plan.profile.profile_hash
            assert 1 <= len(writes) <= 10
            batch = tuple(writes)
            self.writes.append(batch)
            self.persist_count += 1
            for write in batch:
                self.ready_chunks.add(write.chunk_id)
                if write.source == "provider":
                    assert write.embedding is not None
                    self.cache[write.cache_key] = CachedEmbedding(
                        cache_key=write.cache_key,
                        profile_hash=write.profile_hash,
                        embedding_text_sha256=write.embedding_text_sha256,
                        resolved_model=write.resolved_model,
                        dimension=1024,
                    )
            ready_count = len(self.ready_chunks)
            if self.report_incomplete and ready_count:
                ready_count -= 1
            return AssignmentProgress(len(self.plan.items), ready_count)
        finally:
            self._leave()

    def fenced_activate(
        self,
        *,
        lease: JobLease,
        document_version_id: uuid.UUID,
        profile_hash: str,
        expected_count: int,
    ) -> bool:
        self._enter("repo.activate")
        try:
            self.write_leases.append(lease)
            self._maybe_lose("activate")
            assert document_version_id == DOCUMENT_VERSION_ID
            assert profile_hash == self.plan.profile.profile_hash
            return (
                self.activation_allowed
                and expected_count == len(self.plan.items)
                and len(self.ready_chunks) == expected_count
            )
        finally:
            self._leave()


ResultFactory = Callable[[Sequence[str], int], EmbeddingResult]


@dataclass
class SpyEmbeddingProvider:
    repository: SpyRepository
    model: str = "text-embedding-v4"
    result_factory: ResultFactory | None = None
    failures: dict[int, ModelProviderError] = field(default_factory=dict)
    calls: list[tuple[str, ...]] = field(default_factory=list)

    async def embed_documents(self, texts: Sequence[str]) -> EmbeddingResult:
        assert self.repository.in_repository is False
        values = tuple(texts)
        assert 1 <= len(values) <= 10
        self.calls.append(values)
        call_number = len(self.calls)
        self.repository.events.append(f"provider:{len(values)}")
        if call_number in self.failures:
            raise self.failures[call_number]
        if self.result_factory is not None:
            return self.result_factory(values, call_number)
        return EmbeddingResult(
            vectors=tuple(_vector(index) for index in range(len(values))),
            model=self.model,
            request_id=f"embed-request-{call_number}",
            usage=ProviderUsage(input_tokens=len(values), total_tokens=len(values)),
            elapsed_ms=4.0,
        )

    async def embed_query(self, text: str) -> EmbeddingResult:
        return await self.embed_documents((text,))


def _service(
    repository: SpyRepository,
    provider: SpyEmbeddingProvider,
    *,
    synthetic_provider: SpyEmbeddingProvider | None = None,
) -> DocumentIndexingService:
    return DocumentIndexingService(
        repository=repository,
        embedding_provider=provider,
        synthetic_embedding_provider=synthetic_provider,
    )


@pytest.mark.asyncio
async def test_batches_are_ten_and_provider_runs_between_short_repository_calls() -> None:
    repository = SpyRepository(_plan(12))
    provider = SpyEmbeddingProvider(repository)

    result = await _service(repository, provider).index_document_version(
        lease=LEASE,
        document_version_id=DOCUMENT_VERSION_ID,
        trace_id=TRACE_ID,
    )

    assert [len(call) for call in provider.calls] == [10, 2]
    assert [len(batch) for batch in repository.writes] == [10, 2]
    assert repository.events == [
        "repo.load",
        "repo.cache:10",
        "repo.cache:2",
        "repo.begin:10",
        "provider:10",
        "repo.finish:SUCCEEDED",
        "repo.persist:10",
        "repo.begin:2",
        "provider:2",
        "repo.finish:SUCCEEDED",
        "repo.persist:2",
        "repo.activate",
    ]
    assert all(lease == LEASE for lease in repository.write_leases)
    assert result.provider_call_count == 2
    assert result.ready_count == 12
    assert result.activated is True

    provider_writes = [write for batch in repository.writes for write in batch]
    assert [write.batch_index for write in provider_writes[:10]] == list(range(10))
    assert provider_writes[0].embedding == _vector(0)
    assert provider_writes[1].embedding == _vector(1)
    assert all("embedding_text" not in call for call in repository.finish_calls)
    assert all("embedding" not in call for call in repository.finish_calls)


@pytest.mark.asyncio
async def test_cache_key_hits_skip_model_and_only_misses_are_embedded() -> None:
    plan = _plan(3)
    repository = SpyRepository(plan)
    for item in plan.items[:2]:
        key = embedding_cache_key(item.embedding_text_sha256, PROFILE.profile_hash)
        repository.cache[key] = CachedEmbedding(
            cache_key=key,
            profile_hash=PROFILE.profile_hash,
            embedding_text_sha256=item.embedding_text_sha256,
            resolved_model=PROFILE.model,
            dimension=1024,
        )
    provider = SpyEmbeddingProvider(repository)

    result = await _service(repository, provider).index_document_version(
        lease=LEASE,
        document_version_id=DOCUMENT_VERSION_ID,
        trace_id=TRACE_ID,
    )

    expected_key = hashlib.sha256(
        f"{plan.items[0].embedding_text_sha256}{PROFILE.profile_hash}".encode()
    ).hexdigest()
    assert (
        embedding_cache_key(plan.items[0].embedding_text_sha256, PROFILE.profile_hash)
        == expected_key
    )
    assert provider.calls == [(plan.items[2].embedding_text,)]
    assert result.cache_hit_count == 2
    assert result.newly_embedded_count == 1
    assert [write.source for batch in repository.writes for write in batch] == [
        "cache",
        "cache",
        "provider",
    ]


@pytest.mark.asyncio
async def test_partial_batch_resume_never_reembeds_persisted_first_batch() -> None:
    repository = SpyRepository(_plan(12))
    upstream_failure = ModelProviderError(
        operation="embedding.document",
        kind=ProviderErrorKind.UPSTREAM,
        status_code=503,
        retryable=True,
    )
    first_provider = SpyEmbeddingProvider(repository, failures={2: upstream_failure})

    with pytest.raises(ModelProviderError):
        await _service(repository, first_provider).index_document_version(
            lease=LEASE,
            document_version_id=DOCUMENT_VERSION_ID,
            trace_id=TRACE_ID,
        )

    assert len(repository.ready_chunks) == 10
    assert repository.finish_calls[-1]["status"] == "FAILED"
    assert repository.finish_calls[-1]["error_code"] == "EMBEDDING_UPSTREAM"

    repository.events.clear()
    second_provider = SpyEmbeddingProvider(repository)
    result = await _service(repository, second_provider).index_document_version(
        lease=LEASE,
        document_version_id=DOCUMENT_VERSION_ID,
        trace_id=TRACE_ID,
    )

    assert second_provider.calls == [
        tuple(item.embedding_text for item in repository.plan.items[10:])
    ]
    assert result.newly_embedded_count == 2
    assert result.ready_count == 12


def _invalid_result(case: str) -> ResultFactory:
    def factory(texts: Sequence[str], _call_number: int) -> EmbeddingResult:
        vectors = tuple(_vector(index) for index in range(len(texts)))
        model = PROFILE.model
        elapsed_ms = 1.0
        if case == "count":
            vectors = vectors[:-1]
        elif case == "dimension":
            vectors = ((1.0, 0.0),) + vectors[1:]
        elif case == "nonfinite":
            vectors = ((math.nan,) + (0.0,) * 1023,) + vectors[1:]
        elif case == "zero":
            vectors = ((0.0,) * 1024,) + vectors[1:]
        elif case == "model":
            model = "wrong-embedding-model"
        elif case == "elapsed":
            elapsed_ms = -1.0
        return EmbeddingResult(
            vectors=vectors,
            model=model,
            request_id="invalid-response-id",
            usage=ProviderUsage(input_tokens=len(texts), total_tokens=len(texts)),
            elapsed_ms=elapsed_ms,
        )

    return factory


@pytest.mark.asyncio
@pytest.mark.parametrize("case", ["count", "dimension", "nonfinite", "zero", "model", "elapsed"])
async def test_invalid_provider_response_fails_closed_before_vector_persistence(case: str) -> None:
    repository = SpyRepository(_plan(2))
    provider = SpyEmbeddingProvider(repository, result_factory=_invalid_result(case))

    with pytest.raises(InvalidEmbeddingResponseError):
        await _service(repository, provider).index_document_version(
            lease=LEASE,
            document_version_id=DOCUMENT_VERSION_ID,
            trace_id=TRACE_ID,
        )

    assert repository.finish_calls[-1]["status"] == "FAILED"
    assert repository.finish_calls[-1]["error_code"] == "INVALID_EMBEDDING_RESPONSE"
    assert repository.writes == []
    assert "repo.activate" not in repository.events


@pytest.mark.asyncio
@pytest.mark.parametrize(
    ("kind", "expected_status"),
    [
        (ProviderErrorKind.AUTHENTICATION, "FAILED"),
        (ProviderErrorKind.INVALID_REQUEST, "FAILED"),
        (ProviderErrorKind.RATE_LIMITED, "FAILED"),
        (ProviderErrorKind.UPSTREAM, "FAILED"),
        (ProviderErrorKind.TIMEOUT, "UNKNOWN"),
    ],
)
async def test_provider_failures_are_not_retried_by_service_and_never_activate(
    kind: ProviderErrorKind,
    expected_status: InvocationStatus,
) -> None:
    repository = SpyRepository(_plan(1))
    failure = ModelProviderError(
        operation="embedding.document",
        kind=kind,
        status_code=401 if kind is ProviderErrorKind.AUTHENTICATION else 503,
        retryable=kind not in {ProviderErrorKind.AUTHENTICATION, ProviderErrorKind.INVALID_REQUEST},
    )
    provider = SpyEmbeddingProvider(repository, failures={1: failure})

    with pytest.raises(ModelProviderError):
        await _service(repository, provider).index_document_version(
            lease=LEASE,
            document_version_id=DOCUMENT_VERSION_ID,
            trace_id=TRACE_ID,
        )

    assert len(provider.calls) == 1
    assert repository.finish_calls[-1]["status"] == expected_status
    assert repository.writes == []
    assert "repo.activate" not in repository.events


@pytest.mark.asyncio
@pytest.mark.parametrize("lease_loss_operation", ["finish", "persist", "activate"])
async def test_lease_loss_blocks_every_subsequent_write_and_activation(
    lease_loss_operation: str,
) -> None:
    repository = SpyRepository(_plan(1), lose_lease_on=lease_loss_operation)
    provider = SpyEmbeddingProvider(repository)

    with pytest.raises(LeaseLostError):
        await _service(repository, provider).index_document_version(
            lease=LEASE,
            document_version_id=DOCUMENT_VERSION_ID,
            trace_id=TRACE_ID,
        )

    if lease_loss_operation == "finish":
        assert repository.writes == []
        assert "repo.persist:1" not in repository.events
        assert "repo.activate" not in repository.events
    elif lease_loss_operation == "persist":
        assert repository.writes == []
        assert "repo.activate" not in repository.events
    else:
        assert repository.events[-1] == "repo.activate"


@pytest.mark.asyncio
async def test_activation_requires_expected_count_to_equal_ready_count() -> None:
    repository = SpyRepository(_plan(2), report_incomplete=True)
    provider = SpyEmbeddingProvider(repository)

    with pytest.raises(IndexingNotReadyError):
        await _service(repository, provider).index_document_version(
            lease=LEASE,
            document_version_id=DOCUMENT_VERSION_ID,
            trace_id=TRACE_ID,
        )

    assert len(repository.ready_chunks) == 2
    assert "repo.activate" not in repository.events


@pytest.mark.asyncio
async def test_synthetic_profile_uses_only_explicit_local_provider() -> None:
    synthetic_profile = replace(
        PROFILE,
        model="fake-feature-hash-v1",
        synthetic=True,
    )
    repository = SpyRepository(_plan(1, profile=synthetic_profile))
    cloud_provider = SpyEmbeddingProvider(
        repository,
        failures={
            1: ModelProviderError(
                operation="must-not-run",
                kind=ProviderErrorKind.AUTHENTICATION,
            )
        },
    )
    synthetic_provider = SpyEmbeddingProvider(repository, model="fake-feature-hash-v1")

    result = await _service(
        repository,
        cloud_provider,
        synthetic_provider=synthetic_provider,
    ).index_document_version(
        lease=LEASE,
        document_version_id=DOCUMENT_VERSION_ID,
        trace_id=TRACE_ID,
    )

    assert cloud_provider.calls == []
    assert len(synthetic_provider.calls) == 1
    assert result.activated is True