🔎 RAG Patterns Deep Dive — Production Retrieval-Augmented Generation on Fabric¶

# Databricks notebook source
# COMMAND ----------
# MAGIC %md
# MAGIC ## Recursive Chunking with Token-Aware Boundaries

# COMMAND ----------

import tiktoken
from typing import List

ENC = tiktoken.encoding_for_model("text-embedding-3-large")

def recursive_chunk(
    text: str,
    max_tokens: int = 500,
    overlap_tokens: int = 50,
    separators: List[str] = None
) -> List[str]:
    """
    Recursively split text on a hierarchy of separators until each chunk
    is under max_tokens. Adds overlap to reduce boundary loss.
    """
    if separators is None:
        separators = ["\n\n", "\n", ". ", " ", ""]

    def _count(t: str) -> int:
        return len(ENC.encode(t))

    if _count(text) <= max_tokens:
        return [text]

    # Try each separator until we get small-enough pieces
    for sep in separators:
        if sep == "":
            # Last resort: hard split at token boundary
            tokens = ENC.encode(text)
            return [
                ENC.decode(tokens[i:i + max_tokens])
                for i in range(0, len(tokens), max_tokens - overlap_tokens)
            ]
        parts = text.split(sep)
        if len(parts) > 1:
            # Greedily pack parts into chunks under max_tokens
            chunks, current = [], ""
            for part in parts:
                candidate = current + (sep if current else "") + part
                if _count(candidate) <= max_tokens:
                    current = candidate
                else:
                    if current:
                        chunks.append(current)
                    if _count(part) > max_tokens:
                        # Recurse on oversized part
                        chunks.extend(
                            recursive_chunk(part, max_tokens, overlap_tokens, separators)
                        )
                        current = ""
                    else:
                        current = part
            if current:
                chunks.append(current)
            # Add overlap by prepending tail of previous chunk
            return _add_overlap(chunks, overlap_tokens)
    return [text]

def _add_overlap(chunks: List[str], overlap_tokens: int) -> List[str]:
    if overlap_tokens <= 0 or len(chunks) <= 1:
        return chunks
    result = [chunks[0]]
    for prev, curr in zip(chunks[:-1], chunks[1:]):
        prev_tokens = ENC.encode(prev)
        tail = ENC.decode(prev_tokens[-overlap_tokens:])
        result.append(tail + " " + curr)
    return result

# Truncate at embedding time
response = client.embeddings.create(
    input=text,
    model="text-embedding-3-large",
    dimensions=1024  # Matryoshka truncation
)

EXPANSION_PROMPT = """Generate 4 alternative phrasings of this question that would
match the same answer in a knowledge base. Return ONLY the questions, one per line.

Question: {query}

Alternatives:"""

def multi_query_retrieve(query: str, top_k: int = 10) -> list[dict]:
    # Step 1: Generate variants
    variants = chat_complete(EXPANSION_PROMPT.format(query=query)).split("\n")
    variants = [v.strip() for v in variants if v.strip()][:4]
    all_queries = [query] + variants

    # Step 2: Retrieve for each
    seen_ids, results = set(), []
    for q in all_queries:
        for hit in vector_search(q, top_k=top_k):
            if hit["chunk_id"] not in seen_ids:
                seen_ids.add(hit["chunk_id"])
                results.append(hit)

    # Step 3: Rerank merged set
    return rerank(query, results)[:top_k]

HYDE_PROMPT = """Write a concise paragraph (3-5 sentences) that would appear in
a regulatory document and directly answer the following question. Do not say
'I don't know' — write a plausible passage even if you must guess details.

Question: {query}

Passage:"""

def hyde_retrieve(query: str, top_k: int = 10) -> list[dict]:
    fake_passage = chat_complete(HYDE_PROMPT.format(query=query))
    return vector_search(fake_passage, top_k=top_k)

SELF_QUERY_PROMPT = """Extract structured filters from this question. Return JSON:
{{"semantic_query": str, "filters": {{"doc_type": str|null, "date_from": str|null,
"date_to": str|null, "tags": list[str]}}}}

Schema fields: doc_type ∈ {{"CTR","SAR","W2G","policy","regulation"}};
date format YYYY-MM-DD; tags free-form.

Question: {query}
JSON:"""

from sentence_transformers import CrossEncoder

# Load once at app start
reranker = CrossEncoder("BAAI/bge-reranker-v2-m3", max_length=512)

def rerank(query: str, candidates: list[dict], top_k: int = 5) -> list[dict]:
    pairs = [(query, c["chunk_text"]) for c in candidates]
    scores = reranker.predict(pairs, batch_size=32)
    for c, s in zip(candidates, scores):
        c["rerank_score"] = float(s)
    return sorted(candidates, key=lambda x: -x["rerank_score"])[:top_k]

