Retrieval Strategies

# MMR implementation
import numpy as np
from typing import List, Tuple

def maximal_marginal_relevance(
    query_embedding: np.ndarray,
    doc_embeddings: np.ndarray,
    documents: List[str],
    k: int = 5,
    lambda_param: float = 0.5,  # 0 = max diversity, 1 = max relevance
) -> List[Tuple[int, str, float]]:
    """Select documents balancing relevance and diversity."""
    # Calculate query-document similarities
    query_sims = np.dot(doc_embeddings, query_embedding) / (
        np.linalg.norm(doc_embeddings, axis=1) * np.linalg.norm(query_embedding)
    )

    selected = []
    remaining = list(range(len(documents)))

    for _ in range(min(k, len(documents))):
        best_score = -float("inf")
        best_idx = -1

        for idx in remaining:
            # Relevance to query
            relevance = query_sims[idx]

            # Max similarity to already selected documents
            if selected:
                selected_embeddings = doc_embeddings[selected]
                doc_sims = np.dot(selected_embeddings, doc_embeddings[idx]) / (
                    np.linalg.norm(selected_embeddings, axis=1)
                    * np.linalg.norm(doc_embeddings[idx])
                )
                max_sim_to_selected = np.max(doc_sims)
            else:
                max_sim_to_selected = 0

            # MMR score
            mmr_score = lambda_param * relevance - (1 - lambda_param) * max_sim_to_selected

            if mmr_score > best_score:
                best_score = mmr_score
                best_idx = idx

        selected.append(best_idx)
        remaining.remove(best_idx)

    return [(i, documents[i], query_sims[i]) for i in selected]

# Using LangChain's built-in MMR
from langchain_chroma import Chroma

vectorstore = Chroma(persist_directory="./db", embedding_function=embeddings)

# MMR retrieval
retriever = vectorstore.as_retriever(
    search_type="mmr",
    search_kwargs={
        "k": 5,
        "fetch_k": 20,       # Fetch 20 candidates, select 5 diverse ones
        "lambda_mult": 0.5,  # Balance relevance and diversity
    },
)

docs = retriever.invoke("What are the best practices for API design?")

// Multi-query retrieval
import Anthropic from "@anthropic-ai/sdk";

const client = new Anthropic();

async function multiQueryRetrieval(
  question: string,
  vectorStore: any,
  numQueries: number = 3
): Promise<string[]> {
  // Step 1: Generate query variations
  const response = await client.messages.create({
    model: "claude-sonnet-4-20250514",
    max_tokens: 512,
    messages: [
      {
        role: "user",
        content: `Generate ${numQueries} different search queries to find information relevant to this question. Each query should approach the question from a different angle.

Question: ${question}

Return ONLY the queries, one per line, no numbering.`,
      },
    ],
  });

  const queries = (response.content[0] as any).text
    .split("\n")
    .filter((q: string) => q.trim());

  // Step 2: Retrieve for each query
  const allDocs = new Map<string, any>();

  for (const query of [question, ...queries]) {
    const results = await vectorStore.similaritySearch(query, 3);
    for (const doc of results) {
      const key = doc.pageContent.slice(0, 100); // Dedup key
      if (!allDocs.has(key)) {
        allDocs.set(key, doc);
      }
    }
  }

  // Step 3: Return unique documents
  return Array.from(allDocs.values()).map(d => d.pageContent);
}

# HyDE - Hypothetical Document Embeddings
import anthropic
from sentence_transformers import SentenceTransformer

client = anthropic.Anthropic()
embedder = SentenceTransformer("all-MiniLM-L6-v2")

def hyde_retrieval(question: str, collection, n_results: int = 5):
    # Step 1: Generate a hypothetical answer
    response = client.messages.create(
        model="claude-sonnet-4-20250514",
        max_tokens=512,
        messages=[{
            "role": "user",
            "content": f"""Write a short paragraph that would be the perfect answer
to this question, as if it appeared in a document:

Question: {question}

Write ONLY the answer paragraph, no preamble.""",
        }],
    )

    hypothetical_answer = response.content[0].text

    # Step 2: Embed the hypothetical answer (not the question!)
    hyde_embedding = embedder.encode(hypothetical_answer).tolist()

    # Step 3: Search with the hypothetical answer embedding
    results = collection.query(
        query_embeddings=[hyde_embedding],
        n_results=n_results,
    )

    return results

# Comparison: standard vs HyDE
standard_results = collection.query(
    query_texts=["What causes memory leaks in Node.js?"],
    n_results=5,
)

hyde_results = hyde_retrieval(
    "What causes memory leaks in Node.js?",
    collection,
)

// Contextual compression - extract only relevant parts from retrieved docs
async function compressContext(
  question: string,
  documents: string[],
): Promise<string[]> {
  const compressed: string[] = [];

  for (const doc of documents) {
    const response = await client.messages.create({
      model: "claude-sonnet-4-20250514",
      max_tokens: 256,
      messages: [
        {
          role: "user",
          content: `Given this question: "${question}"

Extract ONLY the parts of this document that are relevant to answering the question.
If nothing is relevant, respond with "NOT_RELEVANT".

Document:
${doc}`,
        },
      ],
    });

    const result = (response.content[0] as any).text;
    if (!result.includes("NOT_RELEVANT")) {
      compressed.push(result);
    }
  }

  return compressed;
}