OSOCONSULT/GRITVOICE45

1"""2Vector database integration for semantic memory in GRIT Voice Agent3Cloud-compatible for Hugging Face Spaces4"""5 6import os7import json8import logging9import uuid10from datetime import datetime11from typing import Dict, List, Optional, Union, Any12 13# Configure logging14logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')15logger = logging.getLogger(__name__)16 17# Try to import ChromaDB18try:19    import chromadb20    from chromadb.config import Settings21    CHROMADB_AVAILABLE = True22    logger.info("ChromaDB is available")23except ImportError:24    CHROMADB_AVAILABLE = False25    logger.warning("ChromaDB not available. Install with: pip install chromadb")26 27# Try to import sentence transformers28try:29    from sentence_transformers import SentenceTransformer30    SENTENCE_TRANSFORMERS_AVAILABLE = True31    logger.info("Sentence Transformers is available")32except ImportError:33    SENTENCE_TRANSFORMERS_AVAILABLE = False34    logger.warning("Sentence Transformers not available. Install with: pip install sentence-transformers")35 36# Default paths and settings37DEFAULT_PERSIST_DIR = os.getenv("VECTOR_DB_PATH", "chroma_db")38DEFAULT_COLLECTION = "grit_conversations"39DEFAULT_EMBEDDING_MODEL = "all-MiniLM-L6-v2"  # Small, fast model good for semantic search40 41class VectorMemory:42    """Vector database for semantic memory using ChromaDB"""43    44    def __init__(self, 45                persist_directory: str = DEFAULT_PERSIST_DIR,46                collection_name: str = DEFAULT_COLLECTION,47                embedding_model: str = DEFAULT_EMBEDDING_MODEL):48        """49        Initialize vector memory50        51        Args:52            persist_directory: Directory to persist ChromaDB53            collection_name: Name of the collection54            embedding_model: Name of the sentence transformer model55        """56        self.persist_directory = persist_directory57        self.collection_name = collection_name58        self.embedding_model = embedding_model59        self.client = None60        self.collection = None61        self.model = None62        self.available = False63        64        # Initialize if dependencies are available65        if CHROMADB_AVAILABLE and SENTENCE_TRANSFORMERS_AVAILABLE:66            self.initialize()67    68    def initialize(self):69        """Initialize ChromaDB and embedding model"""70        try:71            # Create persist directory if it doesn't exist72            os.makedirs(self.persist_directory, exist_ok=True)73            74            # Initialize ChromaDB client75            self.client = chromadb.PersistentClient(76                path=self.persist_directory,77                settings=Settings(78                    anonymized_telemetry=False,79                    allow_reset=True80                )81            )82            83            # Get or create collection84            self.collection = self.client.get_or_create_collection(85                name=self.collection_name,86                metadata={"description": "GRIT Voice Agent conversations"}87            )88            89            # Initialize embedding model90            self.model = SentenceTransformer(self.embedding_model)91            92            self.available = True93            logger.info(f"Vector memory initialized with model: {self.embedding_model}")94            95        except Exception as e:96            logger.error(f"Failed to initialize vector memory: {e}")97            self.available = False98    99    def add_memory(self, 100                  user_id: str, 101                  text: str, 102                  metadata: Optional[Dict[str, Any]] = None) -> bool:103        """104        Add text to vector memory105        106        Args:107            user_id: User ID108            text: Text to add109            metadata: Additional metadata110            111        Returns:112            Success status113        """114        if not self.available:115            logger.error("Vector memory not available")116            return False117        118        try:119            # Generate ID120            memory_id = f"{user_id}_{uuid.uuid4().hex}"121            122            # Prepare metadata123            meta = {124                "user_id": user_id,125                "timestamp": datetime.now().isoformat(),126                "type": "conversation"127            }128            129            # Add custom metadata if provided130            if metadata:131                meta.update(metadata)132            133            # Add to collection134            self.collection.add(135                ids=[memory_id],136                documents=[text],137                metadatas=[meta]138            )139            140            logger.info(f"Added memory for {user_id}: {len(text)} chars")141            return True142            143        except Exception as e:144            logger.error(f"Error adding memory: {e}")145            return False146    147    def search_memory(self, 148                     user_id: str, 149                     query: str, 150                     limit: int = 5) -> List[Dict[str, Any]]:151        """152        Search vector memory for relevant context153        154        Args:155            user_id: User ID156            query: Search query157            limit: Maximum number of results158            159        Returns:160            List of relevant memories161        """162        if not self.available:163            logger.error("Vector memory not available")164            return []165        166        try:167            # Query collection168            results = self.collection.query(169                query_texts=[query],170                n_results=limit,171                where={"user_id": user_id}172            )173            174            # Format results175            memories = []176            for i, doc in enumerate(results["documents"][0]):177                if i < len(results["metadatas"][0]):178                    meta = results["metadatas"][0][i]179                    memories.append({180                        "text": doc,181                        "metadata": meta182                    })183                else:184                    memories.append({185                        "text": doc,186                        "metadata": {}187                    })188            189            logger.info(f"Found {len(memories)} relevant memories for {user_id}")190            return memories191            192        except Exception as e:193            logger.error(f"Error searching memory: {e}")194            return []195    196    def clear_user_memory(self, user_id: str) -> bool:197        """198        Clear all memories for a user199        200        Args:201            user_id: User ID202            203        Returns:204            Success status205        """206        if not self.available:207            logger.error("Vector memory not available")208            return False209        210        try:211            # Delete where user_id matches212            self.collection.delete(213                where={"user_id": user_id}214            )215            216            logger.info(f"Cleared all memories for {user_id}")217            return True218            219        except Exception as e:220            logger.error(f"Error clearing user memory: {e}")221            return False222 223# Singleton instance224vector_memory = VectorMemory()225 226def add_to_vector_memory(user_id: str, 227                        text: str, 228                        metadata: Optional[Dict[str, Any]] = None) -> bool:229    """230    Add text to vector memory231    232    Args:233        user_id: User ID234        text: Text to add235        metadata: Additional metadata236        237    Returns:238        Success status239    """240    return vector_memory.add_memory(user_id, text, metadata)241 242def get_relevant_context(user_id: str, 243                        query: str, 244                        limit: int = 5) -> str:245    """246    Get relevant context from vector memory247    248    Args:249        user_id: User ID250        query: Search query251        limit: Maximum number of results252        253    Returns:254        Formatted context string255    """256    memories = vector_memory.search_memory(user_id, query, limit)257    258    if not memories:259        return ""260    261    # Format as context string262    context = "Relevant past conversations:\n\n"263    for memory in memories:264        context += f"{memory['text']}\n\n"265    266    return context.strip()267 268def clear_vector_memory(user_id: str) -> bool:269    """270    Clear all vector memories for a user271    272    Args:273        user_id: User ID274        275    Returns:276        Success status277    """278    return vector_memory.clear_user_memory(user_id)279 280# Example usage281if __name__ == "__main__":282    # Test adding and retrieving memories283    test_user = "test_user_123"284    285    # Add some test memories286    add_to_vector_memory(test_user, "User: What's the best way to improve sales?")287    add_to_vector_memory(test_user, "Assistant: Focus on customer needs, improve your value proposition, and optimize your sales funnel.")288    add_to_vector_memory(test_user, "User: How can I reduce customer churn?")289    add_to_vector_memory(test_user, "Assistant: Improve customer onboarding, gather feedback regularly, and implement a customer success program.")290    291    # Test retrieval292    query = "How do I increase sales?"293    context = get_relevant_context(test_user, query)294    print(f"Query: {query}")295    print(f"Retrieved context:\n{context}")296    297    # Clean up298    clear_vector_memory(test_user)299