RAG vs. prompt stuffing: Do we still need vector retrieval?

pip install ragas==0.2.15 langchain==0.3.25 google-genai==1.11.0

import json
from tqdm import tqdm
from datasets import load_dataset
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.docstore.document import Document as LangchainDocument
from transformers import GPT2TokenizerFast

import os

os.environ["GOOGLE_CLOUD_PROJECT"] = "your-project"
os.environ["GOOGLE_CLOUD_LOCATION"] = "us-central1"
os.environ["GOOGLE_GENAI_USE_VERTEXAI"] = "True"

import json
from tqdm import tqdm
from datasets import load_dataset
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.docstore.document import Document as LangchainDocument
from transformers import GPT2TokenizerFast
from google import genai

client = genai.Client()
MAX_GEMINI_TOKENS = 500_000  # 1/2 million

def gemini_count_tokens(text):
    result = client.models.count_tokens(
        model="gemini-2.5-flash-preview-05-20",
        contents=text
    )
    if isinstance(result, dict):
        return result["total_tokens"]
    return result.total_tokens

# 1. Load dataset
dataset = load_dataset("wikimedia/wikisource", "20231201.en")
train_dataset = dataset['train'].select(range(50000))

# 2. Convert to Langchain Document objects
langchain_docs = [
    LangchainDocument(
        page_content=doc["text"],
        metadata={"id": doc["id"], "url": doc["url"], "title": doc["title"]}
    )
    for doc in tqdm(train_dataset)
]

# 3. Initialize text splitter and tokenizer
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=2000,
    chunk_overlap=200,
    add_start_index=True,
    separators=["\n\n", "\n", ".", " ", ""],
)
tokenizer = GPT2TokenizerFast.from_pretrained("gpt2")

# 4. Process, capping Gemini tokens globally
docs_processed = []
total_gemini_tokens = 0  # <--- THIS IS GLOBAL NOW
done = False 
for doc in tqdm(langchain_docs, desc="Docs"):
    if done:
        break 
    chunks = text_splitter.split_documents([doc])
    for i, chunk in enumerate(chunks):
        chunk.metadata["custom_id"] = f"{chunk.metadata['id']}_chunk_{i}"
        gpt2_tokens = len(tokenizer.encode(chunk.page_content))
        chunk.metadata["gpt2_token_count"] = gpt2_tokens

        gemini_tokens = gemini_count_tokens("ID: " + str(chunk.metadata['id']) +  "\n" + str(chunk.metadata['title']) + "\n" + chunk.page_content)
        print(str(gemini_tokens), flush=True)
        chunk.metadata["gemini_token_count"] = gemini_tokens
        print(total_gemini_tokens, flush=True)
        if total_gemini_tokens + gemini_tokens > MAX_GEMINI_TOKENS:
            print("Global Gemini token limit reached, stopping.")
            done = True 
            break  # Stop processing further chunks (and docs)
        docs_processed.append(chunk)
        total_gemini_tokens += gemini_tokens




# 5. Convert to JSON format
def documents_to_json(docs):
    json_list = []
    for doc in docs:
        json_item = {
            "custom_id": doc.metadata["custom_id"],
            "id": doc.metadata["id"],
            "url": doc.metadata["url"],
            "title": doc.metadata["title"],
            "gpt2_token_count": doc.metadata["gpt2_token_count"],
            "gemini_token_count": doc.metadata["gemini_token_count"],
            "text": doc.page_content
        }
        json_list.append(json_item)
    return json_list

json_data = documents_to_json(docs_processed)

# 6. Save to disk
output_file_path = "./chunked_wikisource_data.json"
with open(output_file_path, 'w', encoding='utf-8') as file:
    json.dump(json_data, file, ensure_ascii=False, indent=4)

print(f"Chunked data has been saved to {output_file_path}")


import os

os.environ["GOOGLE_CLOUD_PROJECT"] = "your_project"
os.environ["GOOGLE_CLOUD_LOCATION"] = "us-central1"
os.environ["GOOGLE_GENAI_USE_VERTEXAI"] = "True"

import json
from typing import List
from langchain.docstore.document import Document as LangchainDocument
from langchain_community.vectorstores import Chroma

