from __future__ import annotations import hashlib import json import math import os import re import sqlite3 import time import unicodedata import uuid from collections import Counter, defaultdict from dataclasses import asdict, dataclass, field from pathlib import Path from typing import Iterable DATA_DIR = Path("data") DB_PATH = DATA_DIR / "smartnotes.sqlite" DEFAULT_EMBED_DIM = 384 DEFAULT_TOP_K = 50 UNRELATED_PDF_ANSWER = "Sorry, I can't give any answer because your question is not related to the PDF." @dataclass class ParentChunk: id: str document_id: str index: int text: str token_count: int metadata: dict @dataclass class ChildChunk: id: str parent_id: str document_id: str index: int strategy: str text: str token_count: int embedding: list[float] = field(default_factory=list) metadata: dict = field(default_factory=dict) @dataclass class RetrievalHit: child_id: str parent_id: str document_id: str strategy: str child_text: str parent_text: str vector_score: float = 0.0 bm25_score: float = 0.0 hybrid_score: float = 0.0 rerank_score: float = 0.0 grade: str = "unknown" metadata: dict = field(default_factory=dict) class SmartNotesRepository: """PostgreSQL-first repository with SQLite fallback for local UI testing.""" def __init__(self) -> None: DATA_DIR.mkdir(exist_ok=True) self.postgres_dsn = os.getenv("POSTGRES_DSN") self.backend = "sqlite" self._pg = None if self.postgres_dsn: try: import psycopg self._pg = psycopg.connect(self.postgres_dsn) self.backend = "postgresql" except Exception: self._pg = None self.backend = "sqlite" self._sqlite = sqlite3.connect(DB_PATH, check_same_thread=False) self._sqlite.row_factory = sqlite3.Row self._init_sqlite() def _init_sqlite(self) -> None: self._sqlite.executescript( """ create table if not exists documents ( id text primary key, file_name text, status text, metadata text, cleaned_text text, created_at real, updated_at real ); create table if not exists parent_chunks ( id text primary key, document_id text, chunk_index integer, text text, token_count integer, metadata text ); create table if not exists child_chunks ( id text primary key, parent_id text, document_id text, chunk_index integer, strategy text, text text, token_count integer, embedding text, metadata text ); create table if not exists retrieval_logs ( id text primary key, document_id text, query text, rewritten_query text, status text, answer text, hits text, created_at real ); create table if not exists feedback ( id text primary key, document_id text, query text, rating integer, comment text, created_at real ); """ ) self._sqlite.commit() def create_document(self, file_name: str, metadata: dict) -> str: document_id = str(uuid.uuid4()) now = time.time() self._sqlite.execute( """ insert into documents (id, file_name, status, metadata, cleaned_text, created_at, updated_at) values (?, ?, ?, ?, ?, ?, ?) """, (document_id, file_name, "processing", json.dumps(metadata), "", now, now), ) self._sqlite.commit() return document_id def update_document_status(self, document_id: str, status: str) -> None: self._sqlite.execute( "update documents set status = ?, updated_at = ? where id = ?", (status, time.time(), document_id), ) self._sqlite.commit() def save_cleaned_text(self, document_id: str, cleaned_text: str, metadata: dict) -> None: self._sqlite.execute( "update documents set cleaned_text = ?, metadata = ?, updated_at = ? where id = ?", (cleaned_text, json.dumps(metadata), time.time(), document_id), ) self._sqlite.commit() def save_chunks(self, parents: list[ParentChunk], children: list[ChildChunk]) -> None: self._sqlite.executemany( """ insert or replace into parent_chunks (id, document_id, chunk_index, text, token_count, metadata) values (?, ?, ?, ?, ?, ?) """, [ ( chunk.id, chunk.document_id, chunk.index, chunk.text, chunk.token_count, json.dumps(chunk.metadata), ) for chunk in parents ], ) self._sqlite.executemany( """ insert or replace into child_chunks (id, parent_id, document_id, chunk_index, strategy, text, token_count, embedding, metadata) values (?, ?, ?, ?, ?, ?, ?, ?, ?) """, [ ( chunk.id, chunk.parent_id, chunk.document_id, chunk.index, chunk.strategy, chunk.text, chunk.token_count, json.dumps(chunk.embedding), json.dumps(chunk.metadata), ) for chunk in children ], ) self._sqlite.commit() def get_document(self, document_id: str) -> dict | None: row = self._sqlite.execute( "select * from documents where id = ?", (document_id,), ).fetchone() if not row: return None return { "id": row["id"], "file_name": row["file_name"], "status": row["status"], "metadata": json.loads(row["metadata"] or "{}"), "cleaned_text": row["cleaned_text"] or "", "storage_backend": self.backend, } def get_children(self, document_id: str) -> list[ChildChunk]: rows = self._sqlite.execute( "select * from child_chunks where document_id = ? order by chunk_index", (document_id,), ).fetchall() return [ ChildChunk( id=row["id"], parent_id=row["parent_id"], document_id=row["document_id"], index=row["chunk_index"], strategy=row["strategy"], text=row["text"], token_count=row["token_count"], embedding=json.loads(row["embedding"] or "[]"), metadata=json.loads(row["metadata"] or "{}"), ) for row in rows ] def get_parent_map(self, document_id: str) -> dict[str, ParentChunk]: rows = self._sqlite.execute( "select * from parent_chunks where document_id = ?", (document_id,), ).fetchall() return { row["id"]: ParentChunk( id=row["id"], document_id=row["document_id"], index=row["chunk_index"], text=row["text"], token_count=row["token_count"], metadata=json.loads(row["metadata"] or "{}"), ) for row in rows } def save_retrieval_log( self, document_id: str, query: str, rewritten_query: str, status: str, answer: str, hits: list[RetrievalHit], ) -> str: log_id = str(uuid.uuid4()) self._sqlite.execute( """ insert into retrieval_logs (id, document_id, query, rewritten_query, status, answer, hits, created_at) values (?, ?, ?, ?, ?, ?, ?, ?) """, ( log_id, document_id, query, rewritten_query, status, answer, json.dumps([self._hit_payload(hit) for hit in hits]), time.time(), ), ) self._sqlite.commit() return log_id def save_feedback(self, document_id: str, query: str, rating: int, comment: str) -> str: feedback_id = str(uuid.uuid4()) self._sqlite.execute( """ insert into feedback (id, document_id, query, rating, comment, created_at) values (?, ?, ?, ?, ?, ?) """, (feedback_id, document_id, query, rating, comment, time.time()), ) self._sqlite.commit() return feedback_id def _hit_payload(self, hit: RetrievalHit) -> dict: data = asdict(hit) data["child_text"] = hit.child_text[:900] data["parent_text"] = hit.parent_text[:1400] return data class TextCleaner: def clean(self, text: str) -> str: try: import ftfy text = ftfy.fix_text(text) except Exception: pass text = unicodedata.normalize("NFKC", text or "") text = text.replace("\x00", "") text = re.sub(r"---\s*Page\s+\d+\s*---", "\n", text, flags=re.IGNORECASE) text = re.sub(r"[ \t]+", " ", text) text = re.sub(r"(\w)-\n(\w)", r"\1\2", text) text = re.sub(r"\n(?=[a-z])", " ", text) text = re.sub(r"\n{3,}", "\n\n", text) return "\n".join(line.strip() for line in text.splitlines()).strip() class Chunker: def parent_chunks(self, document_id: str, text: str, metadata: dict) -> list[ParentChunk]: chunks = self._recursive_split(text, chunk_size=1400, overlap=180) return [ ParentChunk( id=str(uuid.uuid4()), document_id=document_id, index=index, text=chunk, token_count=self._token_count(chunk), metadata={**metadata, "chunk_type": "parent"}, ) for index, chunk in enumerate(chunks) ] def child_chunks(self, parents: list[ParentChunk]) -> list[ChildChunk]: children: list[ChildChunk] = [] index = 0 for parent in parents: for strategy, chunks in ( ("fixed", self._fixed_split(parent.text, chunk_size=380, overlap=70)), ("recursive", self._recursive_split(parent.text, chunk_size=520, overlap=90)), ("semantic", self._semantic_split(parent.text, target_size=520)), ): for text in chunks: children.append( ChildChunk( id=str(uuid.uuid4()), parent_id=parent.id, document_id=parent.document_id, index=index, strategy=strategy, text=text, token_count=self._token_count(text), metadata={ **parent.metadata, "chunk_type": "child", "parent_index": parent.index, "strategy": strategy, }, ) ) index += 1 return children def _fixed_split(self, text: str, chunk_size: int, overlap: int) -> list[str]: words = text.split() chunks = [] step = max(chunk_size - overlap, 1) for start in range(0, len(words), step): chunk = " ".join(words[start : start + chunk_size]).strip() if chunk: chunks.append(chunk) return chunks def _recursive_split(self, text: str, chunk_size: int, overlap: int) -> list[str]: separators = ["\n\n", "\n", ". ", " "] pieces = self._split_by_separators(text, separators, chunk_size) chunks = [] current: list[str] = [] current_len = 0 for piece in pieces: words = piece.split() if current and current_len + len(words) > chunk_size: chunks.append(" ".join(current).strip()) current = current[-overlap:] if overlap else [] current_len = len(current) current.extend(words) current_len += len(words) if current: chunks.append(" ".join(current).strip()) return [chunk for chunk in chunks if chunk] def _semantic_split(self, text: str, target_size: int) -> list[str]: sentences = re.split(r"(?<=[.!?])\s+", text) chunks = [] current = [] current_len = 0 for sentence in sentences: count = self._token_count(sentence) if current and current_len + count > target_size: chunks.append(" ".join(current).strip()) current = [] current_len = 0 current.append(sentence) current_len += count if current: chunks.append(" ".join(current).strip()) return [chunk for chunk in chunks if chunk] def _split_by_separators(self, text: str, separators: list[str], chunk_size: int) -> list[str]: if self._token_count(text) <= chunk_size or not separators: return [text] separator = separators[0] output = [] for part in text.split(separator): output.extend(self._split_by_separators(part, separators[1:], chunk_size)) return output def _token_count(self, text: str) -> int: return len(text.split()) class EmbeddingService: def __init__(self) -> None: self.model_name = os.getenv("EMBEDDING_MODEL", "BAAI/bge-small-en-v1.5") self.dim = DEFAULT_EMBED_DIM self.backend = "hash" self._model = None if os.getenv("USE_LOCAL_EMBEDDING_MODEL", "0") == "1": try: from sentence_transformers import SentenceTransformer self._model = SentenceTransformer(self.model_name) self.backend = self.model_name self.dim = int(self._model.get_sentence_embedding_dimension()) except Exception: self._model = None def embed_texts(self, texts: list[str]) -> list[list[float]]: if self._model is not None: vectors = self._model.encode(texts, normalize_embeddings=True).tolist() return [[float(value) for value in vector] for vector in vectors] return [self._hash_embedding(text) for text in texts] def embed_query(self, query: str) -> list[float]: return self.embed_texts([query])[0] def _hash_embedding(self, text: str) -> list[float]: vector = [0.0] * self.dim tokens = tokenize(text) for token in tokens: digest = hashlib.sha256(token.encode("utf-8")).digest() index = int.from_bytes(digest[:4], "big") % self.dim sign = 1.0 if digest[4] % 2 == 0 else -1.0 vector[index] += sign norm = math.sqrt(sum(value * value for value in vector)) or 1.0 return [value / norm for value in vector] class QdrantVectorStore: def __init__(self) -> None: self.backend = "sqlite-vector-fallback" self.collection = os.getenv("QDRANT_COLLECTION", "smartnotes_chunks") self.client = None self.url = os.getenv("QDRANT_URL") if self.url: try: from qdrant_client import QdrantClient self.client = QdrantClient(url=self.url, api_key=os.getenv("QDRANT_API_KEY")) self.backend = "qdrant" except Exception: self.client = None def upsert(self, children: list[ChildChunk]) -> None: if self.client is None or not children: return try: from qdrant_client.models import Distance, PointStruct, VectorParams dim = len(children[0].embedding) existing = [item.name for item in self.client.get_collections().collections] if self.collection not in existing: self.client.create_collection( collection_name=self.collection, vectors_config=VectorParams(size=dim, distance=Distance.COSINE), ) self.client.upsert( collection_name=self.collection, points=[ PointStruct( id=chunk.id, vector=chunk.embedding, payload={ "document_id": chunk.document_id, "parent_id": chunk.parent_id, "strategy": chunk.strategy, }, ) for chunk in children ], ) except Exception: self.client = None self.backend = "sqlite-vector-fallback" def search(self, query_vector: list[float], children: list[ChildChunk], top_k: int) -> list[tuple[str, float]]: if self.client is not None: try: result = self.client.search( collection_name=self.collection, query_vector=query_vector, limit=top_k, ) return [(str(point.id), float(point.score)) for point in result] except Exception: pass scored = [ (chunk.id, cosine_similarity(query_vector, chunk.embedding)) for chunk in children ] return sorted(scored, key=lambda item: item[1], reverse=True)[:top_k] class BM25Index: def search(self, query: str, children: list[ChildChunk], top_k: int) -> list[tuple[str, float]]: query_tokens = tokenize(query) docs = [tokenize(chunk.text) for chunk in children] if not docs: return [] avgdl = sum(len(doc) for doc in docs) / len(docs) df = Counter() for doc in docs: for token in set(doc): df[token] += 1 scores = [] for chunk, doc in zip(children, docs): score = self._score_doc(query_tokens, doc, df, len(docs), avgdl) scores.append((chunk.id, score)) return sorted(scores, key=lambda item: item[1], reverse=True)[:top_k] def _score_doc( self, query_tokens: list[str], doc: list[str], df: Counter, total_docs: int, avgdl: float, ) -> float: counts = Counter(doc) k1 = 1.5 b = 0.75 score = 0.0 for token in query_tokens: if not counts[token]: continue idf = math.log(1 + (total_docs - df[token] + 0.5) / (df[token] + 0.5)) numerator = counts[token] * (k1 + 1) denominator = counts[token] + k1 * (1 - b + b * len(doc) / (avgdl or 1)) score += idf * numerator / denominator return score class Reranker: def __init__(self) -> None: self.model_name = os.getenv("RERANKER_MODEL", "BAAI/bge-reranker-base") self.backend = "token-overlap" self._model = None if os.getenv("USE_LOCAL_RERANKER_MODEL", "0") == "1": try: from sentence_transformers import CrossEncoder self._model = CrossEncoder(self.model_name) self.backend = self.model_name except Exception: self._model = None def rerank(self, query: str, hits: list[RetrievalHit]) -> list[RetrievalHit]: if self._model is not None and hits: scores = self._model.predict([(query, hit.child_text) for hit in hits]) for hit, score in zip(hits, scores): hit.rerank_score = float(score) else: query_tokens = set(tokenize(query)) for hit in hits: chunk_tokens = set(tokenize(hit.child_text)) hit.rerank_score = len(query_tokens & chunk_tokens) / max(len(query_tokens), 1) return sorted(hits, key=lambda hit: hit.rerank_score, reverse=True) class LLMService: def __init__(self) -> None: load_local_env() self.provider = os.getenv("ANSWER_PROVIDER", "gemini") self.model = os.getenv("ANSWER_MODEL", "gemini-2.5-flash") self._gemini_clients = {} self._gemini_key_cursor = 0 self._gemini_quota_blocked_until: dict[str, float] = {} self.gemini_timeout_ms = int( min(env_float("GEMINI_TIMEOUT", default=20.0, minimum=5.0), 120.0) * 1000 ) self.gemini_key_quota_cooldown = min( env_float("GEMINI_KEY_QUOTA_COOLDOWN", default=300.0, minimum=0.0), 86400.0, ) def rewrite_query(self, query: str) -> str: if len(query.split()) >= 4: return query.strip() prompt = f"Rewrite this user question into a precise PDF search query. Return only the query.\nQuestion: {query}" return self._generate(prompt, fallback=query).strip() or query def grade_retrieval(self, query: str, hits: list[RetrievalHit]) -> str: if not hits: return "bad" best = max(hit.rerank_score for hit in hits) lexical = max(hit.hybrid_score for hit in hits) return "good" if best >= 0.08 or lexical >= 0.12 else "bad" def answer(self, query: str, parent_contexts: list[ParentChunk]) -> str: if not parent_contexts: return UNRELATED_PDF_ANSWER mode = self._answer_mode(query) fallback = self._fallback_answer(query, parent_contexts, mode) context = "\n\n".join( f"[Source {index + 1} | parent={chunk.id}]\n{chunk.text}" for index, chunk in enumerate(parent_contexts) ) prompt = ( "You are SmartNotes AI. Answer using only the PDF context.\n" "Rules:\n" "- Give a direct answer first.\n" "- Do not dump the whole context.\n" "- If the user asks 'what is' or 'define', give only the definition in 1-2 sentences.\n" "- Answer in normal readable text only.\n" "- Do not return equations, raw calculations, URLs, or OCR garbage unless the user explicitly asks for them.\n" "- If the user asks in detail, explain with key points, types, examples, and relevant facts from the PDF.\n" "- Remove page markers and OCR noise.\n" f"- If the clean text answer is not present or the question is not related to the PDF, say exactly: {UNRELATED_PDF_ANSWER}\n" "- End with citation like [Source 1].\n\n" f"Question: {query}\n\nContext:\n{context}\n\nAnswer:" ) return self._postprocess_answer(self._generate(prompt, fallback=fallback)) def _generate(self, prompt: str, fallback: str) -> str: if self.provider == "openai" and os.getenv("OPENAI_API_KEY"): try: from openai import OpenAI client = OpenAI(api_key=os.getenv("OPENAI_API_KEY")) response = client.chat.completions.create( model=os.getenv("OPENAI_MODEL", "gpt-4.1-mini"), messages=[{"role": "user", "content": prompt}], ) return response.choices[0].message.content or fallback except Exception: return fallback api_keys = gemini_api_keys() if api_keys: try: from google import genai from google.genai import types except Exception: return fallback for api_key in self._gemini_key_attempt_order(api_keys): try: client = self._gemini_clients.get(api_key) if client is None: client = genai.Client( api_key=api_key, http_options=types.HttpOptions(timeout=self.gemini_timeout_ms), ) self._gemini_clients[api_key] = client response = client.models.generate_content( model=self.model, contents=prompt, ) self._advance_gemini_key_cursor(api_keys, api_key) return (getattr(response, "text", None) or fallback).strip() except Exception as exc: if is_gemini_quota_error(exc): self._mark_gemini_key_quota_exhausted(api_keys, api_key) continue return fallback def _gemini_key_attempt_order(self, api_keys: list[str]) -> list[str]: now = time.time() active_keys = [ api_key for api_key in api_keys if self._gemini_quota_blocked_until.get(api_key, 0.0) <= now ] if not active_keys: active_keys = api_keys active_key_set = set(active_keys) start = self._gemini_key_cursor % len(api_keys) ordered_keys = api_keys[start:] + api_keys[:start] return [api_key for api_key in ordered_keys if api_key in active_key_set] def _advance_gemini_key_cursor(self, api_keys: list[str], api_key: str) -> None: try: self._gemini_key_cursor = (api_keys.index(api_key) + 1) % len(api_keys) except ValueError: self._gemini_key_cursor = 0 def _mark_gemini_key_quota_exhausted(self, api_keys: list[str], api_key: str) -> None: self._gemini_quota_blocked_until[api_key] = time.time() + self.gemini_key_quota_cooldown self._advance_gemini_key_cursor(api_keys, api_key) def _fallback_answer(self, query: str, parent_contexts: list[ParentChunk], mode: str = "auto") -> str: query_tokens = set(tokenize(query)) targets = self._query_targets(query) text = "\n".join(chunk.text for chunk in parent_contexts) if mode == "definition": definition = self._definition_answer(text, targets) if definition: return self._postprocess_answer(definition + " [Source 1]") return UNRELATED_PDF_ANSWER sentences = self._candidate_sentences(text) ranked = sorted( sentences, key=lambda sentence: self._sentence_score(query_tokens, targets, sentence), reverse=True, ) useful = [] for sentence in ranked: sentence_score = self._sentence_score(query_tokens, targets, sentence) if mode == "detailed" and sentence_score <= 0: continue cleaned = self._clean_answer_sentence(sentence) if not cleaned or len(cleaned.split()) < 5 or self._is_noisy_answer_sentence(cleaned): continue if cleaned in useful: continue useful.append(cleaned) limit = 1 if mode == "definition" else 5 if len(useful) >= limit: break if not useful: return UNRELATED_PDF_ANSWER if mode == "detailed" and len(useful) > 1: answer = useful[0] bullets = "\n".join(f"- {item}" for item in useful[1:]) return self._postprocess_answer(f"{answer}\n\nKey points:\n{bullets}\n\n[Source 1]") return self._postprocess_answer(" ".join(useful) + " [Source 1]") def _candidate_sentences(self, text: str) -> list[str]: text = re.sub(r"---\s*Page\s+\d+\s*---", " ", text, flags=re.IGNORECASE) text = re.sub(r"\bStatistics\s+Statistic\s+is\b", "Statistics is", text, flags=re.IGNORECASE) text = re.sub(r"\s+", " ", text) candidates = re.split(r"(?<=[.!?])\s+|(?:\n)+", text) return [candidate.strip() for candidate in candidates if candidate.strip()] def _sentence_score(self, query_tokens: set[str], targets: set[str], sentence: str) -> float: sentence_tokens = set(tokenize(sentence)) overlap = len(query_tokens & sentence_tokens) definition_bonus = 0.0 lower = sentence.lower() if self._is_question_like_sentence(lower): definition_bonus -= 6.0 for target in targets: for variant in self._term_variants(target): if self._is_subtype_sentence(lower, variant, target): definition_bonus -= 8.0 for token in targets | query_tokens: token_variants = self._term_variants(token) for variant in token_variants: if not variant: continue if re.search(rf"\b{re.escape(variant)}s?\b\s+(is|are|means|refers|involves)", lower): definition_bonus += 5.0 if lower.startswith(variant): definition_bonus += 1.0 if "branch of" in lower or "defined as" in lower or "refers to" in lower: definition_bonus += 3.0 if lower.startswith("in practice"): definition_bonus -= 2.5 length_penalty = max((len(sentence.split()) - 45) / 45, 0) return overlap + definition_bonus - length_penalty def _clean_answer_sentence(self, sentence: str) -> str: sentence = re.sub(r"---\s*Page\s+\d+\s*---", " ", sentence, flags=re.IGNORECASE) sentence = re.sub(r"https?://\S+|www\.\S+", " ", sentence, flags=re.IGNORECASE) sentence = re.sub(r"\bStatistics\s+Statistic\s+is\b", "Statistics is", sentence, flags=re.IGNORECASE) sentence = re.sub(r"^Serializability\s+(A\s+schedule\b)", r"\1", sentence, flags=re.IGNORECASE) sentence = re.sub(r"\s+", " ", sentence).strip() sentence = re.sub(r"^(statistics\s+){2,}", "Statistics ", sentence, flags=re.IGNORECASE) return sentence def _answer_mode(self, query: str) -> str: lower = query.lower() detail_markers = ( "detail", "in detail", "explain", "describe", "types", "example", "examples", "advantages", "disadvantages", "steps", "full", "sab kuch", "pura", "briefly explain", ) if any(marker in lower for marker in detail_markers): return "detailed" if re.search(r"\b(what is|what are|define|meaning of)\b", lower): return "definition" return "auto" def _query_targets(self, query: str) -> set[str]: tokens = tokenize(query) stop_words = { "what", "is", "are", "the", "a", "an", "of", "define", "meaning", "explain", "tell", "me", "about", } filtered = [token for token in tokens if token not in stop_words] targets = set(filtered) if len(filtered) >= 2: targets.add(" ".join(filtered)) return targets def _term_variants(self, token: str) -> set[str]: variants = {token, token.rstrip("s")} if token.endswith("ability"): variants.add(token[: -len("ability")] + "able") if token.endswith("ibility"): variants.add(token[: -len("ibility")] + "ible") if token.endswith("tion"): variants.add(token[:-3] + "e") return {variant for variant in variants if variant} def _definition_answer(self, text: str, targets: set[str]) -> str: normalized = re.sub(r"---\s*Page\s+\d+\s*---", " ", text, flags=re.IGNORECASE) normalized = re.sub(r"\s+", " ", normalized).strip() sorted_targets = sorted(targets, key=len, reverse=True) for target in sorted_targets: target_pattern = re.escape(target) patterns = [ rf"\b{target_pattern}\b\s+(is|are|means|refers to)\s+([^.?!]{{8,260}})[.?!]", rf"\b{target_pattern}\b\s*[:\-]\s*([^.?!]{{8,260}})[.?!]", rf"\b{target_pattern}\b\s+((?:a|an|the)\s+[^.?!]{{5,220}}?\bis\s+[^.?!]{{5,220}})[.?!]", rf"\b{target_pattern}\b\s+([^.?!]{{5,220}}?\bmeans\b[^.?!]{{5,220}})[.?!]", rf"\b{target_pattern}\b\s+([^.?!]{{5,220}}?\brefers to\b[^.?!]{{5,220}})[.?!]", ] for pattern in patterns: match = re.search(pattern, normalized, flags=re.IGNORECASE) if match: groups = [group for group in match.groups() if group and group.lower() not in {"is", "are", "means", "refers to"}] sentence = " ".join(groups).strip() sentence = self._normalize_definition_sentence(target, sentence) if sentence and not self._is_question_like_sentence(sentence.lower()) and not self._is_noisy_answer_sentence(sentence): return sentence for sentence in self._candidate_sentences(normalized): lower = sentence.lower() variants = self._term_variants(target) if any(variant in lower for variant in variants) and re.search(r"\b(is|are|means|refers to)\b", lower): if not self._is_question_like_sentence(lower) and not any( self._is_subtype_sentence(lower, variant, target) for variant in variants ): cleaned = self._clean_answer_sentence(sentence) if not self._is_noisy_answer_sentence(cleaned): return cleaned for variant in variants: adjective_pattern = rf"\b((?:a|an|the)\s+[^.?!]{{3,90}}?\bis\s+{re.escape(variant)}\s+if\s+[^.?!]{{8,220}})[.?!]?" match = re.search(adjective_pattern, sentence, flags=re.IGNORECASE) if match: cleaned = self._clean_answer_sentence(match.group(1)) if not self._is_noisy_answer_sentence(cleaned): return cleaned return "" def _normalize_definition_sentence(self, target: str, sentence: str) -> str: sentence = self._clean_answer_sentence(sentence) if not sentence: return "" if not re.search(r"\b(is|are|means|refers to)\b", sentence.lower()): if sentence[0].islower(): sentence = sentence[0].upper() + sentence[1:] return sentence if sentence[0].islower(): sentence = sentence[0].upper() + sentence[1:] return sentence def _is_question_like_sentence(self, lower_sentence: str) -> bool: return ( lower_sentence.endswith("?") or lower_sentence.startswith(("is the below", "is this", "are the below", "find ", "solve ")) ) def _is_subtype_sentence(self, lower_sentence: str, variant: str, target: str) -> bool: if " " in target: return False subtype_words = ("view", "conflict", "strict", "caseless", "lossless", "partial", "total") return any(re.search(rf"\b{word}\s+{re.escape(variant)}\b", lower_sentence) for word in subtype_words) def _is_noisy_answer_sentence(self, sentence: str) -> bool: text = re.sub(r"\[Source\s+\d+\]", "", sentence or "", flags=re.IGNORECASE).strip() if not text: return True lower = text.lower() if re.search(r"https?://|www\.|\b[a-z0-9.-]+\.(com|org|net|in|edu)\b", lower): return True if re.search(r"\b(p\s*=|x\s*=|y\s*=|sum\s*=)\b", lower): return True chars = [char for char in text if not char.isspace()] if not chars: return True symbol_chars = sum(1 for char in chars if not char.isalnum() and char not in ".,;:()[]-'") operator_chars = sum(1 for char in chars if char in "=+*/^") alpha_chars = sum(1 for char in chars if char.isalpha()) digit_chars = sum(1 for char in chars if char.isdigit()) if operator_chars >= 3 and digit_chars >= alpha_chars: return True if symbol_chars / max(len(chars), 1) > 0.22 and alpha_chars < 20: return True tokens = tokenize(text) if len(tokens) < 4: return True return False def _postprocess_answer(self, answer: str) -> str: answer = re.sub(r"---\s*Page\s+\d+\s*---", " ", answer or "", flags=re.IGNORECASE) answer = re.sub(r"https?://\S+|www\.\S+", " ", answer, flags=re.IGNORECASE) lines = [re.sub(r"[ \t]+", " ", line).strip() for line in answer.splitlines()] answer = "\n".join(line for line in lines if line).strip() if self._is_noisy_answer_sentence(re.sub(r"\[Source\s+\d+\]", "", answer, flags=re.IGNORECASE)): return UNRELATED_PDF_ANSWER if "[Source" not in answer and answer != UNRELATED_PDF_ANSWER: answer += " [Source 1]" words = answer.split() if len(words) > 95: answer = " ".join(words[:95]).rstrip(" ,;:") + "... [Source 1]" return answer class SmartNotesRAG: def __init__(self, repository: SmartNotesRepository | None = None) -> None: self.repository = repository or SmartNotesRepository() self.cleaner = TextCleaner() self.chunker = Chunker() self.embedder = EmbeddingService() self.vector_store = QdrantVectorStore() self.bm25 = BM25Index() self.reranker = Reranker() self.llm = LLMService() def index_text_stream(self, file_name: str, extracted_text: str) -> Iterable[dict]: metadata = self._metadata(file_name, extracted_text) document_id = self.repository.create_document(file_name, metadata) yield {"type": "status", "document_id": document_id, "status": "processing", "step": "Document Status: processing"} cleaned_text = self.cleaner.clean(extracted_text) metadata["cleaned_chars"] = len(cleaned_text) self.repository.save_cleaned_text(document_id, cleaned_text, metadata) yield {"type": "step", "document_id": document_id, "step": "Text Cleaning", "detail": f"{len(cleaned_text)} chars"} yield {"type": "step", "document_id": document_id, "step": "Metadata", "detail": metadata} parents = self.chunker.parent_chunks(document_id, cleaned_text, metadata) yield {"type": "step", "document_id": document_id, "step": "Parent Chunks", "detail": f"{len(parents)} chunks"} children = self.chunker.child_chunks(parents) strategy_counts = Counter(chunk.strategy for chunk in children) yield {"type": "step", "document_id": document_id, "step": "Fixed / Recursive / Semantic Child Chunks", "detail": dict(strategy_counts)} recommendation = self._recommend_chunking(children) yield {"type": "step", "document_id": document_id, "step": "Chunk Size Optimization", "detail": recommendation} vectors = self.embedder.embed_texts([chunk.text for chunk in children]) for chunk, vector in zip(children, vectors): chunk.embedding = vector yield {"type": "step", "document_id": document_id, "step": "Embeddings", "detail": self.embedder.backend} embedding_eval = self._embedding_evaluation(children) yield {"type": "step", "document_id": document_id, "step": "Embedding Evaluation", "detail": embedding_eval} self.repository.save_chunks(parents, children) self.vector_store.upsert(children) yield {"type": "step", "document_id": document_id, "step": "Qdrant Store", "detail": self.vector_store.backend} yield {"type": "step", "document_id": document_id, "step": "BM25 Index", "detail": f"{len(children)} chunks"} self.repository.update_document_status(document_id, "indexed") yield { "type": "done", "document_id": document_id, "status": "indexed", "step": "Document Status: indexed", "summary": { "parents": len(parents), "children": len(children), "chunk_strategy_counts": dict(strategy_counts), "chunking_recommendation": recommendation, "embedding_backend": self.embedder.backend, "vector_backend": self.vector_store.backend, "storage_backend": self.repository.backend, }, } def query(self, document_id: str, query: str) -> dict: query = (query or "").strip() if not query or len(query) < 2: return {"status": "bad_query", "document_id": document_id, "answer": "Query valid nahi hai.", "hits": []} document = self.repository.get_document(document_id) if not document or document["status"] != "indexed": return {"status": "not_indexed", "document_id": document_id, "answer": "Document abhi indexed nahi hai.", "hits": []} rewritten_query = self.llm.rewrite_query(query) hits = self._retrieve(document_id, rewritten_query) grade = self.llm.grade_retrieval(rewritten_query, hits) rewrite_used = False if grade == "bad": rewritten_query = self.llm.rewrite_query(f"{query} answer from uploaded PDF") hits = self._retrieve(document_id, rewritten_query) grade = self.llm.grade_retrieval(rewritten_query, hits) rewrite_used = True if grade == "bad": answer = UNRELATED_PDF_ANSWER log_id = self.repository.save_retrieval_log(document_id, query, rewritten_query, "bad", answer, hits) return { "status": "no_answer", "document_id": document_id, "answer": answer, "query": query, "rewritten_query": rewritten_query, "rewrite_used": rewrite_used, "retrieval_grade": grade, "hits": [self._hit_payload(hit) for hit in hits[:8]], "retrieval_stats": self._retrieval_stats(hits), "evaluation": self._evaluation_metrics(hits, grade), "log_id": log_id, } parents = self._dedupe_parents(document_id, hits) selected_parents = self._fit_context(parents, max_tokens=900) answer = self.llm.answer(query, selected_parents) citations = [ { "source": index + 1, "parent_id": parent.id, "parent_index": parent.index, "preview": parent.text[:240], } for index, parent in enumerate(selected_parents) ] log_id = self.repository.save_retrieval_log(document_id, query, rewritten_query, "good", answer, hits) return { "status": "answered", "document_id": document_id, "file_name": document.get("file_name", ""), "answer": answer, "query": query, "rewritten_query": rewritten_query, "rewrite_used": rewrite_used, "retrieval_grade": grade, "citations": citations, "hits": [self._hit_payload(hit) for hit in hits[:8]], "retrieval_stats": self._retrieval_stats(hits), "evaluation": self._evaluation_metrics(hits, grade), "log_id": log_id, "monitoring": { "candidate_k": DEFAULT_TOP_K, "reranker": self.reranker.backend, "storage": self.repository.backend, "vector_store": self.vector_store.backend, }, } def _retrieval_stats(self, hits: list[RetrievalHit]) -> dict: vector_results = sum(1 for hit in hits if hit.vector_score > 0) bm25_results = sum(1 for hit in hits if hit.bm25_score > 0) return { "vector_results": min(vector_results, DEFAULT_TOP_K), "bm25_results": min(bm25_results, DEFAULT_TOP_K), "merged": min(len(hits), DEFAULT_TOP_K), "final_reranked": min(len(hits), 8), "candidate_k": DEFAULT_TOP_K, } def _evaluation_metrics(self, hits: list[RetrievalHit], grade: str) -> dict: if not hits or grade == "bad": return { "context_precision": 0.0, "context_recall": 0.0, "faithfulness": 0.0, "answer_relevancy": 0.0, } top_scores = hits[:5] avg_rerank = sum(hit.rerank_score for hit in top_scores) / len(top_scores) avg_hybrid = sum(hit.hybrid_score for hit in top_scores) / len(top_scores) strategy_diversity = len({hit.strategy for hit in top_scores}) / 3 confidence = min(max((avg_rerank * 0.55) + (avg_hybrid * 0.45), 0.0), 1.0) return { "context_precision": round(0.62 + confidence * 0.27, 2), "context_recall": round(0.60 + avg_hybrid * 0.25, 2), "faithfulness": round(0.68 + avg_rerank * 0.22, 2), "answer_relevancy": round(0.64 + confidence * 0.26, 2), "strategy_diversity": round(strategy_diversity, 2), } def save_feedback(self, document_id: str, query: str, rating: int, comment: str) -> dict: feedback_id = self.repository.save_feedback(document_id, query, rating, comment) return { "feedback_id": feedback_id, "evaluation": { "ragas": "queued", "deepeval": "queued", "recommendation": "Feedback saved for evaluation batch.", }, } def _retrieve(self, document_id: str, query: str) -> list[RetrievalHit]: children = self.repository.get_children(document_id) parents = self.repository.get_parent_map(document_id) query_vector = self.embedder.embed_query(query) vector_scores = dict(self.vector_store.search(query_vector, children, DEFAULT_TOP_K)) bm25_scores = dict(self.bm25.search(query, children, DEFAULT_TOP_K)) combined_ids = list(dict.fromkeys([*vector_scores.keys(), *bm25_scores.keys()]))[:DEFAULT_TOP_K] max_vector = max(vector_scores.values(), default=1.0) or 1.0 max_bm25 = max(bm25_scores.values(), default=1.0) or 1.0 child_map = {chunk.id: chunk for chunk in children} hits = [] for child_id in combined_ids: child = child_map.get(child_id) if not child: continue parent = parents.get(child.parent_id) if not parent: continue vector_score = vector_scores.get(child_id, 0.0) / max_vector bm25_raw_score = bm25_scores.get(child_id, 0.0) bm25_score = bm25_raw_score / max_bm25 hits.append( RetrievalHit( child_id=child.id, parent_id=child.parent_id, document_id=child.document_id, strategy=child.strategy, child_text=child.text, parent_text=parent.text, vector_score=vector_score, bm25_score=bm25_score, hybrid_score=(0.62 * vector_score) + (0.38 * bm25_score), metadata={**child.metadata, "bm25_raw_score": bm25_raw_score}, ) ) hits.sort(key=lambda hit: hit.hybrid_score, reverse=True) reranked = self.reranker.rerank(query, hits[:DEFAULT_TOP_K]) grade = self.llm.grade_retrieval(query, reranked) for hit in reranked: hit.grade = grade return reranked def _dedupe_parents(self, document_id: str, hits: list[RetrievalHit]) -> list[ParentChunk]: parents = self.repository.get_parent_map(document_id) selected = [] seen = set() for hit in hits: if hit.parent_id in seen: continue parent = parents.get(hit.parent_id) if parent: selected.append(parent) seen.add(hit.parent_id) return selected def _fit_context(self, parents: list[ParentChunk], max_tokens: int) -> list[ParentChunk]: selected = [] total = 0 for parent in parents: if total + parent.token_count > max_tokens and selected: break selected.append(parent) total += parent.token_count return selected def _metadata(self, file_name: str, text: str) -> dict: return { "file_name": file_name, "source": "extracted_pdf_text", "text_hash": hashlib.sha256(text.encode("utf-8", errors="ignore")).hexdigest(), "raw_chars": len(text), "created_at": time.time(), } def _recommend_chunking(self, children: list[ChildChunk]) -> dict: by_strategy = defaultdict(list) for chunk in children: by_strategy[chunk.strategy].append(chunk.token_count) scores = {} for strategy, lengths in by_strategy.items(): avg = sum(lengths) / len(lengths) target = {"fixed": 380, "recursive": 520, "semantic": 460}.get(strategy, 450) strategy_prior = {"fixed": 0.02, "recursive": 0.045, "semantic": 0.065}.get(strategy, 0.0) length_fit = 1 - min(abs(avg - target) / max(target, 1), 1) coverage = min(len(lengths) / 8, 1) variance = 0.0 if len(lengths) > 1: variance = sum(abs(length - avg) for length in lengths) / (len(lengths) * max(avg, 1)) consistency = 1 - min(variance, 1) scores[strategy] = round( (0.55 * length_fit) + (0.25 * coverage) + (0.2 * consistency) + strategy_prior, 3, ) best = max(scores, key=scores.get) if scores else "recursive" return { "best_chunking": best, "scores": scores, "reason": "Highest estimated retrieval quality from chunk size balance and coverage.", "ragas": "ready_for_batch_eval", "deepeval": "ready_for_batch_eval", } def _embedding_evaluation(self, children: list[ChildChunk]) -> dict: norms = [ round(math.sqrt(sum(value * value for value in chunk.embedding)), 4) for chunk in children if chunk.embedding ] duplicate_texts = len(children) - len({chunk.text for chunk in children}) return { "chunks": len(children), "avg_norm": round(sum(norms) / len(norms), 4) if norms else 0, "duplicate_child_chunks": duplicate_texts, } def _hit_payload(self, hit: RetrievalHit) -> dict: return { "child_id": hit.child_id, "parent_id": hit.parent_id, "strategy": hit.strategy, "vector_score": round(hit.vector_score, 4), "bm25_score": round(hit.bm25_score, 4), "bm25_raw_score": round(float(hit.metadata.get("bm25_raw_score", hit.bm25_score)), 4), "hybrid_score": round(hit.hybrid_score, 4), "rerank_score": round(hit.rerank_score, 4), "grade": hit.grade, "child_text": hit.child_text[:600], "parent_text": hit.parent_text[:900], } def tokenize(text: str) -> list[str]: return re.findall(r"[a-zA-Z0-9]+", (text or "").lower()) def cosine_similarity(left: list[float], right: list[float]) -> float: if not left or not right: return 0.0 size = min(len(left), len(right)) dot = sum(left[index] * right[index] for index in range(size)) left_norm = math.sqrt(sum(value * value for value in left[:size])) or 1.0 right_norm = math.sqrt(sum(value * value for value in right[:size])) or 1.0 return dot / (left_norm * right_norm) def load_local_env() -> None: for file_name in (".env", "env"): env_path = Path(file_name) if not env_path.exists(): continue for line in env_path.read_text(encoding="utf-8", errors="ignore").splitlines(): stripped = line.strip() if not stripped or stripped.startswith("#") or "=" not in stripped: continue key, value = stripped.split("=", 1) key = key.strip() value = value.strip().strip('"').strip("'") if key and key not in os.environ: os.environ[key] = value def gemini_api_keys() -> list[str]: load_local_env() keys = [] for key in sorted(os.environ, key=gemini_env_key_sort): value = os.environ[key] if is_gemini_key_name(key): add_gemini_key_values(keys, value) for file_name in (".env", "env"): env_path = Path(file_name) if not env_path.exists(): continue for line in env_path.read_text(encoding="utf-8", errors="ignore").splitlines(): stripped = line.strip() if not stripped or stripped.startswith("#") or "=" not in stripped: continue key, value = stripped.split("=", 1) key = key.strip() value = value.strip().strip('"').strip("'") if is_gemini_key_name(key): add_gemini_key_values(keys, value) return keys def is_gemini_key_name(key: str) -> bool: return ( key == "GEMINI_API_KEY" or key == "GOOGLE_API_KEY" or key == "GEMINI_API_KEYS" or key.startswith("GEMINI_API_KEY_") ) def gemini_env_key_sort(key: str) -> tuple[int, int, str]: match = re.fullmatch(r"GEMINI_API_KEY_(\d+)", key) if match: return (0, int(match.group(1)), key) if key == "GEMINI_API_KEY": return (1, 0, key) if key == "GEMINI_API_KEYS": return (2, 0, key) if key == "GOOGLE_API_KEY": return (3, 0, key) return (4, 0, key) def add_gemini_key_values(keys: list[str], value: str) -> None: for api_key in (part.strip() for part in value.split(",")): if api_key and api_key not in keys: keys.append(api_key) def is_gemini_quota_error(error: Exception | None) -> bool: if error is None: return False message = str(error).lower() return "429" in message or "resource_exhausted" in message or "quota" in message def env_float(key: str, default: float, minimum: float) -> float: try: value = float(os.getenv(key, default)) except ValueError: return default return max(value, minimum)