Retrieval-augmented generation (RAG) end to end

Brew & Bean is open every day from 8 am to 10 pm.

We do not deliver. Orders are pickup only: order on the website, then collect your order from the counter at your pickup time.

You can cancel an order only while its status is still received. Once it is preparing, it cannot be cancelled.

Allergens: the Butter Croissant contains gluten and milk. The Banana Bread contains gluten, eggs and walnuts. The Cappuccino, Filter Coffee and Masala Chai contain milk.

The Espresso and the Cold Brew contain none of the common allergens.

Our kitchen handles nuts, milk and gluten, so we cannot promise that any item is completely free from traces.

Oat milk can be used in any coffee or chai for an extra Rs 40.

We accept UPI, cards and cash at the counter.

The Cold Brew is steeped in cold water for 18 hours and served over ice.

import re
import sys
from pathlib import Path

import numpy as np
from google import genai
from google.genai import types

client = genai.Client()  # reads GEMINI_API_KEY from the environment
EMBED_MODEL = "gemini-embedding-001"
GEN_MODEL = "gemini-flash-latest"  # check the model names on ai.google.dev


def chunk_text(text, source, max_chars=600):
    """Split on blank lines, then pack paragraphs into chunks of about max_chars."""
    paragraphs = [p.strip() for p in re.split(r"\n\s*\n", text) if p.strip()]
    chunks, current = [], ""
    for p in paragraphs:
        if current and len(current) + len(p) > max_chars:
            chunks.append(current)
            current = p
        else:
            current = current + "\n\n" + p if current else p
    if current:
        chunks.append(current)
    return [{"source": source, "text": c} for c in chunks]


def embed(texts, task_type):
    """Turn a list of strings into a matrix of unit-length vectors."""
    vectors = []
    for i in range(0, len(texts), 50):  # send at most 50 texts per request
        resp = client.models.embed_content(
            model=EMBED_MODEL,
            contents=texts[i : i + 50],
            config=types.EmbedContentConfig(task_type=task_type),
        )
        vectors.extend(e.values for e in resp.embeddings)
    arr = np.array(vectors, dtype=np.float32)
    return arr / np.linalg.norm(arr, axis=1, keepdims=True)


def build_index(folder):
    chunks = []
    for path in sorted(Path(folder).glob("*.txt")):
        chunks += chunk_text(path.read_text(encoding="utf-8"), path.name)
    matrix = embed([c["text"] for c in chunks], "RETRIEVAL_DOCUMENT")
    return chunks, matrix


def retrieve(question, chunks, matrix, k=3):
    q = embed([question], "RETRIEVAL_QUERY")[0]
    scores = matrix @ q  # cosine similarity, because all vectors have length 1
    top = np.argsort(-scores)[:k]
    return [(chunks[i], float(scores[i])) for i in top]


SYSTEM = (
    "You answer questions for Brew & Bean, a coffee shop. Use ONLY the numbered "
    "sources in the user message. If the sources do not contain the answer, say "
    "'I don't know based on our documents.' Cite sources like [1] after each claim."
)


def answer(question, chunks, matrix):
    hits = retrieve(question, chunks, matrix)
    context = "\n\n".join(
        f"[{n}] ({c['source']})\n{c['text']}" for n, (c, _) in enumerate(hits, 1)
    )
    prompt = f"Sources:\n{context}\n\nQuestion: {question}"
    resp = client.models.generate_content(
        model=GEN_MODEL,
        contents=prompt,
        config=types.GenerateContentConfig(system_instruction=SYSTEM, temperature=0.2),
    )
    return resp.text, hits


if __name__ == "__main__":
    folder = sys.argv[1] if len(sys.argv) > 1 else "docs"
    chunks, matrix = build_index(folder)
    print(f"Indexed {len(chunks)} chunks. Ask a question (Ctrl+C to quit).")
    while True:
        q = input("\nYou: ")
        text, hits = answer(q, chunks, matrix)
        print("\nAnswer:", text)
        print("\nRetrieved:")
        for chunk, score in hits:
            print(f"  {score:.2f}  {chunk['source']}: {chunk['text'][:60]!r}")

python rag.py docs

1. Does the banana bread have nuts?
2. How long is the cold brew steeped?
3. What is the CEO's name?