elasticdash-sdk 0.2.0

package/src/core/agent-state.ts

@@ -0,0 +1,162 @@
+/**
+ * Agent state serialization, deserialization, and utility functions.
+ *
+ * Enables capturing agent plan state during execution and resuming
+ * agents from any task in the plan without re-executing completed steps.
+ */
+
+import type { AgentPlan, AgentState, AgentTask } from '../types/agent.js'
+import type { WorkflowEvent } from '../capture/event.js'
+
+/**
+ * Serializes an agent plan and its captured trace events into an AgentState
+ * that can be persisted and later used for resumption.
+ *
+ * The resumeFromTaskIndex is automatically determined as the index of the
+ * first non-completed task. If all tasks are completed, it equals tasks.length.
+ */
+export function serializeAgentState(plan: AgentPlan, trace: WorkflowEvent[]): AgentState {
+  const resumeFromTaskIndex = plan.tasks.findIndex(
+    (t) => t.status !== 'completed',
+  )
+  return {
+    plan: JSON.parse(JSON.stringify(plan)) as AgentPlan,
+    trace: JSON.parse(JSON.stringify(trace)) as WorkflowEvent[],
+    resumeFromTaskIndex: resumeFromTaskIndex === -1 ? plan.tasks.length : resumeFromTaskIndex,
+  }
+}
+
+/**
+ * Validates and hydrates an AgentState from a parsed JSON object.
+ * Throws if the state is invalid or cannot be safely used for resumption.
+ */
+export function deserializeAgentState(raw: unknown): AgentState {
+  if (!raw || typeof raw !== 'object') {
+    throw new Error('AgentState must be a non-null object')
+  }
+  const obj = raw as Record<string, unknown>
+
+  if (!obj.plan || typeof obj.plan !== 'object') {
+    throw new Error('AgentState.plan is required')
+  }
+  const plan = obj.plan as AgentPlan
+
+  if (!Array.isArray(plan.tasks)) {
+    throw new Error('AgentState.plan.tasks must be an array')
+  }
+  if (typeof plan.id !== 'string') {
+    throw new Error('AgentState.plan.id must be a string')
+  }
+
+  const trace = Array.isArray(obj.trace) ? (obj.trace as WorkflowEvent[]) : []
+  const resumeFromTaskIndex =
+    typeof obj.resumeFromTaskIndex === 'number' ? obj.resumeFromTaskIndex : 0
+
+  // Validate that all tasks before resumeFromTaskIndex have outputs
+  for (let i = 0; i < resumeFromTaskIndex; i++) {
+    const task = plan.tasks[i]
+    if (!task) continue
+    if (task.status !== 'completed') {
+      throw new Error(
+        `Task at index ${i} (id="${task.id}") has status "${task.status}" but must be "completed" before resumeFromTaskIndex=${resumeFromTaskIndex}`,
+      )
+    }
+    if (task.output === undefined) {
+      throw new Error(
+        `Task at index ${i} (id="${task.id}") is completed but has no output. Cannot resume safely.`,
+      )
+    }
+  }
+
+  return { plan, trace, resumeFromTaskIndex }
+}
+
+/**
+ * Extracts all completed task outputs into a flat map keyed by task ID.
+ * Used for resolving placeholder references in subsequent task inputs.
+ */
+export function extractTaskOutputs(plan: AgentPlan): Record<string, unknown> {
+  const outputs: Record<string, unknown> = {}
+  for (const task of plan.tasks) {
+    if (task.status === 'completed' && task.output !== undefined) {
+      outputs[task.id] = task.output
+    }
+  }
+  return outputs
+}
+
+/**
+ * Resolves placeholder references in a task input.
+ *
+ * Placeholders use the form: `{ $ref: "task-N.output.fieldName" }`
+ * where "task-N" is a task ID and "fieldName" is a dot-separated path
+ * into that task's output.
+ *
+ * Example:
+ *   previousOutputs = { "task-1": { userId: "abc" } }
+ *   input = { $ref: "task-1.output.userId" }
+ *   → returns "abc"
+ *
+ * Works recursively on nested objects and arrays.
+ */
+export function resolveTaskInput(
+  input: unknown,
+  previousOutputs: Record<string, unknown>,
+): unknown {
+  if (input === null || input === undefined) return input
+
+  if (Array.isArray(input)) {
+    return input.map((item) => resolveTaskInput(item, previousOutputs))
+  }
+
+  if (typeof input === 'object') {
+    const obj = input as Record<string, unknown>
+
+    // Check for placeholder: { $ref: "taskId.output.path" }
+    if (typeof obj['$ref'] === 'string') {
+      return resolveRef(obj['$ref'], previousOutputs)
+    }
+
+    // Recursively resolve nested objects
+    const resolved: Record<string, unknown> = {}
+    for (const [k, v] of Object.entries(obj)) {
+      resolved[k] = resolveTaskInput(v, previousOutputs)
+    }
+    return resolved
+  }
+
+  return input
+}
+
+/**
+ * Resolves a dot-separated reference path like "task-1.output.userId"
+ * against the previousOutputs map.
+ */
+function resolveRef(ref: string, previousOutputs: Record<string, unknown>): unknown {
+  const parts = ref.split('.')
+  // Expected format: <taskId>.output.<...path>
+  // We skip the literal "output" segment to navigate into the output object
+  const taskId = parts[0]
+  const pathParts = parts.slice(1) // may start with "output"
+
+  let current: unknown = previousOutputs[taskId]
+  for (const part of pathParts) {
+    if (part === 'output') continue // "output" is implicit — skip this keyword
+    if (current === null || current === undefined) return undefined
+    current = (current as Record<string, unknown>)[part]
+  }
+  return current
+}
+
+/**
+ * Clones a task and marks it as completed with the given output.
+ * Used internally when skipping already-completed tasks during resumption.
+ */
+export function markTaskCompleted(task: AgentTask, output: unknown): AgentTask {
+  return {
+    ...task,
+    status: 'completed',
+    output,
+    completedAt: task.completedAt ?? Date.now(),
+  }
+}

