Spaces:
Sleeping
Sleeping
Download backend/vector_store.py from trigent09/ai-backend: direct link, hf CLI and curl.
- Browser
- Download file 17.2 kB
-
https://huggingface.co/spaces/trigent09/ai-backend/resolve/main/backend/vector_store.py
- Command line
-
hf download hf://spaces/trigent09/ai-backend/backend/vector_store.py
-
curl -L -o vector_store.py https://huggingface.co/spaces/trigent09/ai-backend/resolve/main/backend/vector_store.py
17.2 kB
| """ | |
| Multi-tenant vector store: each school gets its own ChromaDB collection. | |
| Collection naming: school_{school_id}_{board_id} | |
| Within a collection, metadata filters handle class → subject → chapter → topic. | |
| """ | |
| import math | |
| import re | |
| import chromadb | |
| from typing import List, Dict, Any, Optional | |
| from langchain_core.documents import Document | |
| from huggingface_hub import InferenceClient | |
| from backend.config import CHROMA_DIR, EMBEDDING_MODEL_NAME, HINDI_EMBEDDING_MODEL_NAME, HF_TOKEN, USE_LOCAL_EMBEDDINGS | |
| def _normalize_vector(vec: Any) -> List[float]: | |
| """L2 normalize vector for inner-product (ip) similarity.""" | |
| if hasattr(vec, "tolist"): | |
| vec = vec.tolist() | |
| while len(vec) > 0 and isinstance(vec[0], list): | |
| vec = vec[0] | |
| norm = math.sqrt(sum(x * x for x in vec)) | |
| return [x / norm for x in vec] if norm > 0 else vec | |
| class RobustHFEmbeddings: | |
| """ | |
| Lightweight embedding wrapper: | |
| - Default (Cloud): Uses Hugging Face Inference API (<70MB RAM, zero PyTorch load). | |
| - Local: Only loads PyTorch in RAM if USE_LOCAL_EMBEDDINGS=true. | |
| """ | |
| def __init__(self, model_name: str): | |
| self.model_name = model_name | |
| self._local_emb = None | |
| self._client = InferenceClient(token=HF_TOKEN or None) if HF_TOKEN else None | |
| if USE_LOCAL_EMBEDDINGS or not self._client: | |
| try: | |
| from langchain_huggingface import HuggingFaceEmbeddings | |
| self._local_emb = HuggingFaceEmbeddings( | |
| model_name=model_name, | |
| model_kwargs={"device": "cpu"}, | |
| encode_kwargs={"normalize_embeddings": True} | |
| ) | |
| except Exception: | |
| self._local_emb = None | |
| def embed_documents(self, texts: List[str]) -> List[List[float]]: | |
| if self._local_emb: | |
| try: | |
| return [_normalize_vector(v) for v in self._local_emb.embed_documents(texts)] | |
| except Exception: | |
| pass | |
| if not self._client: | |
| raise ValueError("No Hugging Face token (HF_TOKEN) or local embedding model available.") | |
| all_embeddings = [] | |
| batch_size = 32 | |
| for i in range(0, len(texts), batch_size): | |
| batch = texts[i:i + batch_size] | |
| vecs = self._client.feature_extraction(batch, model=self.model_name) | |
| if hasattr(vecs, "tolist"): | |
| vecs = vecs.tolist() | |
| if len(batch) == 1 and isinstance(vecs[0], float): | |
| vecs = [vecs] | |
| all_embeddings.extend([_normalize_vector(v) for v in vecs]) | |
| return all_embeddings | |
| def embed_query(self, text: str) -> List[float]: | |
| if self._local_emb: | |
| try: | |
| return _normalize_vector(self._local_emb.embed_query(text)) | |
| except Exception: | |
| pass | |
| if not self._client: | |
| raise ValueError("No Hugging Face token (HF_TOKEN) or local embedding model available.") | |
| vec = self._client.feature_extraction(text, model=self.model_name) | |
| if hasattr(vec, "tolist"): | |
| vec = vec.tolist() | |
| return _normalize_vector(vec) | |
| def _sanitize_collection_name(name: str) -> str: | |
| """ChromaDB collection names: 3-63 chars, alphanumeric + underscores + hyphens, must start/end with alphanum.""" | |
| name = re.sub(r"[^a-zA-Z0-9_-]", "_", name) | |
| name = name.strip("_-") or "default" | |
| # Pad if too short | |
| if len(name) < 3: | |
| name = name + "_col" | |
| return name[:63] | |
| class VectorStoreManager: | |
| """ | |
| Manages persistent Chroma vector storage with per-tenant collection isolation. | |
| Each school+board combination gets its own collection. | |
| """ | |
| def __init__(self, persist_directory: str = str(CHROMA_DIR)): | |
| self.persist_directory = persist_directory | |
| self.embeddings = RobustHFEmbeddings(EMBEDDING_MODEL_NAME) | |
| self._hindi_embeddings = None | |
| self.client = chromadb.PersistentClient(path=self.persist_directory) | |
| def hindi_embeddings(self): | |
| """Lazy-loaded multilingual model for Hindi & Indian languages (saves RAM until needed).""" | |
| if self._hindi_embeddings is None: | |
| self._hindi_embeddings = RobustHFEmbeddings(HINDI_EMBEDDING_MODEL_NAME) | |
| return self._hindi_embeddings | |
| def get_embeddings(self, subject: Optional[str] = None, text_sample: str = ""): | |
| """Route to Hindi multilingual model if subject is Hindi/Sanskrit or text is Devanagari.""" | |
| is_indic = bool(subject and str(subject).strip().lower() in ("hindi", "हिन्दी", "sanskrit", "संस्कृत")) | |
| has_devanagari = bool(re.search(r"[\u0900-\u097F]", text_sample or "")) | |
| if is_indic or has_devanagari: | |
| return self.hindi_embeddings | |
| return self.embeddings | |
| def _collection_name(self, school_id: str, board_id: str) -> str: | |
| return _sanitize_collection_name(f"school_{school_id}_{board_id}") | |
| def get_tenant_collection(self, school_id: str, board_id: str): | |
| """Get or create the ChromaDB collection for a specific school+board tenant.""" | |
| name = self._collection_name(school_id, board_id) | |
| return self.client.get_or_create_collection( | |
| name=name, | |
| metadata={"hnsw:space": "ip", "hnsw:search_ef": 30} | |
| ) | |
| def add_documents( | |
| self, | |
| documents: List[Document], | |
| school_id: str, | |
| board_id: str | |
| ) -> int: | |
| if not documents: | |
| return 0 | |
| collection = self.get_tenant_collection(school_id, board_id) | |
| texts = [doc.page_content for doc in documents] | |
| metadatas = [doc.metadata for doc in documents] | |
| ids = [ | |
| f"{doc.metadata.get('class_id', 'x')}_{doc.metadata.get('subject', 'x')}_" | |
| f"{doc.metadata.get('chapter', 'x')}_{i}_{hash(doc.page_content)}" | |
| for i, doc in enumerate(documents) | |
| ] | |
| # Dynamically route embeddings based on subject and script | |
| subj = documents[0].metadata.get("subject", "") if documents else "" | |
| emb_engine = self.get_embeddings(subject=subj, text_sample=texts[0] if texts else "") | |
| embeddings = emb_engine.embed_documents(texts) | |
| collection.upsert( | |
| documents=texts, | |
| embeddings=embeddings, | |
| metadatas=metadatas, | |
| ids=ids | |
| ) | |
| return len(documents) | |
| def delete_chapter( | |
| self, | |
| school_id: str, | |
| board_id: str, | |
| class_id: str, | |
| subject: str, | |
| chapter: str, | |
| section_id: Optional[str] = None, | |
| sub_subject: Optional[str] = None, | |
| topic: Optional[str] = None | |
| ) -> int: | |
| """Delete chunks for a specific chapter/lesson, optionally scoped by section, sub_subject, or topic.""" | |
| collection = self.get_tenant_collection(school_id, board_id) | |
| conds = [ | |
| {"class_id": str(class_id)}, | |
| {"subject": subject}, | |
| {"chapter": chapter} | |
| ] | |
| if section_id: | |
| conds.append({"section_id": str(section_id)}) | |
| if sub_subject: | |
| conds.append({"sub_subject": str(sub_subject)}) | |
| if topic: | |
| conds.append({"topic": str(topic)}) | |
| where_filter = {"$and": conds} if len(conds) > 1 else conds[0] | |
| try: | |
| existing = collection.get(where=where_filter) | |
| if existing and existing.get("ids"): | |
| collection.delete(ids=existing["ids"]) | |
| return len(existing["ids"]) | |
| except Exception: | |
| pass | |
| return 0 | |
| def query_context( | |
| self, | |
| query: str, | |
| school_id: str, | |
| board_id: str, | |
| class_id: Optional[str] = None, | |
| section_id: Optional[str] = None, | |
| subject: Optional[str] = None, | |
| sub_subject: Optional[str] = None, | |
| chapter: Optional[str] = None, | |
| lesson: Optional[str] = None, | |
| topic: Optional[str] = None, | |
| n_results: int = 4 | |
| ) -> List[Dict[str, Any]]: | |
| collection = self.get_tenant_collection(school_id, board_id) | |
| effective_lesson = lesson or chapter | |
| clean_topic = str(topic).strip() if topic else "" | |
| is_generic_topic = ( | |
| not clean_topic | |
| or clean_topic.lower() in ("all", "all topics", "all topics (full lesson)", "general", "none") | |
| or (effective_lesson and clean_topic.lower() == str(effective_lesson).strip().lower()) | |
| ) | |
| filter_conditions = [] | |
| if class_id: | |
| filter_conditions.append({"class_id": str(class_id)}) | |
| if section_id and section_id not in ("All", ""): | |
| filter_conditions.append({"section_id": str(section_id)}) | |
| if subject: | |
| filter_conditions.append({"subject": str(subject)}) | |
| if sub_subject and sub_subject not in ("General", ""): | |
| filter_conditions.append({"sub_subject": str(sub_subject)}) | |
| if effective_lesson: | |
| filter_conditions.append({"chapter": str(effective_lesson)}) | |
| if not is_generic_topic: | |
| filter_conditions.append({"topic": clean_topic}) | |
| where_filter = None | |
| if len(filter_conditions) == 1: | |
| where_filter = filter_conditions[0] | |
| elif len(filter_conditions) > 1: | |
| where_filter = {"$and": filter_conditions} | |
| emb_engine = self.get_embeddings(subject=subject, text_sample=query) | |
| query_embedding = emb_engine.embed_query(query) | |
| try: | |
| count = collection.count() | |
| if count == 0: | |
| return [] | |
| results = collection.query( | |
| query_embeddings=[query_embedding], | |
| n_results=min(n_results, count), | |
| where=where_filter | |
| ) | |
| # If no results matched with strict filters, fall back to class/subject/lesson (widest lesson scope) | |
| if (not results or not results.get("documents") or not results["documents"][0]): | |
| fallback_conditions = [] | |
| if class_id: | |
| fallback_conditions.append({"class_id": str(class_id)}) | |
| if subject: | |
| fallback_conditions.append({"subject": str(subject)}) | |
| if effective_lesson: | |
| fallback_conditions.append({"chapter": str(effective_lesson)}) | |
| fb_filter = {"$and": fallback_conditions} if len(fallback_conditions) > 1 else (fallback_conditions[0] if fallback_conditions else None) | |
| if fb_filter: | |
| results = collection.query( | |
| query_embeddings=[query_embedding], | |
| n_results=min(n_results, count), | |
| where=fb_filter | |
| ) | |
| except Exception: | |
| return [] | |
| matched_docs = [] | |
| if results and results.get("documents") and results["documents"][0]: | |
| docs = results["documents"][0] | |
| metas = results["metadatas"][0] if results.get("metadatas") else [{}] * len(docs) | |
| for doc_text, meta in zip(docs, metas): | |
| matched_docs.append({ | |
| "content": doc_text, | |
| "metadata": meta | |
| }) | |
| return matched_docs | |
| def get_curriculum( | |
| self, | |
| school_id: str, | |
| board_id: str, | |
| class_id: Optional[str] = None | |
| ) -> Dict[str, Any]: | |
| """Extracts available hierarchy (Subjects -> Chapters/Lessons -> Topics) for a tenant.""" | |
| collection = self.get_tenant_collection(school_id, board_id) | |
| where_filter = {"class_id": str(class_id)} if class_id else None | |
| try: | |
| data = collection.get(where=where_filter, include=["metadatas"]) | |
| except Exception: | |
| return {} | |
| hierarchy: Dict[str, Dict[str, set]] = {} | |
| if data and data.get("metadatas"): | |
| for meta in data["metadatas"]: | |
| if not meta: | |
| continue | |
| subj = meta.get("subject", "General") | |
| chap = meta.get("lesson") or meta.get("chapter", "General") | |
| top = meta.get("topic", "General") | |
| if subj not in hierarchy: | |
| hierarchy[subj] = {} | |
| if chap not in hierarchy[subj]: | |
| hierarchy[subj][chap] = set() | |
| hierarchy[subj][chap].add(top) | |
| result = {} | |
| for s, chaps in hierarchy.items(): | |
| result[s] = {} | |
| for c, tops in chaps.items(): | |
| clean_tops = [t for t in tops if t.strip().lower() != c.strip().lower()] | |
| result[s][c] = sorted(clean_tops) | |
| return result | |
| def get_erp_curriculum( | |
| self, | |
| school_id: str, | |
| board_id: str, | |
| class_id: Optional[str] = None | |
| ) -> Dict[str, Any]: | |
| """ | |
| Extracts full ERP hierarchy: | |
| Class -> Section -> Subject -> Sub-Subject -> Lesson -> Topics | |
| """ | |
| collection = self.get_tenant_collection(school_id, board_id) | |
| where_filter = {"class_id": str(class_id)} if class_id else None | |
| try: | |
| data = collection.get(where=where_filter, include=["metadatas"]) | |
| except Exception: | |
| return {} | |
| tree: Dict[str, Dict[str, Dict[str, Dict[str, Dict[str, set]]]]] = {} | |
| if data and data.get("metadatas"): | |
| for meta in data["metadatas"]: | |
| if not meta: | |
| continue | |
| cls = str(meta.get("class_id", class_id or "General")) | |
| sec = str(meta.get("section_id", "A")) | |
| subj = meta.get("subject", "General") | |
| sub_subj = meta.get("sub_subject", "General") | |
| lesson = meta.get("lesson") or meta.get("chapter", "General") | |
| top = meta.get("topic", "General") | |
| tree.setdefault(cls, {}).setdefault(sec, {}).setdefault(subj, {}).setdefault(sub_subj, {}).setdefault(lesson, set()).add(top) | |
| # Convert sets to sorted lists, stripping redundant lesson names | |
| clean_tree = {} | |
| for c, secs in tree.items(): | |
| clean_tree[c] = {} | |
| for sec, subjs in secs.items(): | |
| clean_tree[c][sec] = {} | |
| for s, sub_subjs in subjs.items(): | |
| clean_tree[c][sec][s] = {} | |
| for ss, lessons in sub_subjs.items(): | |
| clean_tree[c][sec][s][ss] = {} | |
| for les, tops in lessons.items(): | |
| clean_tops = [t for t in tops if t.strip().lower() != les.strip().lower()] | |
| clean_tree[c][sec][s][ss][les] = sorted(clean_tops) | |
| return clean_tree | |
| def get_training_status( | |
| self, | |
| school_id: str, | |
| board_id: str | |
| ) -> Dict[str, Any]: | |
| """Returns chunk counts grouped by class/subject/chapter for the teacher dashboard.""" | |
| collection = self.get_tenant_collection(school_id, board_id) | |
| try: | |
| data = collection.get(include=["metadatas"]) | |
| except Exception: | |
| return {} | |
| status: Dict[str, Dict[str, Dict[str, Any]]] = {} | |
| topics_map: Dict[str, set] = {} | |
| if data and data.get("metadatas"): | |
| for meta in data["metadatas"]: | |
| if not meta: | |
| continue | |
| cls = meta.get("class_id", "?") | |
| subj = meta.get("subject", "?") | |
| chap = meta.get("chapter", "?") | |
| key = f"Class {cls}" | |
| ckey = f"{key}_{subj}_{chap}" | |
| if key not in status: | |
| status[key] = {} | |
| if subj not in status[key]: | |
| status[key][subj] = {} | |
| if chap not in status[key][subj]: | |
| status[key][subj][chap] = { | |
| "chunk_count": 0, | |
| "source": meta.get("source", ""), | |
| "topics": [] | |
| } | |
| topics_map[ckey] = set() | |
| status[key][subj][chap]["chunk_count"] += 1 | |
| top = meta.get("topic") | |
| if top and top.strip(): | |
| topics_map[ckey].add(top.strip()) | |
| # Finalize topics: convert sets to sorted lists, filtering out chapter title if subtopics exist | |
| for cls_label, subjects in status.items(): | |
| for subj, chapters in subjects.items(): | |
| for chap, info in chapters.items(): | |
| ckey = f"{cls_label}_{subj}_{chap}" | |
| all_tops = topics_map.get(ckey, set()) | |
| # Filter out redundant chapter name if other subtopics are present | |
| clean_tops = [t for t in all_tops if t.lower() != chap.lower()] | |
| info["topics"] = sorted(clean_tops) if clean_tops else sorted(list(all_tops)) | |
| return status | |
| def list_tenants(self) -> List[str]: | |
| """List all tenant collections (admin utility).""" | |
| collections = self.client.list_collections() | |
| return [c.name for c in collections if c.name.startswith("school_")] | |
| vector_manager = VectorStoreManager() | |