""" Persistent per-user memory storage backed by ChromaDB and LlamaIndex. """ from __future__ import annotations import logging import os from typing import Any, Optional from uuid import uuid4 from datetime import datetime, timezone import chromadb import yaml from llama_index.core import Document, Settings, StorageContext, VectorStoreIndex from llama_index.core.vector_stores import MetadataFilter, MetadataFilters from llama_index.vector_stores.chroma import ChromaVectorStore # Configure logging for this module logger = logging.getLogger(__name__) class UserMemoryStore: """Stores and retrieves user-specific memories in a dedicated vector index.""" def __init__( self, config: Optional[dict[str, Any]] = None, config_path: Optional[str] = None, ) -> None: """Initializes the user memory store from a config dict or YAML file.""" # Guardrail: prevent both config and config_path from being provided simultaneously if config is not None and config_path is not None: raise ValueError("Provide either 'config' or 'config_path', not both.") # If a config dict is provided, validate and use it if config is not None: self.config = self._validate_config(config) # Otherwise, load it from the YAML file else: if config_path is None: config_path = os.path.join(os.path.dirname(__file__), "rag_config.yaml") self.config = self._load_config(config_path) self._setup_vector_store() self._index: Optional[VectorStoreIndex] = None logger.info("UserMemoryStore initialized successfully.") def _validate_config(self, config: dict[str, Any]) -> dict[str, Any]: """Validates the config structure required for the user memory store.""" if not isinstance(config, dict): raise ValueError("RAG configuration must be a dictionary.") # Validate the keys needed for memory storage self._require_config_key(config, ("storage",)) self._require_config_key(config, ("storage", "persist_dir")) self._require_config_key(config, ("storage", "user_memory_collection_name")) self._require_config_key(config, ("memory_retrieval",)) self._require_config_key(config, ("memory_retrieval", "similarity_top_k")) return config def _load_config(self, config_path: str) -> dict[str, Any]: """Loads the memory store configuration.""" logger.info("Loading user memory configuration from: %s", config_path) if not os.path.exists(config_path): raise FileNotFoundError(f"Configuration file not found at {config_path}") with open(config_path, "r", encoding="utf-8") as config_file: config = yaml.safe_load(config_file) return self._validate_config(config) def _require_config_key( self, config: dict[str, Any], key_path: tuple[str, ...] ) -> Any: """Returns a required config value or raises an error if the key is missing.""" current: Any = config for key in key_path: if not isinstance(current, dict) or key not in current: dotted_path = ".".join(key_path) raise KeyError(f"Missing required config key: {dotted_path}") current = current[key] return current def _setup_vector_store(self) -> None: """Creates or opens the dedicated Chroma collection for user memory.""" persist_dir = self.config["storage"]["persist_dir"] collection_name = self.config["storage"]["user_memory_collection_name"] # Create the persistence directory if it doesn't exist if not os.path.exists(persist_dir): logger.info("Creating user memory persistence directory: %s", persist_dir) os.makedirs(persist_dir, exist_ok=True) # Create the Chroma client and open/create the user memory collection db = chromadb.PersistentClient(path=persist_dir) chroma_collection = db.get_or_create_collection(collection_name) # Wrap the Chroma collection in a LlamaIndex vector store and storage context self.vector_store = ChromaVectorStore(chroma_collection=chroma_collection) self.storage_context = StorageContext.from_defaults( vector_store=self.vector_store ) logger.info( "User memory vector store ready at %s (collection=%s)", persist_dir, collection_name, ) def _ensure_embed_model(self) -> None: """Fails fast if no global embedding model has been configured.""" if Settings.embed_model is None: raise RuntimeError( "LlamaIndex Settings.embed_model is not configured. " "Configure it before using UserMemoryStore." ) def _get_or_create_index(self) -> VectorStoreIndex: """Loads or creates the vector index for the user memory store.""" if self._index is None: # Check that the embedding model is configured before loading the index self._ensure_embed_model() # Load the index from the existing Chroma vector store logger.info("Loading user memory index from vector store.") self._index = VectorStoreIndex.from_vector_store( self.vector_store, storage_context=self.storage_context, ) return self._index def add_memory( self, user_id: str, memory_text: str, metadata: Optional[dict[str, Any]] = None ) -> None: """Inserts a single memory for a specific user into persistent storage with metadata.""" # Strip whitespace and validate inputs user_id = user_id.strip() memory_text = memory_text.strip() if not user_id: raise ValueError("user_id must be a non-empty string.") if not memory_text: raise ValueError("memory_text must be a non-empty string.") # Combine provided metadata with the required user_id field for filtering memory_metadata: dict[str, Any] = dict(metadata or {}) memory_metadata["user_id"] = user_id memory_metadata.setdefault("created_at", datetime.now(timezone.utc).isoformat()) # Build a LlamaIndex Document with a unique ID for this memory entry document = Document( text=memory_text, metadata=memory_metadata, doc_id=f"user-memory-{user_id}-{uuid4()}", ) logger.info("Adding memory for user_id=%s", user_id) index = self._get_or_create_index() # Insert the document into the vector index index.insert(document) def query_memories( self, user_id: str, query: str, top_k: Optional[int] = None ) -> Any: """Retrieves the most relevant memories for a specific user.""" # Strip whitespace and validate inputs user_id = user_id.strip() query = query.strip() if not user_id: raise ValueError("user_id must be a non-empty string.") if not query: raise ValueError("query must be a non-empty string.") # Use the memory retrieval top_k if not explicitly provided if top_k is None: top_k = self.config["memory_retrieval"]["similarity_top_k"] if top_k <= 0: raise ValueError("top_k must be greater than 0.") # Only retrieve documents that match the user_id in their metadata filters = MetadataFilters( filters=[MetadataFilter(key="user_id", value=user_id)] ) logger.info("Querying memories for user_id=%s with top_k=%s", user_id, top_k) index = self._get_or_create_index() # Build retriever that applies the user_id filter and retrieves most similar memories retriever = index.as_retriever(similarity_top_k=top_k, filters=filters) return retriever.retrieve(query)