File size: 9,380 Bytes
bae15d1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 | """
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})")
|