Vector Databases & RAG Basics

python -m venv .venv && source .venv/bin/activate
pip install numpy
# Optional cloud embeddings (verify package names at build):
# pip install openai
# export OPENAI_API_KEY="sk-..."

# Pattern A: OpenAI-compatible embeddings (verify model id at build)
from openai import OpenAI
 
client = OpenAI()
resp = client.embeddings.create(
    model="text-embedding-3-small",
    input="Refunds are available within 30 days of purchase.",
)
vector = resp.data[0].embedding
print(len(vector), vector[:5])

# Pattern B: offline-friendly stub for tests (not production quality)
import hashlib
import struct
 
def fake_embed(text: str, dim: int = 32) -> list[float]:
    """Deterministic pseudo-embedding for demos when no API is available."""
    out = []
    seed = text.encode("utf-8")
    for i in range(dim):
        h = hashlib.sha256(seed + i.to_bytes(2, "little")).digest()
        out.append(struct.unpack("f", h[:4])[0])
    # L2 normalize
    norm = sum(x * x for x in out) ** 0.5 or 1.0
    return [x / norm for x in out]
 
print(len(fake_embed("hello world")))

import math
 
def cosine(a: list[float], b: list[float]) -> float:
    dot = sum(x * y for x, y in zip(a, b))
    na = math.sqrt(sum(x * x for x in a)) or 1.0
    nb = math.sqrt(sum(y * y for y in b)) or 1.0
    return dot / (na * nb)
 
docs = [
    {"id": "d1", "text": "Refunds are available within 30 days with a receipt."},
    {"id": "d2", "text": "Support hours are Monday to Friday, 9am to 5pm Eastern."},
    {"id": "d3", "text": "Annual plans renew automatically unless cancelled."},
]
 
# Use fake_embed from example 1 or replace with a real embed() call
for d in docs:
    d["vector"] = fake_embed(d["text"])
 
print(f"indexed {len(docs)} documents")

def search(query: str, docs: list[dict], k: int = 2) -> list[dict]:
    q = fake_embed(query)
    ranked = sorted(
        docs,
        key=lambda d: cosine(q, d["vector"]),
        reverse=True,
    )
    return ranked[:k]
 
for hit in search("How long do I have to return a purchase?", docs, k=2):
    print(hit["id"], hit["text"])

def build_rag_prompt(question: str, docs: list[dict], k: int = 2) -> str:
    hits = search(question, docs, k=k)
    blocks = [f"[{h['id']}] {h['text']}" for h in hits]
    context = "\n".join(blocks)
    return (
        "Answer using only the CONTEXT. If missing, say you do not know.\n\n"
        f"CONTEXT:\n{context}\n\n"
        f"QUESTION: {question}\n"
    )
 
prompt = build_rag_prompt("What is the refund window?", docs)
print(prompt)
# llm_answer = chat_model.complete(prompt)  # your LLM client here

docs_meta = [
    {
        "id": "p1",
        "text": "Refunds are available within 30 days with a receipt.",
        "metadata": {"corpus": "policy", "tenant": "acme"},
        "vector": fake_embed("Refunds are available within 30 days with a receipt."),
    },
    {
        "id": "m1",
        "text": "Marketing: spring sale ends Friday.",
        "metadata": {"corpus": "marketing", "tenant": "acme"},
        "vector": fake_embed("Marketing: spring sale ends Friday."),
    },
]
 
def search_filtered(query: str, docs: list[dict], corpus: str, k: int = 3):
    candidates = [d for d in docs if d["metadata"].get("corpus") == corpus]
    return search(query, candidates, k=k)
 
hits = search_filtered("refund window", docs_meta, corpus="policy", k=1)
print(hits[0]["id"], hits[0]["metadata"])

def chunk_text(text: str, size: int = 120, overlap: int = 30) -> list[str]:
    chunks = []
    start = 0
    while start < len(text):
        end = min(len(text), start + size)
        chunks.append(text[start:end])
        if end == len(text):
            break
        start = max(0, end - overlap)
    return chunks
 
policy = (
    "Refunds are available within 30 days of purchase with a receipt. "
    "Digital goods are final sale after download. "
    "Support hours are Monday to Friday, 9am to 5pm Eastern."
)
 
chunk_docs = []
for i, piece in enumerate(chunk_text(policy, size=80, overlap=20)):
    chunk_docs.append(
        {"id": f"policy-{i}", "text": piece, "vector": fake_embed(piece)}
    )
 
print(len(chunk_docs), search("digital goods refund", chunk_docs, k=1)[0]["text"])

import json
from pathlib import Path
 
path = Path("./demo_index.json")
path.write_text(json.dumps(docs, indent=2), encoding="utf-8")
 
loaded = json.loads(path.read_text(encoding="utf-8"))
print("reloaded", len(loaded), "docs; sample id", loaded[0]["id"])

from typing import Any
 
def search_knowledge(query: str, top_k: int = 3, corpus: str = "policy") -> dict[str, Any]:
    """
    Search internal knowledge. Use for refund, support hours, and plan policies.
    Do not use for live order status (call the orders API instead).
    """
    hits = search_filtered(query, docs_meta, corpus=corpus, k=top_k)
    return {
        "ok": True,
        "results": [
            {"id": h["id"], "text": h["text"], "metadata": h.get("metadata", {})}
            for h in hits
        ],
    }
 
print(search_knowledge("When is support available?", top_k=1))