package/src/core/judge-utils.ts

@@ -0,0 +1,232 @@
+/**
+ * judge-utils.ts
+ *
+ * Utilities for preprocessing outputs before sending to LLM-as-a-judge evaluators.
+ * Addresses two problems:
+ * 1. Large outputs cause slow inference
+ * 2. LLMs miss attributes in large JSON payloads due to attention degradation
+ */
+
+const DEFAULT_MAX_CHARS = 8000
+const MAX_STRING_VALUE_LENGTH = 500
+const MAX_ARRAY_EDGE_ITEMS = 5
+
+/**
+ * Prepare an output string for LLM judge evaluation.
+ *
+ * - Small outputs (< maxChars) pass through unchanged.
+ * - Large JSON outputs: extracts subtrees whose keys match keywords from the judge prompt.
+ * - Large non-JSON outputs: truncates with head + tail and a marker in between.
+ *
+ * @param output     - The raw output string to prepare.
+ * @param judgePrompt - The judge/evaluation prompt, used to identify relevant JSON keys.
+ * @param maxChars   - Maximum character budget for the prepared output. Default 8000.
+ * @returns The prepared output string, possibly trimmed.
+ */
+export function prepareOutputForJudge(
+  output: string,
+  judgePrompt: string,
+  maxChars: number = DEFAULT_MAX_CHARS,
+): string {
+  if (!output) return output
+
+  // Always attempt JSON-aware processing (value truncation + key extraction).
+  // Even outputs under maxChars can contain arrays with 100+ items that
+  // drown out the fields the judge actually needs to evaluate.
+  const jsonResult = tryJsonExtract(output, judgePrompt, maxChars)
+  if (jsonResult !== null) return jsonResult
+
+  // Non-JSON: only truncate if over budget
+  if (output.length <= maxChars) return output
+  return truncateHeadTail(output, maxChars)
+}
+
+/**
+ * Attempt to parse the output as JSON and extract only the subtrees
+ * whose keys are relevant to the judge prompt.
+ * Returns null if the output is not valid JSON.
+ */
+function tryJsonExtract(output: string, judgePrompt: string, maxChars: number): string | null {
+  let parsed: unknown
+  try {
+    parsed = JSON.parse(output)
+  } catch {
+    return null
+  }
+
+  if (typeof parsed !== 'object' || parsed === null) return null
+
+  const keywords = extractKeywords(judgePrompt)
+  if (keywords.length === 0) {
+    // No keywords to match — truncate large values then serialize
+    const trimmed = truncateJsonValues(parsed)
+    const trimmedStr = JSON.stringify(trimmed, null, 2)
+    if (trimmedStr.length <= maxChars) return trimmedStr
+    return truncateHeadTail(trimmedStr, maxChars)
+  }
+
+  const extracted = extractRelevantPaths(parsed, keywords)
+
+  // If extraction found relevant content, truncate large values then serialize
+  if (extracted !== undefined && Object.keys(extracted as Record<string, unknown>).length > 0) {
+    const trimmed = truncateJsonValues(extracted)
+    const extractedStr = JSON.stringify(trimmed, null, 2)
+
+    if (extractedStr.length <= maxChars) {
+      const omittedKeys = countOmittedKeys(parsed, extracted)
+      if (omittedKeys > 0) {
+        return `${extractedStr}\n\n[Note: ${omittedKeys} irrelevant key(s) omitted from original output for brevity. Total original size: ${output.length} chars.]`
+      }
+      return extractedStr
+    }
+
+    // Extracted content is still too large — truncate it
+    return truncateHeadTail(extractedStr, maxChars)
+  }
+
+  // Extraction found nothing relevant — truncate values in full object
+  const trimmedFull = truncateJsonValues(parsed)
+  const trimmedFullStr = JSON.stringify(trimmedFull, null, 2)
+  if (trimmedFullStr.length <= maxChars) return trimmedFullStr
+  return truncateHeadTail(trimmedFullStr, maxChars)
+}
+
+/**
+ * Extract lowercase keywords from the judge prompt.
+ * Filters out common stop words and short tokens.
+ */
+function extractKeywords(prompt: string): string[] {
+  const stopWords = new Set([
+    'the', 'a', 'an', 'is', 'are', 'was', 'were', 'be', 'been', 'being',
+    'have', 'has', 'had', 'do', 'does', 'did', 'will', 'would', 'could',
+    'should', 'may', 'might', 'shall', 'can', 'need', 'must',
+    'and', 'or', 'but', 'if', 'then', 'else', 'when', 'where', 'how',
+    'what', 'which', 'who', 'whom', 'this', 'that', 'these', 'those',
+    'it', 'its', 'of', 'in', 'on', 'at', 'to', 'for', 'with', 'by',
+    'from', 'as', 'into', 'about', 'between', 'through', 'after', 'before',
+    'not', 'no', 'nor', 'only', 'also', 'just', 'more', 'most', 'very',
+    'all', 'each', 'every', 'any', 'some', 'such', 'than', 'too',
+    'output', 'evaluate', 'score', 'check', 'whether', 'contains',
+    'following', 'given', 'based', 'respond', 'number', 'scale',
+    'text', 'result', 'response', 'answer', 'return', 'value',
+  ])
+
+  const words = prompt
+    .toLowerCase()
+    .replace(/[^a-z0-9_\s-]/g, ' ')
+    .split(/\s+/)
+    .filter(w => w.length > 2 && !stopWords.has(w))
+
+  return [...new Set(words)]
+}
+
+/**
+ * Recursively extract object entries whose keys match any of the keywords.
+ * For arrays, preserves items that contain matching keys.
+ */
+function extractRelevantPaths(obj: unknown, keywords: string[]): unknown {
+  if (Array.isArray(obj)) {
+    // For arrays: keep items that have relevant keys, limit to first few items
+    const relevant = obj
+      .map(item => extractRelevantPaths(item, keywords))
+      .filter(item => item !== undefined)
+    if (relevant.length === 0) return undefined
+    return relevant
+  }
+
+  if (typeof obj === 'object' && obj !== null) {
+    const result: Record<string, unknown> = {}
+    let hasMatch = false
+
+    for (const [key, value] of Object.entries(obj)) {
+      const keyLower = key.toLowerCase()
+      const keyMatchesDirectly = keywords.some(kw =>
+        keyLower.includes(kw) || kw.includes(keyLower)
+      )
+
+      if (keyMatchesDirectly) {
+        result[key] = value
+        hasMatch = true
+      } else if (typeof value === 'object' && value !== null) {
+        // Recurse into nested objects/arrays
+        const nested = extractRelevantPaths(value, keywords)
+        if (nested !== undefined) {
+          result[key] = nested
+          hasMatch = true
+        }
+      }
+    }
+
+    return hasMatch ? result : undefined
+  }
+
+  return undefined
+}
+
+/**
+ * Recursively truncate large values inside a JSON structure.
+ * - Long strings: keep first/last portions with a marker in between.
+ * - Long arrays: keep first 5 and last 5 items, skip the rest with a marker.
+ */
+function truncateJsonValues(obj: unknown): unknown {
+  if (obj === null || obj === undefined) return obj
+
+  if (typeof obj === 'string') {
+    if (obj.length <= MAX_STRING_VALUE_LENGTH) return obj
+    const headLen = Math.floor(MAX_STRING_VALUE_LENGTH * 0.6)
+    const tailLen = MAX_STRING_VALUE_LENGTH - headLen
+    return `${obj.slice(0, headLen)}...[${obj.length - headLen - tailLen} chars truncated]...${obj.slice(obj.length - tailLen)}`
+  }
+
+  if (Array.isArray(obj)) {
+    if (obj.length <= MAX_ARRAY_EDGE_ITEMS * 2) {
+      return obj.map(item => truncateJsonValues(item))
+    }
+    const head = obj.slice(0, MAX_ARRAY_EDGE_ITEMS).map(item => truncateJsonValues(item))
+    const tail = obj.slice(obj.length - MAX_ARRAY_EDGE_ITEMS).map(item => truncateJsonValues(item))
+    const skipped = obj.length - MAX_ARRAY_EDGE_ITEMS * 2
+    return [...head, `[...${skipped} items skipped...]`, ...tail]
+  }
+
+  if (typeof obj === 'object') {
+    const result: Record<string, unknown> = {}
+    for (const [key, value] of Object.entries(obj)) {
+      result[key] = truncateJsonValues(value)
+    }
+    return result
+  }
+
+  return obj
+}
+
+/**
+ * Count how many top-level keys from the original object are not in the extracted object.
+ */
+function countOmittedKeys(original: unknown, extracted: unknown): number {
+  if (typeof original !== 'object' || original === null) return 0
+  if (typeof extracted !== 'object' || extracted === null) return 0
+  if (Array.isArray(original)) return 0
+
+  const origKeys = Object.keys(original)
+  const extKeys = new Set(Object.keys(extracted as Record<string, unknown>))
+  return origKeys.filter(k => !extKeys.has(k)).length
+}
+
+/**
+ * Truncate a string by keeping the head and tail portions,
+ * inserting a marker in the middle.
+ */
+function truncateHeadTail(text: string, maxChars: number): string {
+  if (text.length <= maxChars) return text
+
+  // Reserve space for the marker
+  const marker = `\n\n[...truncated ${text.length - maxChars} chars out of ${text.length} total...]\n\n`
+  const available = maxChars - marker.length
+  if (available <= 0) return text.slice(0, maxChars)
+
+  // 70% head, 30% tail — head is usually more important
+  const headSize = Math.floor(available * 0.7)
+  const tailSize = available - headSize
+
+  return text.slice(0, headSize) + marker + text.slice(text.length - tailSize)
+}