from google import genai
from google.genai.types import EmbedContentConfig

# ------------- Custom Embedding wrapper for Gemini -------------
class GoogleGeminiEmbeddings:
    def __init__(self, model="text-embedding-004"):
        self.model = model
        self.client = genai.Client()

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        embeddings = []
        total = len(texts)
        print(f"Embedding {total} documents...", flush=True)
        for i, text in enumerate(texts):
            if i % 10 == 0 or i == total - 1:
                print(f"> Embedding doc {i+1}/{total}", flush=True)
            response = self.client.models.embed_content(
                model=self.model,
                contents=text,
                config=EmbedContentConfig(
                    task_type="RETRIEVAL_DOCUMENT",
                    output_dimensionality=768,
                ),
            )
            emb = response.embeddings[0].values
            embeddings.append(emb)
        print("All embeddings complete.")
        return embeddings

    def embed_query(self, text: str) -> List[float]:
        print(f"Embedding query: {text[:60]}{'...' if len(text)>60 else ''}")
        response = self.client.models.embed_content(
            model=self.model,
            contents=text,
            config=EmbedContentConfig(
                task_type="RETRIEVAL_QUERY",
                output_dimensionality=768,
            ),
        )
        print("Query embedding complete.")
        return response.embeddings[0].values

# ----------- Load the chunked data -----------
DATA_PATH = "./chunked_wikisource_data.json"
print(f"Loading chunked data from {DATA_PATH}...")
with open(DATA_PATH, "r", encoding="utf-8") as f:
    data_json = json.load(f)

docs = []
for idx, item in enumerate(data_json):
    if idx < 2:
        print(f"Sample doc {idx+1}: '{item['title'][:60]}...' tokens: {item['gpt2_token_count']}")
    meta = {
        "custom_id": item["custom_id"],
        "id": item["id"],
        "url": item["url"],
        "title": item["title"],
        "gemini_token_count": item["gemini_token_count"],
    }
    docs.append(LangchainDocument(page_content=item["text"], metadata=meta))

print(f"Loaded {len(docs)} documents from {DATA_PATH}")

# ------------ Create Chroma vector db with Gemini Embeddings ----------
persist_dir = "./chroma_wikisource_vector_db_gemini"
os.makedirs(persist_dir, exist_ok=True)

print("Building vector database using Gemini Embeddings...")
embeddings = GoogleGeminiEmbeddings(model="text-embedding-004")

vectorstore = Chroma.from_documents(
    documents=docs,
    embedding=embeddings,
    persist_directory=persist_dir,
)

print(f"Vector DB created and saved at {persist_dir}")

# (Optional) Test a simple retrieval
query = "What is the Magna Carta?"
print(f"\nPerforming retrieval for test query: '{query}'")
retriever = vectorstore.as_retriever()
results = retriever.invoke(query)
print(f"\nExample search results for query: '{query}'\n")
for idx, r in enumerate(results):
    print(f"Result {idx+1}:")
    print("Title:", r.metadata.get('title'))
    print("Tokens:", r.metadata.get('gpt2_token_count'))
    print("Excerpt:", r.page_content[:200], "...\n")


import os
import json
import pandas as pd

# --- CONFIG ---
JSON_PATH = "./chunked_wikisource_data.json"
OUT_CSV   = "./generated_testset.csv"
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
assert OPENAI_API_KEY, "Set your OPENAI_API_KEY variable!"

os.environ["OPENAI_API_KEY"] = OPENAI_API_KEY
TESTSET_SIZE = 30

# --- RAGAS + LANGCHAIN IMPORTS ---
from ragas.llms import LangchainLLMWrapper
from ragas.embeddings import LangchainEmbeddingsWrapper
from ragas.testset.graph import KnowledgeGraph, Node, NodeType
from ragas.testset.transforms import default_transforms, apply_transforms
from ragas.testset import TestsetGenerator
from ragas.testset.synthesizers.single_hop.specific import SingleHopSpecificQuerySynthesizer
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain.docstore.document import Document as LCDocument

# --- 1. LOAD DOCUMENTS ---
with open(JSON_PATH, "r", encoding="utf-8") as f:
    raw_chunks = json.load(f)
