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

86.66% Statements 13/15
81.81% Branches 9/11
100% Functions 2/2
86.66% Lines 13/15

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                                          1x                 3x 3x 3x 2x 2x 2x 1x 1x     1x 1x   1x           1x      
import { SampleGenerationError } from '@metamask/kernel-errors';
import type { Logger } from '@metamask/logger';
import type { SampleCollector } from '@ocap/kernel-agents/types';
 
import { StatementMessage } from './messages.ts';
 
/**
 * A simple sample collector that collects content from a stream until a stop
 * string is encountered and returns the content prior, formatted.
 *
 * XXX This functionality is typically available in the language model service,
 * but a reimplementation appears here to remain agnostic to the service API.
 *
 * @param args - The arguments to make the sample collector.
 * @param args.stop - The stop string to stop collection.
 * @param args.maxChunkCount - The maximum number of chunks to collect before
 *   throwing an error.
 * @param args.logger - The logger to use for the sample collector.
 * @returns A function that collects a delta of a streaming response, returning a
 *   StatementMessage if the stop string is encountered or null otherwise.
 */
export const makeSampleCollector = ({
  stop,
  maxChunkCount = 200,
  logger,
}: {
  stop: string;
  maxChunkCount?: number;
  logger?: Logger;
}): SampleCollector<StatementMessage> => {
  let buffer = '';
  let chunkCount = 0;
  return (delta: string) => {
    chunkCount += 1;
    buffer += delta;
    if (buffer.includes(stop)) {
      const [content] = buffer.split(stop);
      Iif (content === undefined || content.trim() === '') {
        throw new SampleGenerationError(buffer, new Error('Empty content'));
      }
      logger?.info('content:', content);
      return StatementMessage.fromCode(content.trim());
    }
    Iif (maxChunkCount && chunkCount > maxChunkCount) {
      throw new SampleGenerationError(
        buffer,
        new Error('Max chunk count reached'),
      );
    }
    return null;
  };
};