Files
library-desk/src/clients/qdrant_client.py
T
jpmschweitzerandClaude Opus 4.5 37f8e1819e
Build and Push / build (release) Successful in 28s
feat: refactor volatile cache to vector storage with HybridRAG integration
- Migrate volatile backend from Redis to Qdrant for semantic search
- Add natural language conversion for structured data embedding
- Simplify API: /volatile/search, /volatile/store, /{namespace}/{key}
- Integrate volatile into HybridRAG with priority boost in RRF fusion
- Add POST /maintenance/cleanup/volatile for expiry purging
- Update tests for new Qdrant-based architecture (37/37 pass)

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
2025-12-24 20:03:37 +01:00

809 lines
24 KiB
Python

"""
Qdrant client wrapper for Library Desk.
Provides async vector operations with:
- Collection-per-user multi-tenancy
- Document chunk storage with embeddings
- Semantic search
- Similarity queries
"""
from qdrant_client import QdrantClient
from qdrant_client.models import (
Distance, VectorParams, PointStruct,
Filter, FieldCondition, MatchValue, Range
)
from typing import List, Dict, Any, Optional
import uuid
import logging
from src.core.multi_tenancy import get_qdrant_collection_name
logger = logging.getLogger(__name__)
class QdrantClientWrapper:
"""
Qdrant client wrapper with multi-tenancy support.
Pattern: Collection per user (from qdrant_memory.py)
Each user has isolated vector collection for their documents.
"""
def __init__(self, url: str, embedding_dim: int = 768):
"""
Initialize Qdrant client.
Args:
url: Qdrant server URL (e.g., "http://qdrant:6333")
embedding_dim: Vector embedding dimension (default 768 for nomic-embed-text)
"""
self.client = QdrantClient(url=url)
self.embedding_dim = embedding_dim
logger.info(f"Initialized Qdrant client: {url}")
def get_collection_name(self, user: str) -> str:
"""
Get collection name for user.
Args:
user: User identifier
Returns:
Collection name (e.g., "library_desk_jpmschweitzer")
"""
return get_qdrant_collection_name(user)
async def ensure_collection(self, collection_name: str):
"""
Create collection if it doesn't exist.
Args:
collection_name: Collection name
"""
try:
collections = self.client.get_collections()
existing = [c.name for c in collections.collections]
if collection_name not in existing:
logger.info(f"Creating Qdrant collection: {collection_name}")
self.client.create_collection(
collection_name=collection_name,
vectors_config=VectorParams(
size=self.embedding_dim,
distance=Distance.COSINE
)
)
logger.info(f"Created collection: {collection_name}")
except Exception as e:
logger.error(f"Error ensuring collection: {e}", exc_info=True)
raise
async def collection_exists(self, collection_name: str) -> bool:
"""
Check if collection exists.
Args:
collection_name: Collection name
Returns:
True if collection exists
"""
try:
collections = self.client.get_collections()
existing = [c.name for c in collections.collections]
return collection_name in existing
except Exception as e:
logger.error(f"Error checking collection: {e}", exc_info=True)
return False
async def upsert_document_chunks(
self,
user: str,
doc_id: str,
chunks: List[Dict[str, Any]],
embeddings: List[List[float]]
) -> int:
"""
Upsert document chunks with embeddings.
Point structure:
{
id: "doc_123_chunk_0",
vector: [...],
payload: {
doc_id: "doc_123",
chunk_index: 0,
content: "text content",
metadata: {...}
}
}
Args:
user: User identifier
doc_id: Document ID
chunks: List of chunk dictionaries with content
embeddings: List of embedding vectors
Returns:
Number of chunks upserted
Raises:
ValueError: If chunks and embeddings length mismatch
"""
if len(chunks) != len(embeddings):
raise ValueError(
f"Chunks ({len(chunks)}) and embeddings ({len(embeddings)}) length mismatch"
)
collection_name = self.get_collection_name(user)
await self.ensure_collection(user)
points = []
for i, (chunk, embedding) in enumerate(zip(chunks, embeddings)):
# Generate deterministic point ID
point_id = str(uuid.uuid5(
uuid.NAMESPACE_DNS,
f"{doc_id}_chunk_{i}"
))
points.append(PointStruct(
id=point_id,
vector=embedding,
payload={
"doc_id": doc_id,
"chunk_index": i,
"content": chunk.get("content", ""),
"metadata": chunk.get("metadata", {})
}
))
try:
self.client.upsert(
collection_name=collection_name,
points=points
)
logger.info(f"Upserted {len(points)} chunks for {doc_id}")
return len(points)
except Exception as e:
logger.error(f"Failed to upsert chunks: {e}", exc_info=True)
raise
async def search(
self,
user: str,
query_vector: List[float],
limit: int = 10,
score_threshold: float = 0.7,
filter_dict: Optional[Dict[str, Any]] = None
) -> List[Dict[str, Any]]:
"""
Semantic search in user's collection.
Args:
user: User identifier
query_vector: Query embedding vector
limit: Maximum results to return
score_threshold: Minimum similarity score (0.0-1.0)
filter_dict: Optional payload filters
Returns:
List of matches with scores and payloads
Example:
results = await client.search(
user="jpmschweitzer",
query_vector=[0.1, 0.2, ...],
limit=5,
filter_dict={"doc_id": "doc_123"}
)
"""
collection_name = self.get_collection_name(user)
# Build filter if provided
query_filter = None
if filter_dict:
conditions = []
for key, value in filter_dict.items():
conditions.append(
FieldCondition(key=key, match=MatchValue(value=value))
)
query_filter = Filter(must=conditions)
try:
response = self.client.query_points(
collection_name=collection_name,
query=query_vector,
limit=limit,
score_threshold=score_threshold,
query_filter=query_filter,
with_payload=True
)
return [
{
"id": str(point.id),
"score": point.score,
"doc_id": point.payload["doc_id"],
"chunk_index": point.payload["chunk_index"],
"content": point.payload["content"],
"metadata": point.payload.get("metadata", {})
}
for point in response.points
]
except Exception as e:
logger.error(f"Search failed: {e}", exc_info=True)
return []
async def delete_document(self, user: str, doc_id: str) -> bool:
"""
Delete all chunks for a document.
Args:
user: User identifier
doc_id: Document ID
Returns:
True if deleted successfully
"""
collection_name = self.get_collection_name(user)
try:
self.client.delete(
collection_name=collection_name,
points_selector=Filter(
must=[
FieldCondition(
key="doc_id",
match=MatchValue(value=doc_id)
)
]
)
)
logger.info(f"Deleted chunks for {doc_id}")
return True
except Exception as e:
logger.error(f"Failed to delete document: {e}", exc_info=True)
return False
async def get_document_chunks(
self,
user: str,
doc_id: str
) -> List[Dict[str, Any]]:
"""
Get all chunks for a document.
Args:
user: User identifier
doc_id: Document ID
Returns:
List of chunks with content and metadata
"""
collection_name = self.get_collection_name(user)
try:
# Scroll through points with doc_id filter
points, _ = self.client.scroll(
collection_name=collection_name,
scroll_filter=Filter(
must=[
FieldCondition(
key="doc_id",
match=MatchValue(value=doc_id)
)
]
),
limit=1000,
with_payload=True,
with_vectors=False
)
return [
{
"id": str(point.id),
"chunk_index": point.payload["chunk_index"],
"content": point.payload["content"],
"metadata": point.payload.get("metadata", {})
}
for point in points
]
except Exception as e:
logger.error(f"Failed to get document chunks: {e}", exc_info=True)
return []
async def count_documents(self, user: str) -> int:
"""
Count total number of unique documents in user's collection.
Args:
user: User identifier
Returns:
Number of unique documents
"""
collection_name = self.get_collection_name(user)
try:
# Get collection info
collection_info = self.client.get_collection(collection_name)
# This gives total points, not unique docs
# For unique docs, would need to aggregate by doc_id
return collection_info.points_count
except Exception as e:
logger.error(f"Failed to count documents: {e}", exc_info=True)
return 0
async def delete_collection(self, user: str) -> bool:
"""
Delete user's entire collection.
Warning: This removes all data for the user!
Args:
user: User identifier
Returns:
True if deleted successfully
"""
collection_name = self.get_collection_name(user)
try:
self.client.delete_collection(collection_name)
logger.warning(f"Deleted collection: {collection_name}")
return True
except Exception as e:
logger.error(f"Failed to delete collection: {e}", exc_info=True)
return False
async def find_similar_chunks(
self,
user: str,
doc_id: str,
limit: int = 10
) -> List[Dict[str, Any]]:
"""
Find chunks similar to those in a given document.
Strategy: Get all chunks from doc, use their vectors to find similar chunks.
Args:
user: User identifier
doc_id: Source document ID
limit: Maximum results per chunk
Returns:
List of similar chunks from other documents
"""
collection_name = self.get_collection_name(user)
try:
# Get source document chunks with vectors
source_points, _ = self.client.scroll(
collection_name=collection_name,
scroll_filter=Filter(
must=[
FieldCondition(
key="doc_id",
match=MatchValue(value=doc_id)
)
]
),
limit=10, # Sample first 10 chunks
with_payload=True,
with_vectors=True
)
if not source_points:
return []
# Search using first chunk's vector
# (could aggregate multiple chunks for better results)
first_vector = source_points[0].vector
response = self.client.query_points(
collection_name=collection_name,
query=first_vector,
limit=limit * 2, # Get more to filter out same doc
with_payload=True
)
# Filter out chunks from same document
similar = [
{
"id": str(point.id),
"score": point.score,
"doc_id": point.payload["doc_id"],
"content": point.payload["content"]
}
for point in response.points
if point.payload["doc_id"] != doc_id
]
return similar[:limit]
except Exception as e:
logger.error(f"Failed to find similar chunks: {e}", exc_info=True)
return []
async def upsert_vector(
self,
collection_name: str,
vector_id: str,
vector: List[float],
payload: Dict[str, Any]
) -> bool:
"""
Upsert a single vector point.
Args:
collection_name: Collection name
vector_id: Point ID
vector: Embedding vector
payload: Point payload/metadata
Returns:
True if successful
"""
try:
self.client.upsert(
collection_name=collection_name,
points=[PointStruct(
id=vector_id,
vector=vector,
payload=payload
)]
)
return True
except Exception as e:
logger.error(f"Failed to upsert vector: {e}", exc_info=True)
return False
async def delete_by_filter(
self,
collection_name: str,
filter_conditions: Dict[str, Any]
) -> int:
"""
Delete points matching filter conditions.
Args:
collection_name: Collection name
filter_conditions: Filter conditions (e.g., {"page_id": 5})
Returns:
Number of points deleted (approximation)
"""
try:
# Build filter
conditions = []
for key, value in filter_conditions.items():
conditions.append(
FieldCondition(key=key, match=MatchValue(value=value))
)
query_filter = Filter(must=conditions)
# Delete points
result = self.client.delete(
collection_name=collection_name,
points_selector=query_filter
)
# Return operation status (Qdrant doesn't return count directly)
logger.info(f"Deleted points with filter {filter_conditions}")
return 1 # Placeholder, actual count not available from API
except Exception as e:
logger.error(f"Failed to delete by filter: {e}", exc_info=True)
return 0
async def search_vectors(
self,
collection_name: str,
query_vector: List[float],
limit: int = 10,
score_threshold: float = 0.5,
filter_conditions: Optional[Dict[str, Any]] = None
) -> List[Dict[str, Any]]:
"""
Search vectors in collection.
Args:
collection_name: Collection name
query_vector: Query embedding vector
limit: Maximum results
score_threshold: Minimum similarity score
filter_conditions: Optional filter conditions
Returns:
List of search results with scores and payloads
"""
# Build filter if provided
query_filter = None
if filter_conditions:
conditions = []
for key, value in filter_conditions.items():
conditions.append(
FieldCondition(key=key, match=MatchValue(value=value))
)
query_filter = Filter(must=conditions)
try:
response = self.client.query_points(
collection_name=collection_name,
query=query_vector,
limit=limit,
score_threshold=score_threshold,
query_filter=query_filter,
with_payload=True
)
return [
{
"id": str(point.id),
"score": point.score,
"payload": dict(point.payload)
}
for point in response.points
]
except Exception as e:
logger.error(f"Search failed: {e}", exc_info=True)
return []
async def scroll_all_points(
self,
collection_name: str,
batch_size: int = 100,
with_payload: bool = True,
with_vectors: bool = False,
filter_conditions: Optional[Dict[str, Any]] = None
) -> List[Dict[str, Any]]:
"""
Scroll through all points in a collection.
Args:
collection_name: Collection name
batch_size: Number of points per batch
with_payload: Include payload in results
with_vectors: Include vectors in results
filter_conditions: Optional filter conditions
Returns:
List of all points with id and payload
"""
all_points = []
offset = None
# Build filter if provided
scroll_filter = None
if filter_conditions:
conditions = []
for key, value in filter_conditions.items():
conditions.append(
FieldCondition(key=key, match=MatchValue(value=value))
)
scroll_filter = Filter(must=conditions)
try:
while True:
points, next_offset = self.client.scroll(
collection_name=collection_name,
scroll_filter=scroll_filter,
limit=batch_size,
offset=offset,
with_payload=with_payload,
with_vectors=with_vectors
)
for point in points:
all_points.append({
"id": str(point.id),
"payload": dict(point.payload) if point.payload else {}
})
if next_offset is None:
break
offset = next_offset
return all_points
except Exception as e:
logger.error(f"Failed to scroll collection {collection_name}: {e}", exc_info=True)
return []
async def delete_by_ids(
self,
collection_name: str,
point_ids: List[str]
) -> int:
"""
Delete points by their IDs.
Args:
collection_name: Collection name
point_ids: List of point IDs to delete
Returns:
Number of points deleted
"""
if not point_ids:
return 0
try:
self.client.delete(
collection_name=collection_name,
points_selector=point_ids
)
logger.info(f"Deleted {len(point_ids)} points from {collection_name}")
return len(point_ids)
except Exception as e:
logger.error(f"Failed to delete points by IDs: {e}", exc_info=True)
return 0
async def list_collections(self) -> List[Dict[str, Any]]:
"""
List all collections with stats.
Returns:
List of collection info dictionaries
"""
try:
collections = self.client.get_collections()
result = []
for coll in collections.collections:
# Get detailed collection info
try:
info = self.client.get_collection(coll.name)
result.append({
"name": coll.name,
"vectors_count": info.vectors_count or 0,
"points_count": info.points_count or 0,
"segments_count": info.segments_count or 0
})
except Exception as e:
logger.warning(f"Failed to get info for collection {coll.name}: {e}")
result.append({
"name": coll.name,
"vectors_count": 0,
"points_count": 0,
"segments_count": 0
})
return result
except Exception as e:
logger.error(f"Failed to list collections: {e}", exc_info=True)
return []
# ========== Volatile Data Methods ==========
async def search_with_expiry_filter(
self,
collection_name: str,
query_vector: List[float],
current_timestamp: int,
limit: int = 10,
score_threshold: float = 0.7
) -> List[Dict[str, Any]]:
"""
Search vectors filtering out expired records.
Args:
collection_name: Collection name
query_vector: Query embedding vector
current_timestamp: Current time in milliseconds
limit: Maximum results
score_threshold: Minimum similarity score
Returns:
List of non-expired search results
"""
# Filter: ttl_expiry > current_timestamp (not expired)
expiry_filter = Filter(
must=[
FieldCondition(
key="ttl_expiry",
range=Range(gt=current_timestamp)
)
]
)
try:
response = self.client.query_points(
collection_name=collection_name,
query=query_vector,
limit=limit,
score_threshold=score_threshold,
query_filter=expiry_filter,
with_payload=True
)
return [
{
"id": str(point.id),
"score": point.score,
"payload": dict(point.payload)
}
for point in response.points
]
except Exception as e:
logger.error(f"Volatile search failed: {e}", exc_info=True)
return []
async def delete_expired_vectors(
self,
collection_name: str,
current_timestamp: int
) -> int:
"""
Delete all vectors where ttl_expiry < current_timestamp.
Args:
collection_name: Collection name
current_timestamp: Current time in milliseconds
Returns:
Number of points deleted (approximate)
"""
# Filter: ttl_expiry < current_timestamp (expired)
expiry_filter = Filter(
must=[
FieldCondition(
key="ttl_expiry",
range=Range(lt=current_timestamp)
)
]
)
try:
# First count how many will be deleted (scroll to count)
count = 0
offset = None
while True:
points, next_offset = self.client.scroll(
collection_name=collection_name,
scroll_filter=expiry_filter,
limit=100,
offset=offset,
with_payload=False
)
count += len(points)
if next_offset is None:
break
offset = next_offset
if count == 0:
return 0
# Delete expired points
self.client.delete(
collection_name=collection_name,
points_selector=expiry_filter
)
logger.info(f"Deleted {count} expired vectors from {collection_name}")
return count
except Exception as e:
logger.error(f"Failed to delete expired vectors: {e}", exc_info=True)
return 0
async def get_volatile_collections(self) -> List[str]:
"""
Get all volatile collections (prefixed with 'volatile_').
Returns:
List of volatile collection names
"""
try:
collections = self.client.get_collections()
return [
c.name for c in collections.collections
if c.name.startswith("volatile_")
]
except Exception as e:
logger.error(f"Failed to list volatile collections: {e}", exc_info=True)
return []