""" Vector database integration for semantic memory in GRIT Voice Agent Cloud-compatible for Hugging Face Spaces """ import os import json import logging import uuid from datetime import datetime from typing import Dict, List, Optional, Union, Any # Configure logging logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s') logger = logging.getLogger(__name__) # Try to import ChromaDB try: import chromadb from chromadb.config import Settings CHROMADB_AVAILABLE = True logger.info("ChromaDB is available") except ImportError: CHROMADB_AVAILABLE = False logger.warning("ChromaDB not available. Install with: pip install chromadb") # Try to import sentence transformers try: from sentence_transformers import SentenceTransformer SENTENCE_TRANSFORMERS_AVAILABLE = True logger.info("Sentence Transformers is available") except ImportError: SENTENCE_TRANSFORMERS_AVAILABLE = False logger.warning("Sentence Transformers not available. Install with: pip install sentence-transformers") # Default paths and settings DEFAULT_PERSIST_DIR = os.getenv("VECTOR_DB_PATH", "chroma_db") DEFAULT_COLLECTION = "grit_conversations" DEFAULT_EMBEDDING_MODEL = "all-MiniLM-L6-v2" # Small, fast model good for semantic search class VectorMemory: """Vector database for semantic memory using ChromaDB""" def __init__(self, persist_directory: str = DEFAULT_PERSIST_DIR, collection_name: str = DEFAULT_COLLECTION, embedding_model: str = DEFAULT_EMBEDDING_MODEL): """ Initialize vector memory Args: persist_directory: Directory to persist ChromaDB collection_name: Name of the collection embedding_model: Name of the sentence transformer model """ self.persist_directory = persist_directory self.collection_name = collection_name self.embedding_model = embedding_model self.client = None self.collection = None self.model = None self.available = False # Initialize if dependencies are available if CHROMADB_AVAILABLE and SENTENCE_TRANSFORMERS_AVAILABLE: self.initialize() def initialize(self): """Initialize ChromaDB and embedding model""" try: # Create persist directory if it doesn't exist os.makedirs(self.persist_directory, exist_ok=True) # Initialize ChromaDB client self.client = chromadb.PersistentClient( path=self.persist_directory, settings=Settings( anonymized_telemetry=False, allow_reset=True ) ) # Get or create collection self.collection = self.client.get_or_create_collection( name=self.collection_name, metadata={"description": "GRIT Voice Agent conversations"} ) # Initialize embedding model self.model = SentenceTransformer(self.embedding_model) self.available = True logger.info(f"Vector memory initialized with model: {self.embedding_model}") except Exception as e: logger.error(f"Failed to initialize vector memory: {e}") self.available = False def add_memory(self, user_id: str, text: str, metadata: Optional[Dict[str, Any]] = None) -> bool: """ Add text to vector memory Args: user_id: User ID text: Text to add metadata: Additional metadata Returns: Success status """ if not self.available: logger.error("Vector memory not available") return False try: # Generate ID memory_id = f"{user_id}_{uuid.uuid4().hex}" # Prepare metadata meta = { "user_id": user_id, "timestamp": datetime.now().isoformat(), "type": "conversation" } # Add custom metadata if provided if metadata: meta.update(metadata) # Add to collection self.collection.add( ids=[memory_id], documents=[text], metadatas=[meta] ) logger.info(f"Added memory for {user_id}: {len(text)} chars") return True except Exception as e: logger.error(f"Error adding memory: {e}") return False def search_memory(self, user_id: str, query: str, limit: int = 5) -> List[Dict[str, Any]]: """ Search vector memory for relevant context Args: user_id: User ID query: Search query limit: Maximum number of results Returns: List of relevant memories """ if not self.available: logger.error("Vector memory not available") return [] try: # Query collection results = self.collection.query( query_texts=[query], n_results=limit, where={"user_id": user_id} ) # Format results memories = [] for i, doc in enumerate(results["documents"][0]): if i < len(results["metadatas"][0]): meta = results["metadatas"][0][i] memories.append({ "text": doc, "metadata": meta }) else: memories.append({ "text": doc, "metadata": {} }) logger.info(f"Found {len(memories)} relevant memories for {user_id}") return memories except Exception as e: logger.error(f"Error searching memory: {e}") return [] def clear_user_memory(self, user_id: str) -> bool: """ Clear all memories for a user Args: user_id: User ID Returns: Success status """ if not self.available: logger.error("Vector memory not available") return False try: # Delete where user_id matches self.collection.delete( where={"user_id": user_id} ) logger.info(f"Cleared all memories for {user_id}") return True except Exception as e: logger.error(f"Error clearing user memory: {e}") return False # Singleton instance vector_memory = VectorMemory() def add_to_vector_memory(user_id: str, text: str, metadata: Optional[Dict[str, Any]] = None) -> bool: """ Add text to vector memory Args: user_id: User ID text: Text to add metadata: Additional metadata Returns: Success status """ return vector_memory.add_memory(user_id, text, metadata) def get_relevant_context(user_id: str, query: str, limit: int = 5) -> str: """ Get relevant context from vector memory Args: user_id: User ID query: Search query limit: Maximum number of results Returns: Formatted context string """ memories = vector_memory.search_memory(user_id, query, limit) if not memories: return "" # Format as context string context = "Relevant past conversations:\n\n" for memory in memories: context += f"{memory['text']}\n\n" return context.strip() def clear_vector_memory(user_id: str) -> bool: """ Clear all vector memories for a user Args: user_id: User ID Returns: Success status """ return vector_memory.clear_user_memory(user_id) # Example usage if __name__ == "__main__": # Test adding and retrieving memories test_user = "test_user_123" # Add some test memories add_to_vector_memory(test_user, "User: What's the best way to improve sales?") add_to_vector_memory(test_user, "Assistant: Focus on customer needs, improve your value proposition, and optimize your sales funnel.") add_to_vector_memory(test_user, "User: How can I reduce customer churn?") add_to_vector_memory(test_user, "Assistant: Improve customer onboarding, gather feedback regularly, and implement a customer success program.") # Test retrieval query = "How do I increase sales?" context = get_relevant_context(test_user, query) print(f"Query: {query}") print(f"Retrieved context:\n{context}") # Clean up clear_vector_memory(test_user)