""" rag.py — a complete, runnable, OFFLINE Retrieval-Augmented Generation pipeline that faithfully mirrors the Azure AI Search architecture from Knowledge Module 02. Read this top to bottom: it is the "book chapter" in code form. Every stage that Azure performs for you in the cloud is implemented here in a few lines of pure Python so you can SEE what hybrid search, RRF, and reranking actually do. THE PIPELINE (offline mirror of the Azure one): documents --[chunk]--> chunks --[embed]--> vector index (AI Search vector field) chunks --[index]--> BM25 index (AI Search searchable field) query --> [ BM25 keyword ] --+ [ vector cosine ] --+--> RRF fuse --> semantic rerank --> top-k | grounding prompt (cite + abstain) --> LLM --> answer HOW TO GO LIVE ON AZURE: Replace three classes and nothing else changes: Embedder -> AzureOpenAIEmbedder (text-embedding-3-large) LocalLLM -> AzureOpenAILLM (gpt-4o deployment) LocalHybridIndex -> AzureAISearchIndex (vector + hybrid + semantic ranker) The retrieval/grounding LOGIC is identical; only the provider changes. Each Azure class below is stubbed with the real SDK call shape in its docstring. """ from __future__ import annotations import math import re from collections import Counter, defaultdict from dataclasses import dataclass, field from data import Document # =========================================================================== # 0. Tokenization & chunking # =========================================================================== # Latin/number tokens + Arabic LETTERS and Eastern-Arabic DIGITS only. We must # exclude Arabic punctuation (؟ ، ؛, U+060C/061B/061F) and diacritics, or a query # word like "الجاري؟" would never match the document's "الجاري". _TOKEN_RE = re.compile(r"[A-Za-z0-9.%]+|[ء-ي٠-٩]+") def _stem(tok: str) -> str: """Minimal plural stemmer so 'statements' matches 'statement', 'fees' matches 'fee', 'transfers' matches 'transfer'. Real Azure AI Search uses a full language analyzer (Lucene) that stems for you; this is the toy version. Only strips a single trailing 's' (never 'ss'), and leaves non-ASCII (Arabic) tokens untouched.""" if tok.isascii() and tok.isalpha() and len(tok) > 3 \ and tok.endswith("s") and not tok.endswith("ss"): return tok[:-1] return tok def tokenize(text: str) -> list[str]: """Lowercase, stem word/number tokenizer. Keeps numbers and percents (e.g. '5.99', '25', '60') because in banking the EXACT figures are what a keyword query must match. Also matches Arabic Unicode ranges so Arabic docs index.""" # Drop single-character tokens ('s' from "bank's", stray letters) — they # carry no retrieval signal and would otherwise inflate term overlap. return [_stem(t.lower()) for t in _TOKEN_RE.findall(text) if len(t) > 1 or t.isdigit()] @dataclass class Chunk: """A retrievable passage. In Azure AI Search this is one *document* in the index: an id, the searchable `content`, a `contentVector`, and filter fields (source/section/language/acl). We carry the same shape here.""" id: str # e.g. "fees-transfers#0" doc_id: str # source document id -> citation section: int # chunk index within the doc -> citation page/section text: str language: str acl: list[str] vector: dict[str, float] = field(default_factory=dict) # sparse TF-IDF vector tokens: list[str] = field(default_factory=list) def chunk_text(text: str, max_tokens: int = 60, overlap: int = 12) -> list[str]: """Fixed-size chunking with overlap — the simplest of the four strategies in K02 §3. Real corpora use recursive/structure-aware or layout-aware chunking, but the trade-off (too small loses context, too large dilutes the embedding) is identical. Overlap keeps a fact that straddles a boundary intact in at least one chunk. We chunk on sentence boundaries first, then pack to size.""" sentences = re.split(r"(?<=[.!?])\s+", text.strip()) chunks, cur, cur_len = [], [], 0 for s in sentences: n = len(tokenize(s)) if cur and cur_len + n > max_tokens: chunks.append(" ".join(cur)) # carry the tail `overlap` tokens of words into the next chunk carry = " ".join(cur).split()[-overlap:] cur, cur_len = [" ".join(carry)], len(carry) cur.append(s) cur_len += n if cur: chunks.append(" ".join(cur)) return chunks # =========================================================================== # 1. Embeddings (offline stand-in for text-embedding-3-large) # =========================================================================== # A SMALL banking synonym map. Real embedding models learn that "declined" and # "failed", or "waived" and "free", point in similar directions. Pure TF-IDF # can't — it only sees shared tokens. So we expand tokens with these synonyms # before vectorizing: a deterministic, dependency-free way to demonstrate that # VECTOR/semantic retrieval matches *meaning*, not just surface words. In # production this whole trick disappears — a real embedding model handles it. SYNONYMS: dict[str, list[str]] = { "refused": ["declined", "rejected", "failed", "denied"], "declined": ["refused", "rejected", "failed", "denied"], "failed": ["refused", "declined", "rejected", "denied"], "waived": ["free", "waiver", "exempt", "removed"], "free": ["waived", "complimentary", "exempt"], "fee": ["charge", "cost", "price"], "charge": ["fee", "cost", "price"], "transfer": ["wire", "payment", "remittance", "send"], "interest": ["rate", "profit", "apr"], "dispute": ["claim", "chargeback", "contest"], "balance": ["funds", "money", "amount"], "document": ["paper", "paperwork", "id"], "early": ["prepay", "settlement", "prepayment"], "penalty": ["fee", "charge", "settlement"], "freeze": ["block", "lock", "suspend"], "statement": ["report"], } def _expand(tokens: list[str]) -> list[str]: out = list(tokens) for t in tokens: out.extend(SYNONYMS.get(t, [])) return out class Embedder: """Deterministic TF-IDF embedder with synonym expansion. Stands in for `text-embedding-3-large`. Vectors are sparse dicts {term: weight}; similarity is cosine. Same interface a real embedder would expose: `.fit`, `.embed`.""" def __init__(self) -> None: self.idf: dict[str, float] = {} self.n_docs = 0 def fit(self, corpus_tokens: list[list[str]]) -> None: """Compute inverse-document-frequency so rare, discriminative terms (an account code, '5.99') weigh more than common ones ('the', 'fee').""" self.n_docs = len(corpus_tokens) df: Counter[str] = Counter() for toks in corpus_tokens: for term in set(_expand(toks)): df[term] += 1 self.idf = {t: math.log((self.n_docs + 1) / (c + 1)) + 1.0 for t, c in df.items()} def embed(self, text: str) -> dict[str, float]: toks = _expand(tokenize(text)) tf = Counter(toks) vec = {t: (c / len(toks)) * self.idf.get(t, math.log(self.n_docs + 1) + 1.0) for t, c in tf.items()} return _l2_normalize(vec) def _l2_normalize(vec: dict[str, float]) -> dict[str, float]: norm = math.sqrt(sum(v * v for v in vec.values())) or 1.0 return {t: v / norm for t, v in vec.items()} _STOPWORDS = { "the", "a", "an", "is", "are", "was", "were", "be", "of", "to", "for", "and", "or", "my", "i", "you", "it", "at", "in", "on", "do", "does", "did", "how", "what", "when", "where", "why", "who", "with", "may", "can", "could", "would", "should", "this", "that", "there", "your", "me", "about", "tell", "get", "have", "has", "from", "by", "as", "if", "so", "not", "no", } def cosine(a: dict[str, float], b: dict[str, float]) -> float: """Cosine similarity of two sparse vectors = dot product (already normalized). This is exactly what an ANN index (HNSW) approximates at scale — K02 §4.""" if len(a) > len(b): a, b = b, a return sum(w * b.get(t, 0.0) for t, w in a.items()) # =========================================================================== # 2. BM25 keyword search # =========================================================================== class BM25: """Classic BM25 ranking — the keyword half of hybrid search (K02 §5). It scores by term frequency x inverse document frequency with length normalization. BM25 nails EXACT tokens (account numbers, '5.99', 'SWIFT') that vector search blurs together. k1 and b are the standard knobs.""" def __init__(self, k1: float = 1.5, b: float = 0.75) -> None: self.k1, self.b = k1, b self.docs: list[list[str]] = [] self.ids: list[str] = [] self.df: Counter[str] = Counter() self.avgdl = 0.0 def index(self, ids: list[str], docs_tokens: list[list[str]]) -> None: self.ids, self.docs = ids, docs_tokens self.avgdl = sum(len(d) for d in docs_tokens) / max(len(docs_tokens), 1) self.df = Counter() for toks in docs_tokens: for term in set(toks): self.df[term] += 1 def _idf(self, term: str) -> float: n, df = len(self.docs), self.df.get(term, 0) return math.log((n - df + 0.5) / (df + 0.5) + 1.0) def search(self, query: str, top_k: int = 10) -> list[tuple[str, float]]: q = tokenize(query) scores: list[tuple[str, float]] = [] for cid, toks in zip(self.ids, self.docs): tf = Counter(toks) dl = len(toks) s = 0.0 for term in q: if term not in tf: continue num = tf[term] * (self.k1 + 1) den = tf[term] + self.k1 * (1 - self.b + self.b * dl / self.avgdl) s += self._idf(term) * num / den if s > 0: scores.append((cid, s)) scores.sort(key=lambda x: x[1], reverse=True) return scores[:top_k] # =========================================================================== # 3. The hybrid index (vector + BM25 + RRF + semantic rerank) # =========================================================================== def rrf_fuse(rankings: list[list[tuple[str, float]]], k: int = 60) -> list[tuple[str, float]]: """Reciprocal Rank Fusion (K02 §6). Combine several ranked lists by summing 1/(k+rank). Rank-based, so it needs NO score calibration between BM25 and cosine (their scales differ wildly). A doc ranked high in EITHER list scores well; high in BOTH wins. This is exactly how Azure AI Search fuses hybrid.""" fused: dict[str, float] = defaultdict(float) for ranking in rankings: for rank, (cid, _score) in enumerate(ranking): fused[cid] += 1.0 / (k + rank) return sorted(fused.items(), key=lambda x: x[1], reverse=True) class LocalHybridIndex: """Offline stand-in for an Azure AI Search index. Holds chunks with both a vector field and a BM25 index, and supports four retrieval modes so you can A/B them: 'vector', 'bm25', 'hybrid' (RRF), 'hybrid_semantic' (RRF + rerank). AZURE EQUIVALENT — see AzureAISearchIndex below. There you would create an index with a `contentVector` field (HNSW/cosine) + a searchable `content` field + a semantic configuration, and a single `search()` call with `query_type=SEMANTIC` does hybrid+RRF+rerank server-side.""" def __init__(self) -> None: self.chunks: dict[str, Chunk] = {} self.embedder = Embedder() self.bm25 = BM25() def ingest(self, documents: list[Document], max_tokens: int = 60, overlap: int = 12) -> None: # 1. chunk every document, attaching metadata to each chunk chunks: list[Chunk] = [] for doc in documents: for i, piece in enumerate(chunk_text(doc.text, max_tokens, overlap)): chunks.append(Chunk( id=f"{doc.id}#{i}", doc_id=doc.id, section=i, text=piece, language=doc.language, acl=list(doc.acl), tokens=tokenize(piece))) # 2. fit the embedder on the corpus, then embed each chunk (the vector field) self.embedder.fit([c.tokens for c in chunks]) for c in chunks: c.vector = self.embedder.embed(c.text) self.chunks[c.id] = c # 3. build the BM25 index (the searchable field) self.bm25.index([c.id for c in chunks], [c.tokens for c in chunks]) # --- filtering: the banking-critical part --- def _allowed(self, chunk: Chunk, user_groups: list[str], language: str | None) -> bool: """SECURITY TRIMMING (K02 §9). A user may only retrieve chunks whose ACL intersects their groups. This is enforced HERE, in retrieval — never by asking the LLM to 'not reveal' something. Optional language filter too.""" if not (set(chunk.acl) & set(user_groups)): return False if language and chunk.language != language: return False return True def retrieve(self, query: str, mode: str = "hybrid_semantic", top_k: int = 5, user_groups: list[str] | None = None, language: str | None = None ) -> list[tuple[Chunk, float]]: user_groups = user_groups or ["retail"] cand_ids = [cid for cid, c in self.chunks.items() if self._allowed(c, user_groups, language)] # vector ranking (dense) qv = self.embedder.embed(query) vec_rank = sorted(((cid, cosine(qv, self.chunks[cid].vector)) for cid in cand_ids), key=lambda x: x[1], reverse=True) vec_rank = [r for r in vec_rank if r[1] > 0][:25] # bm25 ranking (sparse) — restricted to allowed candidates bm = [(cid, s) for cid, s in self.bm25.search(query, top_k=50) if cid in set(cand_ids)][:25] if mode == "vector": ranked = vec_rank elif mode == "bm25": ranked = bm else: # hybrid / hybrid_semantic ranked = rrf_fuse([vec_rank, bm]) ranked = ranked[: max(top_k * 4, 12)] # over-fetch for reranking if mode == "hybrid_semantic": ranked = self._semantic_rerank(query, [cid for cid, _ in ranked]) return [(self.chunks[cid], score) for cid, score in ranked[:top_k]] def _semantic_rerank(self, query: str, cand_ids: list[str]) -> list[tuple[str, float]]: """Cross-encoder stand-in (K02 §7). A real semantic ranker reads the query and each passage TOGETHER (not as separate vectors) for far more accurate relevance, then reorders. We approximate that 'joint reading' with a lexical+semantic coverage score: how many query terms (and their synonyms) the passage actually covers, weighted by IDF. Cheap, but it demonstrates why retrieve-then-rerank lifts the right chunk into top-k.""" q_terms = set(_expand(tokenize(query))) scored: list[tuple[str, float]] = [] for cid in cand_ids: c = self.chunks[cid] c_terms = set(_expand(c.tokens)) covered = q_terms & c_terms coverage = sum(self.embedder.idf.get(t, 1.0) for t in covered) denom = sum(self.embedder.idf.get(t, 1.0) for t in q_terms) or 1.0 scored.append((cid, coverage / denom)) scored.sort(key=lambda x: x[1], reverse=True) return scored # =========================================================================== # 4. The LLM (offline grounded generator) # =========================================================================== class LocalLLM: """A DETERMINISTIC, extractive 'generator'. It does NOT invent text — it selects the sentence(s) from the retrieved SOURCES that best answer the question, and prepends citations. This is a faithful stand-in for the grounding contract you'd give gpt-4o (K02 §10, K06 §6): answer only from sources, cite, and ABSTAIN when the answer isn't there. Crucially, because it can only emit text present in the sources, it is grounded by construction — exactly the property you want a bank's bot to have, and what you'd verify with a groundedness evaluator (Lab 06).""" ABSTAIN = ("I don't have that information in the available sources. " "I can connect you with a human agent who can help.") # If the best retrieved sentence covers less than this fraction of the # query's IDF "mass", we abstain rather than return a loosely-related snippet. RELEVANCE_THRESHOLD = 0.34 def answer(self, question: str, sources: list[tuple[Chunk, float]], idf: dict[str, float] | None = None, idf_default: float = 3.0) -> dict: if not sources: return self._abstain() idf = idf or {} # Coverage is measured over CONTENT terms only (drop stopwords), so that # a sentence sharing only 'the'/'is'/'bank' with the query can't look # relevant — the distinctive uncovered terms ('stock', 'today') dominate. content = [t for t in tokenize(question) if t not in _STOPWORDS] q_terms = set(_expand(content)) # Weight each query term by IDF so rare/specific terms ('5.99', 'swift', # 'murabaha') dominate and common ones ('the', 'bank') barely count. def w(t: str) -> float: return idf.get(t, idf_default) denom = sum(w(t) for t in q_terms) or 1.0 best = (0.0, None, None) # (coverage_ratio, sentence, chunk) for chunk, _score in sources: for sent in re.split(r"(?<=[.!?])\s+", chunk.text): s_terms = set(_expand(tokenize(sent))) covered = q_terms & s_terms ratio = sum(w(t) for t in covered) / denom if ratio > best[0]: best = (ratio, sent, chunk) ratio, best_sent, best_chunk = best if ratio < self.RELEVANCE_THRESHOLD or best_sent is None: return self._abstain() cite = f"[{best_chunk.doc_id}:{best_chunk.section}]" return { "answer": f"{best_sent.strip()} {cite}", "citations": [cite], "grounded": True, "abstained": False, "source_doc": best_chunk.doc_id, "relevance": round(ratio, 3), } def _abstain(self) -> dict: return {"answer": self.ABSTAIN, "citations": [], "grounded": True, "abstained": True, "relevance": 0.0} # =========================================================================== # 5. The pipeline (ties it together) # =========================================================================== class RagPipeline: """End-to-end RAG. `ask()` runs retrieve -> ground -> generate and returns a structured, cited, possibly-abstaining answer. The default providers are the offline ones; pass Azure providers to go live (see __init__ signature).""" def __init__(self, index: LocalHybridIndex | None = None, llm: LocalLLM | None = None) -> None: self.index = index or LocalHybridIndex() self.llm = llm or LocalLLM() def ingest(self, documents: list[Document]) -> None: self.index.ingest(documents) def ask(self, question: str, mode: str = "hybrid_semantic", top_k: int = 5, user_groups: list[str] | None = None, language: str | None = None) -> dict: sources = self.index.retrieve(question, mode=mode, top_k=top_k, user_groups=user_groups, language=language) import math as _m idf_default = _m.log(self.index.embedder.n_docs + 1) + 1.0 result = self.llm.answer(question, sources, idf=self.index.embedder.idf, idf_default=idf_default) result["retrieved"] = [(c.id, round(s, 4)) for c, s in sources] result["mode"] = mode return result # =========================================================================== # 6. Azure provider stubs — the ONE-CLASS-SWAP to go live (documented, not run) # =========================================================================== class AzureOpenAIEmbedder: """PRODUCTION embedder. Drop-in for `Embedder`. Requires `openai` + `azure-identity`. Keyless auth via Managed Identity (K00 §6). from openai import AzureOpenAI from azure.identity import DefaultAzureCredential, get_bearer_token_provider tok = get_bearer_token_provider(DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default") client = AzureOpenAI(azure_endpoint=ENDPOINT, api_version="2024-10-21", azure_ad_token_provider=tok) def embed(self, text): v = client.embeddings.create(model="text-embedding-3-large", input=text) return v.data[0].embedding # a 3072-dim dense list[float] """ class AzureOpenAILLM: """PRODUCTION generator. Drop-in for `LocalLLM`. The grounding PROMPT is the same contract LocalLLM enforces in code: SYSTEM = ("You are a banking assistant. Answer ONLY from the SOURCES. " "If the answer isn't there, say you don't have it and offer a " "human agent. Never invent figures, rates, or policies. " "Cite each fact as [source:section].") resp = client.chat.completions.create( model="gpt-4o-prod", temperature=0, messages=[{"role":"system","content":SYSTEM}, {"role":"user","content": build_prompt(question, sources)}]) """ class AzureAISearchIndex: """PRODUCTION retriever. Drop-in for `LocalHybridIndex`. ONE call does hybrid + RRF + semantic rerank server-side: from azure.search.documents import SearchClient results = client.search( search_text=query, # BM25 half vector_queries=[VectorizedQuery(vector=embed(query), fields="contentVector", k_nearest_neighbors=50)], query_type="semantic", semantic_configuration_name="default", filter=f"acl/any(g: search.in(g, '{','.join(user_groups)}'))", # security trim top=5) Use INTEGRATED VECTORIZATION (indexer + embedding skill) to let Azure embed chunks and the query for you (K02 §8)."""