""" Vector Storage Manager - Traditional vector storage backend for dual storage comparison. Provides vector embeddings storage with local fallback and future Modal integration. """ import os import json import time import logging from typing import Dict, List, Any, Optional from pathlib import Path import numpy as np try: from sentence_transformers import SentenceTransformer import faiss VECTOR_DEPS_AVAILABLE = True except ImportError: logging.warning( "Vector storage dependencies not available (sentence-transformers, faiss)" ) SentenceTransformer = None faiss = None VECTOR_DEPS_AVAILABLE = False class VectorStorageManager: """ Vector storage backend for dual storage comparison. Provides traditional embedding-based storage with local FAISS index. Future: Modal integration for production scaling. """ def __init__( self, data_dir: str = "data", model_name: str = "all-MiniLM-L6-v2", storage_handler=None, ): """ Initialize vector storage manager. Args: data_dir (str): Base directory for storage model_name (str): Sentence transformer model name storage_handler: HF Dataset storage handler for persistence """ self.logger = logging.getLogger(__name__) self.data_dir = Path(data_dir) self.model_name = model_name self.storage_handler = storage_handler # For HF Dataset persistence # Initialize embedding model self.encoder = None if VECTOR_DEPS_AVAILABLE: try: self.encoder = SentenceTransformer(model_name) self.logger.info(f"Vector storage initialized with model: {model_name}") except Exception as e: self.logger.error(f"Failed to load embedding model: {e}") else: self.logger.warning("Vector storage not available - missing dependencies") # Client indices self.client_indices = {} # client_id -> faiss index self.client_texts = {} # client_id -> list of texts self.client_metadata = {} # client_id -> list of metadata def store_embedding( self, text: str, client_id: str, metadata: Dict[str, Any] = None ) -> str: """ Store text as vector embedding. Args: text (str): Text to store client_id (str): Client identifier metadata (dict): Additional metadata Returns: str: Storage result message """ try: if not VECTOR_DEPS_AVAILABLE: return "Error: Vector storage dependencies not available (sentence-transformers, faiss)" if not self.encoder: return "Error: Embedding model not loaded" # Generate embedding start_time = time.time() embedding = self.encoder.encode([text]) embedding_time = time.time() - start_time # Initialize client storage if needed if client_id not in self.client_indices: self._init_client_storage(client_id, embedding.shape[1]) # Add to client index self.client_indices[client_id].add(embedding) self.client_texts[client_id].append(text) self.client_metadata[client_id].append(metadata or {}) # Save to disk self._save_client_index(client_id) # Auto-backup to HF Dataset for persistence on HF Spaces self.auto_backup_after_store(client_id, self.storage_handler) total_embeddings = len(self.client_texts[client_id]) return f"Vector embedding stored for client {client_id}. Embedding time: {embedding_time:.3f}s. Total embeddings: {total_embeddings}" except Exception as e: error_msg = f"Error storing vector embedding: {str(e)}" self.logger.error(error_msg) return error_msg def search_embeddings(self, query: str, client_id: str, top_k: int = 5) -> str: """ Search embeddings using vector similarity. Args: query (str): Search query client_id (str): Client identifier top_k (int): Number of results Returns: str: JSON string with search results """ try: if not VECTOR_DEPS_AVAILABLE: return json.dumps( {"error": "Vector storage dependencies not available"} ) if not self.encoder: return json.dumps({"error": "Embedding model not loaded"}) if client_id not in self.client_indices: return json.dumps( {"error": f"No embeddings found for client {client_id}"} ) # Generate query embedding query_embedding = self.encoder.encode([query]) # Search index scores, indices = self.client_indices[client_id].search( query_embedding, top_k ) # Prepare results results = [] for i, (score, idx) in enumerate(zip(scores[0], indices[0])): if idx < len(self.client_texts[client_id]): result = { "text": self.client_texts[client_id][idx], "score": float(score), "rank": i + 1, "metadata": self.client_metadata[client_id][idx], } results.append(result) return json.dumps( { "query": query, "client_id": client_id, "total_results": len(results), "results": results, "backend": "vector_storage", }, indent=2, ) except Exception as e: error_msg = f"Error searching vector embeddings: {str(e)}" self.logger.error(error_msg) return json.dumps({"error": error_msg}) def delete_memory(self, client_id: str, memory_name: str = "") -> str: """ Delete embeddings for a client. Args: client_id (str): Client identifier memory_name (str): Memory name (not used in vector storage) Returns: str: Deletion result """ try: if client_id in self.client_indices: # Clear client data del self.client_indices[client_id] del self.client_texts[client_id] del self.client_metadata[client_id] # Remove saved files client_dir = self._get_client_dir(client_id) if client_dir.exists(): import shutil shutil.rmtree(client_dir) return f"Vector embeddings deleted for client {client_id}" else: return f"No vector embeddings found for client {client_id}" except Exception as e: error_msg = f"Error deleting vector embeddings: {str(e)}" self.logger.error(error_msg) return error_msg def get_stats(self, client_id: str) -> str: """ Get vector storage statistics. Args: client_id (str): Client identifier Returns: str: JSON string with statistics """ try: if client_id not in self.client_indices: return json.dumps( { "client_id": client_id, "total_embeddings": 0, "storage_backend": "vector_storage", "status": "no_data", } ) total_embeddings = len(self.client_texts[client_id]) total_text_size = sum(len(text) for text in self.client_texts[client_id]) # Calculate storage size client_dir = self._get_client_dir(client_id) storage_size = 0 if client_dir.exists(): storage_size = sum( f.stat().st_size for f in client_dir.rglob("*") if f.is_file() ) return json.dumps( { "client_id": client_id, "total_embeddings": total_embeddings, "total_text_size_bytes": total_text_size, "storage_size_bytes": storage_size, "storage_backend": "vector_storage", "embedding_model": self.model_name, "status": "active", }, indent=2, ) except Exception as e: error_msg = f"Error getting vector storage stats: {str(e)}" self.logger.error(error_msg) return json.dumps({"error": error_msg}) def _init_client_storage(self, client_id: str, embedding_dim: int) -> None: """Initialize storage for a new client.""" # Create FAISS index self.client_indices[client_id] = faiss.IndexFlatIP( embedding_dim ) # Inner product similarity self.client_texts[client_id] = [] self.client_metadata[client_id] = [] # Create client directory client_dir = self._get_client_dir(client_id) client_dir.mkdir(parents=True, exist_ok=True) def _get_client_dir(self, client_id: str) -> Path: """Get client-specific directory for vector storage.""" return self.data_dir / f"{client_id}_vector" def _save_client_index(self, client_id: str) -> None: """Save client index and data to disk.""" try: client_dir = self._get_client_dir(client_id) # Save FAISS index faiss.write_index( self.client_indices[client_id], str(client_dir / "vector_index.faiss") ) # Save texts and metadata with open(client_dir / "texts.json", "w", encoding="utf-8") as f: json.dump(self.client_texts[client_id], f, indent=2) with open(client_dir / "metadata.json", "w", encoding="utf-8") as f: json.dump(self.client_metadata[client_id], f, indent=2) except Exception as e: self.logger.error(f"Error saving client index for {client_id}: {e}") def _load_client_index(self, client_id: str) -> bool: """Load client index and data from disk.""" try: client_dir = self._get_client_dir(client_id) if not (client_dir / "vector_index.faiss").exists(): return False # Load FAISS index self.client_indices[client_id] = faiss.read_index( str(client_dir / "vector_index.faiss") ) # Load texts and metadata with open(client_dir / "texts.json", "r", encoding="utf-8") as f: self.client_texts[client_id] = json.load(f) with open(client_dir / "metadata.json", "r", encoding="utf-8") as f: self.client_metadata[client_id] = json.load(f) return True except Exception as e: self.logger.error(f"Error loading client index for {client_id}: {e}") return False def load_client_data(self, client_id: str) -> str: """ Load client data from disk. Args: client_id (str): Client identifier Returns: str: Load result message """ try: if self._load_client_index(client_id): total_embeddings = len(self.client_texts[client_id]) return f"Vector storage loaded for client {client_id}: {total_embeddings} embeddings" else: return f"No vector storage data found for client {client_id}" except Exception as e: error_msg = f"Error loading client data: {str(e)}" self.logger.error(error_msg) return error_msg # Future Modal integration methods (placeholders) def enable_modal_backend(self, modal_token: str) -> str: """ Enable Modal backend for production scaling. Args: modal_token (str): Modal API token Returns: str: Activation result """ # TODO: Implement Modal integration return ( "Modal backend integration not yet implemented. Using local FAISS storage." ) def migrate_to_modal(self, client_id: str) -> str: """ Migrate client data to Modal backend. Args: client_id (str): Client identifier Returns: str: Migration result """ # TODO: Implement Modal migration return "Modal migration not yet implemented. Data remains in local storage." # HF Dataset Integration for Persistence on HF Spaces def backup_to_hf_dataset(self, client_id: str, storage_handler) -> str: """ Backup vector storage to HuggingFace Dataset for persistence. Args: client_id (str): Client identifier storage_handler: HF Dataset storage handler Returns: str: Backup result """ try: if not storage_handler or not storage_handler.hf_enabled: return "HF Dataset backup not available - no storage handler or HF not enabled" client_dir = self._get_client_dir(client_id) if not client_dir.exists(): return f"No vector data found for client {client_id}" # Use storage handler to backup vector files success = storage_handler.backup_client_data(client_id, client_dir) if success: return f"Successfully backed up vector storage for client {client_id} to HF Dataset" else: return f"Failed to backup vector storage for client {client_id}" except Exception as e: error_msg = f"Error backing up vector storage: {str(e)}" self.logger.error(error_msg) return error_msg def restore_from_hf_dataset(self, client_id: str, storage_handler) -> str: """ Restore vector storage from HuggingFace Dataset. Args: client_id (str): Client identifier storage_handler: HF Dataset storage handler Returns: str: Restore result """ try: if not storage_handler or not storage_handler.hf_enabled: return "HF Dataset restore not available - no storage handler or HF not enabled" client_dir = self._get_client_dir(client_id) # Use storage handler to restore vector files success = storage_handler.restore_client_data(client_id, client_dir) if success: # Load the restored data into memory if self._load_client_index(client_id): total_embeddings = len(self.client_texts[client_id]) return f"Successfully restored vector storage for client {client_id}: {total_embeddings} embeddings" else: return f"Vector files restored but failed to load into memory for client {client_id}" else: return f"Failed to restore vector storage for client {client_id}" except Exception as e: error_msg = f"Error restoring vector storage: {str(e)}" self.logger.error(error_msg) return error_msg def auto_backup_after_store(self, client_id: str, storage_handler) -> None: """ Automatically backup after storing embeddings (for HF Spaces persistence). Args: client_id (str): Client identifier storage_handler: HF Dataset storage handler """ try: if storage_handler and storage_handler.hf_enabled: # Auto-backup in background (non-blocking) self.backup_to_hf_dataset(client_id, storage_handler) except Exception as e: self.logger.warning(f"Auto-backup failed for client {client_id}: {e}")