Aigenthix/Graph_RAG

1"""Vector database service wrapper"""2 3from typing import List, Dict, Any, Optional4from ..vector_dbs.factory import VectorDBFactory5from ..vector_dbs.base import VectorDBProvider6import logging7 8logger = logging.getLogger(__name__)9 10 11class VectorDBService:12    """Service for managing vector database operations"""13 14    def __init__(self, db_type: str, config: Dict[str, Any]):15        self.db_type = db_type16        self.provider: VectorDBProvider = VectorDBFactory.create(db_type, config)17 18    def add_chunks(19        self,20        chunks: List[Dict[str, Any]],21    ) -> bool:22        """Add document chunks to vector database"""23        try:24            vectors = [chunk["embedding"] for chunk in chunks]25            ids = [26                f"{chunk['metadata'].get('doc_id', 'unknown')}_{chunk['chunk_id']}"27                for chunk in chunks28            ]29            metadatas = [30                {31                    "text": chunk["text"],32                    "doc_id": chunk["metadata"].get("doc_id", ""),33                    "chunk_id": chunk["chunk_id"],34                    **chunk["metadata"],35                }36                for chunk in chunks37            ]38 39            return self.provider.add_vectors(vectors, ids, metadatas)40        except Exception as e:41            logger.error(f"Failed to add chunks to vector database: {e}")42            return False43 44    def search(45        self,46        query_embedding: List[float],47        top_k: int = 5,48    ) -> List[Dict[str, Any]]:49        """Search for similar chunks"""50        try:51            results = self.provider.search(query_embedding, top_k)52            return results53        except Exception as e:54            logger.error(f"Search failed: {e}")55            return []56 57    def delete_document(self, doc_id: str) -> bool:58        """Delete all chunks for a document"""59        try:60            return self.provider.delete(doc_id)61        except Exception as e:62            logger.error(f"Failed to delete document: {e}")63            return False64 65    def health_check(self) -> bool:66        """Check if vector database is healthy"""67        return self.provider.health_check()68 69    def get_stats(self) -> Dict[str, Any]:70        """Get vector database statistics"""71        return self.provider.get_stats()72