import cohere

co = cohere.Client(api_key=os.environ["COHERE_API_KEY"])

def cohere_rerank(query: str, candidates: list[dict], top_k: int = 5) -> list[dict]:
    docs = [c["chunk_text"] for c in candidates]
    resp = co.rerank(
        model="rerank-3.5",
        query=query,
        documents=docs,
        top_n=top_k,
        return_documents=False
    )
    return [
        {**candidates[r.index], "rerank_score": r.relevance_score}
        for r in resp.results
    ]

LLM_RERANK_PROMPT = """Rate each passage's relevance to the query on a 0-10 scale.
Return ONLY a JSON array of integers, one per passage, in order.

Query: {query}

Passages:
{passages}

Scores (JSON array):"""

def rrf_fuse(rankings: list[list[str]], k: int = 60) -> list[str]:
    """Each ranking is an ordered list of chunk_ids. Returns merged ranking."""
    scores = {}
    for ranking in rankings:
        for rank, chunk_id in enumerate(ranking, start=1):
            scores[chunk_id] = scores.get(chunk_id, 0) + 1.0 / (k + rank)
    return [cid for cid, _ in sorted(scores.items(), key=lambda x: -x[1])]

STUFF_TEMPLATE = """You are a compliance analyst assistant. Answer the question using
ONLY the passages below. Cite each claim with the passage id like [P3]. If the
passages do not contain the answer, say "I don't have that information."

Passages:
{passages}

Question: {question}

Answer (with citations):"""

def build_stuffed_prompt(query: str, chunks: list[dict]) -> str:
    passages = "\n\n".join(
        f"[P{i+1}] (source: {c['document_title']} > {c['section_path']})\n{c['chunk_text']}"
        for i, c in enumerate(chunks)
    )
    return STUFF_TEMPLATE.format(passages=passages, question=query)

MAP_PROMPT = """Extract any facts from this passage that help answer the question.
Return only the facts, one per line. If none, return 'NONE'.

Question: {question}
Passage: {passage}

Facts:"""

REDUCE_PROMPT = """Synthesize a final answer using these extracted facts. Cite
sources by [P#].

Question: {question}
Facts:
{facts}

Answer:"""

import re

CITATION_RE = re.compile(r"\[P(\d+)\]")

def extract_citations(answer: str, chunks: list[dict]) -> list[dict]:
    """Map [P#] markers in answer to chunk metadata."""
    cited_idxs = {int(m.group(1)) - 1 for m in CITATION_RE.finditer(answer)}
    return [
        {
            "marker": f"P{i+1}",
            "chunk_id": chunks[i]["chunk_id"],
            "document_title": chunks[i]["document_title"],
            "document_uri": chunks[i]["document_uri"],
            "section_path": chunks[i]["section_path"]
        }
        for i in cited_idxs if i < len(chunks)
    ]

from ragas import evaluate
from ragas.metrics import (
    faithfulness, answer_relevancy,
    context_precision, context_recall
)
from datasets import Dataset

eval_data = Dataset.from_dict({
    "question":    [r["query"]              for r in eval_set],
    "answer":      [r["generated_answer"]   for r in eval_set],
    "contexts":    [r["retrieved_chunks"]   for r in eval_set],
    "ground_truth":[r["expected_answer"]    for r in eval_set],
})

scores = evaluate(
    eval_data,
    metrics=[faithfulness, answer_relevancy, context_precision, context_recall]
)

print(scores)
# {'faithfulness': 0.92, 'answer_relevancy': 0.88,
#  'context_precision': 0.84, 'context_recall': 0.86}

# Notebook: 18_bronze_rag_chunking.py
# COMMAND ----------
# MAGIC %md
# MAGIC ## Bronze: Document Ingestion + Chunking
# MAGIC Reads documents from OneLake, parses, chunks, writes Delta Bronze table.

# COMMAND ----------

import uuid
from pyspark.sql import Row
from pyspark.sql.types import (StructType, StructField, StringType,
                               IntegerType, TimestampType)
from datetime import datetime
import notebookutils as mssparkutils

BRONZE_PATH = "abfss://<workspace>@onelake.dfs.fabric.microsoft.com/lh_bronze.Lakehouse/Tables/bronze_rag_chunks"
SOURCE_DIR  = "abfss://<workspace>@onelake.dfs.fabric.microsoft.com/lh_bronze.Lakehouse/Files/rag_corpus"

# COMMAND ----------

# Load chunker (defined in shared utils)
from rag_utils.chunker import recursive_chunk
from rag_utils.parser  import parse_document  # PDF/DOCX/MD/HTML

