RAG and Vector Search

Query → Embed query → Vector search → Top-k relevant chunks
      → Augment prompt with chunks → LLM generates answer grounded in retrieved context

from sentence_transformers import SentenceTransformer
import numpy as np

model = SentenceTransformer("BAAI/bge-large-en-v1.5")  # 1024-dim embeddings

# Embed documents
docs = [
    "Python is a high-level programming language",
    "Photosynthesis converts sunlight to glucose",
    "Neural networks are inspired by the brain"
]
embeddings = model.encode(docs, normalize_embeddings=True)
# shape: (3, 1024)

# Semantic similarity via dot product (after L2 normalization = cosine sim)
query_emb = model.encode(["What is Python?"], normalize_embeddings=True)
scores = (query_emb @ embeddings.T)[0]
print(f"Similarities: {scores}")

import faiss
import numpy as np

d = 1024    # embedding dimension
n = 100000  # number of documents

embeddings = np.random.randn(n, d).astype("float32")
faiss.normalize_L2(embeddings)   # L2 normalize in-place

# Flat index (exact, brute force) — best for <100k
index = faiss.IndexFlatIP(d)   # Inner Product = cosine after normalization
index.add(embeddings)

# Search
query = np.random.randn(1, d).astype("float32")
faiss.normalize_L2(query)
distances, indices = index.search(query, k=10)

# IVF index (approximate, for millions of vectors)
quantizer  = faiss.IndexFlatIP(d)
index_ivf  = faiss.IndexIVFFlat(quantizer, d, 1000, faiss.METRIC_INNER_PRODUCT)  # nlist=1000
index_ivf.train(embeddings)     # cluster centroids
index_ivf.add(embeddings)
index_ivf.nprobe = 20           # how many clusters to check at query time
dist, idx = index_ivf.search(query, k=10)

# HNSW (best single-machine ANN, no training needed)
index_hnsw = faiss.IndexHNSWFlat(d, 32)  # 32 = M parameter
index_hnsw.add(embeddings)
dist, idx = index_hnsw.search(query, k=10)

# Chroma (minimal setup)
import chromadb

client = chromadb.PersistentClient(path="./chroma_db")
collection = client.get_or_create_collection("documents", metadata={"hnsw:space": "cosine"})

# Add documents
collection.add(
    documents=docs,
    embeddings=embeddings.tolist(),
    ids=[f"doc_{i}" for i in range(len(docs))],
    metadatas=[{"source": "wiki", "year": 2024} for _ in docs]
)

# Query
results = collection.query(
    query_embeddings=query_emb.tolist(),
    n_results=5,
    where={"year": {"$gte": 2023}},   # metadata filtering
    include=["documents", "distances", "metadatas"]
)

from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import PyPDFLoader
from sentence_transformers import SentenceTransformer
import chromadb

# Load documents
loader = PyPDFLoader("document.pdf")
pages  = loader.load()

# Chunk text (important: chunk size matches embedding model context)
splitter = RecursiveCharacterTextSplitter(
    chunk_size=512,          # chars per chunk
    chunk_overlap=50,        # overlap for continuity
    separators=["\n\n", "\n", ". ", " ", ""]
)
chunks = splitter.split_documents(pages)

# Embed and store
embed_model = SentenceTransformer("BAAI/bge-large-en-v1.5")
texts = [c.page_content for c in chunks]
embeddings = embed_model.encode(texts, normalize_embeddings=True)

collection = client.get_or_create_collection("docs")
collection.add(
    documents=texts,
    embeddings=embeddings.tolist(),
    ids=[f"chunk_{i}" for i in range(len(texts))],
    metadatas=[{"page": c.metadata.get("page", 0)} for c in chunks]
)

from openai import OpenAI

openai_client = OpenAI()

def rag_query(question: str, n_results: int = 5) -> str:
    # Embed the question
    q_emb = embed_model.encode([question], normalize_embeddings=True)

    # Retrieve relevant chunks
    results = collection.query(query_embeddings=q_emb.tolist(), n_results=n_results)
    context = "\n\n---\n\n".join(results["documents"][0])

    # Generate answer
    response = openai_client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content":
             "Answer the question using only the provided context. "
             "If the context doesn't contain the answer, say 'I don't know'."},
            {"role": "user", "content":
             f"Context:\n{context}\n\nQuestion: {question}"}
        ],
        temperature=0
    )
    return response.choices[0].message.content

answer = rag_query("What is the company's revenue growth target?")

def hyde_retrieval(query: str) -> list:
    # Generate hypothetical answer
    hyp_answer = openai_client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content":
                   f"Write a short paragraph that would answer: {query}"}]
    ).choices[0].message.content

    # Embed the hypothetical answer
    hyp_emb = embed_model.encode([hyp_answer], normalize_embeddings=True)
    return collection.query(query_embeddings=hyp_emb.tolist(), n_results=5)

from rank_bm25 import BM25Okapi

# Sparse BM25 index
tokenized = [doc.split() for doc in texts]
bm25 = BM25Okapi(tokenized)

def hybrid_search(query, alpha=0.5, k=10):
    # Dense retrieval
    q_emb = embed_model.encode([query], normalize_embeddings=True)
    dense_res = collection.query(query_embeddings=q_emb.tolist(), n_results=k)

    # Sparse BM25 retrieval
    sparse_scores = bm25.get_scores(query.split())
    sparse_top = np.argsort(sparse_scores)[::-1][:k]

    # Reciprocal Rank Fusion
    scores = {}
    for rank, idx in enumerate(dense_res["ids"][0]):
        scores[idx] = scores.get(idx, 0) + alpha / (rank + 60)
    for rank, idx in enumerate(sparse_top):
        doc_id = f"chunk_{idx}"
        scores[doc_id] = scores.get(doc_id, 0) + (1-alpha) / (rank + 60)
    return sorted(scores, key=scores.get, reverse=True)[:k]

from sentence_transformers import CrossEncoder

reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")

query = "What is photosynthesis?"
candidates = retrieved_docs[:20]   # initial retrieval

pairs  = [[query, doc] for doc in candidates]
scores = reranker.predict(pairs)
ranked = sorted(zip(scores, candidates), reverse=True)
top5   = [doc for _, doc in ranked[:5]]

from ragas import evaluate
from ragas.metrics import (faithfulness, answer_relevancy,
                            context_recall, context_precision)
from datasets import Dataset

# Prepare evaluation dataset
data = Dataset.from_dict({
    "question": questions,
    "answer": generated_answers,
    "contexts": retrieved_contexts,
    "ground_truth": reference_answers
})

result = evaluate(data, metrics=[faithfulness, answer_relevancy,
                                   context_recall, context_precision])
print(result)

# Simple iterative RAG with OpenAI tools
tools = [{
    "type": "function",
    "function": {
        "name": "search_documents",
        "description": "Search the knowledge base for relevant information",
        "parameters": {
            "type": "object",
            "properties": {"query": {"type": "string"}},
            "required": ["query"]
        }
    }
}]

response = openai_client.chat.completions.create(
    model="gpt-4o",
    messages=messages,
    tools=tools,
    tool_choice="auto"
)