RAGcore-Public/tests/integration/qdrant/test_index_adapter.py

from __future__ import annotations

import inspect
import os
from collections.abc import AsyncIterator
from dataclasses import replace
from uuid import uuid4

import pytest
from qdrant_client import AsyncQdrantClient

from ragcore.infrastructure.qdrant import (
    CollectionKind,
    CollectionProfile,
    CollectionProfileMismatchError,
    IndexPayload,
    IndexPoint,
    QdrantIndexAdapter,
)


@pytest.fixture
def profile() -> CollectionProfile:
    return CollectionProfile(
        kind=CollectionKind.DENSE,
        profile="test",
        version=uuid4().hex[:12],
        dimensions=4,
        model_digest="a" * 64,
    )


@pytest.fixture
async def adapter(
    profile: CollectionProfile,
) -> AsyncIterator[tuple[QdrantIndexAdapter, CollectionProfile]]:
    client = AsyncQdrantClient(
        url=os.getenv("RAGCORE_QDRANT_URL", "http://qdrant:6333"),
        api_key=os.getenv("RAGCORE_QDRANT_API_KEY", "replace-me"),
        timeout=60,
    )
    index = QdrantIndexAdapter(client)
    try:
        yield index, profile
    finally:
        await client.delete_collection(profile.collection_name)
        await client.close()


def make_point(profile: CollectionProfile) -> IndexPoint:
    document_id = uuid4()
    version_id = uuid4()
    chunk_id = uuid4()
    return IndexPoint(
        point_id=chunk_id,
        dense_vector=(0.1, 0.2, 0.3, 0.4),
        sparse_indices=(2, 9),
        sparse_values=(0.8, 0.3),
        payload=IndexPayload(
            workspace_id=uuid4(),
            space_id=uuid4(),
            document_id=document_id,
            document_version_id=version_id,
            chunk_id=chunk_id,
            parent_chunk_id=None,
            source_id=uuid4(),
            status="published",
            content_type="text/markdown",
            language="en",
            trust_level="source",
            location={"page": 1, "block": "heading"},
            parser_version="text-v1",
            chunk_profile_version=profile.version,
            embedding_model_digest=profile.model_digest,
            embedding_space_id=profile.embedding_space_id,
            text="A safe indexed passage.",
            excerpt="A safe indexed passage.",
        ),
    )


def test_profile_rejects_invalid_dimensions_and_digest() -> None:
    with pytest.raises(ValueError, match="dimensions"):
        CollectionProfile(CollectionKind.DENSE, "test", "v1", 0, "a" * 64)
    with pytest.raises(ValueError, match="16384"):
        CollectionProfile(CollectionKind.DENSE, "test", "v1", 16_385, "a" * 64)
    with pytest.raises(ValueError, match="SHA-256"):
        CollectionProfile(CollectionKind.DENSE, "test", "v1", 4, "not-a-digest")


def test_point_payload_rejects_unbounded_or_nonfinite_values() -> None:
    profile = CollectionProfile(CollectionKind.DENSE, "test", "v1", 4, "a" * 64)
    point = make_point(profile)
    with pytest.raises(ValueError, match="payload limit"):
        replace(point.payload, text="x" * (1_048_576 + 1))
    with pytest.raises(ValueError, match="finite"):
        IndexPoint(point_id=point.point_id, dense_vector=(float("nan"),), payload=point.payload)


@pytest.mark.integration
async def test_collection_profile_indexes_alias_and_idempotent_setup(
    adapter: tuple[QdrantIndexAdapter, CollectionProfile],
) -> None:
    index, profile = adapter
    await index.ensure_collection(profile)
    await index.ensure_collection(profile)

    client = index._client
    info = await client.get_collection(profile.collection_name)
    assert info.config.params.vectors["dense"].size == profile.dimensions
    assert "sparse" in info.config.params.sparse_vectors
    assert set(
        (
            "workspace_id",
            "space_id",
            "document_id",
            "document_version_id",
            "status",
            "content_type",
            "language",
            "trust_level",
        )
    ).issubset(info.payload_schema)
    aliases = await client.get_aliases()
    assert any(
        alias.alias_name == profile.alias_name and alias.collection_name == profile.collection_name
        for alias in aliases.aliases
    )


@pytest.mark.integration
async def test_point_upsert_and_explicit_delete_are_idempotent(
    adapter: tuple[QdrantIndexAdapter, CollectionProfile],
) -> None:
    index, profile = adapter
    await index.ensure_collection(profile)
    point = make_point(profile)

    await index.upsert_points(profile, (point,))
    await index.upsert_points(profile, (point,))
    stored = await index._client.retrieve(
        collection_name=profile.collection_name,
        ids=[str(point.point_id)],
    )
    assert len(stored) == 1
    with pytest.raises(CollectionProfileMismatchError):
        await index.upsert_points(
            profile,
            (
                IndexPoint(
                    point_id=point.point_id,
                    dense_vector=(0.1, 0.2),
                    payload=point.payload,
                ),
            ),
        )

    await index.delete_points(profile, (point.point_id,))
    await index.delete_points(profile, (point.point_id,))
    assert (
        await index._client.retrieve(
            collection_name=profile.collection_name,
            ids=[str(point.point_id)],
        )
        == []
    )


def test_public_point_operations_do_not_accept_raw_filters() -> None:
    for method_name in ("upsert_points", "delete_points"):
        parameters = inspect.signature(getattr(QdrantIndexAdapter, method_name)).parameters
        assert "filter" not in parameters
        assert "raw_filter" not in parameters


@pytest.mark.asyncio
async def test_point_upsert_rejects_a_different_embedding_space_before_qdrant_io() -> None:
    profile = CollectionProfile(
        kind=CollectionKind.DENSE,
        profile="test",
        version="v2",
        dimensions=4,
        model_digest="a" * 64,
        embedding_space_id=uuid4(),
    )
    point = make_point(profile)
    mismatched_point = replace(point, payload=replace(point.payload, embedding_space_id=uuid4()))
    adapter = QdrantIndexAdapter(None)  # type: ignore[arg-type]

    with pytest.raises(CollectionProfileMismatchError, match="embedding space"):
        await adapter.upsert_points(profile, (mismatched_point,))