ai-rpg/app/core/qdrant_client.py

"""Qdrant client wrapper (singleton) with health check."""

from __future__ import annotations

from typing import Any

from qdrant_client import AsyncQdrantClient
from qdrant_client.http.models import (
    Distance,
    PayloadSchemaType,
    VectorParams,
)

from app.config import get_settings
from app.core.logging import get_logger

_logger = get_logger(__name__)

_client: AsyncQdrantClient | None = None


def get_qdrant_client() -> AsyncQdrantClient:
    """Return the singleton AsyncQdrantClient."""
    global _client
    if _client is None:
        cfg = get_settings()
        _client = AsyncQdrantClient(
            url=cfg.qdrant_url,
            api_key=cfg.qdrant_api_key or None,
            timeout=cfg.qdrant_timeout,
        )
    return _client


async def dispose_qdrant_client() -> None:
    """Close the Qdrant client (on shutdown)."""
    global _client
    if _client is not None:
        try:
            await _client.close()
        except Exception as e:  # noqa: BLE001
            _logger.warning("qdrant_close_failed", error=str(e))
        _client = None


async def ping_qdrant() -> bool:
    """Health-check: returns True if Qdrant responds."""
    try:
        client = get_qdrant_client()
        await client.get_collections()
        return True
    except Exception as e:  # noqa: BLE001
        _logger.warning("qdrant_ping_failed", error=str(e))
        return False


async def init_qdrant_collections(dimension: int) -> dict[str, Any]:
    """Create collections `entities` and `story_entries` if missing.

    Returns a dict with the list of created collection names and the dimension used.
    """
    cfg = get_settings()
    prefix = cfg.qdrant_collection_prefix or ""
    client = get_qdrant_client()

    existing = {c.name for c in (await client.get_collections()).collections}
    created: list[str] = []

    collections_config = {
        f"{prefix}entities": [
            ("world_id", PayloadSchemaType.KEYWORD),
            ("entity_type", PayloadSchemaType.KEYWORD),
            ("deleted", PayloadSchemaType.BOOL),
        ],
        f"{prefix}story_entries": [
            ("world_id", PayloadSchemaType.KEYWORD),
            ("entry_type", PayloadSchemaType.KEYWORD),
            ("created_at", PayloadSchemaType.INTEGER),
        ],
    }

    for name, indexes in collections_config.items():
        if name in existing:
            continue
        await client.create_collection(
            collection_name=name,
            vectors_config=VectorParams(size=dimension, distance=Distance.COSINE),
        )
        for field, schema_type in indexes:
            await client.create_payload_index(name, field, schema_type)
        created.append(name)
        _logger.info("qdrant_collection_created", name=name, dimension=dimension)

    return {"created": created, "dimension": dimension, "existing": sorted(existing)}


async def cleanup_world_points(world_id: str) -> None:
    """Best-effort delete of all Qdrant points for a given world_id."""
    from qdrant_client.http.models import (
        FieldCondition,
        Filter,
        FilterSelector,
        MatchValue,
    )

    cfg = get_settings()
    prefix = cfg.qdrant_collection_prefix or ""
    client = get_qdrant_client()
    for collection in (f"{prefix}entities", f"{prefix}story_entries"):
        try:
            await client.delete(
                collection_name=collection,
                points_selector=FilterSelector(
                    filter=Filter(
                        must=[
                            FieldCondition(
                                key="world_id",
                                match=MatchValue(value=str(world_id)),
                            )
                        ]
                    )
                ),
            )
        except Exception as e:  # noqa: BLE001
            _logger.error(
                "qdrant_cleanup_failed",
                collection=collection,
                world_id=str(world_id),
                error=str(e),
            )