RAG in Production: Typed Inputs, Elastic Storage, and the Latency Tradeoffs That Actually Matter

// TypeScript: Typed context assembler with per-category budget enforcement
const BUDGET: Record<ContextType, number> = {
  system: 512,
  documents: 2048,
  conversation: 1024,
  tools: 512,
};

type ContextType = 'system' | 'documents' | 'conversation' | 'tools';

interface TypedContextInput {
  type: ContextType;
  content: string;
  sourceId: string;  // for provenance logging
  tokenCount: number;
}

function assembleContext(
  inputs: TypedContextInput[],
  tokenizer: (s: string) => number
): { prompt: string; provenance: Record<ContextType, string[]>; overflow: ContextType[] } {
  const grouped = inputs.reduce((acc, input) => {
    acc[input.type] = [...(acc[input.type] ?? []), input];
    return acc;
  }, {} as Record<ContextType, TypedContextInput[]>);

  const overflow: ContextType[] = [];
  const provenance: Record<ContextType, string[]> = { system: [], documents: [], conversation: [], tools: [] };
  const parts: string[] = [];

  for (const [type, items] of Object.entries(grouped) as [ContextType, TypedContextInput[]][]) {
    let budget = BUDGET[type];
    for (const item of items) {
      if (item.tokenCount > budget) {
        overflow.push(type);
        continue;
      }
      parts.push(item.content);
      provenance[type].push(item.sourceId);
      budget -= item.tokenCount;
    }
  }

  return { prompt: parts.join('\n\n'), provenance, overflow };
}

// TypeScript: Retrieval adapter that sources from LTAP-backed query endpoint
// No CDC polling, no vector store TTL management — queries live columnar data directly

interface LTAPRetrievalConfig {
  endpoint: string;
  maxRows: number;
  freshnessWindowMs: number; // still useful for rate-limiting staleness assertions
}

async function retrieveFromLTAP(
  query: string,
  entityId: string,
  config: LTAPRetrievalConfig
): Promise<{ rows: Record<string, unknown>[]; retrievedAtMs: number; lagMs: number }> {
  const start = Date.now();

  const response = await fetch(config.endpoint, {
    method: 'POST',
    headers: { 'Content-Type': 'application/json' },
    body: JSON.stringify({
      sql: query,
      bindings: { entity_id: entityId },
      limit: config.maxRows,
    }),
  });

  const { rows, committedAtMs } = await response.json();
  const retrievedAtMs = Date.now();
  const lagMs = retrievedAtMs - committedAtMs;

  if (lagMs > config.freshnessWindowMs) {
    // Emit to observability — not throw; agent should degrade gracefully
    console.warn({ event: 'ltap_freshness_breach', lagMs, entityId, retrievedAtMs });
  }

  return { rows, retrievedAtMs, lagMs };
}

// TypeScript: Local inference sidecar client for RAG preprocessing
// Wraps llama.cpp HTTP server — same interface whether local or hosted

const localInference = {
  async expandQuery(query: string): Promise<string[]> {
    const res = await fetch('http://localhost:8080/completion', {
      method: 'POST',
      headers: { 'Content-Type': 'application/json' },
      body: JSON.stringify({
        prompt: `Generate 3 search query variants for: "${query}"\nReturn as JSON array.`,
        max_tokens: 256,
        temperature: 0.3,
      }),
    });
    const { content } = await res.json();
    try {
      return JSON.parse(content) as string[];
    } catch {
      return [query]; // degrade gracefully — original query still works
    }
  },

  async rerankChunks(
    query: string,
    chunks: { id: string; text: string }[]
  ): Promise<{ id: string; score: number }[]> {
    const res = await fetch('http://localhost:8080/rerank', {
      method: 'POST',
      headers: { 'Content-Type': 'application/json' },
      body: JSON.stringify({ query, documents: chunks.map(c => c.text) }),
    });
    const { scores } = await res.json();
    return chunks
      .map((c, i) => ({ id: c.id, score: scores[i] as number }))
      .sort((a, b) => b.score - a.score);
  },
};

// TypeScript: Policy-gated retrieval with provenance tagging
// Policy gate runs before VectorDB query, not after — fail fast, not fail quietly

interface AgentContext {
  agentRole: 'customer_facing' | 'internal_ops' | 'admin';
  sessionId: string;
}

const ROLE_ALLOWED_COLLECTIONS: Record<AgentContext['agentRole'], string[]> = {
  customer_facing: ['product_docs', 'public_kb', 'order_history'],
  internal_ops: ['product_docs', 'public_kb', 'order_history', 'internal_runbooks'],
  admin: ['product_docs', 'public_kb', 'order_history', 'internal_runbooks', 'cost_data'],
};