docs = [
    LCDocument(
        page_content=chunk["text"],
        metadata={
            "custom_id": chunk.get("custom_id"),
            "id": chunk.get("id"),
            "url": chunk.get("url"),
            "title": chunk.get("title"),
            "gpt2_token_count": chunk.get("gpt2_token_count"),
        },
    )
    for chunk in raw_chunks
]
print(f"Loaded {len(docs)} document chunks.")

# --- 2. CREATE KNOWLEDGE GRAPH & ADD DOCS ---
kg = KnowledgeGraph()
for doc in docs:  
    kg.nodes.append(Node(
        type=NodeType.DOCUMENT,
        properties={
            "page_content": doc.page_content,
            "document_metadata": doc.metadata
        }
    ))
print(f"KnowledgeGraph: {len(kg.nodes)} nodes, {len(kg.relationships)} relationships (before enrichment)")

# --- 3. ENRICH KNOWLEDGEGRAPH with TRANSFORMS ---
generator_llm      = LangchainLLMWrapper(ChatOpenAI(model="gpt-4o-mini", temperature=0))
generator_emb      = LangchainEmbeddingsWrapper(OpenAIEmbeddings())
transforms         = default_transforms(
    documents=docs,
    llm=generator_llm,
    embedding_model=generator_emb
)
apply_transforms(kg, transforms)
print(f"Enriched KnowledgeGraph: {len(kg.nodes)} nodes, {len(kg.relationships)} relationships (after enrichment)")

# Optionally, save & reload the knowledge graph
kg.save("knowledge_graph.json")
kg = KnowledgeGraph.load("knowledge_graph.json")

# --- 4. GENERATE TESTSET using the enriched KG! ---
synthesizers = [
    (SingleHopSpecificQuerySynthesizer(llm=generator_llm), 1.0),
]
generator = TestsetGenerator(
    llm=generator_llm,
    embedding_model=generator_emb,
    knowledge_graph=kg,
)
testset = generator.generate(
    testset_size=TESTSET_SIZE,
    query_distribution=synthesizers
)
df = testset.to_pandas()
df.to_csv(OUT_CSV, index=False)
print(f"\nDONE. Saved testset to {OUT_CSV} ({len(df)} rows)")

# ---- 5. SHOW SAMPLES ----
print(df[["user_input", "reference"]].head(5).to_markdown(index=False))


import os
import json
import pandas as pd
import ast
from tqdm import tqdm
import warnings
from typing import List, Dict, Any

from ragas import evaluate
from ragas.llms import LangchainLLMWrapper
from ragas.dataset_schema import EvaluationDataset
from ragas.metrics import (
    LLMContextRecall, Faithfulness, FactualCorrectness,
    LLMContextPrecisionWithoutReference, NoiseSensitivity,
    ResponseRelevancy, ContextEntityRecall
)
from ragas.run_config import RunConfig
import plotly.graph_objects as go
import time

from langchain.docstore.document import Document as LangchainDocument
from langchain_community.vectorstores import Chroma
from google import genai
from google.genai.types import GenerateContentConfig
from langchain_core.language_models.llms import LLM
import wandb
from google.genai.types import EmbedContentConfig


# SETUP
os.environ["GOOGLE_CLOUD_PROJECT"] = "your_project"
os.environ["GOOGLE_CLOUD_LOCATION"] = "us-central1"
os.environ["GOOGLE_GENAI_USE_VERTEXAI"] = "True"

class GoogleGeminiLLM(LLM):
    model_name: str = "gemini-2.5-flash-preview-05-20"
    temperature: float = 0.0
    api_version: str = "v1"
    api_key: Any = None
    max_tokens: Any = None

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        object.__setattr__(self, "_client", genai.Client(api_key=self.api_key, http_options={"api_version": self.api_version}))


        
    def _call(
        self,
        prompt: str,
        stop: Any = None,
        run_manager: Any = None,
        **kwargs: Any,
    ) -> str:
        config = GenerateContentConfig(
            temperature=self.temperature,
            max_output_tokens=self.max_tokens if self.max_tokens else None
        )
        response = self._client.models.generate_content(
            model=self.model_name,
            contents=prompt,
            config=config
        )
        return response.text if hasattr(response, 'text') else str(response)

    @property
    def _llm_type(self) -> str:
        return "google-gemini"

    @property
    def _identifying_params(self) -> Dict[str, Any]:
        return {"model_name": self.model_name,"temperature": self.temperature}