def process_document(file_path: str, doc_type: str, tenant_id: str) -> list[Row]:
    title, sections = parse_document(file_path)
    rows = []
    for sec_path, sec_text in sections:
        for idx, chunk_text in enumerate(
            recursive_chunk(sec_text, max_tokens=500, overlap_tokens=50)
        ):
            rows.append(Row(
                chunk_id=str(uuid.uuid4()),
                document_id=file_path.split("/")[-1],
                document_title=title,
                document_uri=file_path,
                chunk_index=idx,
                chunk_text=f"{title}\n{sec_path}\n\n{chunk_text}",  # contextual prefix
                token_count=len(chunk_text.split()) * 4 // 3,       # rough
                section_path=sec_path,
                doc_type=doc_type,
                tenant_id=tenant_id,
                sensitivity_label="internal",
                language="en",
                created_at=datetime.utcnow(),
                updated_at=datetime.utcnow(),
                embedding_model=None,
                embedding=None
            ))
    return rows

# COMMAND ----------

files = mssparkutils.fs.ls(SOURCE_DIR)
all_rows = []
for f in files:
    all_rows.extend(process_document(f.path, doc_type="policy", tenant_id="casino-prod"))

df = spark.createDataFrame(all_rows)
df.write.format("delta").mode("append").save(BRONZE_PATH)
print(f"Wrote {df.count()} chunks to Bronze")

# Notebook: 18_silver_rag_embed.py
# COMMAND ----------
# MAGIC %md
# MAGIC ## Silver: Embed Bronze Chunks → Eventhouse Vector Table

# COMMAND ----------

from openai import AzureOpenAI
from azure.identity import DefaultAzureCredential, get_bearer_token_provider

AOAI_ENDPOINT   = "https://your-aoai.openai.azure.com"
AOAI_DEPLOYMENT = "text-embedding-3-large"
EMBED_DIMS      = 1024  # Matryoshka truncation

token_provider = get_bearer_token_provider(
    DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default"
)
client = AzureOpenAI(
    azure_endpoint=AOAI_ENDPOINT,
    azure_ad_token_provider=token_provider,
    api_version="2024-10-21"
)

def embed_batch(texts: list[str]) -> list[list[float]]:
    resp = client.embeddings.create(
        input=texts, model=AOAI_DEPLOYMENT, dimensions=EMBED_DIMS
    )
    return [d.embedding for d in resp.data]

# COMMAND ----------

bronze = spark.read.format("delta").load(BRONZE_PATH).filter("embedding IS NULL")
print(f"To embed: {bronze.count()}")

# Process in batches of 100 (AOAI rate-limit friendly)
BATCH = 100
rows = bronze.collect()
embedded = []
for i in range(0, len(rows), BATCH):
    batch = rows[i:i+BATCH]
    vectors = embed_batch([r.chunk_text for r in batch])
    for r, v in zip(batch, vectors):
        d = r.asDict()
        d["embedding"] = v
        d["embedding_model"] = f"{AOAI_DEPLOYMENT}-{EMBED_DIMS}d"
        d["updated_at"] = datetime.utcnow()
        embedded.append(d)

# COMMAND ----------

# Write to Eventhouse via Kusto SDK
from azure.kusto.data import KustoConnectionStringBuilder
from azure.kusto.ingest import QueuedIngestClient, IngestionProperties, DataFormat

KUSTO_URI    = "https://<eventhouse>.kusto.fabric.microsoft.com"
KUSTO_DB     = "db_rag"
KUSTO_TABLE  = "rag_chunks"

kcsb = KustoConnectionStringBuilder.with_az_cli_authentication(KUSTO_URI)
ingest = QueuedIngestClient(kcsb)
props  = IngestionProperties(database=KUSTO_DB, table=KUSTO_TABLE,
                             data_format=DataFormat.JSON)

import json, tempfile
with tempfile.NamedTemporaryFile(suffix=".json", delete=False, mode="w") as fh:
    for r in embedded:
        fh.write(json.dumps(r, default=str) + "\n")
    fh.flush()
    ingest.ingest_from_file(fh.name, ingestion_properties=props)

print(f"Ingested {len(embedded)} chunks into Eventhouse")

# Notebook: 18_gold_rag_query.py
# COMMAND ----------
# MAGIC %md
# MAGIC ## Gold: End-to-End RAG Query with Hybrid Retrieval, Reranking, Logging

# COMMAND ----------

import time, uuid, json
from azure.kusto.data import KustoClient
from sentence_transformers import CrossEncoder

reranker = CrossEncoder("BAAI/bge-reranker-v2-m3", max_length=512)
kusto    = KustoClient(kcsb)
chat     = AzureOpenAI(azure_endpoint=AOAI_ENDPOINT,
                      azure_ad_token_provider=token_provider,
                      api_version="2024-10-21")
CHAT_DEPLOYMENT = "gpt-4o"

# COMMAND ----------

