""" Vector Store for Kalim Chatbot Uses ChromaDB for storing and retrieving major embeddings """ import os from typing import List, Dict, Any, Optional from pathlib import Path try: import chromadb from chromadb.config import Settings HAS_CHROMADB = True except ImportError: HAS_CHROMADB = False from config.settings import VECTORSTORE_DIR, TOP_K_RESULTS from data.loader import load_majors from rag.embeddings import ( get_embedding_model, create_major_document, create_major_metadata ) class KalimVectorStore: """Vector store for Kalim major recommendations.""" COLLECTION_NAME = "kalim_majors" def __init__(self, persist_directory: Optional[Path] = None): """Initialize the vector store. Args: persist_directory: Directory to persist the database """ if not HAS_CHROMADB: raise ImportError( "chromadb is required. " "Install with: pip install chromadb" ) self.persist_directory = persist_directory or VECTORSTORE_DIR self.persist_directory.mkdir(parents=True, exist_ok=True) # Initialize ChromaDB client with persistence self._client = chromadb.PersistentClient( path=str(self.persist_directory), settings=Settings( anonymized_telemetry=False, allow_reset=True ) ) self._collection = None self._embedding_model = None @property def embedding_model(self): """Lazy load embedding model.""" if self._embedding_model is None: self._embedding_model = get_embedding_model() return self._embedding_model @property def collection(self): """Get or create the collection.""" if self._collection is None: self._collection = self._client.get_or_create_collection( name=self.COLLECTION_NAME, metadata={"description": "Lebanese University majors for Kalim chatbot"} ) return self._collection def build_index(self, majors: Optional[List[Dict]] = None, force_rebuild: bool = False): """Build the vector index from majors data. Args: majors: List of major dictionaries. If None, loads from JSON. force_rebuild: If True, rebuilds even if collection exists """ if majors is None: majors = load_majors() # Check if already built if self.collection.count() > 0 and not force_rebuild: print(f"Collection already has {self.collection.count()} documents. " "Use force_rebuild=True to rebuild.") return # Clear existing data if rebuilding if force_rebuild and self.collection.count() > 0: print("Clearing existing collection...") self._client.delete_collection(self.COLLECTION_NAME) self._collection = self._client.create_collection( name=self.COLLECTION_NAME, metadata={"description": "Lebanese University majors for Kalim chatbot"} ) print(f"Building index for {len(majors)} majors...") # Prepare documents, embeddings, and metadata documents = [] embeddings = [] metadatas = [] ids = [] for major in majors: doc = create_major_document(major) metadata = create_major_metadata(major) major_id = str(major.get('id', len(ids) + 1)) documents.append(doc) metadatas.append(metadata) ids.append(major_id) # Generate embeddings in batch print("Generating embeddings...") embeddings = self.embedding_model.embed_texts(documents) # Add to collection print("Adding to collection...") self.collection.add( documents=documents, embeddings=embeddings, metadatas=metadatas, ids=ids ) print(f"Index built successfully with {self.collection.count()} documents") def search( self, query: str, n_results: int = TOP_K_RESULTS, where: Optional[Dict] = None, where_document: Optional[Dict] = None ) -> List[Dict[str, Any]]: """Search for similar majors. Args: query: Search query text n_results: Number of results to return where: Metadata filter conditions where_document: Document content filter Returns: List of matching majors with scores """ # Generate query embedding query_embedding = self.embedding_model.embed_query(query) # Search results = self.collection.query( query_embeddings=[query_embedding], n_results=n_results, where=where, where_document=where_document, include=["documents", "metadatas", "distances"] ) # Format results formatted_results = [] if results and results['ids'] and results['ids'][0]: for i, doc_id in enumerate(results['ids'][0]): result = { 'id': doc_id, 'document': results['documents'][0][i] if results['documents'] else '', 'metadata': results['metadatas'][0][i] if results['metadatas'] else {}, 'distance': results['distances'][0][i] if results['distances'] else 0.0, 'score': 1 - results['distances'][0][i] if results['distances'] else 1.0 } formatted_results.append(result) return formatted_results def search_by_holland( self, holland_codes: List[str], n_results: int = TOP_K_RESULTS ) -> List[Dict[str, Any]]: """Search for majors matching Holland codes. Args: holland_codes: List of Holland codes (e.g., ['E', 'C', 'S']) n_results: Number of results Returns: List of matching majors """ if not holland_codes: return [] # Build query from Holland descriptions from config.settings import HOLLAND_CODES query_parts = [] for code in holland_codes: if code in HOLLAND_CODES: query_parts.append(HOLLAND_CODES[code]['description_ar']) query = " ".join(query_parts) if query_parts else " ".join(holland_codes) # Build filter for primary Holland code primary_code = holland_codes[0] if holland_codes else None where_filter = None if primary_code: where_filter = {"holland_primary": primary_code} return self.search( query=query, n_results=n_results, where=where_filter ) def get_all_majors(self) -> List[Dict[str, Any]]: """Get all majors from the collection. Returns: List of all major metadata """ results = self.collection.get( include=["documents", "metadatas"] ) formatted = [] if results and results['ids']: for i, doc_id in enumerate(results['ids']): formatted.append({ 'id': doc_id, 'document': results['documents'][i] if results['documents'] else '', 'metadata': results['metadatas'][i] if results['metadatas'] else {} }) return formatted def get_major_by_id(self, major_id: str) -> Optional[Dict[str, Any]]: """Get a specific major by ID. Args: major_id: Major ID string Returns: Major data or None """ results = self.collection.get( ids=[major_id], include=["documents", "metadatas"] ) if results and results['ids']: return { 'id': results['ids'][0], 'document': results['documents'][0] if results['documents'] else '', 'metadata': results['metadatas'][0] if results['metadatas'] else {} } return None def reset(self): """Reset the collection (delete all data).""" self._client.delete_collection(self.COLLECTION_NAME) self._collection = None print("Collection reset successfully") def get_vectorstore() -> KalimVectorStore: """Get the vector store instance.""" return KalimVectorStore() def build_vectorstore(force_rebuild: bool = False) -> KalimVectorStore: """Build and return the vector store. Args: force_rebuild: If True, rebuilds the index Returns: Built vector store """ store = get_vectorstore() store.build_index(force_rebuild=force_rebuild) return store if __name__ == "__main__": # Build the vector store when run directly print("Building Kalim Vector Store...") store = build_vectorstore(force_rebuild=True) print(f"Vector store built with {store.collection.count()} documents") # Test search print("\nTesting search for 'هندسة البرمجيات'...") results = store.search("هندسة البرمجيات", n_results=5) for r in results: print(f" - {r['metadata'].get('name_ar')} (score: {r['score']:.3f})")