257 lines
9.3 KiB
Python
257 lines
9.3 KiB
Python
from __future__ import annotations
|
|
|
|
from collections.abc import AsyncIterator, Sequence
|
|
from dataclasses import replace
|
|
from uuid import UUID, uuid4
|
|
|
|
import pytest
|
|
|
|
from ragcore.application.indexes.models import (
|
|
IndexBuildError,
|
|
MalformedIndexPoint,
|
|
StoredIndexPoint,
|
|
index_point_fingerprint,
|
|
)
|
|
from ragcore.application.indexes.shadow_migration import (
|
|
ShadowEmbeddingMigrationService,
|
|
ShadowMigrationCheckpoint,
|
|
ShadowMigrationStatus,
|
|
ShadowPreflight,
|
|
)
|
|
from ragcore.infrastructure.qdrant.models import (
|
|
CollectionKind,
|
|
CollectionProfile,
|
|
IndexPayload,
|
|
IndexPoint,
|
|
)
|
|
|
|
|
|
class MemoryCheckpoints:
|
|
def __init__(self, *, cancel_after_batch: bool = False) -> None:
|
|
self.item: ShadowMigrationCheckpoint | None = None
|
|
self.cancel_after_batch = cancel_after_batch
|
|
|
|
async def load(self, _migration_id: UUID) -> ShadowMigrationCheckpoint | None:
|
|
return self.item
|
|
|
|
async def save(self, checkpoint: ShadowMigrationCheckpoint) -> None:
|
|
if self.cancel_after_batch and checkpoint.request_batches == 1:
|
|
checkpoint = replace(checkpoint, cancel_requested=True)
|
|
self.item = checkpoint
|
|
|
|
|
|
class Source:
|
|
def __init__(self, points: Sequence[IndexPoint]) -> None:
|
|
self.points = points
|
|
|
|
async def iter_points(self, _profile: CollectionProfile) -> AsyncIterator[IndexPoint]:
|
|
for point in self.points:
|
|
yield point
|
|
|
|
|
|
class Store:
|
|
def __init__(self, *, fail_once: bool = False, omit_last: bool = False) -> None:
|
|
self.points: dict[UUID, StoredIndexPoint] = {}
|
|
self.fail_once = fail_once
|
|
self.omit_last = omit_last
|
|
self.alias = "rag_dense_legacy_v1"
|
|
self.alias_mutations = 0
|
|
self.upsert_calls = 0
|
|
|
|
async def create_shadow(self, _profile: CollectionProfile) -> None:
|
|
return None
|
|
|
|
async def delete_shadow(self, _profile: CollectionProfile) -> None:
|
|
raise AssertionError("validated shadow is retained")
|
|
|
|
async def upsert_points(
|
|
self, _profile: CollectionProfile, points: Sequence[IndexPoint]
|
|
) -> None:
|
|
self.upsert_calls += 1
|
|
if self.fail_once:
|
|
self.fail_once = False
|
|
raise TimeoutError("transient")
|
|
for point in points:
|
|
self.points[point.point_id] = StoredIndexPoint(
|
|
point.point_id, index_point_fingerprint(point)
|
|
)
|
|
|
|
async def delete_points(self, _profile: CollectionProfile, _point_ids: Sequence[UUID]) -> None:
|
|
raise AssertionError("no destructive cleanup")
|
|
|
|
async def delete_raw_points(self, _profile: CollectionProfile, _raw_ids: Sequence[str]) -> None:
|
|
raise AssertionError("no destructive cleanup")
|
|
|
|
async def count_points(self, _profile: CollectionProfile) -> int:
|
|
return len(self.points)
|
|
|
|
async def list_points(
|
|
self, _profile: CollectionProfile
|
|
) -> tuple[tuple[StoredIndexPoint, ...], tuple[MalformedIndexPoint, ...]]:
|
|
values = tuple(self.points.values())
|
|
return (values[:-1] if self.omit_last and values else values), ()
|
|
|
|
async def current_alias_target(self, _alias_name: str) -> str:
|
|
return self.alias
|
|
|
|
async def activate(self, _profile: CollectionProfile) -> None:
|
|
self.alias_mutations += 1
|
|
|
|
async def rollback(self, _alias_name: str, _previous_collection: str | None) -> None:
|
|
self.alias_mutations += 1
|
|
|
|
|
|
def profile() -> CollectionProfile:
|
|
return CollectionProfile(
|
|
kind=CollectionKind.DENSE,
|
|
profile="modelforge-shadow",
|
|
version="v1",
|
|
dimensions=3,
|
|
model_digest="a" * 64,
|
|
embedding_space_id=uuid4(),
|
|
)
|
|
|
|
|
|
def points(item: CollectionProfile, count: int = 3) -> tuple[IndexPoint, ...]:
|
|
workspace, space, document, version, source = (uuid4() for _ in range(5))
|
|
return tuple(
|
|
IndexPoint(
|
|
point_id=(chunk := uuid4()),
|
|
dense_vector=(1.0, 0.0, 0.0),
|
|
payload=IndexPayload(
|
|
workspace_id=workspace,
|
|
space_id=space,
|
|
document_id=document,
|
|
document_version_id=version,
|
|
chunk_id=chunk,
|
|
parent_chunk_id=None,
|
|
source_id=source,
|
|
status="published",
|
|
content_type="text/plain",
|
|
language="en",
|
|
trust_level="internal",
|
|
location={"structure_path": f"/{index}"},
|
|
parser_version="parser-v1",
|
|
chunk_profile_version=item.version,
|
|
embedding_model_digest=item.model_digest,
|
|
text=f"chunk {index}",
|
|
excerpt=f"chunk {index}",
|
|
embedding_space_id=item.embedding_space_id,
|
|
),
|
|
)
|
|
for index in range(count)
|
|
)
|
|
|
|
|
|
def checkpoint(item: CollectionProfile, total: int = 3) -> ShadowMigrationCheckpoint:
|
|
assert item.embedding_space_id is not None
|
|
return ShadowMigrationCheckpoint(
|
|
migration_id=uuid4(),
|
|
source_index="rag_dense_legacy_v1",
|
|
target_index=item.collection_name,
|
|
source_embedding_space="legacy-observed",
|
|
target_embedding_space_id=item.embedding_space_id,
|
|
corpus_revision="corpus-sha256",
|
|
total_chunks=total,
|
|
)
|
|
|
|
|
|
async def preflight(*, passed: bool = True) -> ShadowPreflight:
|
|
return ShadowPreflight(
|
|
production_healthy=passed,
|
|
source_available=True,
|
|
modelforge_healthy=True,
|
|
credential_valid=True,
|
|
embedding_space_matches=True,
|
|
dimension_matches=True,
|
|
distance_metric_matches=True,
|
|
qdrant_reachable=True,
|
|
storage_sufficient=True,
|
|
target_conflict_free=True,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_preflight_failure_writes_no_vectors() -> None:
|
|
item = profile()
|
|
store, state = Store(), MemoryCheckpoints()
|
|
service = ShadowEmbeddingMigrationService(store, state)
|
|
with pytest.raises(IndexBuildError, match="preflight"):
|
|
await service.run(
|
|
checkpoint(item), item, Source(points(item)), lambda: preflight(passed=False)
|
|
)
|
|
assert not store.points
|
|
assert state.item is not None and state.item.status is ShadowMigrationStatus.FAILED
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_backfill_is_bounded_validated_and_never_mutates_source_alias() -> None:
|
|
item, store, state = profile(), Store(), MemoryCheckpoints()
|
|
service = ShadowEmbeddingMigrationService(store, state, batch_size=2)
|
|
result = await service.run(checkpoint(item), item, Source(points(item)), preflight)
|
|
assert result.status is ShadowMigrationStatus.READY_FOR_EVALUATION
|
|
assert result.completed_chunks == 3 and result.failed_chunks == 0
|
|
assert result.request_batches == 2
|
|
assert result.validation["passed"] is True
|
|
assert result.validation["source_alias_unchanged"] is True
|
|
assert store.alias_mutations == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_resume_skips_integrity_valid_existing_chunk_and_is_idempotent() -> None:
|
|
item, store, states = profile(), Store(), MemoryCheckpoints()
|
|
source_points = points(item)
|
|
existing = source_points[0]
|
|
store.points[existing.point_id] = StoredIndexPoint(
|
|
existing.point_id, index_point_fingerprint(existing)
|
|
)
|
|
initial = replace(
|
|
checkpoint(item), status=ShadowMigrationStatus.BACKFILLING, completed_chunks=1
|
|
)
|
|
states.item = initial
|
|
service = ShadowEmbeddingMigrationService(store, states, batch_size=1)
|
|
result = await service.run(initial, item, Source(source_points), preflight)
|
|
assert result.completed_chunks == 3
|
|
assert store.upsert_calls == 2
|
|
repeated = await service.run(initial, item, Source(source_points), preflight)
|
|
assert repeated.status is ShadowMigrationStatus.READY_FOR_EVALUATION
|
|
assert store.upsert_calls == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancel_finishes_current_batch_and_remains_resumable() -> None:
|
|
item, store, states = profile(), Store(), MemoryCheckpoints(cancel_after_batch=True)
|
|
source_points = points(item)
|
|
service = ShadowEmbeddingMigrationService(store, states, batch_size=1)
|
|
cancelled = await service.run(checkpoint(item), item, Source(source_points), preflight)
|
|
assert cancelled.status is ShadowMigrationStatus.CANCELLED
|
|
assert cancelled.completed_chunks == 1
|
|
states.cancel_after_batch = False
|
|
resumed = await service.run(cancelled, item, Source(source_points), preflight)
|
|
assert resumed.status is ShadowMigrationStatus.READY_FOR_EVALUATION
|
|
assert resumed.completed_chunks == 3
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_transient_retry_is_bounded_and_counted() -> None:
|
|
item, store, states = profile(), Store(fail_once=True), MemoryCheckpoints()
|
|
service = ShadowEmbeddingMigrationService(
|
|
store,
|
|
states,
|
|
batch_size=3,
|
|
max_retries=1,
|
|
retryable=lambda error: isinstance(error, TimeoutError),
|
|
)
|
|
result = await service.run(checkpoint(item), item, Source(points(item)), preflight)
|
|
assert result.retried_chunks == 3
|
|
assert store.upsert_calls == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_missing_point_blocks_evaluation_readiness() -> None:
|
|
item, store, states = profile(), Store(omit_last=True), MemoryCheckpoints()
|
|
service = ShadowEmbeddingMigrationService(store, states, batch_size=3)
|
|
with pytest.raises(IndexBuildError, match="validation"):
|
|
await service.run(checkpoint(item), item, Source(points(item)), preflight)
|
|
assert states.item is not None and states.item.status is ShadowMigrationStatus.FAILED
|