All files / kernel-test/src/vats lms-queue-vat.ts

0% Statements 0/17
100% Branches 0/0
0% Functions 0/10
0% Lines 0/17

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                                                                                                                                 
import { makeDefaultExo } from '@metamask/kernel-utils/exo';
import { makeQueueService } from '@ocap/kernel-language-model-service/test-utils';
import { makeExoGenerator } from '@ocap/remote-iterables';
 
type QueueModel = {
  getInfo: () => unknown;
  load: () => Promise<void>;
  unload: () => Promise<void>;
  sample: (prompt: string) => Promise<{
    stream: AsyncIterable<unknown>;
    abort: () => void;
  }>;
  push: (text: string) => void;
};
 
/**
 * An envatted @ocap/kernel-language-model-service package.
 *
 * @returns A QueueLanguageModelService instance.
 */
// eslint-disable-next-line @typescript-eslint/explicit-function-return-type
export function buildRootObject() {
  const queueService = makeQueueService();
  return makeDefaultExo('root', {
    async makeInstance(config: unknown) {
      const model = (await queueService.makeInstance(config)) as QueueModel;
      return makeDefaultExo('queueLanguageModel', {
        async getInfo() {
          return model.getInfo();
        },
        async load() {
          return model.load();
        },
        async unload() {
          return model.unload();
        },
        async sample(prompt: string) {
          const result = await model.sample(prompt);
          // Convert the async iterable stream to an async generator and make it remotable
          const streamGenerator = async function* (): AsyncGenerator<unknown> {
            for await (const chunk of result.stream) {
              yield chunk;
            }
          };
          const streamRef = makeExoGenerator(streamGenerator());
          // Store abort function for later use
          const abortFn = result.abort;
          // Return a remotable object with getStream and abort as methods
          return makeDefaultExo('sampleResult', {
            getStream() {
              return streamRef;
            },
            async abort() {
              return abortFn();
            },
          });
        },
        push(text: string) {
          return model.push(text);
        },
      });
    },
  });
}