feat(trace-analyst): add Ax RLM trace analyst (#25)

Author: drewstone

src/index.ts

@@ -908,3 +908,6 @@ export type {
 
 // Pareto extensions (paretoFrontier + dominates already exported above)
 export { scalarScore, crowdingDistance, paretoFrontierWithCrowding } from './pareto'
+
+// Ax RLM trace analyst.
+export * from './trace-analyst'

src/trace-analyst/analyst.test.ts

@@ -0,0 +1,203 @@
+import { describe, expect, it, vi } from 'vitest'
+
+import { analyzeTraces } from './analyst'
+import type { TraceAnalysisStore } from './store'
+import type { DatasetOverview } from './types'
+
+const axMock = vi.hoisted(() => ({
+  agentCalls: [] as Array<{ signature: string; options: Record<string, unknown> }>,
+  forwardCalls: [] as Array<{ ai: unknown; values: unknown }>,
+}))
+
+vi.mock('@ax-llm/ax', () => {
+  const field = {
+    optional() {
+      return field
+    },
+    array() {
+      return field
+    },
+  }
+  const f = Object.assign(
+    () => ({
+      input() {
+        return this
+      },
+      output() {
+        return this
+      },
+      build() {
+        return {}
+      },
+    }),
+    {
+      string: () => field,
+      number: () => field,
+      boolean: () => field,
+      json: () => field,
+    },
+  )
+  const fn = (name: string) => {
+    const tool: Record<string, unknown> = { name }
+    const builder = {
+      description(value: string) {
+        tool.description = value
+        return builder
+      },
+      namespace(value: string) {
+        tool.namespace = value
+        return builder
+      },
+      arg() {
+        return builder
+      },
+      returns() {
+        return builder
+      },
+      handler(value: unknown) {
+        tool.handler = value
+        return builder
+      },
+      example() {
+        return builder
+      },
+      build() {
+        return tool
+      },
+    }
+    return builder
+  }
+  return {
+    f,
+    fn,
+    AxJSRuntime: class {
+      readonly options: unknown
+      constructor(options?: unknown) {
+        this.options = options
+      }
+    },
+    agent: (signature: string, options: Record<string, unknown>) => {
+      axMock.agentCalls.push({ signature, options })
+      return {
+        async forward(ai: unknown, values: unknown) {
+          axMock.forwardCalls.push({ ai, values })
+          const onTurn = options.actorTurnCallback
+          if (typeof onTurn === 'function') {
+            await onTurn({
+              turn: 1,
+              actionLogEntryCount: 1,
+              guidanceLogEntryCount: 0,
+              actorResult: {},
+              code: 'const overview = await traces.getDatasetOverview({})',
+              result: {},
+              output: 'overview loaded',
+              isError: false,
+              thought: 'inspect first',
+            })
+          }
+          return {
+            answer: 'publish_finding hits MaxTurnsExceeded in t000000000001/s004',
+            findings: ['t000000000001/s004: publish_finding hit MaxTurnsExceeded'],
+          }
+        },
+        getUsage() {
+          return { actor: [{ tokens: { totalTokens: 10 } }], responder: [] }
+        },
+        getChatLog() {
+          return { actor: [{ role: 'assistant' }], responder: [] }
+        },
+        resetUsage() {},
+      }
+    },
+  }
+})
+
+describe('analyzeTraces', () => {
+  it('constructs an Ax RLM analyst with bounded trace tools and returns run telemetry', async () => {
+    const overview: DatasetOverview = {
+      total_traces: 1,
+      raw_jsonl_bytes: 100,
+      services: ['bench'],
+      agents: ['driver'],
+      models: ['model-a'],
+      tool_names: ['publish_finding'],
+      sample_trace_ids: ['t000000000001'],
+      errors: { trace_count: 1, span_count: 1 },
+      time_range: null,
+    }
+    const store: TraceAnalysisStore = {
+      async getOverview() {
+        return overview
+      },
+      async queryTraces() {
+        return {
+          traces: [],
+          total: 0,
+          has_more: false,
+        }
+      },
+      async countTraces() {
+        return 1
+      },
+      async viewTrace() {
+        return { trace_id: 't000000000001', spans: [] }
+      },
+      async viewSpans() {
+        return {
+          trace_id: 't000000000001',
+          spans: [],
+          missing_span_ids: [],
+          truncated_attribute_count: 0,
+        }
+      },
+      async searchTrace() {
+        return {
+          trace_id: 't000000000001',
+          hits: [],
+          total_matches: 0,
+          has_more: false,
+        }
+      },
+      async searchSpan() {
+        return {
+          trace_id: 't000000000001',
+          span_id: 's004',
+          hits: [],
+          total_matches: 0,
+          has_more: false,
+        }
+      },
+    }
+
+    const ai = { provider: 'test' }
+    const result = await analyzeTraces(
+      { question: 'Which harness failure mode blocks success?' },
+      { source: store, ai, model: 'rlm-test', maxDepth: 1 },
+    )
+
+    expect(axMock.agentCalls).toHaveLength(1)
+    expect(axMock.agentCalls[0].signature).toBe('question:string -> answer:string, findings:string[]')
+    expect(axMock.agentCalls[0].options.mode).toBe('advanced')
+    expect(axMock.agentCalls[0].options.functions).toMatchObject({
+      local: expect.arrayContaining([
+        expect.objectContaining({ namespace: 'traces', name: 'getDatasetOverview' }),
+        expect.objectContaining({ namespace: 'traces', name: 'searchSpan' }),
+      ]),
+    })
+    expect(axMock.forwardCalls).toEqual([
+      {
+        ai,
+        values: { question: 'Which harness failure mode blocks success?' },
+      },
+    ])
+    expect(result.answer).toContain('MaxTurnsExceeded')
+    expect(result.findings).toEqual(['t000000000001/s004: publish_finding hit MaxTurnsExceeded'])
+    expect(result.turnCount).toBe(1)
+    expect(result.turns[0]).toMatchObject({
+      turn: 1,
+      isError: false,
+      output: 'overview loaded',
+    })
+    expect(result.actorPromptVersion).toMatch(/^trace-analyst-actor-v1-/)
+  })
+})

src/trace-analyst/analyst.ts

@@ -0,0 +1,175 @@
+import {
+  AxJSRuntime,
+  agent,
+  type AxActorTurn,
+  type AxAIService,
+  type AxFunction,
+} from '@ax-llm/ax'
+
+import { TraceFileMissingError } from './store-otlp'
+import {
+  TRACE_ANALYST_ACTOR_DESCRIPTION,
+  TRACE_ANALYST_ACTOR_DESCRIPTION_VERSION,
+  TRACE_ANALYST_SUBAGENT_DESCRIPTION,
+} from './prompts'
+import { buildTraceAnalystTools } from './tools'
+import type { TraceAnalysisStore } from './store'
+import { OtlpFileTraceStore } from './store-otlp'
+
+export interface AnalyzeTracesInput {
+  /** The user-facing question. Domain framing belongs here, not in the
+   *  actor description. */
+  question: string
+}
+
+export interface AnalyzeTracesResult {
+  /** The responder's prose answer. */
+  answer: string
+  /** Bulleted findings extracted from the responder's structured output. */
+  findings: string[]
+  /** Per-actor-turn snapshots captured via `actorTurnCallback`. */
+  turns: AnalyzeTracesTurnSnapshot[]
+  /** Total turns the actor took. */
+  turnCount: number
+  /** Token usage by role. */
+  usage: { actor: unknown[]; responder: unknown[] }
+  /** Full system + assistant + tool message log by role. */
+  chatLog: { actor: unknown[]; responder: unknown[] }
+  /** Prompt version that produced this run. */
+  actorPromptVersion: string
+}
+
+export interface AnalyzeTracesTurnSnapshot {
+  turn: number
+  isError: boolean
+  /** The JS code the actor produced for this turn. */
+  code: string
+  /** The formatted action-log entry the actor sees on the next turn. */
+  output: string
+  /** Provider thought (when `actorOptions.showThoughts` is true and the
+   *  provider returns it). */
+  thought?: string
+}
+
+export interface AnalyzeTracesOptions {
+  /** Trace data source. Pass either an OTLP-JSONL path or a custom store. */
+  source: string | TraceAnalysisStore
+  /** Caller-provided AxAIService. */
+  ai: AxAIService
+  /** Model id forwarded to actor + responder. */
+  model?: string
+  /** Recursion depth. 0 = no sub-agent dispatch. Default 1. */
+  maxDepth?: number
+  /** Maximum actor turns. Default 12. */
+  maxTurns?: number
+  /** Maximum parallel sub-agent calls in batched llmQuery. Default 2. */
+  maxParallelSubagents?: number
+  /** Override the actor description. */
+  actorDescription?: string
+  /** Override the subagent description. */
+  subagentDescription?: string
+  /** Per-turn observability hook. */
+  onTurn?: (turn: AnalyzeTracesTurnSnapshot) => void | Promise<void>
+  /** Override max runtime characters per turn. Default 6000. */
+  maxRuntimeChars?: number
+}
+
+/**
+ * Run the trace analyst.
+ *
+ * Throws:
+ *   - `TraceFileMissingError` if `source` is a path and doesn't exist.
+ *   - `AxAgentClarificationError` if the analyst asks for clarification.
+ *   - Provider errors (auth, rate limits) propagate from the AI service.
+ */
+export async function analyzeTraces(
+  input: AnalyzeTracesInput,
+  options: AnalyzeTracesOptions,
+): Promise<AnalyzeTracesResult> {
+  if (!input.question || typeof input.question !== 'string') {
+    throw new TypeError('analyzeTraces: input.question must be a non-empty string')
+  }
+
+  const store: TraceAnalysisStore =
+    typeof options.source === 'string'
+      ? new OtlpFileTraceStore({ path: options.source })
+      : options.source
+
+  // Pre-warm file stores so missing inputs fail before the RLM starts.
+  if (store instanceof OtlpFileTraceStore) {
+    await store.ensureIndexed()
+  }
+
+  const tools: AxFunction[] = buildTraceAnalystTools({ store })
+  const turns: AnalyzeTracesTurnSnapshot[] = []
+
+  const actorTurnCallback = async (turn: AxActorTurn): Promise<void> => {
+    const snap: AnalyzeTracesTurnSnapshot = {
+      turn: turn.turn,
+      isError: turn.isError,
+      code: turn.code,
+      output: turn.output,
+      thought: turn.thought,
+    }
+    turns.push(snap)
+    if (options.onTurn) await options.onTurn(snap)
+  }
+
+  const maxDepth = options.maxDepth ?? 1
+  const maxTurns = options.maxTurns ?? 12
+  const maxParallelSubagents = options.maxParallelSubagents ?? 2
+  const maxRuntimeChars = options.maxRuntimeChars ?? 6000
+
+  const analyst = agent<{ question: string }, { answer: string; findings: string[] }>(
+    'question:string -> answer:string, findings:string[]',
+    {
+      agentIdentity: {
+        name: 'TraceAnalyst',
+        description:
+          'Analyzes OTLP-shaped JSONL traces using bounded discovery tools to identify systemic failure modes.',
+      },
+      contextFields: ['question'],
+      runtime: new AxJSRuntime({
+        permissions: [],
+        blockDynamicImport: true,
+        allowedModules: [],
+        freezeIntrinsics: true,
+        blockShadowRealm: true,
+        preventGlobalThisExtensions: true,
+      }),
+      mode: maxDepth > 0 ? 'advanced' : 'simple',
+      recursionOptions: maxDepth > 0 ? { maxDepth } : undefined,
+      maxTurns,
+      maxRuntimeChars,
+      maxBatchedLlmQueryConcurrency: maxParallelSubagents,
+      promptLevel: 'detailed',
+      contextPolicy: { preset: 'checkpointed', budget: 'balanced' },
+      functions: { local: tools },
+      actorOptions: {
+        description: options.actorDescription ?? TRACE_ANALYST_ACTOR_DESCRIPTION,
+        ...(options.model ? { model: options.model } : {}),
+      },
+      responderOptions: {
+        ...(options.model ? { model: options.model } : {}),
+        description:
+          options.subagentDescription ?? TRACE_ANALYST_SUBAGENT_DESCRIPTION,
+      },
+      actorTurnCallback,
+      bubbleErrors: [TraceFileMissingError],
+    },
+  )
+
+  const result = await analyst.forward(options.ai, { question: input.question })
+
+  return {
+    answer: typeof result.answer === 'string' ? result.answer : String(result.answer ?? ''),
+    findings: Array.isArray(result.findings)
+      ? result.findings.filter((s): s is string => typeof s === 'string')
+      : [],
+    turns,
+    turnCount: turns.length,
+    usage: analyst.getUsage(),
+    chatLog: analyst.getChatLog(),
+    actorPromptVersion: TRACE_ANALYST_ACTOR_DESCRIPTION_VERSION,
+  }
+}