async function policyGatedRetrieval(
  query: string,
  requestedCollections: string[],
  agentCtx: AgentContext,
  vectorClient: VectorClient
): Promise<{ chunks: ProvenanceTaggedChunk[]; blockedCollections: string[] }> {
  const allowed = ROLE_ALLOWED_COLLECTIONS[agentCtx.agentRole];
  const permitted = requestedCollections.filter(c => allowed.includes(c));
  const blocked = requestedCollections.filter(c => !allowed.includes(c));

  if (blocked.length > 0) {
    // Emit governance event — not a throw; agent should proceed with permitted scope
    emit('retrieval_policy_block', { agentRole: agentCtx.agentRole, blocked, sessionId: agentCtx.sessionId });
  }

  const rawChunks = await vectorClient.search(query, { collections: permitted, topK: 10 });

  const chunks: ProvenanceTaggedChunk[] = rawChunks.map(c => ({
    ...c,
    sourceCollection: c.metadata.collection,
    retrievedAtMs: Date.now(),
    sessionId: agentCtx.sessionId,
  }));

  return { chunks, blockedCollections: blocked };
}

function emit(event: string, payload: Record<string, unknown>): void {
  console.log(JSON.stringify({ event, ...payload, ts: new Date().toISOString() }));
}

// TypeScript: Explicit retrieval-action boundary in a LangGraph-style node graph
// Retrieval nodes are typed as ReadOnlyNode — action nodes require explicit promotion

type NodeType = 'retrieval' | 'synthesis' | 'action';

interface AgentNode {
  id: string;
  type: NodeType;
  execute: (state: AgentState) => Promise<Partial<AgentState>>;
}

const retrieveCustomerHistory: AgentNode = {
  id: 'retrieve_customer_history',
  type: 'retrieval',  // safe to retry, safe to cache, safe to replay
  execute: async (state) => {
    const history = await policyGatedRetrieval(
      state.query,
      ['order_history', 'customer_kb'],
      state.agentContext,
      vectorClient
    );
    return { retrievedChunks: history.chunks, retrievalProvenance: history };
  },
};

const issueRefund: AgentNode = {
  id: 'issue_refund',
  type: 'action',  // irreversible — requires human checkpoint before execution
  execute: async (state) => {
    if (!state.humanApproval?.approved) {
      throw new Error('Action node requires explicit human approval in state');
    }
    const result = await refundService.process(state.orderId, state.refundAmount);
    return { actionResult: result, actionCompletedAt: Date.now() };
  },
};

function buildGraph(nodes: AgentNode[]): CompiledGraph {
  const actionNodes = nodes.filter(n => n.type === 'action');
  // Automatically inject human-in-the-loop checkpoint before first action node
  return graph.addCheckpointBefore(actionNodes[0].id, humanApprovalNode);
}

// TypeScript: Retrieval quality vs model size cost estimator
// Helps calibrate infrastructure decisions before running full benchmarks

interface ModelConfig {
  name: string;
  paramsBillions: number;
  costPerMillionTokens: number; // USD
  avgContextDegradationFactor: number; // 1.0 = no degradation, <1.0 = degrades with noisy context
}

interface RetrievalConfig {
  precisionScore: number; // 0-1, measured from your eval set
  avgChunksReturned: number;
  avgTokensPerChunk: number;
}

function estimateQualityPerDollar(
  model: ModelConfig,
  retrieval: RetrievalConfig,
  queriesPerHour: number
): { qualityScore: number; costPerHour: number; qualityPerDollar: number } {
  const tokensPerQuery = retrieval.avgChunksReturned * retrieval.avgTokensPerChunk;
  const effectiveQuality = retrieval.precisionScore * model.avgContextDegradationFactor;
  const costPerQuery = (tokensPerQuery / 1_000_000) * model.costPerMillionTokens;
  const costPerHour = costPerQuery * queriesPerHour;

  return {
    qualityScore: effectiveQuality,
    costPerHour,
    qualityPerDollar: effectiveQuality / costPerHour,
  };
}

// Usage: compare models to find compute-optimal frontier for your retrieval config
const configs: ModelConfig[] = [
  { name: 'llama-70b-hosted', paramsBillions: 70, costPerMillionTokens: 0.90, avgContextDegradationFactor: 0.95 },
  { name: 'qwen-13b-hosted',  paramsBillions: 13, costPerMillionTokens: 0.20, avgContextDegradationFactor: 0.88 },
  { name: 'gemma-7b-local',   paramsBillions: 7,  costPerMillionTokens: 0.02, avgContextDegradationFactor: 0.80 },
];

const retrieval: RetrievalConfig = { precisionScore: 0.87, avgChunksReturned: 6, avgTokensPerChunk: 320 };

configs
  .map(m => ({ model: m.name, ...estimateQualityPerDollar(m, retrieval, 5000) }))
  .sort((a, b) => b.qualityPerDollar - a.qualityPerDollar)
  .forEach(r => console.log(r));