All files / kernel-language-model-service/src client.ts

100% Statements 15/15
100% Branches 10/10
100% Functions 4/4
100% Lines 15/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 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110                                                      4x                                                                       4x 4x 1x   3x 3x 1x   2x     4x                               4x               3x       3x 3x 1x   2x      
import { E } from '@endo/eventual-send';
import type { ERef } from '@endo/eventual-send';
 
import type {
  ChatParams,
  ChatResult,
  ChatService,
  ChatStreamChunk,
  SampleParams,
  SampleResult,
  SampleService,
} from './types.ts';
 
/**
 * Wraps a remote service reference with Open /v1-style chat completion ergonomics.
 *
 * Usage:
 * ```ts
 * const client = makeChatClient(lmsRef, 'gpt-4o');
 * const result = await client.chat.completions.create({ messages });
 * const stream = await client.chat.completions.create({ messages, stream: true });
 * ```
 *
 * @param lmsRef - Reference to a service with a `chat` method.
 * @param defaultModel - Default model name used when params do not specify one.
 * @returns A client object with `chat.completions.create`.
 */
export const makeChatClient = (
  lmsRef: ERef<ChatService>,
  defaultModel?: string,
): {
  chat: {
    completions: {
      create(
        params: Omit<ChatParams, 'model'> & { model?: string; stream: true },
      ): Promise<AsyncIterable<ChatStreamChunk>>;
      create(
        params: Omit<ChatParams, 'model'> & { model?: string; stream?: false },
      ): Promise<ChatResult>;
    };
  };
} => {
  type BaseParams = Omit<ChatParams, 'model'> & { model?: string };
 
  /**
   * @param params - Chat completion parameters with `stream: true`.
   * @returns A promise resolving to an async iterable of stream chunks.
   */
  function create(
    params: BaseParams & { stream: true },
  ): Promise<AsyncIterable<ChatStreamChunk>>;
  /**
   * @param params - Chat completion parameters.
   * @returns A promise resolving to the full chat result.
   */
  function create(params: BaseParams & { stream?: false }): Promise<ChatResult>;
  /**
   * @param params - Chat completion parameters.
   * @returns A promise resolving to a stream or full result depending on `stream`.
   */
  async function create(
    params: BaseParams,
  ): Promise<AsyncIterable<ChatStreamChunk> | ChatResult> {
    const model = params.model ?? defaultModel;
    if (!model) {
      throw new Error('model is required');
    }
    const fullParams = harden({ ...params, model });
    if (fullParams.stream === true) {
      return E(lmsRef).chat(fullParams as ChatParams & { stream: true });
    }
    return E(lmsRef).chat(fullParams as ChatParams & { stream?: false });
  }
 
  return harden({ chat: { completions: { create } } });
};
 
/**
 * Wraps a remote service reference with raw token-prediction ergonomics.
 *
 * Usage:
 * ```ts
 * const client = makeSampleClient(lmsRef, 'llama3');
 * const result = await client.sample({ prompt: 'Once upon' });
 * ```
 *
 * @param lmsRef - Reference to a service with a `sample` method.
 * @param defaultModel - Default model name used when params do not specify one.
 * @returns A client object with `sample`.
 */
export const makeSampleClient = (
  lmsRef: ERef<SampleService>,
  defaultModel?: string,
): {
  sample: (
    params: Omit<SampleParams, 'model'> & { model?: string },
  ) => Promise<SampleResult>;
} =>
  harden({
    async sample(
      params: Omit<SampleParams, 'model'> & { model?: string },
    ): Promise<SampleResult> {
      const model = params.model ?? defaultModel;
      if (!model) {
        throw new Error('model is required');
      }
      return E(lmsRef).sample(harden({ ...params, model }));
    },
  });