RAG (Retrieval-Augmented Generation)

from openai import OpenAI
import numpy as np
from typing import List, Dict

client = OpenAI()

# ══════════════════════════════════════════════════════════
#  PHASE 1: INDEXING (runs once when you upload a document)
# ══════════════════════════════════════════════════════════

def chunk_text(text: str, chunk_size: int = 500, overlap: int = 100) -> List[str]:
    """Split text into overlapping chunks."""
    words = text.split()
    chunks = []
    step = chunk_size - overlap

    for i in range(0, len(words), step):
        chunk = " ".join(words[i : i + chunk_size])
        if chunk:
            chunks.append(chunk)

    return chunks


def embed_texts(texts: List[str]) -> np.ndarray:
    """Embed a list of texts using OpenAI's embedding model."""
    response = client.embeddings.create(
        model="text-embedding-3-small",
        input=texts
    )
    embeddings = [r.embedding for r in response.data]
    return np.array(embeddings)


# Your document (in practice, extracted from PDF)
document = """
Photosynthesis is the process by which plants convert sunlight into glucose.
It occurs in two stages: the light-dependent reactions in the thylakoids,
and the Calvin cycle in the stroma of the chloroplast.
The overall equation is: 6CO2 + 6H2O + light → C6H12O6 + 6O2.
...
"""

# Split into chunks and embed
chunks = chunk_text(document)
chunk_embeddings = embed_texts(chunks)     # shape: (n_chunks, 1536)

# Normalize for fast cosine similarity
norms = np.linalg.norm(chunk_embeddings, axis=1, keepdims=True)
chunk_embeddings_norm = chunk_embeddings / norms

# Store in "vector database" (in-memory here; use pgvector/Pinecone in production)
index = {"chunks": chunks, "embeddings": chunk_embeddings_norm}

print(f"Indexed {len(chunks)} chunks")

# ══════════════════════════════════════════════════════════
#  PHASE 2: QUERY (runs on every user question)
# ══════════════════════════════════════════════════════════

def retrieve(query: str, index: Dict, top_k: int = 3) -> List[str]:
    """Find the most relevant chunks for a query."""
    query_emb = embed_texts([query])[0]
    query_emb = query_emb / np.linalg.norm(query_emb)     # normalize

    # Cosine similarity = dot product with normalized vectors
    scores = index["embeddings"] @ query_emb               # (n_chunks,)
    top_indices = np.argsort(scores)[::-1][:top_k]

    return [index["chunks"][i] for i in top_indices]


def rag_answer(question: str, index: Dict) -> str:
    """Retrieve relevant context, then generate a grounded answer."""

    # Step 1: Retrieve
    context_chunks = retrieve(question, index, top_k=3)
    context = "\n\n".join(f"[Chunk {i+1}]: {c}" for i, c in enumerate(context_chunks))

    # Step 2: Generate (grounded answer)
    system_prompt = """You are a precise AI assistant.
Answer ONLY from the provided context below.
If the answer is not in the context, say "I don't have that information in the provided documents."
Always cite which chunk your answer comes from."""

    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user",   "content": f"Context:\n{context}\n\nQuestion: {question}"}
        ],
        temperature=0    # deterministic for factual Q&A
    )

    return response.choices[0].message.content


answer = rag_answer("What is the equation for photosynthesis?", index)
print(answer)
# "According to Chunk 1, the equation for photosynthesis is:
#  6CO2 + 6H2O + light → C6H12O6 + 6O2"

SYSTEM: You are a precise study assistant. Answer questions based ONLY on 
the provided context passages below. 

Rules:
1. If the answer is directly stated in the context, cite it with [Chunk N].  
2. If the answer is implied but not directly stated, say "Based on the context..."
3. If the answer is NOT in the context, respond: "This isn't covered in the 
   provided material. Try asking a more specific question or uploading a 
   document that covers this topic."
4. NEVER use your general knowledge to fill gaps — only use the context.

This is critical for academic use: fabricating information could harm students.

# pip install graphrag
# GraphRAG builds a knowledge graph from your documents — enabling 
# "what are the main themes across all documents?" style queries

import os
from graphrag.index import run_pipeline
from graphrag.query import GraphRagConfig, run_local_search, run_global_search

# Initialize GraphRAG workspace
os.makedirs("./graphrag_workspace/input", exist_ok=True)

# Put your documents in ./graphrag_workspace/input/
# Then initialize and run the indexing pipeline:
# $ graphrag init --root ./graphrag_workspace
# $ graphrag index --root ./graphrag_workspace
# (Indexing builds entities, relationships, communities, and summaries)

# --- Local search: specific entity questions ---
local_answer = run_local_search(
    root_dir="./graphrag_workspace",
    query="What did the research find about transformer attention mechanisms?",
    # Returns entity-focused, attribution-rich answer
)

# --- Global search: corpus-wide synthesis ---
global_answer = run_global_search(
    root_dir="./graphrag_workspace",
    query="What are the three most important themes across all research papers?",
    # Synthesizes community summaries — impossible with naive RAG
)

print("LOCAL:", local_answer.response[:300])
print("GLOBAL:", global_answer.response[:300])
# GraphRAG global search: particularly powerful for research synthesis,
# legal corpus analysis, and large enterprise knowledge bases