@od-oneapp/ai-platform 0.1.0

package/src/grounding/rag/trace.ts

@@ -0,0 +1,503 @@
+/**
+ * @fileoverview Research trace storage with cost & latency tracking
+ *
+ * Extends the artifact pattern for multi-model research runs:
+ * - Per-step artifact storage (context expansion → reasoning → verification)
+ * - Cost tracking per model (configurable via PipelineModelConfig.costRates)
+ * - Latency profiling
+ * - Fast indexing for dashboard queries
+ *
+ * Storage pattern:
+ * - File-based for prototype (like artifacts.ts)
+ * - Production: Neon Postgres + Upstash Redis + R2
+ *
+ * @module @od-oneapp/ai-platform/research/trace
+ */
+
+import { createHash, randomUUID } from 'node:crypto';
+import { promises as fs } from 'node:fs';
+import { join, resolve } from 'node:path';
+
+import { logError } from '@repo/shared/logs';
+
+import type { ResearchTier } from './types';
+
+/** Directory for storing research traces. */
+const TRACES_DIR = process.env.NEXT_AI_TRACES_PATH ?? '/tmp/next-ai-traces';
+
+/**
+ * Validates run ID to prevent path traversal attacks.
+ * Format: {timestamp}-{12-char-hex}
+ * Example: 1733875200000-a1b2c3d4e5f6
+ *
+ * @param runId - Run ID to validate.
+ * @returns `true` if run ID is valid, `false` otherwise.
+ * @internal
+ */
+function validateRunId(runId: string): boolean {
+  // runId format: timestamp-hash (e.g., 1733875200000-a1b2c3d4e5f6)
+  const runIdRegex = /^[0-9]+-[0-9a-f]{12}$/i;
+  return runIdRegex.test(runId);
+}
+
+/**
+ * Safely constructs run path with validation.
+ * Prevents path traversal attacks by ensuring paths stay within TRACES_DIR.
+ *
+ * @param runId - Run ID (validated).
+ * @param filename - Optional filename to append.
+ * @returns Resolved path or `null` if validation fails.
+ * @internal
+ */
+function getRunPath(runId: string, filename?: string): string | null {
+  if (!validateRunId(runId)) {
+    logError('Invalid run ID format', { runId });
+    return null;
+  }
+
+  const runPath = filename
+    ? resolve(join(TRACES_DIR, runId, filename))
+    : resolve(join(TRACES_DIR, runId));
+  const baseDir = resolve(TRACES_DIR);
+
+  // Ensure path is within TRACES_DIR
+  if (!runPath.startsWith(baseDir)) {
+    logError('Path traversal attempt detected', { runId });
+    return null;
+  }
+
+  return runPath;
+}
+
+/**
+ * Pipeline stage kind for trace tracking.
+ * Uses role-based names instead of provider-specific names.
+ */
+export type TraceStepKind = 'context-expansion' | 'reasoning' | 'verification';
+
+export interface TraceStep {
+  /** Step identifier. */
+  id: string;
+  /** Pipeline stage kind. */
+  kind: TraceStepKind;
+  /** Timestamp when step started (Unix epoch milliseconds). */
+  startAt: number;
+  /** Timestamp when step ended (Unix epoch milliseconds). */
+  endAt: number;
+  /** Step latency in milliseconds. */
+  latencyMs: number;
+  /** Model used for this step. */
+  model?: string;
+  /** Input tokens consumed. */
+  inputTokens?: number;
+  /** Output tokens generated. */
+  outputTokens?: number;
+  /** Estimated cost in USD. */
+  costUSD?: number;
+  /** Step artifacts (in-memory). */
+  artifacts: Record<string, unknown>;
+  /** Artifact keys (file paths or R2 keys). */
+  artifactKeys: string[];
+}
+
+/**
+ * Complete research run trace information.
+ * Tracks full pipeline execution with all steps, cost, and latency.
+ */
+export interface TraceRun {
+  /** Unique run identifier. */
+  runId: string;
+  /** Application/client identifier. */
+  app: string;
+  /** Research question. */
+  question: string;
+  /** Research tier used. */
+  tier: ResearchTier;
+  /** Timestamp when run started (Unix epoch milliseconds). */
+  startedAt: number;
+  /** Timestamp when run ended (Unix epoch milliseconds). */
+  endedAt?: number;
+  /** Total latency in milliseconds. */
+  totalLatencyMs?: number;
+  /** Total cost in USD. */
+  totalCostUSD?: number;
+  /** Array of flag strings. */
+  flags?: string[];
+  /** Array of research steps. */
+  steps: TraceStep[];
+  /** Optional metadata. */
+  metadata?: Record<string, unknown>;
+}
+
+/**
+ * Default token-based cost rates ($/1k tokens).
+ * Override via `PipelineModelConfig.costRates` or the `customRates` param on `estimateCost`.
+ */
+export const DEFAULT_COST_RATES: Record<
+  string,
+  { input?: number; output?: number; flat?: number }
+> = {
+  'perplexity-latest': { flat: 0.002 },
+  'llama-3.1-sonar-large-128k-online': { flat: 0.002 },
+  'gpt-4-turbo': { input: 0.01, output: 0.03 },
+  'gpt-5': { input: 0.01, output: 0.03 },
+  'claude-3-5-sonnet-20241022': { input: 0.003, output: 0.015 },
+};
+
+/**
+ * Estimates cost for a model call based on token usage.
+ *
+ * @param model - Model identifier.
+ * @param inputTokens - Input tokens consumed (defaults to 0).
+ * @param outputTokens - Output tokens generated (defaults to 0).
+ * @returns Estimated cost in USD.
+ * @remarks Cost rates ($/1k tokens):
+ * - Perplexity: $0.002 flat per call
+ * - Claude 3.5 Sonnet: $0.003 input, $0.015 output
+ * - GPT-4 Turbo: $0.01 input, $0.03 output
+ * Returns 0 if model is not in cost rates table.
+ * @example
+ * ```ts
+ * const cost = estimateCost('claude-3-5-sonnet-20241022', 1000, 500);
+ * // Returns: (1000/1000 * 0.003) + (500/1000 * 0.015) = 0.0105
+ * ```
+ */
+export function estimateCost(
+  model: string,
+  inputTokens = 0,
+  outputTokens = 0,
+  customRates?: Record<string, { input?: number; output?: number; flat?: number }>,
+): number {
+  const rates = customRates ? { ...DEFAULT_COST_RATES, ...customRates } : DEFAULT_COST_RATES;
+  const rate = rates[model];
+  if (!rate) return 0;
+
+  if (rate.flat) return rate.flat;
+
+  return (inputTokens / 1000) * (rate.input ?? 0) + (outputTokens / 1000) * (rate.output ?? 0);
+}
+
+/**
+ * Ensures traces directory exists.
+ * Creates directory recursively if it doesn't exist.
+ * @internal
+ */
+async function ensureTracesDir(): Promise<void> {
+  try {
+    // eslint-disable-next-line security/detect-non-literal-fs-filename -- TRACES_DIR is validated via environment variable with safe default
+    await fs.mkdir(TRACES_DIR, { recursive: true });
+  } catch (error) {
+    logError('Failed to create traces directory', { error });
+  }
+}
+
+/**
+ * Starts a new research run trace.
+ * Creates a new TraceRun with generated runId and initial timestamp.
+ *
+ * @param question - Research question.
+ * @param tier - Research tier.
+ * @param app - Application/client identifier (defaults to 'research-agent').
+ * @returns New TraceRun object.
+ * @example
+ * ```ts
+ * const run = startRun('What is quantum computing?', 'balanced', 'my-app');
+ * ```
+ */
+export function startRun(question: string, tier: ResearchTier, app = 'research-agent'): TraceRun {
+  const timestamp = Date.now();
+  const hash = createHash('sha256').update(`${timestamp}-${question}`).digest('hex').slice(0, 12);
+
+  return {
+    runId: `${timestamp}-${hash}`,
+    app,
+    question,
+    tier,
+    startedAt: timestamp,
+    steps: [],
+  };
+}
+
+/**
+ * Starts a new research step trace.
+ * Creates a new TraceStep with generated ID and start timestamp.
+ *
+ * @param kind - Pipeline stage kind ('context-expansion', 'reasoning', or 'verification').
+ * @param model - Optional model identifier.
+ * @returns New TraceStep object.
+ * @example
+ * ```ts
+ * const step = startStep('context-expansion', 'llama-3.1-sonar-large-128k-online');
+ * ```
+ */
+export function startStep(kind: TraceStepKind, model?: string): TraceStep {
+  return {
+    id: randomUUID(),
+    kind,
+    startAt: Date.now(),
+    endAt: 0,
+    latencyMs: 0,
+    model,
+    artifacts: {},
+    artifactKeys: [],
+  };
+}
+
+/**
+ * Ends a research step with usage metrics.
+ * Calculates latency and cost, updates step with token usage.
+ *
+ * @param step - Step to end.
+ * @param usage - Optional usage metrics (inputTokens, outputTokens, costUSD).
+ * @returns Updated step with end timestamp, latency, and cost.
+ * @example
+ * ```ts
+ * const completedStep = endStep(step, {
+ *   inputTokens: 1000,
+ *   outputTokens: 500
+ * });
+ * ```
+ */
+export function endStep(
+  step: TraceStep,
+  usage?: { inputTokens?: number; outputTokens?: number; costUSD?: number },
+): TraceStep {
+  step.endAt = Date.now();
+  step.latencyMs = step.endAt - step.startAt;
+
+  if (usage?.inputTokens) step.inputTokens = usage.inputTokens;
+  if (usage?.outputTokens) step.outputTokens = usage.outputTokens;
+  if (usage?.costUSD !== undefined) {
+    step.costUSD = usage.costUSD;
+  } else if (step.model) {
+    step.costUSD = estimateCost(step.model, step.inputTokens ?? 0, step.outputTokens ?? 0);
+  }
+
+  return step;
+}
+
+/**
+ * Ends a research run and calculates totals.
+ * Computes total latency and cost from all steps.
+ *
+ * @param run - Run to end.
+ * @returns Updated run with end timestamp, total latency, and total cost.
+ * @example
+ * ```ts
+ * const completedRun = endRun(run);
+ * console.log(`Total cost: $${completedRun.totalCostUSD}`);
+ * ```
+ */
+export function endRun(run: TraceRun): TraceRun {
+  run.endedAt = Date.now();
+  run.totalLatencyMs = run.endedAt - run.startedAt;
+  run.totalCostUSD = run.steps.reduce((acc, s) => acc + (s.costUSD ?? 0), 0);
+  return run;
+}
+
+/**
+ * Persists step artifacts to storage.
+ * Writes step artifacts as JSON files in the run directory.
+ *
+ * @param run - Research run.
+ * @param step - Research step.
+ * @param payloads - Artifact payloads keyed by name.
+ * @throws {Error} If run ID is invalid or file write fails.
+ * @remarks Production: Replace with R2 writes.
+ * Artifact keys are added to step.artifactKeys array for reference.
+ * Files are named as: `{step.kind}-{step.id}.{name}.json`
+ * @example
+ * ```ts
+ * await persistStep(run, step, {
+ *   p_notes: { notes: '...', cites: [...] },
+ *   p_index: { sources: [...], fetchedAt: '...' }
+ * });
+ * ```
+ */
+export async function persistStep(
+  run: TraceRun,
+  step: TraceStep,
+  payloads: Record<string, unknown>,
+): Promise<void> {
+  await ensureTracesDir();
+
+  const runDir = getRunPath(run.runId);
+  if (!runDir) {
+    throw new Error(`Invalid run ID: ${run.runId}`);
+  }
+
+  // eslint-disable-next-line security/detect-non-literal-fs-filename -- runDir is validated by getRunPath which checks for path traversal
+  await fs.mkdir(runDir, { recursive: true });
+
+  for (const [name, body] of Object.entries(payloads)) {
+    const filename = `${step.kind}-${step.id}.${name}.json`;
+    const filepath = join(runDir, filename);
+    // eslint-disable-next-line security/detect-non-literal-fs-filename -- filepath is validated via getRunPath and filename sanitization
+    await fs.writeFile(filepath, JSON.stringify(body, null, 2), 'utf-8');
+    step.artifactKeys.push(filename);
+  }
+}
+
+/**
+ * Persists run index for fast listing.
+ * Creates an index file with run metadata for quick queries.
+ *
+ * @param run - Research run to index.
+ * @remarks Production: Write to Upstash Redis ZSET.
+ * @example
+ * ```ts
+ * await persistRunIndex(run);
+ * ```
+ */
+export async function persistRunIndex(run: TraceRun): Promise<void> {
+  await ensureTracesDir();
+
+  const indexPath = join(TRACES_DIR, `${run.runId}.index.json`);
+  const index = {
+    runId: run.runId,
+    app: run.app,
+    question: run.question,
+    tier: run.tier,
+    startedAt: run.startedAt,
+    endedAt: run.endedAt ?? 0,
+    totalLatencyMs: run.totalLatencyMs ?? 0,
+    totalCostUSD: run.totalCostUSD ?? 0,
+    flags: run.flags ?? [],
+  };
+
+  // eslint-disable-next-line security/detect-non-literal-fs-filename -- indexPath uses validated runId format
+  await fs.writeFile(indexPath, JSON.stringify(index, null, 2), 'utf-8');
+}
+
+/**
+ * Persists full run metadata.
+ * Writes complete run object as JSON file.
+ *
+ * @param run - Research run to persist.
+ * @throws {Error} If run ID is invalid or file write fails.
+ * @remarks Writes to `{runId}/run.meta.json` in traces directory.
+ * Contains complete run state including all steps and metadata.
+ * @example
+ * ```ts
+ * await persistRunMeta(run);
+ * ```
+ */
+export async function persistRunMeta(run: TraceRun): Promise<void> {
+  await ensureTracesDir();
+
+  const runDir = getRunPath(run.runId);
+  if (!runDir) {
+    throw new Error(`Invalid run ID: ${run.runId}`);
+  }
+
+  // eslint-disable-next-line security/detect-non-literal-fs-filename -- runDir is validated by getRunPath which checks for path traversal
+  await fs.mkdir(runDir, { recursive: true });
+
+  const metaPath = join(runDir, 'run.meta.json');
+  // eslint-disable-next-line security/detect-non-literal-fs-filename -- metaPath uses validated runDir
+  await fs.writeFile(metaPath, JSON.stringify(run, null, 2), 'utf-8');
+}
+
+/**
+ * Lists recent research runs.
+ * Returns runs sorted by start time (most recent first).
+ *
+ * @param limit - Maximum number of runs to return (defaults to 100).
+ * @returns Array of TraceRun objects sorted by startedAt descending.
+ * @remarks Production: Query from Redis ZSET.
+ * Reads index files from traces directory and parses JSON.
+ * Returns empty array on error (logs error to console).
+ * @example
+ * ```ts
+ * const runs = await listRuns(50);
+ * ```
+ */
+export async function listRuns(limit = 100): Promise<TraceRun[]> {
+  await ensureTracesDir();
+
+  try {
+    // eslint-disable-next-line security/detect-non-literal-fs-filename -- TRACES_DIR is validated via environment variable with safe default
+    const files = await fs.readdir(TRACES_DIR);
+    const indexFiles = files.filter(f => f.endsWith('.index.json'));
+
+    const indices = await Promise.all(
+      indexFiles.map(async file => {
+        // eslint-disable-next-line security/detect-non-literal-fs-filename -- file is filtered to .index.json files only
+        const content = await fs.readFile(join(TRACES_DIR, file), 'utf-8');
+        return JSON.parse(content);
+      }),
+    );
+
+    // Sort by startedAt desc
+    return indices.sort((a, b) => b.startedAt - a.startedAt).slice(0, limit);
+  } catch (error) {
+    logError('Failed to list runs', { error });
+    return [];
+  }
+}
+
+/**
+ * Loads full run metadata from storage.
+ *
+ * @param runId - Run identifier.
+ * @returns TraceRun object or `null` if not found or invalid.
+ * @remarks Reads from `{runId}/run.meta.json` file.
+ * Returns `null` if file doesn't exist, runId is invalid, or parse fails.
+ * Logs errors to console.
+ * @example
+ * ```ts
+ * const run = await loadRun('1733875200000-a1b2c3d4e5f6');
+ * ```
+ */
+export async function loadRun(runId: string): Promise<TraceRun | null> {
+  try {
+    const metaPath = getRunPath(runId, 'run.meta.json');
+    if (!metaPath) {
+      return null;
+    }
+    // eslint-disable-next-line security/detect-non-literal-fs-filename -- metaPath is validated by getRunPath which checks for path traversal
+    const content = await fs.readFile(metaPath, 'utf-8');
+    return JSON.parse(content) as TraceRun;
+  } catch (error) {
+    logError('Failed to load run', { error, runId });
+    return null;
+  }
+}
+
+/**
+ * Loads a step artifact from storage.
+ * Validates artifact key to prevent path traversal attacks.
+ *
+ * @param runId - Run identifier.
+ * @param artifactKey - Artifact key (filename).
+ * @returns Artifact data or `null` if not found or validation fails.
+ * @remarks Rejects artifact keys containing '..' or '/' to prevent path traversal.
+ * Returns `null` on any error (file not found, invalid key, parse error).
+ * @example
+ * ```ts
+ * const artifact = await loadStepArtifact(runId, 'perplexity-123.p_notes.json');
+ * ```
+ */
+export async function loadStepArtifact(
+  runId: string,
+  artifactKey: string,
+): Promise<unknown | null> {
+  try {
+    // Validate artifactKey doesn't contain path traversal
+    if (artifactKey.includes('..') || artifactKey.includes('/')) {
+      logError('Invalid artifact key', { artifactKey });
+      return null;
+    }
+
+    const artifactPath = getRunPath(runId, artifactKey);
+    if (!artifactPath) {
+      return null;
+    }
+    // eslint-disable-next-line security/detect-non-literal-fs-filename -- artifactPath is validated by getRunPath and artifactKey sanitization
+    const content = await fs.readFile(artifactPath, 'utf-8');
+    return JSON.parse(content);
+  } catch (error) {
+    logError('Failed to load artifact', { error, artifactKey, runId });
+    return null;
+  }
+}