ai-backend / backend /vector_store.py
arinsujal's picture
feat: clean project deployment without cache or binaries
a4a8767
Raw History Blame Contribute Delete
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)
@property
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()