Initial curated RAGcore source from c16667f0c9397e2bfc8dd0255b114fba23c8d05e
This commit is contained in:
Public Source Publisher committed
commit
7ca726caf2
659 files changed
+98171
No files matched your search
@@ -0,0 +1,241 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from collections.abc import AsyncIterator, Iterable
|
||||
from dataclasses import replace
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
from qdrant_client import AsyncQdrantClient
|
||||
from sqlalchemy.ext.asyncio import AsyncEngine
|
||||
|
||||
from ragcore.application.indexes import IndexBuildService, IndexBuildState
|
||||
from ragcore.config import Settings
|
||||
from ragcore.infrastructure.database import create_database_engine, create_session_factory
|
||||
from ragcore.infrastructure.database.indexes import PostgresIndexOperationStore
|
||||
from ragcore.infrastructure.qdrant import (
|
||||
CollectionKind,
|
||||
CollectionProfile,
|
||||
IndexPayload,
|
||||
IndexPoint,
|
||||
QdrantIndexAdapter,
|
||||
)
|
||||
from ragcore.infrastructure.qdrant.rebuild import QdrantRebuildStore
|
||||
|
||||
|
||||
class FixtureSource:
|
||||
def __init__(self, points: Iterable[IndexPoint]) -> None:
|
||||
self.points = tuple(points)
|
||||
|
||||
async def iter_points(self, profile: CollectionProfile) -> AsyncIterator[IndexPoint]:
|
||||
for point in self.points:
|
||||
yield point
|
||||
|
||||
|
||||
class QdrantGoldenSmoke:
|
||||
def __init__(self, client: AsyncQdrantClient) -> None:
|
||||
self.client = client
|
||||
|
||||
async def run(self, profile: CollectionProfile) -> None:
|
||||
response = await self.client.query_points(
|
||||
collection_name=profile.collection_name,
|
||||
query=[0.1, 0.2, 0.3, 0.4],
|
||||
using="dense",
|
||||
limit=1,
|
||||
with_payload=True,
|
||||
timeout=30,
|
||||
)
|
||||
if not response.points:
|
||||
raise AssertionError("golden smoke query returned no points")
|
||||
|
||||
|
||||
def make_profile() -> CollectionProfile:
|
||||
return CollectionProfile(
|
||||
kind=CollectionKind.DENSE,
|
||||
profile="rebuild-test",
|
||||
version=uuid4().hex[:12],
|
||||
dimensions=4,
|
||||
model_digest="a" * 64,
|
||||
)
|
||||
|
||||
|
||||
def make_point(profile: CollectionProfile, index: int) -> IndexPoint:
|
||||
point_id = uuid4()
|
||||
return IndexPoint(
|
||||
point_id=point_id,
|
||||
dense_vector=(0.1 + index, 0.2, 0.3, 0.4),
|
||||
sparse_indices=(index + 1, index + 9),
|
||||
sparse_values=(0.8, 0.3),
|
||||
payload=IndexPayload(
|
||||
workspace_id=uuid4(),
|
||||
space_id=uuid4(),
|
||||
document_id=uuid4(),
|
||||
document_version_id=uuid4(),
|
||||
chunk_id=point_id,
|
||||
parent_chunk_id=None,
|
||||
source_id=uuid4(),
|
||||
status="published",
|
||||
content_type="text/markdown",
|
||||
language="en",
|
||||
trust_level="source",
|
||||
location={"page": index + 1},
|
||||
parser_version="text-v1",
|
||||
chunk_profile_version=profile.version,
|
||||
embedding_model_digest=profile.model_digest,
|
||||
text=f"A safe indexed passage {index}.",
|
||||
excerpt=f"A safe indexed passage {index}.",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def qdrant() -> AsyncIterator[AsyncQdrantClient]:
|
||||
client = AsyncQdrantClient(
|
||||
url=os.getenv("RAGCORE_QDRANT_URL", "http://qdrant:6333"),
|
||||
api_key=os.getenv("RAGCORE_QDRANT_API_KEY", "replace-me"),
|
||||
timeout=60,
|
||||
)
|
||||
try:
|
||||
yield client
|
||||
finally:
|
||||
await client.close()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
async def db_engine() -> AsyncIterator[AsyncEngine]:
|
||||
engine = create_database_engine(Settings(env="test"))
|
||||
try:
|
||||
yield engine
|
||||
finally:
|
||||
await engine.dispose()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def operations(db_engine: AsyncEngine) -> PostgresIndexOperationStore:
|
||||
return PostgresIndexOperationStore(create_session_factory(db_engine))
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
async def test_rebuild_from_empty_validates_and_activates(
|
||||
qdrant: AsyncQdrantClient,
|
||||
operations: PostgresIndexOperationStore,
|
||||
) -> None:
|
||||
profile = make_profile()
|
||||
points = (make_point(profile, 0), make_point(profile, 1))
|
||||
store = QdrantRebuildStore(qdrant)
|
||||
service = IndexBuildService(store, operations)
|
||||
try:
|
||||
report = await service.build(profile, FixtureSource(points), QdrantGoldenSmoke(qdrant))
|
||||
|
||||
assert report.state is IndexBuildState.ACTIVE
|
||||
assert report.previous_collection is None
|
||||
assert report.validation.expected_count == 2
|
||||
assert report.validation.actual_count == 2
|
||||
assert report.validation.passed
|
||||
assert await store.current_alias_target(profile.alias_name) == profile.collection_name
|
||||
finally:
|
||||
await store.rollback(profile.alias_name, None)
|
||||
await qdrant.delete_collection(profile.collection_name)
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
async def test_alias_rollback_restores_previous_active_collection(
|
||||
qdrant: AsyncQdrantClient,
|
||||
operations: PostgresIndexOperationStore,
|
||||
) -> None:
|
||||
first = make_profile()
|
||||
second = CollectionProfile(
|
||||
kind=first.kind,
|
||||
profile=first.profile,
|
||||
version=uuid4().hex[:12],
|
||||
dimensions=first.dimensions,
|
||||
model_digest=first.model_digest,
|
||||
)
|
||||
first_points = (make_point(first, 0),)
|
||||
second_points = (make_point(second, 1),)
|
||||
store = QdrantRebuildStore(qdrant)
|
||||
service = IndexBuildService(store, operations)
|
||||
try:
|
||||
await service.build(first, FixtureSource(first_points), QdrantGoldenSmoke(qdrant))
|
||||
second_report = await service.build(
|
||||
second, FixtureSource(second_points), QdrantGoldenSmoke(qdrant)
|
||||
)
|
||||
assert second_report.previous_collection == first.collection_name
|
||||
|
||||
await service.rollback(second_report)
|
||||
|
||||
assert await store.current_alias_target(first.alias_name) == first.collection_name
|
||||
finally:
|
||||
await store.rollback(first.alias_name, None)
|
||||
await qdrant.delete_collection(first.collection_name)
|
||||
await qdrant.delete_collection(second.collection_name)
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
async def test_rebuild_rejects_direct_payload_tampering(
|
||||
qdrant: AsyncQdrantClient,
|
||||
operations: PostgresIndexOperationStore,
|
||||
) -> None:
|
||||
profile = make_profile()
|
||||
point = make_point(profile, 0)
|
||||
store = QdrantRebuildStore(qdrant)
|
||||
service = IndexBuildService(store, operations)
|
||||
try:
|
||||
await service.build(profile, FixtureSource((point,)), QdrantGoldenSmoke(qdrant))
|
||||
await qdrant.set_payload(
|
||||
collection_name=profile.collection_name,
|
||||
payload={"text": "tampered indexed content"},
|
||||
points=[str(point.point_id)],
|
||||
wait=True,
|
||||
timeout=30,
|
||||
)
|
||||
valid, malformed = await store.list_points(profile)
|
||||
assert valid == ()
|
||||
assert len(malformed) == 1
|
||||
assert malformed[0].point_id == point.point_id
|
||||
assert "wire integrity" in malformed[0].reason
|
||||
finally:
|
||||
await store.rollback(profile.alias_name, None)
|
||||
await qdrant.delete_collection(profile.collection_name)
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
async def test_reconciliation_repairs_missing_changed_and_orphan_points(
|
||||
qdrant: AsyncQdrantClient,
|
||||
operations: PostgresIndexOperationStore,
|
||||
) -> None:
|
||||
profile = make_profile()
|
||||
points = (make_point(profile, 0), make_point(profile, 1))
|
||||
rogue = make_point(profile, 99)
|
||||
store = QdrantRebuildStore(qdrant)
|
||||
adapter = QdrantIndexAdapter(qdrant)
|
||||
service = IndexBuildService(store, operations)
|
||||
try:
|
||||
await service.build(profile, FixtureSource(points), QdrantGoldenSmoke(qdrant))
|
||||
await adapter.delete_points(profile, (points[0].point_id,))
|
||||
await adapter.upsert_points(profile, (rogue,))
|
||||
changed = replace(
|
||||
points[1],
|
||||
payload=replace(
|
||||
points[1].payload,
|
||||
text="A changed indexed passage.",
|
||||
excerpt="A changed indexed passage.",
|
||||
),
|
||||
)
|
||||
await adapter.upsert_points(profile, (changed,))
|
||||
|
||||
report = await service.reconcile(profile, FixtureSource(points))
|
||||
|
||||
assert report.missing_point_ids == (points[0].point_id,)
|
||||
assert report.orphan_point_ids == (rogue.point_id,)
|
||||
assert report.changed_point_ids == (points[1].point_id,)
|
||||
# repaired_point_ids is deterministically sorted by str(uuid), not by
|
||||
# fixture insertion order, since both ids are independently random.
|
||||
assert set(report.repaired_point_ids) == {points[0].point_id, points[1].point_id}
|
||||
assert report.repaired_point_ids == tuple(
|
||||
sorted((points[0].point_id, points[1].point_id), key=str)
|
||||
)
|
||||
assert await store.count_points(profile) == 2
|
||||
finally:
|
||||
await store.rollback(profile.alias_name, None)
|
||||
await qdrant.delete_collection(profile.collection_name)
|
||||
Reference in new issue
Block a user