# ------------- Custom Embedding wrapper for Gemini -------------
class GoogleGeminiEmbeddings:
    def __init__(self, model="text-embedding-004"):
        self.model = model
        self.client = genai.Client()

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        embeddings = []
        total = len(texts)
        print(f"Embedding {total} documents...", flush=True)
        for i, text in enumerate(texts):
            if i % 10 == 0 or i == total - 1:
                print(f"> Embedding doc {i+1}/{total}", flush=True)
            response = self.client.models.embed_content(
                model=self.model,
                contents=text,
                config=EmbedContentConfig(
                    task_type="RETRIEVAL_DOCUMENT",
                    output_dimensionality=768,
                ),
            )
            emb = response.embeddings[0].values
            embeddings.append(emb)
        print("All embeddings complete.")
        return embeddings

    def embed_query(self, text: str) -> List[float]:
        print(f"Embedding query: {text[:60]}{'...' if len(text)>60 else ''}")
        response = self.client.models.embed_content(
            model=self.model,
            contents=text,
            config=EmbedContentConfig(
                task_type="RETRIEVAL_QUERY",
                output_dimensionality=768,
            ),
        )
        print("Query embedding complete.")
        return response.embeddings[0].values






llm = GoogleGeminiLLM(
    model_name="gemini-2.5-flash-preview-05-20",
    temperature=0.0,
    # api_version="v1"
)

# DATA LOAD
testset_path = "generated_testset.csv"
data_path = "./chunked_wikisource_data.json"
# df = pd.read_csv(testset_path)

df = pd.read_csv(testset_path)#.head(3)


with open(data_path, "r", encoding="utf-8") as f:
    docs_all_json = json.load(f)

def safe_parse(val):
    try:
        return ast.literal_eval(val) if isinstance(val, str) else val
    except:
        return []
        
df["reference_contexts"] = df["reference_contexts"].apply(safe_parse)

docs_all = [
    LangchainDocument(
        page_content=docj["text"],
        metadata={
            "custom_id": docj.get("custom_id"),
            "id": docj.get("id"),
            "url": docj.get("url"),
            "title": docj.get("title"),
            "gpt2_token_count": docj.get("gpt2_token_count"),
        }
    )
    for docj in docs_all_json
]
docs_by_id = {doc["custom_id"]: doc for doc in docs_all_json}

# (A) Vanilla RAG (vector DB)
persist_dir = "./chroma_wikisource_vector_db_gemini"
if not os.path.exists(persist_dir):
    raise ValueError(f"Chroma vector DB '{persist_dir}' not found!")



embeddings = GoogleGeminiEmbeddings(model="text-embedding-004")

vectorstore = Chroma(
    persist_directory=persist_dir,
    embedding_function=embeddings # Already built with Gemini embeddings
)
retriever = vectorstore.as_retriever(search_kwargs={"k": 20})

def format_context(docs):
    context = ""
    for i, doc in enumerate(docs):
        context += f"Source [{i+1}] (Title: {doc.metadata.get('title', 'n/a')} | ID: {doc.metadata.get('custom_id')})\n"
        context += doc.page_content.strip() + "\n---\n"
    return context

def rag_vanilla(query: str) -> (str, List[str]):
    docs = retriever.invoke(query)
    context = format_context(docs)
    prompt = (
        f"Given the following documents and a user question, write a detailed, accurate answer using ONLY the information in the sources. "
        f"Do not use outside knowledge or hallucinate details. Cite your supporting source chunk(s) as [1], [2], etc. as appropriate.\n\n"
        f"=== Sources ===\n{context}"
        f"=== User Question ===\n{query}\n\n"
        f"=== Answer ==="
    )
    answer = llm.invoke(prompt)
    return answer, [doc.page_content for doc in docs]

# (B) In-Context All-Docs Prompt
def make_doc_summary(doc):
    snippet = " ".join(doc["text"].split())
    return f'{doc["custom_id"]} | {doc["title"]}: {snippet}'

