""" ByteAstra — ChromaDB RAG service. Handles: - Initialising the ChromaDB client and per-domain collections - Embedding queries using sentence-transformers - Retrieving the top-K most relevant chunks for a query - Indexing new documents (used by ingest.py) """ from __future__ import annotations import logging from dataclasses import dataclass from functools import lru_cache import chromadb from chromadb import Settings as ChromaSettings from chromadb.api.types import Documents, EmbeddingFunction, Embeddings import torch from transformers import AutoTokenizer, AutoModel from app.config import get_settings from app.schemas import Citation logger = logging.getLogger(__name__) settings = get_settings() class TransformersEmbeddingFunction(EmbeddingFunction[Documents]): def __init__(self, model_name: str = "sentence-transformers/all-MiniLM-L6-v2", device: str = "cpu"): self.device = device if "/" not in model_name: model_name = f"sentence-transformers/{model_name}" # Pin embedding CPU threads to 1 to save the second core for the LLM torch.set_num_threads(1) try: import os os.environ["HF_HUB_OFFLINE"] = "1" os.environ["TRANSFORMERS_OFFLINE"] = "1" self.tokenizer = AutoTokenizer.from_pretrained(model_name, local_files_only=True) self.model = AutoModel.from_pretrained(model_name, local_files_only=True).to(self.device) logger.info("Loaded embedding model %s from local files.", model_name) except Exception: logger.info("Local files not found for %s. Enabling online download...", model_name) import os os.environ["HF_HUB_OFFLINE"] = "0" os.environ["TRANSFORMERS_OFFLINE"] = "0" self.tokenizer = AutoTokenizer.from_pretrained(model_name, local_files_only=False) self.model = AutoModel.from_pretrained(model_name, local_files_only=False).to(self.device) logger.info("Successfully downloaded and loaded embedding model %s.", model_name) def __call__(self, input: Documents) -> Embeddings: encoded_input = self.tokenizer(input, padding=True, truncation=True, return_tensors='pt').to(self.device) with torch.no_grad(): model_output = self.model(**encoded_input) token_embeddings = model_output[0] attention_mask = encoded_input['attention_mask'] input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float() sum_embeddings = torch.sum(token_embeddings * input_mask_expanded, 1) sum_mask = torch.clamp(input_mask_expanded.sum(1), min=1e-9) embeddings = sum_embeddings / sum_mask import torch.nn.functional as F embeddings = F.normalize(embeddings, p=2, dim=1) return embeddings.cpu().numpy().tolist() @dataclass class RetrievedChunk: chunk_id: str content: str source: str chapter: str | None section: str | None relevance_score: float # cosine similarity (0–1) def to_citation(self) -> Citation: return Citation( chunk_id=self.chunk_id, source=self.source, chapter=self.chapter, section=self.section, relevance_score=self.relevance_score, content=self.content, ) @lru_cache(maxsize=1) def _get_embedding_fn() -> TransformersEmbeddingFunction: logger.info("Loading embedding model: %s", settings.embedding_model) return TransformersEmbeddingFunction(model_name=settings.embedding_model, device="cpu") @lru_cache(maxsize=1) def _get_chroma_client() -> chromadb.ClientAPI: logger.info("Initialising ChromaDB at: %s", settings.chroma_persist_path) return chromadb.PersistentClient( path=settings.chroma_persist_path, settings=ChromaSettings(anonymized_telemetry=False), ) def get_or_create_collection(collection_name: str) -> chromadb.Collection: """Get or create a ChromaDB collection for a specific domain.""" client = _get_chroma_client() return client.get_or_create_collection( name=collection_name, metadata={"hnsw:space": "cosine"}, ) def retrieve( query: str, collection_name: str, top_k: int = 5, relevance_threshold: float = 0.45, ) -> list[RetrievedChunk]: """ Retrieve the most relevant chunks from the domain's collection. Only returns chunks with similarity score >= relevance_threshold. Returns an empty list if nothing meets the threshold — the engine will treat this as an out-of-syllabus query. """ import re collection = get_or_create_collection(collection_name) # Compute query embedding using the custom function embedding_fn = _get_embedding_fn() query_embeddings = embedding_fn([query]) # Fetch candidates for lexical re-ranking (reduced for speed on CPU) candidate_k = max(top_k * 2, 6) results = collection.query( query_embeddings=query_embeddings, n_results=candidate_k, include=["documents", "metadatas", "distances"], ) candidates: list[tuple[RetrievedChunk, float]] = [] if not results["ids"] or not results["ids"][0]: return [] # Extract significant query words for lexical boosting q_words = re.findall(r"\b[a-zA-Z]{3,}\b", query.lower()) stop_words = {"how", "does", "that", "relate", "to", "what", "is", "are", "the", "and", "in", "of", "a", "an", "with", "about", "for", "its", "why", "who"} sig_words = [w for w in q_words if w not in stop_words] for i, chunk_id in enumerate(results["ids"][0]): # ChromaDB cosine distance → similarity: similarity = 1 - distance distance = results["distances"][0][i] similarity = 1.0 - distance if similarity < relevance_threshold: logger.debug("Chunk %s below threshold (%.3f < %.3f) — skipped", chunk_id, similarity, relevance_threshold) continue content = results["documents"][0][i] meta = results["metadatas"][0][i] or {} # Calculate lexical boost based on query term frequency/overlap boost = 0.0 if sig_words: content_lower = content.lower() for w in sig_words: if re.search(r"\b" + re.escape(w) + r"\b", content_lower): boost += 0.03 boosted_score = similarity + boost chunk = RetrievedChunk( chunk_id=chunk_id, content=content, source=meta.get("source", "Unknown Source"), chapter=meta.get("chapter"), section=meta.get("section"), relevance_score=round(similarity, 4), ) candidates.append((chunk, boosted_score)) # Sort candidates by boosted score descending and take top_k candidates.sort(key=lambda x: x[1], reverse=True) selected_chunks = [x[0] for x in candidates[:top_k]] logger.info("Retrieved %d relevant chunks for query (candidates=%d, threshold=%.2f)", len(selected_chunks), len(candidates), relevance_threshold) return selected_chunks def index_chunks( collection_name: str, chunks: list[dict], ) -> None: """ Index a batch of text chunks into the domain's ChromaDB collection. Each chunk dict must have: - id: str (unique, e.g. "ayurveda_charaka_ch1_000") - text: str - source: str - chapter: str | None - section: str | None """ collection = get_or_create_collection(collection_name) ids = [c["id"] for c in chunks] documents = [c["text"] for c in chunks] metadatas = [ { "source": c.get("source", ""), "chapter": c.get("chapter") or "", "section": c.get("section") or "", } for c in chunks ] # Compute embeddings manually to bypass ChromaDB internal conflicts embedding_fn = _get_embedding_fn() embeddings = embedding_fn(documents) collection.upsert(ids=ids, embeddings=embeddings, documents=documents, metadatas=metadatas) logger.info("Indexed %d chunks into collection '%s'", len(chunks), collection_name)