Your RAG Pipeline Is Answering the Wrong Question

// naive-rag.ts — The fixed embed-search-generate pipeline
import OpenAI from "openai";
 
const openai = new OpenAI();
 
interface Chunk {
  id: string;
  text: string;
  embedding: number[];
  metadata: { source: string; section: string };
}
 
const chunks: Chunk[] = []; // populated from your knowledge base
 
async function embed(text: string): Promise<number[]> {
  const response = await openai.embeddings.create({
    model: "text-embedding-3-small",
    input: text,
  });
  return response.data[0].embedding;
}
 
function cosineSimilarity(a: number[], b: number[]): number {
  let dot = 0, normA = 0, normB = 0;
  for (let i = 0; i < a.length; i++) {
    dot += a[i] * b[i];
    normA += a[i] * a[i];
    normB += b[i] * b[i];
  }
  return dot / (Math.sqrt(normA) * Math.sqrt(normB));
}
 
// This is the problem: one query, one retrieval pass, no evaluation
async function naiveRag(query: string): Promise<string> {
  const queryEmbedding = await embed(query);
 
  // Rank all chunks by cosine similarity — no strategy selection
  const retrieved = chunks
    .map((chunk) => ({ chunk, score: cosineSimilarity(queryEmbedding, chunk.embedding) }))
    .sort((a, b) => b.score - a.score)
    .slice(0, 3) // top-3, regardless of whether 3 is enough
    .map((r) => r.chunk);
 
  const context = retrieved.map((c) => c.text).join("\n\n---\n\n");
 
  const response = await openai.chat.completions.create({
    model: "gpt-4o-mini",
    messages: [
      {
        role: "system",
        content: `Answer based ONLY on the provided context. If the context doesn't contain the answer, say so.\n\nContext:\n${context}`,
      },
      { role: "user", content: query },
    ],
    temperature: 0.1,
  });
 
  return response.choices[0].message.content ?? "No answer generated";
}

# naive_rag.py — The fixed embed-search-generate pipeline
import numpy as np
from openai import OpenAI
 
client = OpenAI()
chunks: list[dict] = []  # populated from your knowledge base
 
def embed(text: str) -> list[float]:
    response = client.embeddings.create(model="text-embedding-3-small", input=text)
    return response.data[0].embedding
 
def cosine_similarity(a: list[float], b: list[float]) -> float:
    a_arr, b_arr = np.array(a), np.array(b)
    return float(np.dot(a_arr, b_arr) / (np.linalg.norm(a_arr) * np.linalg.norm(b_arr)))
 
# Same problem: one query, one pass, no evaluation
def naive_rag(query: str) -> str:
    query_embedding = embed(query)
    scored = [(c, cosine_similarity(query_embedding, c["embedding"])) for c in chunks]
    scored.sort(key=lambda x: x[1], reverse=True)
    retrieved = [c for c, _ in scored[:3]]  # top-3, hope for the best
 
    context = "\n\n---\n\n".join(c["text"] for c in retrieved)
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": f"Answer based ONLY on the provided context.\n\nContext:\n{context}"},
            {"role": "user", "content": query},
        ],
        temperature=0.1,
    )
    return response.choices[0].message.content or "No answer generated"

// agentic-rag.ts — Retrieval as tools, not a pipeline
import OpenAI from "openai";
import type { ChatCompletionTool } from "openai/resources/chat/completions";
 
const openai = new OpenAI();
 
// Each tool represents a different retrieval STRATEGY
const retrievalTools: ChatCompletionTool[] = [
  {
    type: "function",
    function: {
      name: "semantic_search",
      // Agent sees this description and decides when to use it
      description: "Search by meaning. Best for concepts and explanations.",
      parameters: {
        type: "object",
        properties: {
          query: { type: "string", description: "Search query optimized for semantic retrieval." },
          top_k: { type: "number", description: "Results to return. 3 focused, 5-7 broad.", default: 3 },
        },
        required: ["query"],
      },
    },
  },
  {
    type: "function",
    function: {
      name: "keyword_search",
      // For exact terms the agent knows appear in the docs
      description: "Exact keyword matching. Best for product names, policy numbers, specific terms.",
      parameters: {
        type: "object",
        properties: {
          keywords: { type: "string", description: "Space-separated keywords." },
        },
        required: ["keywords"],
      },
    },
  },
  {
    type: "function",
    function: {
      name: "metadata_filter",
      // Surgical retrieval: go straight to a known document section
      description: "Filter by document source or section before searching.",
      parameters: {
        type: "object",
        properties: {
          source: { type: "string", description: "Document source name." },
          section: { type: "string", description: "Section within a document." },
        },
        required: [],
      },
    },
  },
  {
    type: "function",
    function: {
      name: "no_retrieval_needed",
      // Prevents unnecessary retrieval on math/greeting queries
      description: "Use when the question needs no document search (math, greetings, clarifications).",
      parameters: {
        type: "object",
        properties: {
          reason: { type: "string", description: "Why no retrieval is needed." },
        },
        required: ["reason"],
      },
    },
  },
];