# *** HERE IS THE FULL-UNLIMITED DOC ID LIST ***
ALL_DOC_IDS = [doc["custom_id"] for doc in docs_all_json]

def incontext_retrieval_and_response(query, num_to_retrieve=20):
    doc_summaries = [make_doc_summary(docs_by_id[doc_id]) for doc_id in ALL_DOC_IDS]
    retrieval_prompt = (
        f"You are a helpful retrieval engine. Given a user question and a list of document summaries (each with a unique ID and title), "
        f"return a comma-separated list of exactly {num_to_retrieve} document IDs most relevant for answering the question. "
        f"Only use IDs from the list provided - no made-up IDs!\n\n"
        f"Question: {query}\n"
        f"DOCUMENT SUMMARIES:\n" +
        "\n".join(doc_summaries) +
        "\n\nYour response: (just the comma-separated list of IDs)\n"
    )
    ids_raw = llm.invoke(retrieval_prompt)
    ids = []
    for piece in ids_raw.replace("\n",",").split(","):
        x = piece.strip()
        if x in docs_by_id:
            ids.append(x)
        if len(ids) == num_to_retrieve:
            break
    if not ids:
        print("DOC ID PARSING FAILED!!!!!!!")
        ids = ALL_DOC_IDS[:num_to_retrieve]
    context_chunks = [docs_by_id[i] for i in ids]
    context_text = ""
    for i, doc in enumerate(context_chunks):
        context_text += f"[Doc {i+1} | ID: {doc['custom_id']}] {doc['title']}:\n{doc['text']}\n---\n"
    answer_prompt = (
        f"You are a helpful question-answering assistant. Using ONLY the information in the provided documents, answer the user's question as specifically and factually as possible."
        f" If the answer is not contained, say so. Cite sources as [1], [2], etc. as appropriate.\n\n"
        f"=== DOCUMENTS ===\n{context_text}\n"
        f"=== QUESTION ===\n{query}\n\n"
        f"=== ANSWER ==="
    )
    answer = llm.invoke(answer_prompt)
    return answer, [doc["text"] for doc in context_chunks]

# Run both modes over testset
def run_and_save(mode_func, mode_name, **kwargs):
    sys_answers = []
    sys_contexts = []
    print(f"\nRunning {mode_name} over testset...")
    for i, row in tqdm(df.iterrows(), total=len(df)):
        query = row['user_input']
        try:
            ans, contexts = mode_func(query, **kwargs)
        except Exception as e:
            print(f"Error at idx {i}: {e}")
            ans, contexts = "", []
        sys_answers.append(ans)
        sys_contexts.append(contexts)
    col_prefix = "incontext" if mode_name.lower().startswith("incontext") else "vanilla"
    df[f"{col_prefix}_response"] = sys_answers
    df[f"{col_prefix}_contexts"] = sys_contexts
    out = f"results_{col_prefix}_gemini.csv"
    df.to_csv(out, index=False)
    print(f"Saved {mode_name} system answers to {out}")
    return out, f"{col_prefix}_response", f"{col_prefix}_contexts"

# Run Vanilla RAG
out_vanilla_filename, resp_col_vanilla, ctx_col_vanilla = run_and_save(rag_vanilla, "Vanilla")
# Run In-Context RAG, FULL DOCS
out_incontext_filename, resp_col_incontext, ctx_col_incontext = run_and_save(incontext_retrieval_and_response, "InContext")



