rebase

app/core/qdrant_client.py

@@ -0,0 +1,130 @@
+"""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),
+            )