# agentic_rag.py — Retrieval as tools, not a pipeline
from openai import OpenAI
 
client = OpenAI()
 
# Each tool = a different retrieval strategy the agent can choose
retrieval_tools = [
    {
        "type": "function",
        "function": {
            "name": "semantic_search",
            "description": "Search by meaning. Best for concepts and explanations.",
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {"type": "string", "description": "Query optimized for semantic retrieval."},
                    "top_k": {"type": "number", "description": "Results count. 3 focused, 5-7 broad.", "default": 3},
                },
                "required": ["query"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "keyword_search",
            "description": "Exact keyword matching. Best for product names, policy numbers.",
            "parameters": {
                "type": "object",
                "properties": {
                    "keywords": {"type": "string", "description": "Space-separated keywords."},
                },
                "required": ["keywords"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "metadata_filter",
            "description": "Filter by document source or section before searching.",
            "parameters": {
                "type": "object",
                "properties": {
                    "source": {"type": "string", "description": "Document source name."},
                    "section": {"type": "string", "description": "Section within a document."},
                },
                "required": [],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "no_retrieval_needed",
            "description": "Use when the question needs no document search.",
            "parameters": {
                "type": "object",
                "properties": {
                    "reason": {"type": "string", "description": "Why no retrieval is needed."},
                },
                "required": ["reason"],
            },
        },
    },
]

// The agentic loop: plan → retrieve → evaluate → repeat or answer
import type {
  ChatCompletionMessageParam,
  ChatCompletionToolMessageParam,
} from "openai/resources/chat/completions";
 
const SYSTEM_PROMPT = `You are a research assistant with access to a knowledge base.
 
Strategy:
1. Analyze what information the question needs.
2. Choose the best retrieval tool(s) for the question type.
3. For comparison questions, search for EACH item separately.
4. After retrieval, evaluate: do you have enough context?
5. If not, reformulate and try again.
6. Once sufficient, answer directly. Never fabricate beyond retrieved docs.`;
 
async function agenticRag(
  query: string,
  maxIterations = 3 // hard cap prevents runaway retrieval loops
): Promise<{ answer: string; toolCalls: string[]; iterations: number }> {
  const messages: ChatCompletionMessageParam[] = [
    { role: "system", content: SYSTEM_PROMPT },
    { role: "user", content: query },
  ];
  const toolCallLog: string[] = [];
  let iterations = 0;
 
  while (iterations < maxIterations) {
    iterations++;
 
    const response = await openai.chat.completions.create({
      model: "gpt-4o-mini",
      messages,
      tools: retrievalTools,
      tool_choice: "auto", // agent decides whether to call a tool or answer
      temperature: 0.1,
    });
 
    const message = response.choices[0].message;
    messages.push(message);
 
    // No tool calls = agent is ready to answer with what it has
    if (!message.tool_calls || message.tool_calls.length === 0) {
      return { answer: message.content ?? "No answer generated", toolCalls: toolCallLog, iterations };
    }
 
    // Execute each tool call and feed results back into conversation
    for (const toolCall of message.tool_calls) {
      const args = JSON.parse(toolCall.function.arguments);
      const result = await executeTool(toolCall.function.name, args);
      toolCallLog.push(`${toolCall.function.name}(${JSON.stringify(args)})`);
 
      // Tool result becomes part of the conversation — agent sees it next iteration
      const toolMessage: ChatCompletionToolMessageParam = {
        role: "tool",
        tool_call_id: toolCall.id,
        content: result,
      };
      messages.push(toolMessage);
    }
  }
 
  // Safety net: force an answer if we hit the iteration cap
  const finalResponse = await openai.chat.completions.create({
    model: "gpt-4o-mini",
    messages: [
      ...messages,
      { role: "user", content: "Based on all information retrieved, provide your best answer now." },
    ],
    temperature: 0.1,
  });
 
  return {
    answer: finalResponse.choices[0].message.content ?? "No answer generated",
    toolCalls: toolCallLog,
    iterations,
  };
}

# The agentic loop: plan → retrieve → evaluate → repeat or answer
import json
 
SYSTEM_PROMPT = """You are a research assistant with access to a knowledge base.
 
Strategy:
1. Analyze what information the question needs.
2. Choose the best retrieval tool(s) for the question type.
3. For comparison questions, search for EACH item separately.
4. After retrieval, evaluate: do you have enough context?
5. If not, reformulate and try again.
6. Once sufficient, answer directly. Never fabricate beyond retrieved docs."""
 
def agentic_rag(query: str, max_iterations: int = 3) -> dict:
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": query},
    ]
    tool_call_log = []
    iterations = 0
 
    while iterations < max_iterations:
        iterations += 1
        response = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=messages,
            tools=retrieval_tools,
            tool_choice="auto",  # agent decides: tool call or answer
            temperature=0.1,
        )
        message = response.choices[0].message
        messages.append(message)
 
        # No tool calls = agent is ready to answer
        if not message.tool_calls:
            return {"answer": message.content or "No answer", "tool_calls": tool_call_log, "iterations": iterations}
 
        # Execute tools, feed results back into conversation history
        for tool_call in message.tool_calls:
            args = json.loads(tool_call.function.arguments)
            result = execute_tool(tool_call.function.name, args)
            tool_call_log.append(f"{tool_call.function.name}({json.dumps(args)})")
            messages.append({"role": "tool", "tool_call_id": tool_call.id, "content": result})
 
    # Safety net: force answer at iteration cap
    messages.append({"role": "user", "content": "Based on all information retrieved, provide your best answer now."})
    final = client.chat.completions.create(model="gpt-4o-mini", messages=messages, temperature=0.1)
    return {"answer": final.choices[0].message.content or "No answer", "tool_calls": tool_call_log, "iterations": iterations}

// Evaluation tool — agent calls this to assess its own retrieval quality
const evaluationTool: ChatCompletionTool = {
  type: "function",
  function: {
    name: "evaluate_retrieval",
    description: "Assess whether retrieved context is sufficient. Call after each retrieval.",
    parameters: {
      type: "object",
      properties: {
        retrieved_context_summary: { type: "string", description: "What was retrieved so far." },
        missing_information: { type: "string", description: "What's still needed. 'none' if sufficient." },
        confidence: { type: "number", description: "0-1 confidence the context can answer the question." },
        // This is the key decision: does the agent search more or answer now?
        next_action: {
          type: "string",
          enum: ["answer_now", "search_more", "refine_query"],
          description: "What to do next.",
        },
      },
      required: ["retrieved_context_summary", "missing_information", "confidence", "next_action"],
    },
  },
};
 
// When the agent calls evaluate_retrieval, parse its self-assessment
function handleEvaluation(args: {
  confidence: number;
  next_action: "answer_now" | "search_more" | "refine_query";
  missing_information: string;
}): { shouldContinue: boolean; feedback: string } {
  console.log(`[Eval] confidence=${args.confidence} action=${args.next_action}`);
 
  // High confidence + answer_now = stop retrieving
  if (args.confidence >= 0.8 && args.next_action === "answer_now") {
    return { shouldContinue: false, feedback: "Context sufficient. Generate your answer." };
  }
 
  // Agent identified a gap — tell it to target that gap specifically
  if (args.next_action === "refine_query") {
    return { shouldContinue: true, feedback: `Missing: ${args.missing_information}. Reformulate to target this gap.` };
  }
 
  return { shouldContinue: true, feedback: `Confidence: ${args.confidence}. Try a different strategy.` };
}

# When the agent calls evaluate_retrieval, parse its self-assessment
def handle_evaluation(args: dict) -> tuple[bool, str]:
    """Returns (should_continue, feedback_to_agent)."""
    confidence = args["confidence"]
    action = args["next_action"]
    missing = args["missing_information"]
 
    print(f"[Eval] confidence={confidence} action={action}")
 
    # High confidence = stop retrieving, generate answer
    if confidence >= 0.8 and action == "answer_now":
        return False, "Context sufficient. Generate your answer."
 
    # Agent found a gap — direct it to fill that specific gap
    if action == "refine_query":
        return True, f"Missing: {missing}. Reformulate to target this gap."
 
    return True, f"Confidence: {confidence}. Try a different retrieval strategy."

// Stage 2: Query router — cheapest path to agentic gains
async function hybridRag(query: string): Promise<string> {
  // One cheap LLM call to classify the query type
  const classification = await openai.chat.completions.create({
    model: "gpt-4o-mini",
    messages: [
      {
        role: "system",
        content: `Classify this query into one category:
- "simple": single fact, one document
- "multi_hop": needs 2+ documents or comparison
- "no_retrieval": answerable from general knowledge
Respond with just the category name.`,
      },
      { role: "user", content: query },
    ],
    temperature: 0,
    max_tokens: 20,
  });
 
  const queryType = classification.choices[0].message.content?.trim() ?? "simple";
 
  // Only complex queries pay the agentic cost
  switch (queryType) {
    case "no_retrieval": return directLlmAnswer(query);
    case "multi_hop":    return (await agenticRag(query)).answer;
    default:             return naiveRag(query);
  }
}