""" Qdrant client wrapper for memory vector storage. Provides async operations for storing and retrieving memory embeddings: - Collection management (per-user collections) - Memory upsert/search/delete - Filtering by memory type Adapted from library-desk patterns. """ from typing import Any from uuid import uuid4, uuid5, NAMESPACE_DNS from qdrant_client import QdrantClient from qdrant_client.http import models as qdrant_models from .config import config from .logging_config import get_logger from .multi_tenancy import get_memory_collection_name logger = get_logger(__name__) class MemoryQdrantClient: """ Qdrant client wrapper for memory storage. Manages per-user collections with the pattern: memories_{user} Stores memory embeddings with metadata (type, content, timestamps). Usage: client = MemoryQdrantClient() await client.ensure_collection("jpmschweitzer") await client.upsert_memory( user="jpmschweitzer", memory_id="mem_123", vector=[0.1, 0.2, ...], payload={"type": "fact", "content": "User prefers dark mode"} ) """ def __init__( self, url: str | None = None, embedding_dim: int | None = None, ): """ Initialize Qdrant client. Args: url: Qdrant server URL (defaults to config.qdrant_url) embedding_dim: Vector dimension (defaults to config.QDRANT_EMBEDDING_DIM) """ self.url = url or config.qdrant_url self.embedding_dim = embedding_dim or config.QDRANT_EMBEDDING_DIM self._client = QdrantClient(url=self.url) logger.info( "qdrant_client_initialized", url=self.url, embedding_dim=self.embedding_dim, ) def close(self) -> None: """Close Qdrant client.""" if self._client is not None: self._client.close() async def ensure_collection(self, user: str) -> bool: """ Ensure collection exists for user, create if not. Args: user: User identifier Returns: True if collection exists or was created successfully Example: >>> await client.ensure_collection("jpmschweitzer") True """ collection_name = get_memory_collection_name(user) try: # Check if collection exists collections = self._client.get_collections() existing = [c.name for c in collections.collections] if collection_name in existing: logger.debug( "qdrant_collection_exists", collection=collection_name, ) return True # Create collection with cosine distance self._client.create_collection( collection_name=collection_name, vectors_config=qdrant_models.VectorParams( size=self.embedding_dim, distance=qdrant_models.Distance.COSINE, ), ) logger.info( "qdrant_collection_created", collection=collection_name, embedding_dim=self.embedding_dim, ) return True except Exception as e: logger.error( "qdrant_ensure_collection_failed", collection=collection_name, error=str(e), ) return False async def upsert_memory( self, user: str, memory_id: str | None, vector: list[float], payload: dict[str, Any], ) -> str | None: """ Upsert a memory point. Args: user: User identifier memory_id: Memory ID (generated if None) vector: Embedding vector payload: Memory metadata (should include 'type', 'content', etc.) Returns: Memory ID if successful, None on failure Example: >>> memory_id = await client.upsert_memory( ... user="jpmschweitzer", ... memory_id=None, ... vector=[0.1, 0.2, ...], ... payload={ ... "type": "fact", ... "content": "User prefers dark mode", ... "created_at": "2024-01-01T00:00:00Z" ... } ... ) """ collection_name = get_memory_collection_name(user) # Generate deterministic UUID from memory_id (or random if not provided) # Qdrant requires UUID or integer IDs, not arbitrary strings if memory_id: # Deterministic UUID from string - same memory_id = same UUID point_id = str(uuid5(NAMESPACE_DNS, f"{user}:{memory_id}")) else: point_id = str(uuid4()) memory_id = point_id # Use UUID as the memory_id too try: # Ensure collection exists await self.ensure_collection(user) # Create point (store original memory_id in payload for reference) payload["memory_id"] = memory_id point = qdrant_models.PointStruct( id=point_id, vector=vector, payload=payload, ) # Upsert self._client.upsert( collection_name=collection_name, points=[point], ) logger.debug( "qdrant_memory_upserted", collection=collection_name, memory_id=memory_id, memory_type=payload.get("type"), ) return memory_id except Exception as e: logger.error( "qdrant_upsert_memory_failed", collection=collection_name, memory_id=memory_id, error=str(e), ) return None async def search_memories( self, user: str, query_vector: list[float], limit: int = 10, memory_type: str | None = None, score_threshold: float = 0.5, ) -> list[dict[str, Any]]: """ Search memories by vector similarity. Args: user: User identifier query_vector: Query embedding vector limit: Maximum results memory_type: Filter by memory type (e.g., "fact", "preference", "profile") score_threshold: Minimum similarity score (0-1) Returns: List of matching memories with scores Example: >>> memories = await client.search_memories( ... user="jpmschweitzer", ... query_vector=[0.1, 0.2, ...], ... limit=5, ... memory_type="fact" ... ) >>> memories[0] {"id": "mem_123", "score": 0.89, "type": "fact", "content": "..."} """ collection_name = get_memory_collection_name(user) try: # Build filter if memory_type specified query_filter = None if memory_type: query_filter = qdrant_models.Filter( must=[ qdrant_models.FieldCondition( key="type", match=qdrant_models.MatchValue(value=memory_type), ) ] ) # Search using new Query API (qdrant-client >= 1.10) results = self._client.query_points( collection_name=collection_name, query=query_vector, limit=limit, query_filter=query_filter, score_threshold=score_threshold, ).points # Format results memories = [] for hit in results: memory = { "id": hit.id, "score": hit.score, **hit.payload, } memories.append(memory) logger.debug( "qdrant_search_memories", collection=collection_name, results_count=len(memories), memory_type=memory_type, ) return memories except Exception as e: logger.error( "qdrant_search_memories_failed", collection=collection_name, error=str(e), ) return [] async def get_memory(self, user: str, memory_id: str) -> dict[str, Any] | None: """ Get a specific memory by ID. Args: user: User identifier memory_id: Memory ID Returns: Memory data or None if not found """ collection_name = get_memory_collection_name(user) # Convert memory_id to UUID point_id point_id = str(uuid5(NAMESPACE_DNS, f"{user}:{memory_id}")) try: points = self._client.retrieve( collection_name=collection_name, ids=[point_id], ) if not points: return None point = points[0] return { "id": point.id, **point.payload, } except Exception as e: logger.error( "qdrant_get_memory_failed", collection=collection_name, memory_id=memory_id, error=str(e), ) return None async def delete_memory(self, user: str, memory_id: str) -> bool: """ Delete a memory by ID. Args: user: User identifier memory_id: Memory ID to delete Returns: True if deleted successfully, False otherwise Example: >>> await client.delete_memory("jpmschweitzer", "mem_123") True """ collection_name = get_memory_collection_name(user) # Convert memory_id to UUID point_id point_id = str(uuid5(NAMESPACE_DNS, f"{user}:{memory_id}")) try: self._client.delete( collection_name=collection_name, points_selector=qdrant_models.PointIdsList( points=[point_id], ), ) logger.debug( "qdrant_memory_deleted", collection=collection_name, memory_id=memory_id, ) return True except Exception as e: logger.error( "qdrant_delete_memory_failed", collection=collection_name, memory_id=memory_id, error=str(e), ) return False async def delete_memories_by_type(self, user: str, memory_type: str) -> int: """ Delete all memories of a specific type. Args: user: User identifier memory_type: Type of memories to delete Returns: Number of memories deleted (approximate) """ collection_name = get_memory_collection_name(user) try: # Delete by filter self._client.delete( collection_name=collection_name, points_selector=qdrant_models.FilterSelector( filter=qdrant_models.Filter( must=[ qdrant_models.FieldCondition( key="type", match=qdrant_models.MatchValue(value=memory_type), ) ] ) ), ) logger.info( "qdrant_memories_deleted_by_type", collection=collection_name, memory_type=memory_type, ) return -1 # Qdrant doesn't return count for filter deletes except Exception as e: logger.error( "qdrant_delete_memories_by_type_failed", collection=collection_name, memory_type=memory_type, error=str(e), ) return 0 async def count_memories(self, user: str) -> int: """ Count total memories for a user. Args: user: User identifier Returns: Number of memories in user's collection """ collection_name = get_memory_collection_name(user) try: info = self._client.get_collection(collection_name) return info.points_count except Exception as e: logger.error( "qdrant_count_memories_failed", collection=collection_name, error=str(e), ) return 0 async def health_check(self) -> bool: """ Check if Qdrant server is reachable. Returns: True if healthy, False otherwise """ try: self._client.get_collections() return True except Exception as e: logger.error("qdrant_health_check_failed", error=str(e)) return False # Global client instance (lazy initialization) _qdrant_client: MemoryQdrantClient | None = None def get_qdrant_client() -> MemoryQdrantClient: """ Get global Qdrant client instance. Returns: MemoryQdrantClient instance """ global _qdrant_client if _qdrant_client is None: _qdrant_client = MemoryQdrantClient() return _qdrant_client