def eval_with_ragas(
    response_col,
    contexts_col,
    mode_label,
    project="rag-eval",
    run=None,
):


    if df is None:
        raise ValueError("Input dataframe 'df' must be provided.")

    eval_records = []
    missing_keys = set()

    # Safely build the evaluation records
    for i, row in df.iterrows():
        record = {}
        try:
            record["user_input"] = row["user_input"]
            record["reference"] = row["reference"]
            record["response"] = row[response_col]
            record["retrieved_contexts"] = row[contexts_col]
        except KeyError as e:
            print(f"Row {i}: missing expected column ({e}). Skipping row.")
            missing_keys.add(str(e))
            continue
        eval_records.append(record)

    if not eval_records:
        raise ValueError("No valid records to evaluate. Check your dataframe columns.")

    try:
        eval_dataset = EvaluationDataset.from_list(eval_records)
    except Exception as err:
        print(f"Failed to create EvaluationDataset: {err}")
        return

    try:
        eval_llm = LangchainLLMWrapper(
            GoogleGeminiLLM(
                model_name="gemini-2.5-flash-preview-05-20",
                temperature=0.0,
                api_version="v1"
            )
        )
    except Exception as err:
        print(f"Failed to initialize LLM: {err}")
        return

    run_config = RunConfig(max_workers=1, timeout=180)

    print(f"\nEvaluating {mode_label} with RAGAS/Gemini judge (can take a while)...")
    
    metrics = [
        LLMContextRecall(), Faithfulness(),
        LLMContextPrecisionWithoutReference(),
        ResponseRelevancy(), ContextEntityRecall()
    ]

    try:
        results = evaluate(
            dataset=eval_dataset,
            metrics=metrics,
            llm=eval_llm,
            run_config=run_config,
        )
        results_df = results.to_pandas()
    except Exception as err:
        print(f"Evaluation failed: {err}")
        return

    eval_csv_out = f"evaluation_results_{mode_label.lower().replace('-','_')}_gemini.csv"
    try:
        results_df.to_csv(eval_csv_out, index=False)
        print(f"Saved RAGAS evaluation to {eval_csv_out}")
    except Exception as err:
        print(f"Failed to save results CSV: {err}")

    # Compute radar metrics (with KeyError protection)
    metric_map = {
        "Context Recall": "context_recall",
        "Faithfulness": "faithfulness",
        "Context Precision": "llm_context_precision_without_reference",
        "Answer Relevancy": "answer_relevancy",
        "Context Entity Recall": "context_entity_recall",
    }
    radar_metrics = {}
    for display, col in metric_map.items():
        try:
            radar_metrics[display] = results_df[col].mean()
        except KeyError:
            print(f"Warning: Metric '{col}' missing in results. Setting to 0.")
            radar_metrics[display] = 0.0

    for k, v in radar_metrics.items():
        print(f"{mode_label} | {k}: {v:.3f}")

    # Radar plot with error handling
    fig = go.Figure()
    metrics_list = list(radar_metrics.keys())
    vals = list(radar_metrics.values())
    if metrics_list and vals:
        # close the radar chart loop
        metrics_list.append(metrics_list[0])
        vals.append(vals[0])
        fig.add_trace(go.Scatterpolar(
            r=vals, theta=metrics_list, fill='toself', name=mode_label
        ))
        fig.update_layout(
            polar=dict(radialaxis=dict(visible=True, range=[0, 1], tickvals=[0, 0.5, 1])),
            title=dict(text=f"RAGAS Results ({mode_label})", x=0.5),
            showlegend=True
        )
        try:
            path = f"./radar_results_{mode_label.lower()}_{int(time.time())}.html"
            fig.write_html(path, auto_play=False)
            print(f"Radar plot for {mode_label} saved: {path}")
        except Exception as err:
            print(f"Failed to save radar plot: {err}")
    else:
        print("Radar metrics are empty: unable to create plot.")

    # --- WANDB Logging ---
    try:
        if run is None:
            wandb_run = wandb.init(project=project, name=mode_label)
        else:
            wandb_run = run

        wandb_run.log({f"{mode_label}/" + k: v for k, v in radar_metrics.items()})

        if 'fig' in locals():
            wandb_run.log({f"{mode_label}_radar_chart": wandb.Plotly(fig)})

        if run is None:
            wandb_run.finish()
    except Exception as err:
        print(f"WANDB logging failed: {err}")

    if missing_keys:
        print(f"Completed with missing columns: {missing_keys}")



# -------- Example Usage (run just once at start) --------

wandb_run = wandb.init(project="rag-eval", name="rag_eval_run")


# InContext RAG
eval_with_ragas(resp_col_incontext, ctx_col_incontext, "InContext RAG", run=wandb_run)


# Vanilla RAG
eval_with_ragas(resp_col_vanilla, ctx_col_vanilla, "Vanilla RAG", run=wandb_run)



wandb_run.finish()
print("\nEvaluation completed for BOTH Vanilla and In-Context modes and logged to wandb.")