Master Retrieval-Augmented Generation

# ── Step 02: Text Preprocessing ──────────────────
import re
import unicodedata

def clean_text(text: str) -> str:
    """Normalize text for embedding quality."""
    # Normalize unicode (e.g. curly quotes → straight)
    text = unicodedata.normalize("NFKC", text)
    # Fix mid-sentence line breaks
    text = re.sub(r'(?<![.!?])\n(?=[a-z])', ' ', text)
    # Collapse multiple spaces / tabs
    text = re.sub(r'[ \t]+', ' ', text)
    # Normalize multiple newlines → double
    text = re.sub(r'\n{3,}', '\n\n', text)
    # Strip non-printable characters
    text = re.sub(r'[^\x20-\x7E\n]', '', text)
    # Remove page numbers / standalone numbers
    text = re.sub(r'^\d+\s*$', '', text, flags=re.MULTILINE)
    return text.strip()

def preprocess_documents(docs: list) -> list:
    """Clean all documents in pipeline."""
    cleaned = []
    for doc in docs:
        clean = clean_text(doc["content"])
        if len(clean) > 30:  # skip near-empty pages
            doc["content"] = clean
            cleaned.append(doc)
    return cleaned

# ── Usage ──────────────────────────────────────────
cleaned_docs = preprocess_documents(docs)
print(f"Cleaned: {len(cleaned_docs)} docs remain")
print(cleaned_docs[0]["content"][:200])
# → Vikash Innovative Tech Warranty Policy
#   All laptops come with a 1-year limited warranty...

# ── Step 03: Chunking (all 3 strategies) ─────────
import re

# ── Strategy 1: Fixed Size ────────────────────────
def chunk_fixed(text, size=300):
    return [text[i:i+size].strip()
            for i in range(0,len(text),size)
            if text[i:i+size].strip()]

# ── Strategy 2: Sliding Window ────────────────────
def chunk_sliding(text, size=300, overlap=60):
    chunks, i = [], 0
    while i < len(text):
        end = min(i+size, len(text))
        t = text[i:end].strip()
        if t: chunks.append(t)
        if end == len(text): break
        i += size - overlap
    return chunks

# ── Strategy 3: Recursive Splitter ───────────────
def chunk_recursive(text, size=300,
                     seps=['\n\n','\n','. ','! ',' ']):
    if len(text) <= size: return [text.strip()]
    sep = next((s for s in seps if s in text), ' ')
    parts = text.split(sep)
    chunks, cur = [], ""
    for p in parts:
        candidate = cur + (sep if cur else "") + p
        if len(candidate) <= size: cur = candidate
        else:
            if cur.strip(): chunks.append(cur.strip())
            cur = p
    if cur.strip(): chunks.append(cur.strip())
    return [c for c in chunks if len(c) >= 15]

# ── Wrap chunks with metadata ─────────────────────
def chunk_documents(docs, strategy="recursive", **kw):
    fn_map = {
        "fixed": chunk_fixed,
        "sliding": chunk_sliding,
        "recursive": chunk_recursive
    }
    fn = fn_map[strategy]
    all_chunks = []
    for doc in docs:
        for j, txt in enumerate(fn(doc["content"], **kw)):
            all_chunks.append({
                "content": txt,
                "metadata": {**doc["metadata"], "chunk": j}
            })
    return all_chunks

chunks = chunk_documents(cleaned_docs, strategy="recursive", size=300)
print(f"Created {len(chunks)} chunks")
# → Created 22 chunks

# ── Step 04: Embeddings ──────────────────────────
from openai import OpenAI
from sentence_transformers import SentenceTransformer
import numpy as np

# ── Option A: OpenAI (paid, production) ───────────
oai = OpenAI()

def embed_openai(texts: list[str]) -> list:
    resp = oai.embeddings.create(
        input=texts,
        model="text-embedding-3-small"  # 1536 dims
    )
    return [d.embedding for d in resp.data]

# ── Option B: Sentence Transformers (free) ─────────
model = SentenceTransformer("all-MiniLM-L6-v2")

def embed_local(texts: list[str]) -> np.ndarray:
    return model.encode(texts, batch_size=32,
                         show_progress_bar=True)

# ── Embed all chunks in batches ───────────────────
def embed_chunks(chunks, batch_size=100):
    texts = [c["content"] for c in chunks]
    all_emb = []
    for i in range(0, len(texts), batch_size):
        batch = texts[i:i+batch_size]
        all_emb.extend(embed_openai(batch))
    return all_emb