package/src/core/registry.ts

@@ -0,0 +1,92 @@
+import type { AITestContext } from '../trace-adapter/context.js'
+
+export type TestFunction = (ctx: AITestContext) => Promise<void> | void
+
+export interface TestEntry {
+  name: string
+  fn: TestFunction
+}
+
+export interface Registry {
+  tests: TestEntry[]
+  beforeAllHooks: Array<() => Promise<void> | void>
+  afterAllHooks: Array<() => Promise<void> | void>
+  beforeEachHooks: Array<() => Promise<void> | void>
+  afterEachHooks: Array<() => Promise<void> | void>
+}
+
+const REGISTRY_KEY = '__elasticdash_registry__'
+
+function getGlobalRegistry(): Registry {
+  if (!(globalThis as any)[REGISTRY_KEY]) {
+    (globalThis as any)[REGISTRY_KEY] = createEmptyRegistry()
+  }
+  return (globalThis as any)[REGISTRY_KEY] as Registry
+}
+
+function createEmptyRegistry(): Registry {
+  return {
+    tests: [],
+    beforeAllHooks: [],
+    afterAllHooks: [],
+    beforeEachHooks: [],
+    afterEachHooks: [],
+  }
+}
+
+export function clearRegistry(): void {
+  (globalThis as any)[REGISTRY_KEY] = createEmptyRegistry()
+  console.log('[elasticdash] clearRegistry called. Registry reset.')
+}
+
+export function getRegistry(): Registry {
+  const registry = getGlobalRegistry()
+  console.log('[elasticdash] getRegistry called. Current tests:', registry.tests.map(t => t.name))
+  return registry
+}
+
+export function aiTest(name: string, fn: TestFunction): void {
+  const registry = getGlobalRegistry()
+  registry.tests.push({ name, fn })
+  console.log(`[elasticdash] Registered test: ${name}`)
+}
+
+export function beforeAll(fn: () => Promise<void> | void): void {
+  const registry = getGlobalRegistry()
+  registry.beforeAllHooks.push(fn)
+}
+
+export function afterAll(fn: () => Promise<void> | void): void {
+  const registry = getGlobalRegistry()
+  registry.afterAllHooks.push(fn)
+}
+
+export function beforeEach(fn: () => Promise<void> | void): void {
+  const registry = getGlobalRegistry()
+  registry.beforeEachHooks.push(fn)
+}
+
+export function afterEach(fn: () => Promise<void> | void): void {
+  const registry = getGlobalRegistry()
+  registry.afterEachHooks.push(fn)
+}
+
+// Expose globally so test files can use without importing
+declare global {
+  // eslint-disable-next-line no-var
+  var aiTest: (name: string, fn: TestFunction) => void
+  // eslint-disable-next-line no-var
+  var beforeAll: (fn: () => Promise<void> | void) => void
+  // eslint-disable-next-line no-var
+  var afterAll: (fn: () => Promise<void> | void) => void
+  // eslint-disable-next-line no-var
+  var beforeEach: (fn: () => Promise<void> | void) => void
+  // eslint-disable-next-line no-var
+  var afterEach: (fn: () => Promise<void> | void) => void
+}
+
+globalThis.aiTest = aiTest
+globalThis.beforeAll = beforeAll
+globalThis.afterAll = afterAll
+globalThis.beforeEach = beforeEach
+globalThis.afterEach = afterEach