HYBRID_KQL = """
let user_query = '{q}';
let query_vec = toscalar(
    evaluate ai_embed_text(user_query, '{ep}', '{m}') | project embedding
);
let k = 60; let top_n = 50;
let vec_r = rag_chunks
    | where tenant_id == '{tenant}'
    | extend sim = series_cosine_similarity(embedding, query_vec)
    | top top_n by sim desc
    | extend vec_rank = row_number()
    | project chunk_id, vec_rank;
let bm_r = rag_chunks
    | where tenant_id == '{tenant}'
    | where chunk_text has_any (split(user_query, ' '))
    | extend kw_score = countof(chunk_text, user_query, 'normal')
    | top top_n by kw_score desc
    | extend kw_rank = row_number()
    | project chunk_id, kw_rank;
vec_r
| join kind=fullouter bm_r on chunk_id
| extend chunk_id = coalesce(chunk_id, chunk_id1)
| extend rrf = iff(isnotnull(vec_rank), 1.0/(k+vec_rank), 0.0)
              + iff(isnotnull(kw_rank),  1.0/(k+kw_rank),  0.0)
| top 30 by rrf desc
| join kind=inner rag_chunks on chunk_id
| project chunk_id, chunk_text, document_title, section_path,
          document_uri, rrf
"""

def hybrid_retrieve(query: str, tenant: str = "casino-prod") -> list[dict]:
    kql = HYBRID_KQL.format(
        q=query.replace("'", "''"),
        ep=AOAI_ENDPOINT, m="text-embedding-3-large", tenant=tenant
    )
    resp = kusto.execute(KUSTO_DB, kql)
    return [r.to_dict() for r in resp.primary_results[0]]

def rerank(query: str, candidates: list[dict], top_k: int = 5) -> list[dict]:
    pairs  = [(query, c["chunk_text"]) for c in candidates]
    scores = reranker.predict(pairs, batch_size=32)
    for c, s in zip(candidates, scores):
        c["rerank_score"] = float(s)
    return sorted(candidates, key=lambda x: -x["rerank_score"])[:top_k]

# COMMAND ----------

ANSWER_PROMPT = """You are a casino compliance analyst. Answer ONLY from the passages
below. Cite each claim as [P#]. If passages don't cover the question, say
"I don't have that information."

Passages:
{passages}

Question: {q}

Answer (cite as [P#]):"""

def generate(query: str, chunks: list[dict]) -> dict:
    passages = "\n\n".join(
        f"[P{i+1}] ({c['document_title']} > {c['section_path']})\n{c['chunk_text']}"
        for i, c in enumerate(chunks)
    )
    resp = chat.chat.completions.create(
        model=CHAT_DEPLOYMENT,
        messages=[{"role": "user",
                   "content": ANSWER_PROMPT.format(passages=passages, q=query)}],
        temperature=0.1,
        max_tokens=600
    )
    return {
        "answer": resp.choices[0].message.content,
        "tokens_in":  resp.usage.prompt_tokens,
        "tokens_out": resp.usage.completion_tokens
    }

# COMMAND ----------

def rag(query: str, tenant: str = "casino-prod", user_id: str = "anonymous") -> dict:
    qid = str(uuid.uuid4())
    t0  = time.time()
    candidates = hybrid_retrieve(query, tenant)
    t1  = time.time()
    top = rerank(query, candidates, top_k=5)
    t2  = time.time()
    out = generate(query, top)
    t3  = time.time()

    log_row = {
        "query_id": qid, "query_text": query,
        "retrieved_ids": [c["chunk_id"] for c in top],
        "rerank_scores": [c["rerank_score"] for c in top],
        "generated_answer": out["answer"],
        "latency_ms": int((t3 - t0) * 1000),
        "total_tokens": out["tokens_in"] + out["tokens_out"],
        "cost_usd": (out["tokens_in"]/1e6)*2.50 + (out["tokens_out"]/1e6)*10.00,
        "user_id": user_id,
        "ts": datetime.utcnow().isoformat()
    }
    # Async log to Eventhouse rag_query_log
    log_to_kusto(log_row)
    return {**out, **log_row,
            "retrieve_ms": int((t1-t0)*1000),
            "rerank_ms":   int((t2-t1)*1000),
            "generate_ms": int((t3-t2)*1000)}

# COMMAND ----------

result = rag("What triggers a CTR for cash transactions over $10,000?")
print(result["answer"])

def detect_hallucination(answer: str, chunks: list[dict]) -> dict:
    has_citations = bool(re.search(r"\[P\d+\]", answer))
    refuses_when_uncertain = "I don't have that information" in answer
    # Optional: run Ragas faithfulness on (answer, chunks)
    return {
        "has_citations": has_citations,
        "appears_to_refuse": refuses_when_uncertain,
        "warning": (not has_citations and not refuses_when_uncertain)
    }