embeddings = embed_chunks(chunks)
print(f"Shape: {np.array(embeddings).shape}")
# → Shape: (22, 1536)

# ── Manual cosine similarity ──────────────────────
def cosine_sim(a, b):
    a, b = np.array(a), np.array(b)
    return float(np.dot(a,b) / (np.linalg.norm(a)*np.linalg.norm(b)))

# ── Step 05: Vector Store (ChromaDB) ─────────────
import chromadb
from chromadb.config import Settings

# ── Persistent client (survives restarts) ─────────
client = chromadb.PersistentClient(
    path="./vikash_rag_db"
)

# ── Create or load collection ─────────────────────
collection = client.get_or_create_collection(
    name="vikash_innovative_tech",
    metadata={"hnsw:space": "cosine"}  # cosine dist
)

def store_chunks(chunks, embeddings):
    """Add chunks + embeddings to ChromaDB."""
    collection.add(
        documents=[c["content"] for c in chunks],
        embeddings=embeddings,
        metadatas=[c["metadata"] for c in chunks],
        ids=[str(i) for i in range(len(chunks))]
    )
    print(f"✓ Stored {collection.count()} chunks")

store_chunks(chunks, embeddings)
# → ✓ Stored 22 chunks in vikash_rag_db

# ── Metadata filtering (powerful feature) ─────────
# Find only chunks from warranty policy pages:
results = collection.query(
    query_embeddings=[query_emb],
    n_results=3,
    where={"source": "warranty_policy.pdf"}
)

# ── Check if collection already built ─────────────
if collection.count() > 0:
    print("Loaded existing index — skip rebuild")
else:
    store_chunks(chunks, embeddings)

# ── Step 06: Retrieval ───────────────────────────
import numpy as np

def retrieve(query: str, k=3) -> list[dict]:
    """Find k most relevant chunks for a query."""
    # 1. Embed the query (same model as chunks!)
    q_emb = embed_openai([query])[0]

    # 2. Query vector store
    results = collection.query(
        query_embeddings=[q_emb],
        n_results=k
    )

    # 3. Format results with similarity scores
    retrieved = []
    for i in range(len(results["documents"][0])):
        sim = 1 - results["distances"][0][i]
        if sim > 0.5:  # Filter low-relevance
            retrieved.append({
                "text": results["documents"][0][i],
                "score": round(sim, 4),
                "source": results["metadatas"][0][i],
            })
    return sorted(retrieved,
                   key=lambda x:x["score"],
                   reverse=True)

# ── Example ───────────────────────────────────────
results = retrieve("What does the warranty cover?")
for r in results:
    print(f"Score: {r['score']:.3f}")
    print(r["text"][:100] + "...")
# → Score: 0.940
#   All laptops come with a 1-year limited...
# → Score: 0.872
#   Warranty covers: manufacturing defects...

# ── Step 07: LLM + Prompt Engineering ────────────
from openai import OpenAI

client = OpenAI()

def build_context(retrieved_chunks):
    """Format retrieved chunks as numbered context."""
    return "\n\n".join(
        f"[{i+1}] (Source: {c['source'].get('source','?')}, "
        f"page {c['source'].get('page','?')})\n{c['text']}"
        for i, c in enumerate(retrieved_chunks)
    )

def rag_answer(question: str, k=3) -> str:
    """Full RAG: retrieve → prompt → answer."""
    # Retrieve relevant chunks
    chunks = retrieve(question, k=k)
    context = build_context(chunks)

    # Build system prompt with strict grounding
    system = """You are the AI assistant for Vikash Innovative Tech.
Answer ONLY based on the retrieved context provided.
If the answer is not in the context, say exactly:
'I don't have enough information to answer this.'
Always cite sources as [1], [2], etc.
Be concise and accurate."""

    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content":
             f"Context:\n{context}\n\nQuestion: {question}"}
        ],
        temperature=0.1
    )
    return response.choices[0].message.content

# ── Complete pipeline in 3 lines ──────────────────
answer = rag_answer("What does the warranty cover?")
print(answer)
# → Based on [1], Vikash Innovative Tech's 1-year
#   warranty covers manufacturing defects, hardware
#   malfunctions, and faulty components...