All files / src/sdk json-output.ts

68.57% Statements 24/35
60% Branches 18/30
80% Functions 4/5
68.57% Lines 24/35

Press n or j to go to the next uncovered block, b, p or k for the previous block.

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142              18x 18x 18x                                               2x 2x                     83x 83x           83x 83x       83x                           3x 3x 1x           2x                                           2x               2x                       86x 86x 3x 3x 3x         83x 83x 83x                  
import type { z } from 'zod';
import type { UsageStats } from '../types/index.js';
import { extractJson } from './haiku.js';
import { canUseRuntimeAuth } from './extract.js';
import { buildJsonOutputSection, buildTaggedSection, joinPromptSections } from './prompt-sections.js';
import { getRuntime, type Runtime, type RuntimeName } from './runtimes/index.js';
 
const JSON_REPAIR_MAX_CHARS = 60_000;
const JSON_REPAIR_MAX_TOKENS = 16_384;
const JSON_REPAIR_TIMEOUT_MS = 30_000;
 
export type ParseJsonFromOutputResult<T> =
  | { success: true; data: T; json: string; repaired: boolean; usage?: UsageStats }
  | { success: false; error: string; json?: string; usage?: UsageStats };
 
export interface JsonOutputRepairOptions {
  apiKey?: string;
  agentName?: string;
  runtime?: Runtime;
  runtimeName?: RuntimeName;
  model?: string;
  maxRetries?: number;
  maxTokens?: number;
  timeout?: number;
}
 
export interface ParseJsonFromOutputOptions<T> {
  output: string;
  schema: z.ZodType<T>;
  repair?: JsonOutputRepairOptions;
}
 
function truncateForRepair(output: string): string {
  Eif (output.length <= JSON_REPAIR_MAX_CHARS) {
    return output;
  }
  return `${output.slice(0, JSON_REPAIR_MAX_CHARS)}\n[... truncated]`;
}
 
function validationError(error: z.ZodError): string {
  return `validation_failed: ${error.message}`;
}
 
function parseExtractedJson<T>(json: string, schema: z.ZodType<T>): ParseJsonFromOutputResult<T> {
  let parsed: unknown;
  try {
    parsed = JSON.parse(json);
  } catch (error) {
    const message = error instanceof Error ? error.message : String(error);
    return { success: false, error: `invalid_json: ${message}`, json };
  }
 
  const validated = schema.safeParse(parsed);
  Iif (!validated.success) {
    return { success: false, error: validationError(validated.error), json };
  }
 
  return {
    success: true,
    data: validated.data,
    json,
    repaired: false,
  };
}
 
async function repairJsonOutput<T>(
  output: string,
  schema: z.ZodType<T>,
  reason: string,
  repair: JsonOutputRepairOptions,
): Promise<ParseJsonFromOutputResult<T>> {
  const runtime = repair.runtime ?? getRuntime(repair.runtimeName ?? 'claude');
  if (!canUseRuntimeAuth({ apiKey: repair.apiKey, runtime: runtime.name })) {
    return {
      success: false,
      error: `${reason}; repair_skipped: missing_api_key`,
    };
  }
 
  const result = await runtime.runAuxiliary({
    task: 'extraction',
    agentName: repair.agentName,
    apiKey: repair.apiKey,
    model: repair.model,
    maxRetries: repair.maxRetries,
    maxTokens: repair.maxTokens ?? JSON_REPAIR_MAX_TOKENS,
    timeout: repair.timeout ?? JSON_REPAIR_TIMEOUT_MS,
    schema,
    prompt: joinPromptSections([
      `<task>
Extract and repair the JSON value from this model output.
</task>`,
      buildJsonOutputSection(`Return JSON accepted by the provided schema.
Preserve the model's structured content as much as possible.
If the output contains markdown fences, escaped newlines, or prose around JSON, remove only the wrapper/prose and repair JSON escaping.
Do not summarize or invent new content.`),
      buildTaggedSection('parse_error', reason),
      buildTaggedSection('model_output', truncateForRepair(output)),
    ]),
  });
 
  Iif (!result.success) {
    return {
      success: false,
      error: `${reason}; repair_failed: ${result.error}`,
      usage: result.usage,
    };
  }
 
  return {
    success: true,
    data: result.data,
    json: JSON.stringify(result.data),
    repaired: true,
    usage: result.usage,
  };
}
 
export async function parseJsonFromOutput<T>(
  options: ParseJsonFromOutputOptions<T>,
): Promise<ParseJsonFromOutputResult<T>> {
  const json = extractJson(options.output);
  if (!json) {
    const reason = 'no_json';
    Eif (options.repair) {
      return repairJsonOutput(options.output, options.schema, reason, options.repair);
    }
    return { success: false, error: reason };
  }
 
  const parsed = parseExtractedJson(json, options.schema);
  Eif (parsed.success || !options.repair) {
    return parsed;
  }
 
  const repaired = await repairJsonOutput(options.output, options.schema, parsed.error, options.repair);
  if (!repaired.success && parsed.json) {
    return { ...repaired, json: parsed.json };
  }
  return repaired;
}