Practical AI: Building a Robust Data Correction System with JSON Schema and LLMs

async function fixValueWithAI<T>(
  schema: JSONSchema,
  value: unknown,
  options?: {
    maxAttempts?: number;
    model?: string;
  }
): Promise<T> {
  // Try validating first - maybe it's already valid
  const validationResult = safeValidate(schema, value);
  if (validationResult.valid) {
    return validationResult.value as T;
  }

  // Not valid, let's try AI correction
  const attempts = options?.maxAttempts ?? 2;
  let currentValue = value;

  for (let i = 0; i < attempts; i++) {
    try {
      // Prepare instructions for the LLM
      const prompt = buildCorrectionPrompt({
        schema: stripInternalProperties(schema),
        originalValue: currentValue,
        errors: validationResult.error,
      });

      // Get AI correction
      const aiOutput = await askAI(prompt, options?.model ?? "gpt-4o-mini");

      // Check if AI signals it cannot fix
      const escapeHatchCheck = EscapeHatchSchema.safeParse(aiOutput);
      if (escapeHatchCheck.success) {
        throw new Error("AI indicates it cannot fix the value");
      }

      // Validate AI's proposed fix
      const newValidationResult = safeValidate(schema, aiOutput);
      if (newValidationResult.valid) {
        return newValidationResult.value as T;
      }

      // Still not valid, but use this as the starting point for next attempt
      currentValue = aiOutput;
    } catch (err) {
      // Log the error but continue to next attempt
      logger.warn(`AI correction attempt ${i + 1} failed`, { error: err });
    }
  }

  // All attempts failed
  throw new ValidationError("Could not correct value after multiple attempts");
}

function buildCorrectionPrompt({ schema, originalValue, errors }) {
  return `
You are a data correction expert. Your task is to fix a JSON value that fails validation.

THE JSON SCHEMA:
${JSON.stringify(schema, null, 2)}

THE ORIGINAL VALUE:
${JSON.stringify(originalValue, null, 2)}

VALIDATION ERRORS:
${JSON.stringify(errors, null, 2)}

INSTRUCTIONS:
1. Analyze the schema requirements and validation errors
2. Transform the original value to make it conform to the schema
3. ONLY fix what's wrong - preserve all other data
4. DO NOT invent values if you don't have enough information
5. If you cannot fix the value without guessing, respond with {"cannotFix": true}
6. Respond ONLY with the fixed JSON object or the cannotFix object

FIXED VALUE:
`;
}

{
  "type": "object",
  "properties": {
    "tenant": {
      "type": "string",
      "description": "Name of the tenant company"
    },
    "annual_rent": {
      "type": "number",
      "minimum": 0,
      "description": "Annual rent in euros"
    },
    "effective_date": {
      "type": "string",
      "format": "date",
      "description": "Start date of the lease"
    },
    "duration": {
      "type": "integer",
      "minimum": 1,
      "description": "Duration of the lease in years"
    }
  },
  "required": ["tenant", "annual_rent", "effective_date"]
}

{
  "tenant": "SCI Les Oliviers",
  "annual_rent": "25 000€ excl. tax",
  "effective_date": "first of January 2023",
  "duration": "3+6+9"
}

[
  {
    "path": "/annual_rent",
    "message": "must be number, found: string",
    "value": "25 000€ excl. tax"
  },
  {
    "path": "/effective_date",
    "message": "invalid date format",
    "value": "first of January 2023"
  },
  {
    "path": "/duration",
    "message": "must be integer, found: string",
    "value": "3+6+9"
  }
]

{
  "tenant": "SCI Les Oliviers",
  "annual_rent": 25000,
  "effective_date": "2023-01-01",
  "duration": 3
}

function safeValidate(schema: JSONSchema, value: unknown): ValidationResult {
  try {
    // Use AJV for validation
    const valid = ajv.validate(schema, value);

    if (valid) {
      return { valid: true, value };
    } else {
      return {
        valid: false,
        error: transformAjvErrors(ajv.errors),
      };
    }
  } catch (err) {
    // Handle errors in the validation process itself
    return {
      valid: false,
      error: [
        {
          path: "",
          message: "Validation process failed",
          value,
        },
      ],
    };
  }
}

// Transform AJV's verbose errors into a more concise format
function transformAjvErrors(errors: Ajv.ErrorObject[]): ErrorDetail[] {
  return errors.map((err) => ({
    path: err.instancePath,
    message: err.message,
    value: getValueAtPath(originalValue, err.instancePath),
  }));
}