Aigenthix/Graph_RAG

1"""Chroma vector database implementation"""2 3from typing import List, Dict, Any4import chromadb5import logging6import os7 8from .base import VectorDBProvider9 10logger = logging.getLogger(__name__)11 12 13class ChromaDB(VectorDBProvider):14    """Chroma vector database implementation"""15 16    def __init__(self):17        self.client = None18        self.collection = None19 20    def initialize(self, config: Dict[str, Any]) -> None:21        """Initialize Chroma database"""22        try:23            path = config.get("storage_path", "./data/chroma_data")24            os.makedirs(path, exist_ok=True)25 26            self.client = chromadb.PersistentClient(path=path)27            self.collection = self.client.get_or_create_collection(28                name="rag_documents",29                metadata={"hnsw:space": "cosine"},30            )31 32            logger.info(f"Chroma database initialized at {path}")33        except Exception as e:34            logger.error(f"Failed to initialize Chroma: {e}")35            raise36 37    def add_vectors(38        self,39        vectors: List[List[float]],40        ids: List[str],41        metadata: List[Dict[str, Any]],42    ) -> bool:43        """Add vectors to Chroma collection"""44        try:45            self.collection.upsert(46                ids=ids,47                embeddings=vectors,48                metadatas=metadata,49            )50            logger.info(f"Added {len(ids)} vectors to Chroma")51            return True52        except Exception as e:53            logger.error(f"Failed to add vectors to Chroma: {e}")54            return False55 56    def search(57        self,58        query_vector: List[float],59        top_k: int = 5,60    ) -> List[Dict[str, Any]]:61        """Search for similar vectors in Chroma"""62        try:63            results = self.collection.query(64                query_embeddings=[query_vector],65                n_results=top_k,66                include=["embeddings", "metadatas", "distances"],67            )68 69            search_results = []70            if results and results["ids"] and len(results["ids"]) > 0:71                for idx, doc_id in enumerate(results["ids"][0]):72                    distance = results["distances"][0][idx]73                    similarity = 1 - distance  # Convert distance to similarity74 75                    metadata = results["metadatas"][0][idx]76 77                    search_results.append({78                        "id": doc_id,79                        "content": metadata.get("text", ""),80                        "metadata": metadata,81                        "similarity_score": similarity,82                    })83 84            return search_results85        except Exception as e:86            logger.error(f"Search failed in Chroma: {e}")87            return []88 89    def delete(self, doc_id: str) -> bool:90        """Delete document from Chroma"""91        try:92            self.collection.delete(ids=[doc_id])93            logger.info(f"Deleted {doc_id} from Chroma")94            return True95        except Exception as e:96            logger.error(f"Failed to delete from Chroma: {e}")97            return False98 99    def health_check(self) -> bool:100        """Check Chroma database health"""101        try:102            if self.collection is None:103                return False104            count = self.collection.count()105            return True106        except Exception as e:107            logger.error(f"Chroma health check failed: {e}")108            return False109 110    def get_stats(self) -> Dict[str, Any]:111        """Get Chroma database statistics"""112        try:113            count = self.collection.count()114            return {115                "type": "chroma",116                "documents": count,117                "collection_name": "rag_documents",118            }119        except Exception as e:120            logger.error(f"Failed to get Chroma stats: {e}")121            return {}122