Extract JSON array from LLM text

extract-json.ts

function matchBalancedBrackets(inputText) {
  const matches = []
  let stack = 0
  let startIndex

  for (let i = 0; i < inputText.length; i++) {
    const char = inputText[i]
    if (char === '[') {
      if (stack === 0) {
        startIndex = i
      }
      stack++
    } else if (char === ']') {
      stack--
      if (stack === 0) {
        matches.push(inputText.slice(startIndex, i + 1))
      }
    }
  }

  return matches
}

export function extractJson(inputText: string): WordsArray {
  // Find all JSON-like substrings
  const candidates: string[] = matchBalancedBrackets(inputText)

  // Iterate through the candidates and check if they conform to the WordsArray type
  for (const candidate of candidates) {
    try {
      const jsonArray = JSON.parse(candidate)
      const parsedArray = ExpectedSchema.safeParse(jsonArray)
      console.log(parsedArray, null, 2)

      if (parsedArray.success) {
        return parsedArray.data
      }
    } catch (error) {
      console.log('Error:', error)
    }
  }

  return []
}