All files / kernel-agents/src/strategies/json sample-collector.ts

100% Statements 25/25
90% Branches 9/10
100% Functions 2/2
100% Lines 25/25

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                              2x                 9x 9x 9x 9x 9x 12x 12x 12x 12x 12x 12x 12x 12x 12x 12x 12x 12x       5x 1x                 4x 4x 4x 1x         3x      
import { SampleGenerationError } from '@metamask/kernel-errors';
import type { Logger } from '@metamask/logger';
 
import type { SampleCollector } from '../../types.ts';
 
/**
 * A quick and dirty sample collector for a streaming response.
 *
 * @param args - The arguments to make the sample collector.
 * @param args.prefix - The prefix to prepend to the response
 * @param args.maxChunkCount - The maximum number of chunks to parse
 * @param args.logger - The logger to use for the sample collector
 * @returns A function that collects a delta of a streaming response,
 *   returning the result value if collecting is complete or null otherwise.
 */
export const makeSampleCollector = <Result = unknown>({
  prefix = '',
  maxChunkCount = 200,
  logger,
}: {
  prefix?: string;
  maxChunkCount?: number;
  logger?: Logger;
}): SampleCollector<Result> => {
  let response = prefix;
  let chunkCount = 0;
  let leftBracketCount = prefix.split('{').length - 1;
  let rightBracketCount = prefix.split('}').length - 1;
  return (delta: string) => {
    chunkCount += 1;
    const subchunks = delta.split('}');
    const lastSubchunk = subchunks.pop() as string;
    for (const subchunk of subchunks) {
      rightBracketCount += 1;
      leftBracketCount += subchunk.split('{').length - 1;
      response += `${subchunk}}`;
      logger?.info('toParse:', response);
      try {
        const result = JSON.parse(response);
        logger?.info('parsed:', result);
        return result;
      } catch (cause) {
        // XXX There are other ways to detect an irrecoverable state.
        // This is the simplest.
        if (leftBracketCount === rightBracketCount) {
          throw new SampleGenerationError(
            response,
            cause instanceof Error
              ? cause
              : new Error('Invalid JSON', { cause }),
          );
        }
      }
    }
    leftBracketCount += lastSubchunk.split('{').length - 1;
    response += lastSubchunk;
    if (maxChunkCount && chunkCount > maxChunkCount) {
      throw new SampleGenerationError(
        response,
        new Error('Max chunk count reached'),
      );
    }
    return null;
  };
};