jpmschweitzerandClaude Opus 4.5
4c6ac89808
Add multi-tenancy support and memory storage infrastructure:
- Add ContextVar-based request context (src/core/context.py)
- Async-safe user/conversation tracking via contextvars
- RequestContext manager for clean setup/teardown
- get_user(), get_conversation_id() helpers
- Add multi-tenancy utilities (src/core/multi_tenancy.py)
- User ID sanitization for collection/key names
- get_memory_collection_name(), get_session_key() helpers
- Add Ollama embedding client (src/core/embeddings.py)
- nomic-embed-text model (768 dimensions)
- embed(), embed_batch(), health_check() methods
- Add Qdrant client wrapper (src/core/qdrant.py)
- Per-user collection pattern: memories_{user}
- upsert_memory(), search_memories(), delete_memory()
- Type-based filtering support
- Add Redis memory cache (src/core/memory_cache.py)
- Session context with 24h TTL
- Recent entities tracking
- Separate from benchmarks (db=2)
- Update config with memory settings
- QDRANT_HOST, QDRANT_PORT, QDRANT_EMBEDDING_DIM
- OLLAMA_EMBEDDING_MODEL
- REDIS_MEMORY_DB, REDIS_MEMORY_TTL_HOURS
- Add user field to ResponseRequest (OpenAI standard)
- Set context in router, reset in finally block
- Update librarian client to use get_user() (12 methods)
All 333 unit tests pass.
🤖 Generated with [Claude Code](https://claude.com/claude-code)
Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
447 lines
13 KiB
Python
447 lines
13 KiB
Python
"""
|
|
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
|
|
|
|
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)
|
|
memory_id = memory_id or f"mem_{uuid4().hex[:16]}"
|
|
|
|
try:
|
|
# Ensure collection exists
|
|
await self.ensure_collection(user)
|
|
|
|
# Create point
|
|
point = qdrant_models.PointStruct(
|
|
id=memory_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
|
|
results = self._client.search(
|
|
collection_name=collection_name,
|
|
query_vector=query_vector,
|
|
limit=limit,
|
|
query_filter=query_filter,
|
|
score_threshold=score_threshold,
|
|
)
|
|
|
|
# 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)
|
|
|
|
try:
|
|
points = self._client.retrieve(
|
|
collection_name=collection_name,
|
|
ids=[memory_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)
|
|
|
|
try:
|
|
self._client.delete(
|
|
collection_name=collection_name,
|
|
points_selector=qdrant_models.PointIdsList(
|
|
points=[memory_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
|