diff --git a/packages/kernel-language-model-service/package.json b/packages/kernel-language-model-service/package.json index 30e3c4fea..48a23ea2a 100644 --- a/packages/kernel-language-model-service/package.json +++ b/packages/kernel-language-model-service/package.json @@ -104,6 +104,7 @@ "node": ">=22" }, "dependencies": { + "@endo/eventual-send": "^1.3.4", "@metamask/superstruct": "^3.2.1", "ollama": "^0.5.16", "ses": "^1.14.0" diff --git a/packages/kernel-language-model-service/src/client.test.ts b/packages/kernel-language-model-service/src/client.test.ts new file mode 100644 index 000000000..22dd6febe --- /dev/null +++ b/packages/kernel-language-model-service/src/client.test.ts @@ -0,0 +1,107 @@ +import { describe, it, expect, vi } from 'vitest'; + +import { makeChatClient, makeSampleClient } from './client.ts'; +import type { ChatResult, SampleResult } from './types.ts'; + +const MODEL = 'glm-4.7-flash'; + +vi.mock('@endo/eventual-send', () => ({ + E: vi.fn((obj: unknown) => obj), +})); + +const makeChatResult = (): ChatResult => ({ + id: 'chat-1', + model: MODEL, + choices: [ + { + message: { role: 'assistant', content: 'hello' }, + index: 0, + finish_reason: 'stop', + }, + ], + usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 }, +}); + +const makeSampleResult = (): SampleResult => ({ text: 'hi there' }); + +describe('makeChatClient', () => { + it('calls chat on the lmsRef with merged model', async () => { + const chatResult = makeChatResult(); + const lmsRef = { chat: vi.fn().mockResolvedValue(chatResult) }; + + const client = makeChatClient(lmsRef, MODEL); + const result = await client.chat.completions.create({ + messages: [{ role: 'user', content: 'hello' }], + }); + + expect(lmsRef.chat).toHaveBeenCalledWith( + expect.objectContaining({ + model: MODEL, + messages: [{ role: 'user', content: 'hello' }], + }), + ); + expect(result).toStrictEqual(chatResult); + }); + + it('params.model overrides defaultModel', async () => { + const lmsRef = { chat: vi.fn().mockResolvedValue(makeChatResult()) }; + const client = makeChatClient(lmsRef, 'gpt-3.5'); + + await client.chat.completions.create({ + messages: [{ role: 'user', content: 'hi' }], + model: MODEL, + }); + + expect(lmsRef.chat).toHaveBeenCalledWith( + expect.objectContaining({ model: MODEL }), + ); + }); + + it('throws when no model is available', async () => { + const lmsRef = { chat: vi.fn() }; + const client = makeChatClient(lmsRef); + + await expect( + client.chat.completions.create({ + messages: [{ role: 'user', content: 'hi' }], + }), + ).rejects.toThrow('model is required'); + }); +}); + +describe('makeSampleClient', () => { + it('calls sample on the lmsRef with merged model', async () => { + const rawResult = makeSampleResult(); + const lmsRef = { sample: vi.fn().mockResolvedValue(rawResult) }; + + const client = makeSampleClient(lmsRef, 'llama3'); + const result = await client.sample({ prompt: 'Once upon' }); + + expect(lmsRef.sample).toHaveBeenCalledWith( + expect.objectContaining({ model: 'llama3', prompt: 'Once upon' }), + ); + expect(result).toStrictEqual(rawResult); + }); + + it('params.model overrides defaultModel', async () => { + const lmsRef = { + sample: vi.fn().mockResolvedValue(makeSampleResult()), + }; + const client = makeSampleClient(lmsRef, 'llama3'); + + await client.sample({ prompt: 'hi', model: 'mistral' }); + + expect(lmsRef.sample).toHaveBeenCalledWith( + expect.objectContaining({ model: 'mistral' }), + ); + }); + + it('throws when no model is available', async () => { + const lmsRef = { sample: vi.fn() }; + const client = makeSampleClient(lmsRef); + + await expect(client.sample({ prompt: 'hi' })).rejects.toThrow( + 'model is required', + ); + }); +}); diff --git a/packages/kernel-language-model-service/src/client.ts b/packages/kernel-language-model-service/src/client.ts new file mode 100644 index 000000000..24973d6da --- /dev/null +++ b/packages/kernel-language-model-service/src/client.ts @@ -0,0 +1,85 @@ +import { E } from '@endo/eventual-send'; +import type { ERef } from '@endo/eventual-send'; + +import type { + ChatParams, + ChatResult, + ChatService, + 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 }); + * ``` + * + * @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, + defaultModel?: string, +): { + chat: { + completions: { + create: ( + params: Omit & { model?: string }, + ) => Promise; + }; + }; +} => + harden({ + chat: harden({ + completions: harden({ + async create( + params: Omit & { model?: string }, + ): Promise { + const model = params.model ?? defaultModel; + if (!model) { + throw new Error('model is required'); + } + return E(lmsRef).chat(harden({ ...params, model })); + }, + }), + }), + }); + +/** + * 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, + defaultModel?: string, +): { + sample: ( + params: Omit & { model?: string }, + ) => Promise; +} => + harden({ + async sample( + params: Omit & { model?: string }, + ): Promise { + const model = params.model ?? defaultModel; + if (!model) { + throw new Error('model is required'); + } + return E(lmsRef).sample(harden({ ...params, model })); + }, + }); diff --git a/packages/kernel-language-model-service/src/index.ts b/packages/kernel-language-model-service/src/index.ts index c20e32944..1b31cdf67 100644 --- a/packages/kernel-language-model-service/src/index.ts +++ b/packages/kernel-language-model-service/src/index.ts @@ -1 +1,7 @@ export type * from './types.ts'; +export { + LANGUAGE_MODEL_SERVICE_NAME, + makeKernelLanguageModelService, +} from './kernel-service.ts'; +export { makeOpenV1NodejsService } from './open-v1/nodejs.ts'; +export { makeChatClient, makeSampleClient } from './client.ts'; diff --git a/packages/kernel-language-model-service/src/kernel-service.test.ts b/packages/kernel-language-model-service/src/kernel-service.test.ts new file mode 100644 index 000000000..a6ff8712f --- /dev/null +++ b/packages/kernel-language-model-service/src/kernel-service.test.ts @@ -0,0 +1,100 @@ +import { describe, it, expect, vi } from 'vitest'; + +import { + LANGUAGE_MODEL_SERVICE_NAME, + makeKernelLanguageModelService, +} from './kernel-service.ts'; +import type { + ChatParams, + ChatResult, + SampleParams, + SampleResult, +} from './types.ts'; + +const makeChatResult = (): ChatResult => ({ + id: 'chat-1', + model: 'test-model', + choices: [ + { + message: { role: 'assistant', content: 'hi' }, + index: 0, + finish_reason: 'stop', + }, + ], + usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 }, +}); + +const makeSampleResult = (): SampleResult => ({ text: 'hello' }); + +describe('LANGUAGE_MODEL_SERVICE_NAME', () => { + it('equals languageModelService', () => { + expect(LANGUAGE_MODEL_SERVICE_NAME).toBe('languageModelService'); + }); +}); + +describe('makeKernelLanguageModelService', () => { + it('returns object with correct name and a service', () => { + const chat = vi.fn(); + const result = makeKernelLanguageModelService(chat); + expect(result).toMatchObject({ + name: LANGUAGE_MODEL_SERVICE_NAME, + service: expect.any(Object), + }); + }); + + it('service has chat and sample methods', () => { + const chat = vi.fn(); + const { service } = makeKernelLanguageModelService(chat); + expect(service).toMatchObject({ + chat: expect.any(Function), + sample: expect.any(Function), + }); + }); + + it('chat delegates to underlying function and returns hardened result', async () => { + const chatResult = makeChatResult(); + const chat = vi.fn().mockResolvedValue(chatResult); + const { service } = makeKernelLanguageModelService(chat); + + const params: ChatParams = { + model: 'test', + messages: [{ role: 'user', content: 'hi' }], + }; + const result = await ( + service as { chat: (p: ChatParams) => Promise } + ).chat(params); + + expect(chat).toHaveBeenCalledWith(params); + expect(result).toStrictEqual(chatResult); + }); + + it('sample delegates to provided function and returns hardened result', async () => { + const rawResult = makeSampleResult(); + const chat = vi.fn(); + const sample = vi.fn().mockResolvedValue(rawResult); + const { service } = makeKernelLanguageModelService(chat, sample); + + const params: SampleParams = { model: 'test', prompt: 'hello' }; + const result = await ( + service as { + sample: (p: SampleParams) => Promise; + } + ).sample(params); + + expect(sample).toHaveBeenCalledWith(params); + expect(result).toStrictEqual(rawResult); + }); + + it('sample throws when no sample function provided', async () => { + const chat = vi.fn(); + const { service } = makeKernelLanguageModelService(chat); + + await expect( + ( + service as { + sample: (p: SampleParams) => Promise; + } + ).sample({ model: 'test', prompt: 'hello' }), + ).rejects.toThrow('raw sampling not supported by this backend'); + }); +}); diff --git a/packages/kernel-language-model-service/src/kernel-service.ts b/packages/kernel-language-model-service/src/kernel-service.ts new file mode 100644 index 000000000..9b6e2b661 --- /dev/null +++ b/packages/kernel-language-model-service/src/kernel-service.ts @@ -0,0 +1,41 @@ +import type { + ChatParams, + ChatResult, + SampleParams, + SampleResult, +} from './types.ts'; + +/** + * Canonical service name for the language model service in `ClusterConfig.services`. + */ +export const LANGUAGE_MODEL_SERVICE_NAME = 'languageModelService'; + +/** + * Wraps `chat` and optional `sample` functions into a flat, stateless kernel service object. + * Use the returned `{ name, service }` with `kernel.registerKernelServiceObject(name, service)`. + * + * Return values are plain hardened data — no exos — so they are safely serializable + * across the kernel marshal boundary. + * + * @param chat - Function that performs a chat completion request. + * @param sample - Optional function that performs a raw token-prediction request. + * If not provided, `service.sample()` throws "raw sampling not supported by this backend". + * @returns An object with `name` and `service` fields for use with the kernel. + */ +export const makeKernelLanguageModelService = ( + chat: (params: ChatParams & { stream?: true & false }) => Promise, + sample?: (params: SampleParams) => Promise, +): { name: string; service: object } => { + const service = harden({ + async chat(params: ChatParams): Promise { + return harden(await chat(params as ChatParams & { stream?: never })); + }, + async sample(params: SampleParams): Promise { + if (!sample) { + throw new Error('raw sampling not supported by this backend'); + } + return harden(await sample(params)); + }, + }); + return harden({ name: LANGUAGE_MODEL_SERVICE_NAME, service }); +}; diff --git a/packages/kernel-language-model-service/src/ollama/base.ts b/packages/kernel-language-model-service/src/ollama/base.ts index 37b923422..064c63107 100644 --- a/packages/kernel-language-model-service/src/ollama/base.ts +++ b/packages/kernel-language-model-service/src/ollama/base.ts @@ -1,6 +1,13 @@ import type { GenerateResponse, ListResponse } from 'ollama'; -import type { LanguageModelService } from '../types.ts'; +import type { + ChatParams, + ChatResult, + ChatRole, + LanguageModelService, + SampleParams, + SampleResult, +} from '../types.ts'; import { parseModelConfig } from './parse.ts'; import type { OllamaInstanceConfig, @@ -46,6 +53,88 @@ export class OllamaBaseService return await client.list(); } + /** + * Performs a chat completion request via the Ollama chat API. + * + * @param params - The chat parameters. + * @returns A hardened chat result. + */ + async chat(params: ChatParams): Promise { + const { model, messages, temperature, seed, stop } = params; + const ollama = await this.#makeClient(); + let stopArr: string[] | undefined; + if (stop !== undefined) { + stopArr = Array.isArray(stop) ? stop : [stop]; + } + const response = await ollama.chat({ + model, + messages, + stream: false, + options: { + ...(temperature !== undefined && { temperature }), + ...(params.top_p !== undefined && { top_p: params.top_p }), + ...(seed !== undefined && { seed }), + ...(params.max_tokens !== undefined && { + num_predict: params.max_tokens, + }), + ...(stopArr !== undefined && { stop: stopArr }), + }, + }); + const promptTokens = response.prompt_eval_count ?? 0; + const completionTokens = response.eval_count ?? 0; + return harden({ + id: 'ollama-chat', + model: response.model, + choices: [ + { + message: { + role: response.message.role as ChatRole, + content: response.message.content, + }, + index: 0, + finish_reason: response.done_reason ?? 'stop', + }, + ], + usage: { + prompt_tokens: promptTokens, + completion_tokens: completionTokens, + total_tokens: promptTokens + completionTokens, + }, + }); + } + + /** + * Performs a raw token-prediction request via Ollama's generate API with raw=true, + * bypassing the model's chat template. + * + * @param params - The raw sample parameters. + * @returns A hardened raw sample result. + */ + async sample(params: SampleParams): Promise { + const { model, prompt, temperature, seed, stop } = params; + const ollama = await this.#makeClient(); + let stopArr: string[] | undefined; + if (stop !== undefined) { + stopArr = Array.isArray(stop) ? stop : [stop]; + } + const response = await ollama.generate({ + model, + prompt, + raw: true, + stream: false, + options: { + ...(temperature !== undefined && { temperature }), + ...(params.top_p !== undefined && { top_p: params.top_p }), + ...(seed !== undefined && { seed }), + ...(params.max_tokens !== undefined && { + num_predict: params.max_tokens, + }), + ...(stopArr !== undefined && { stop: stopArr }), + }, + }); + return harden({ text: response.response }); + } + /** * Creates a new language model instance with the specified configuration. * The returned instance is hardened for object capability security. @@ -62,7 +151,7 @@ export class OllamaBaseService }; const mandatoryOptions = { model, - stream: true, + stream: true as const, raw: true, }; diff --git a/packages/kernel-language-model-service/src/ollama/nodejs.ts b/packages/kernel-language-model-service/src/ollama/nodejs.ts index f0bdc2e14..ed7d3c37a 100644 --- a/packages/kernel-language-model-service/src/ollama/nodejs.ts +++ b/packages/kernel-language-model-service/src/ollama/nodejs.ts @@ -1,5 +1,11 @@ import { Ollama } from 'ollama'; +import type { + ChatParams, + ChatResult, + SampleParams, + SampleResult, +} from '../types.ts'; import { OllamaBaseService } from './base.ts'; import { defaultClientConfig } from './constants.ts'; import type { OllamaClient, OllamaNodejsConfig } from './types.ts'; @@ -34,3 +40,23 @@ export class OllamaNodejsService extends OllamaBaseService { ); } } + +/** + * Creates a hardened kernel service backend backed by a local Ollama instance. + * + * @param config - Configuration for the Ollama Node.js service. + * @returns An object with `chat` and `sample` methods for use with + * `makeKernelLanguageModelService`. + */ +export const makeOllamaNodejsKernelService = ( + config: OllamaNodejsConfig, +): { + chat: (params: ChatParams) => Promise; + sample: (params: SampleParams) => Promise; +} => { + const service = new OllamaNodejsService(config); + return harden({ + chat: async (params: ChatParams) => service.chat(params), + sample: async (params: SampleParams) => service.sample(params), + }); +}; diff --git a/packages/kernel-language-model-service/src/ollama/types.ts b/packages/kernel-language-model-service/src/ollama/types.ts index 65f7e8289..3630932d2 100644 --- a/packages/kernel-language-model-service/src/ollama/types.ts +++ b/packages/kernel-language-model-service/src/ollama/types.ts @@ -6,21 +6,33 @@ import type { ListResponse, AbortableAsyncIterator, Config, + ChatRequest, + ChatResponse, } from 'ollama'; import type { LanguageModel } from '../types.ts'; /** - * Interface for an Ollama client that can list models and generate responses. + * Interface for an Ollama client that can list models, generate responses, and chat. * Provides the minimal interface required for Ollama operations. */ type OllamaClient = { list: () => Promise; - generate: ( - request: GenerateRequest, - ) => Promise>; + generate( + request: GenerateRequest & { stream: true }, + ): Promise>; + generate( + request: GenerateRequest & { stream?: false }, + ): Promise; + chat(request: ChatRequest & { stream?: false }): Promise; +}; +export type { + GenerateRequest, + GenerateResponse, + OllamaClient, + ChatRequest, + ChatResponse, }; -export type { GenerateRequest, GenerateResponse, OllamaClient }; /** * Configuration for creating an Ollama service in a Node.js environment. diff --git a/packages/kernel-language-model-service/src/open-v1/base.test.ts b/packages/kernel-language-model-service/src/open-v1/base.test.ts new file mode 100644 index 000000000..6ecb36c8b --- /dev/null +++ b/packages/kernel-language-model-service/src/open-v1/base.test.ts @@ -0,0 +1,246 @@ +import { describe, it, expect, vi, beforeEach } from 'vitest'; + +import type { ChatResult, ChatStreamChunk } from '../types.ts'; +import { OpenV1BaseService } from './base.ts'; + +const MODEL = 'glm-4.7-flash'; + +const makeChatResult = (): ChatResult => ({ + id: 'chat-1', + model: MODEL, + choices: [ + { + message: { role: 'assistant', content: 'hi there' }, + index: 0, + finish_reason: 'stop', + }, + ], + usage: { prompt_tokens: 5, completion_tokens: 3, total_tokens: 8 }, +}); + +const makeMockFetch = (json: unknown): typeof globalThis.fetch => + vi.fn().mockResolvedValue({ json: vi.fn().mockResolvedValue(json) }); + +const makeSSEStream = ( + chunks: ChatStreamChunk[], + // eslint-disable-next-line n/no-unsupported-features/node-builtins +): ReadableStream => { + const encoder = new TextEncoder(); + const lines = chunks + .map((chunk) => `data: ${JSON.stringify(chunk)}\n\n`) + .join(''); + const body = `${lines}data: [DONE]\n\n`; + // eslint-disable-next-line n/no-unsupported-features/node-builtins + return new ReadableStream({ + start(controller) { + controller.enqueue(encoder.encode(body)); + controller.close(); + }, + }); +}; + +const makeStreamChunk = (content: string): ChatStreamChunk => ({ + id: 'chat-1', + model: MODEL, + choices: [{ delta: { content }, index: 0, finish_reason: null }], +}); + +const makeMockStreamFetch = ( + chunks: ChatStreamChunk[], +): typeof globalThis.fetch => + vi.fn().mockResolvedValue({ body: makeSSEStream(chunks) }); + +describe('OpenV1BaseService', () => { + let service: OpenV1BaseService; + let mockFetch: ReturnType; + + beforeEach(() => { + mockFetch = makeMockFetch(makeChatResult()); + service = new OpenV1BaseService( + mockFetch, + 'http://localhost:11434', + 'sk-test', + ); + }); + + describe('chat', () => { + it('pOSTs to /v1/chat/completions with serialized params', async () => { + const params = { + model: MODEL, + messages: [{ role: 'user' as const, content: 'hello' }], + }; + await service.chat(params); + + expect(mockFetch).toHaveBeenCalledWith( + 'http://localhost:11434/v1/chat/completions', + expect.objectContaining({ + method: 'POST', + body: JSON.stringify({ ...params, stream: false }), + }), + ); + }); + + it('sends stream: false when stream is not set', async () => { + await service.chat({ + model: MODEL, + messages: [{ role: 'user', content: 'hi' }], + }); + + const [, init] = (mockFetch as ReturnType).mock + .calls[0] as [string, RequestInit]; + expect(JSON.parse(init.body as string)).toMatchObject({ stream: false }); + }); + + it('includes Authorization header when apiKey is set', async () => { + await service.chat({ + model: MODEL, + messages: [{ role: 'user', content: 'hi' }], + }); + + const [, init] = (mockFetch as ReturnType).mock + .calls[0] as [string, RequestInit]; + expect((init.headers as Record).Authorization).toBe( + 'Bearer sk-test', + ); + }); + + it('omits Authorization header when no apiKey', async () => { + const noKeyFetch = makeMockFetch(makeChatResult()); + const noKeyService = new OpenV1BaseService( + noKeyFetch, + 'http://localhost:11434', + ); + await noKeyService.chat({ + model: MODEL, + messages: [{ role: 'user', content: 'hi' }], + }); + + const [, init] = (noKeyFetch as ReturnType).mock + .calls[0] as [string, RequestInit]; + expect( + (init.headers as Record).Authorization, + ).toBeUndefined(); + }); + + it('returns the parsed JSON response', async () => { + const expected = makeChatResult(); + mockFetch = makeMockFetch(expected); + service = new OpenV1BaseService(mockFetch, 'http://localhost:11434'); + + const result = await service.chat({ + model: MODEL, + messages: [{ role: 'user', content: 'hi' }], + }); + + expect(result).toStrictEqual(expected); + }); + + it('throws on invalid params (empty model)', () => { + expect(() => { + // eslint-disable-next-line no-void + void service.chat({ + model: '', + messages: [{ role: 'user', content: 'hi' }], + }); + }).toThrow('Expected a string with a length between'); + }); + + it('throws on invalid params (invalid role)', () => { + expect(() => { + // eslint-disable-next-line no-void + void service.chat({ + model: MODEL, + messages: [{ role: 'unknown' as never, content: 'hi' }], + }); + }).toThrow('Expected the value to satisfy a union'); + }); + + it('uses custom baseUrl', async () => { + const customFetch = makeMockFetch(makeChatResult()); + const customService = new OpenV1BaseService( + customFetch, + 'https://my-llm.internal', + ); + await customService.chat({ + model: 'my-model', + messages: [{ role: 'user', content: 'hi' }], + }); + + expect(customFetch).toHaveBeenCalledWith( + 'https://my-llm.internal/v1/chat/completions', + expect.any(Object), + ); + }); + }); + + describe('chat with stream: true', () => { + it('pOSTs to /v1/chat/completions with stream: true in body', async () => { + const streamFetch = makeMockStreamFetch([makeStreamChunk('hi')]); + const streamService = new OpenV1BaseService( + streamFetch, + 'http://localhost:11434', + ); + // eslint-disable-next-line @typescript-eslint/no-unused-vars + for await (const _ of streamService.chat({ + model: MODEL, + messages: [{ role: 'user', content: 'hi' }], + stream: true, + })) { + // drain + } + + expect(streamFetch).toHaveBeenCalledWith( + 'http://localhost:11434/v1/chat/completions', + expect.objectContaining({ + method: 'POST', + body: JSON.stringify({ + model: MODEL, + messages: [{ role: 'user', content: 'hi' }], + stream: true, + }), + }), + ); + }); + + it('yields parsed chunks and stops at [DONE]', async () => { + const expected = [makeStreamChunk('Hello'), makeStreamChunk(', world!')]; + const streamFetch = makeMockStreamFetch(expected); + const streamService = new OpenV1BaseService( + streamFetch, + 'http://localhost:11434', + ); + + const received: ChatStreamChunk[] = []; + for await (const chunk of streamService.chat({ + model: MODEL, + messages: [{ role: 'user', content: 'hi' }], + stream: true, + })) { + received.push(chunk); + } + + expect(received).toStrictEqual(expected); + }); + + it('throws when response body is null', async () => { + const nullBodyFetch: typeof globalThis.fetch = vi + .fn() + .mockResolvedValue({ body: null }); + const streamService = new OpenV1BaseService( + nullBodyFetch, + 'http://localhost:11434', + ); + + await expect(async () => { + // eslint-disable-next-line @typescript-eslint/no-unused-vars + for await (const _ of streamService.chat({ + model: MODEL, + messages: [{ role: 'user', content: 'hi' }], + stream: true, + })) { + // drain + } + }).rejects.toThrow('No response body for streaming'); + }); + }); +}); diff --git a/packages/kernel-language-model-service/src/open-v1/base.ts b/packages/kernel-language-model-service/src/open-v1/base.ts new file mode 100644 index 000000000..54a396843 --- /dev/null +++ b/packages/kernel-language-model-service/src/open-v1/base.ts @@ -0,0 +1,136 @@ +import { assert } from '@metamask/superstruct'; + +import type { ChatParams, ChatResult, ChatStreamChunk } from '../types.ts'; +import { ChatParamsStruct } from './types.ts'; + +/** + * Base service for any Open /v1-compatible HTTP endpoint. + * + * Accepts an injected `fetch` endowment so it runs safely under lockdown. + * Pass `stream: true` in params for SSE streaming; omit for a single JSON response. + */ +export class OpenV1BaseService { + readonly #fetch: typeof globalThis.fetch; + + readonly #baseUrl: string; + + readonly #apiKey: string | undefined; + + /** + * @param fetchFn - The fetch implementation to use for HTTP requests. + * @param baseUrl - Base URL of the API (e.g. `'https://api.openai.com'`). + * @param apiKey - Optional API key sent as a Bearer token. + */ + constructor( + fetchFn: typeof globalThis.fetch, + baseUrl: string, + apiKey?: string, + ) { + this.#fetch = fetchFn; + this.#baseUrl = baseUrl; + this.#apiKey = apiKey; + harden(this); + } + + /** + * Performs a chat completion request against `/v1/chat/completions`. + * + * When `params.stream` is `true`, returns an async iterable of + * {@link ChatStreamChunk}s, one per SSE event. + * When `params.stream` is `false` or omitted, awaits and returns the full + * {@link ChatResult}. + * + * @param params - The chat parameters. + * @returns An async iterable of stream chunks when `stream: true`. + */ + chat(params: ChatParams & { stream: true }): AsyncIterable; + + /** + * @param params - The chat parameters. + * @returns A promise resolving to the full chat result. + */ + chat(params: ChatParams & { stream?: false }): Promise; + + /** + * @param params - The chat parameters. + * @returns An async iterable or promise depending on `params.stream`. + */ + chat( + params: ChatParams, + ): AsyncIterable | Promise { + assert(params, ChatParamsStruct); + if (params.stream === true) { + return this.#streamingChat(params); + } + return this.#nonStreamingChat(params); + } + + /** + * @param params - The chat parameters. + * @returns A promise resolving to the full chat result. + */ + async #nonStreamingChat(params: ChatParams): Promise { + const response = await this.#fetch(`${this.#baseUrl}/v1/chat/completions`, { + method: 'POST', + headers: this.#makeHeaders(), + body: JSON.stringify({ ...params, stream: false }), + }); + const result = (await response.json()) as ChatResult; + return harden(result); + } + + /** + * @param params - The chat parameters. + * @yields One {@link ChatStreamChunk} per SSE event until `[DONE]`. + */ + async *#streamingChat(params: ChatParams): AsyncGenerator { + const response = await this.#fetch(`${this.#baseUrl}/v1/chat/completions`, { + method: 'POST', + headers: this.#makeHeaders(), + body: JSON.stringify({ ...params, stream: true }), + }); + if (!response.body) { + throw new Error('No response body for streaming'); + } + const reader = response.body.getReader(); + const decoder = new TextDecoder(); + let buffer = ''; + try { + while (true) { + const { done, value } = await reader.read(); + if (value) { + buffer += decoder.decode(value, { stream: !done }); + } + let newlineIdx: number; + while ((newlineIdx = buffer.indexOf('\n')) !== -1) { + const line = buffer.slice(0, newlineIdx).trimEnd(); + buffer = buffer.slice(newlineIdx + 1); + if (line.startsWith('data: ')) { + const data = line.slice(6); + if (data === '[DONE]') { + return; + } + if (data) { + yield harden(JSON.parse(data) as ChatStreamChunk); + } + } + } + if (done) { + break; + } + } + } finally { + reader.releaseLock(); + } + } + + /** + * @returns Headers for the request, including Authorization if an API key is set. + */ + #makeHeaders(): Record { + return harden({ + 'Content-Type': 'application/json', + ...(this.#apiKey ? { Authorization: `Bearer ${this.#apiKey}` } : {}), + }); + } +} diff --git a/packages/kernel-language-model-service/src/open-v1/nodejs.ts b/packages/kernel-language-model-service/src/open-v1/nodejs.ts new file mode 100644 index 000000000..a1eb20ac0 --- /dev/null +++ b/packages/kernel-language-model-service/src/open-v1/nodejs.ts @@ -0,0 +1,39 @@ +import type { ChatParams, ChatResult, ChatStreamChunk } from '../types.ts'; +import { OpenV1BaseService } from './base.ts'; + +/** + * Creates an Open /v1-compatible service for Node.js environments. + * + * Requires `fetch` to be explicitly endowed for object-capability security. + * + * Pass `stream: true` in params for SSE streaming; omit for a single JSON response. + * + * @param config - Configuration for the service. + * @param config.endowments - Required endowments. + * @param config.endowments.fetch - The fetch implementation to use for HTTP requests. + * @param config.baseUrl - Base URL of the API (e.g. `'https://api.openai.com'`). + * @param config.apiKey - Optional API key sent as a Bearer token. + * @returns An object with a `chat` method. Raw sampling is not supported by this backend. + */ +export const makeOpenV1NodejsService = (config: { + endowments: { fetch: typeof globalThis.fetch }; + baseUrl: string; + apiKey?: string; +}): { + chat: { + (params: ChatParams & { stream: true }): AsyncIterable; + (params: ChatParams & { stream?: false }): Promise; + }; +} => { + const { endowments, baseUrl, apiKey } = config; + if (!endowments?.fetch) { + throw new Error('Must endow a fetch implementation.'); + } + const service = new OpenV1BaseService(endowments.fetch, baseUrl, apiKey); + return harden({ + chat: service.chat.bind(service) as { + (params: ChatParams & { stream: true }): AsyncIterable; + (params: ChatParams & { stream?: false }): Promise; + }, + }); +}; diff --git a/packages/kernel-language-model-service/src/open-v1/types.ts b/packages/kernel-language-model-service/src/open-v1/types.ts new file mode 100644 index 000000000..997dd8c99 --- /dev/null +++ b/packages/kernel-language-model-service/src/open-v1/types.ts @@ -0,0 +1,50 @@ +import { + array, + boolean, + literal, + number, + object, + optional, + size, + string, + union, +} from '@metamask/superstruct'; + +export type { + ChatChoice, + ChatMessage, + ChatParams, + ChatResult, + ChatRole, + ChatStreamChunk, + ChatStreamDelta, + Usage, +} from '../types.ts'; + +const ChatRoleStruct = union([ + literal('system'), + literal('user'), + literal('assistant'), +]); + +const ChatMessageStruct = object({ + role: ChatRoleStruct, + content: string(), +}); + +const StopStruct = optional(union([string(), array(string())])); + +/** + * Superstruct schema for chat completion request parameters. + */ +export const ChatParamsStruct = object({ + model: size(string(), 1, Infinity), + messages: array(ChatMessageStruct), + max_tokens: optional(number()), + temperature: optional(number()), + top_p: optional(number()), + stop: StopStruct, + seed: optional(number()), + n: optional(number()), + stream: optional(boolean()), +}); diff --git a/packages/kernel-language-model-service/src/test-utils/index.ts b/packages/kernel-language-model-service/src/test-utils/index.ts index 9fc946e73..8fcf6100e 100644 --- a/packages/kernel-language-model-service/src/test-utils/index.ts +++ b/packages/kernel-language-model-service/src/test-utils/index.ts @@ -1,4 +1,2 @@ -export { makeQueueService } from './queue/service.ts'; -export { makeQueueModel } from './queue/model.ts'; -export type { QueueLanguageModel } from './queue/model.ts'; -export type { QueueLanguageModelService } from './queue/service.ts'; +export { makeMockOpenV1Fetch } from './mock-fetch.ts'; +export { makeMockSample } from './mock-sample.ts'; diff --git a/packages/kernel-language-model-service/src/test-utils/mock-fetch.ts b/packages/kernel-language-model-service/src/test-utils/mock-fetch.ts new file mode 100644 index 000000000..a122eb8a4 --- /dev/null +++ b/packages/kernel-language-model-service/src/test-utils/mock-fetch.ts @@ -0,0 +1,31 @@ +/** + * Returns a fetch implementation that responds to Open /v1 chat completion requests + * with a sequence of non-streaming JSON responses (one content string per request). + * + * @param responses - Content strings to return, in order, for each request. + * @param model - Model name to include in the response (default `'test-model'`). + * @returns A fetch function suitable for use as an endowment. + */ +export const makeMockOpenV1Fetch = ( + responses: string[], + model = 'test-model', +): typeof globalThis.fetch => { + let idx = 0; + return async (_url, _init) => { + const content = responses[idx] ?? ''; + idx += 1; + const result = harden({ + id: `chat-${idx}`, + model, + choices: [ + { + message: { role: 'assistant', content }, + index: 0, + finish_reason: 'stop', + }, + ], + usage: { prompt_tokens: 0, completion_tokens: 0, total_tokens: 0 }, + }); + return { json: async () => result } as unknown as globalThis.Response; + }; +}; diff --git a/packages/kernel-language-model-service/src/test-utils/mock-sample.ts b/packages/kernel-language-model-service/src/test-utils/mock-sample.ts new file mode 100644 index 000000000..454b33a0c --- /dev/null +++ b/packages/kernel-language-model-service/src/test-utils/mock-sample.ts @@ -0,0 +1,18 @@ +import type { SampleParams, SampleResult } from '../types.ts'; + +/** + * Returns a sample function that returns a sequence of result texts (one per call). + * + * @param responses - Text strings to return, in order, for each call. + * @returns A function matching the sample service signature. + */ +export const makeMockSample = ( + responses: string[], +): ((params: SampleParams) => Promise) => { + let idx = 0; + return async (_params) => { + const text = responses[idx] ?? ''; + idx += 1; + return harden({ text }); + }; +}; diff --git a/packages/kernel-language-model-service/src/test-utils/queue/README.md b/packages/kernel-language-model-service/src/test-utils/queue/README.md deleted file mode 100644 index 9b16967a8..000000000 --- a/packages/kernel-language-model-service/src/test-utils/queue/README.md +++ /dev/null @@ -1,49 +0,0 @@ -# Queue-based Language Model Service (Testing Utility) - -[`makeQueueService`](./service.ts) is a testing utility that creates a `LanguageModelService` implementation for use in tests. It provides a queue-based language model where responses are manually queued using the `push()` method and consumed by `sample()` calls. - -## Usage - -1. Create a service using `makeQueueService()` -2. Create a model instance using `makeInstance()` -3. Queue responses using `push()` on the model instance -4. Consume responses by calling `sample()` - -Note that `makeInstance` and `sample` ignore their arguments, but expect them nonetheless. - -## Examples - -### Basic Example - -```typescript -import { makeQueueService } from '@ocap/kernel-language-model-service/test-utils'; - -const service = makeQueueService(); -const model = await service.makeInstance({ model: 'test' }); - -// Queue a response -model.push('Hello, world!'); - -// Consume the response -const result = await model.sample({ prompt: 'Say hello' }); -for await (const chunk of result.stream) { - console.log(chunk.response); // 'Hello, world!' -} -``` - -### Multiple Queued Responses - -```typescript -const service = makeQueueService(); -const model = await service.makeInstance({ model: 'test' }); - -// Queue multiple responses -model.push('First response'); -model.push('Second response'); - -// Each sample() call consumes the next queued response -const first = await model.sample({ prompt: 'test' }); -const second = await model.sample({ prompt: 'test' }); - -// Process streams... -``` diff --git a/packages/kernel-language-model-service/src/test-utils/queue/model.test.ts b/packages/kernel-language-model-service/src/test-utils/queue/model.test.ts deleted file mode 100644 index 479fd54ba..000000000 --- a/packages/kernel-language-model-service/src/test-utils/queue/model.test.ts +++ /dev/null @@ -1,371 +0,0 @@ -import '@ocap/repo-tools/test-utils/mock-endoify'; -import { describe, it, expect, vi, beforeEach } from 'vitest'; - -import { makeQueueModel } from './model.ts'; -import type { ResponseFormatter } from './response.ts'; -import type { Tokenizer } from './tokenizer.ts'; -import type { StreamWithAbort } from './utils.ts'; -import * as utils from './utils.ts'; - -vi.mock('./utils.ts', () => ({ - makeAbortableAsyncIterable: vi.fn(), - makeEmptyStreamWithAbort: vi.fn(), - mapAsyncIterable: vi.fn(), - normalizeToAsyncIterable: vi.fn(), -})); - -describe('makeQueueModel', () => { - let mockTokenizer: ReturnType>; - let mockResponseFormatter: ReturnType< - typeof vi.fn> - >; - let mockMakeAbortableAsyncIterable: ReturnType; - let mockMakeEmptyStreamWithAbort: ReturnType; - let mockMapAsyncIterable: ReturnType; - let mockNormalizeToAsyncIterable: ReturnType; - - beforeEach(() => { - vi.clearAllMocks(); - - mockTokenizer = vi.fn(); - mockResponseFormatter = - vi.fn>(); - mockMakeAbortableAsyncIterable = vi.mocked( - utils.makeAbortableAsyncIterable, - ); - mockMakeEmptyStreamWithAbort = vi.mocked(utils.makeEmptyStreamWithAbort); - mockMapAsyncIterable = vi.mocked(utils.mapAsyncIterable); - mockNormalizeToAsyncIterable = vi.mocked(utils.normalizeToAsyncIterable); - }); - - it('creates model with default parameters', () => { - const model = makeQueueModel(); - expect(model).toMatchObject({ - getInfo: expect.any(Function), - load: expect.any(Function), - unload: expect.any(Function), - sample: expect.any(Function), - push: expect.any(Function), - }); - }); - - it('creates model with custom tokenizer', () => { - const model = makeQueueModel({ tokenizer: mockTokenizer }); - expect(model).toBeDefined(); - }); - - it('creates model with custom responseFormatter', () => { - const model = makeQueueModel({ responseFormatter: mockResponseFormatter }); - expect(model).toBeDefined(); - }); - - it('creates model with custom responseQueue', () => { - const mockStream: StreamWithAbort<{ response: string; done: boolean }> = { - stream: (async function* () { - // Empty stream for testing - })() as AsyncIterable<{ - response: string; - done: boolean; - }>, - abort: vi.fn<() => Promise>(), - }; - const responseQueue = [mockStream]; - const model = makeQueueModel({ responseQueue }); - expect(model).toBeDefined(); - }); - - describe('getInfo', () => { - it('returns model info', async () => { - const model = makeQueueModel(); - const info = await model.getInfo(); - expect(info).toStrictEqual({ model: 'test' }); - }); - }); - - describe('load', () => { - it('resolves without error', async () => { - const model = makeQueueModel(); - expect(await model.load()).toBeUndefined(); - }); - }); - - describe('unload', () => { - it('resolves without error', async () => { - const model = makeQueueModel(); - expect(await model.unload()).toBeUndefined(); - }); - }); - - describe('sample', () => { - it('returns stream from queue when available', async () => { - const mockStream: StreamWithAbort<{ response: string; done: boolean }> = { - stream: (async function* () { - yield { response: 'test', done: false }; - })(), - abort: vi.fn<() => Promise>(), - }; - const responseQueue = [mockStream]; - const model = makeQueueModel({ responseQueue }); - - const result = await model.sample(''); - const values: { response: string; done: boolean }[] = []; - for await (const value of result.stream) { - values.push(value); - } - - expect(values).toStrictEqual([{ response: 'test', done: false }]); - expect(responseQueue).toHaveLength(0); - }); - - it('returns empty stream when queue is empty', async () => { - const emptyStream: StreamWithAbort<{ response: string; done: boolean }> = - { - stream: (async function* () { - // Empty stream for testing - })() as AsyncIterable<{ - response: string; - done: boolean; - }>, - abort: vi.fn<() => Promise>(), - }; - mockMakeEmptyStreamWithAbort.mockReturnValue(emptyStream); - - const model = makeQueueModel(); - const result = await model.sample(''); - - expect(mockMakeEmptyStreamWithAbort).toHaveBeenCalledTimes(1); - expect(result).toBe(emptyStream); - }); - }); - - describe('push', () => { - it('pushes stream to queue', () => { - const responseQueue: StreamWithAbort<{ - response: string; - done: boolean; - }>[] = []; - const mockStream: StreamWithAbort<{ response: string; done: boolean }> = { - stream: (async function* () { - // Empty stream for testing - })() as AsyncIterable<{ - response: string; - done: boolean; - }>, - abort: vi.fn<() => Promise>(), - }; - - mockTokenizer.mockReturnValue(['token1', 'token2']); - mockNormalizeToAsyncIterable.mockReturnValue( - (async function* () { - yield 'token1'; - yield 'token2'; - })(), - ); - mockMapAsyncIterable.mockReturnValue( - (async function* () { - yield { response: 'token1', done: false }; - yield { response: 'token2', done: true }; - })(), - ); - mockMakeAbortableAsyncIterable.mockReturnValue(mockStream); - - const model = makeQueueModel({ - tokenizer: mockTokenizer, - responseFormatter: mockResponseFormatter, - responseQueue, - }); - - model.push('test text'); - - expect(mockTokenizer).toHaveBeenCalledWith('test text'); - expect(mockNormalizeToAsyncIterable).toHaveBeenCalledWith([ - 'token1', - 'token2', - ]); - expect(mockMapAsyncIterable).toHaveBeenCalledWith( - expect.anything(), - mockResponseFormatter, - ); - expect(mockMakeAbortableAsyncIterable).toHaveBeenCalledTimes(1); - expect(responseQueue).toHaveLength(1); - expect(responseQueue[0]).toBe(mockStream); - }); - - it('pushes multiple streams to queue', () => { - const responseQueue: StreamWithAbort<{ - response: string; - done: boolean; - }>[] = []; - const mockStream1: StreamWithAbort<{ response: string; done: boolean }> = - { - stream: (async function* () { - // Empty stream for testing - })() as AsyncIterable<{ - response: string; - done: boolean; - }>, - abort: vi.fn<() => Promise>(), - }; - const mockStream2: StreamWithAbort<{ response: string; done: boolean }> = - { - stream: (async function* () { - // Empty stream for testing - })() as AsyncIterable<{ - response: string; - done: boolean; - }>, - abort: vi.fn<() => Promise>(), - }; - - mockTokenizer.mockReturnValue(['token']); - mockNormalizeToAsyncIterable.mockReturnValue( - (async function* () { - yield 'token'; - })(), - ); - mockMapAsyncIterable.mockReturnValue( - (async function* () { - yield { response: 'token', done: true }; - })(), - ); - mockMakeAbortableAsyncIterable - .mockReturnValueOnce(mockStream1) - .mockReturnValueOnce(mockStream2); - - const model = makeQueueModel({ - tokenizer: mockTokenizer, - responseFormatter: mockResponseFormatter, - responseQueue, - }); - - model.push('text1'); - model.push('text2'); - - expect(responseQueue).toHaveLength(2); - expect(responseQueue[0]).toBe(mockStream1); - expect(responseQueue[1]).toBe(mockStream2); - }); - - it('handles async iterable tokenizer', () => { - const responseQueue: StreamWithAbort<{ - response: string; - done: boolean; - }>[] = []; - const mockStream: StreamWithAbort<{ response: string; done: boolean }> = { - stream: (async function* () { - // Empty stream for testing - })() as AsyncIterable<{ - response: string; - done: boolean; - }>, - abort: vi.fn<() => Promise>(), - }; - - const asyncIterable = (async function* () { - yield 'async'; - yield 'token'; - })(); - mockTokenizer.mockReturnValue(asyncIterable); - mockNormalizeToAsyncIterable.mockReturnValue(asyncIterable); - mockMapAsyncIterable.mockReturnValue( - (async function* () { - yield { response: 'async', done: false }; - yield { response: 'token', done: true }; - })(), - ); - mockMakeAbortableAsyncIterable.mockReturnValue(mockStream); - - const model = makeQueueModel({ - tokenizer: mockTokenizer, - responseFormatter: mockResponseFormatter, - responseQueue, - }); - - model.push('test'); - - expect(mockTokenizer).toHaveBeenCalledWith('test'); - expect(mockNormalizeToAsyncIterable).toHaveBeenCalledWith(asyncIterable); - }); - }); - - describe('integration', () => { - it('pushes and samples from queue in order', async () => { - const responseQueue: StreamWithAbort<{ - response: string; - done: boolean; - }>[] = []; - const mockStream1: StreamWithAbort<{ response: string; done: boolean }> = - { - stream: (async function* () { - yield { response: 'first', done: false }; - yield { response: ' stream', done: true }; - })(), - abort: vi.fn<() => Promise>(), - }; - const mockStream2: StreamWithAbort<{ response: string; done: boolean }> = - { - stream: (async function* () { - yield { response: 'second', done: false }; - yield { response: ' stream', done: true }; - })(), - abort: vi.fn<() => Promise>(), - }; - - mockTokenizer.mockReturnValue(['token']); - mockNormalizeToAsyncIterable.mockReturnValue( - (async function* () { - yield 'token'; - })(), - ); - mockMapAsyncIterable.mockReturnValue( - (async function* () { - yield { response: 'token', done: true }; - })(), - ); - mockMakeAbortableAsyncIterable - .mockReturnValueOnce(mockStream1) - .mockReturnValueOnce(mockStream2); - - const model = makeQueueModel({ - tokenizer: mockTokenizer, - responseFormatter: mockResponseFormatter, - responseQueue, - }); - - model.push('first'); - model.push('second'); - - const [result1, result2] = await Promise.all([ - model.sample(''), - model.sample(''), - ]); - - const [values1, values2] = await Promise.all([ - (async () => { - const values: { response: string; done: boolean }[] = []; - for await (const value of result1.stream) { - values.push(value); - } - return values; - })(), - (async () => { - const values: { response: string; done: boolean }[] = []; - for await (const value of result2.stream) { - values.push(value); - } - return values; - })(), - ]); - - expect(values1).toStrictEqual([ - { response: 'first', done: false }, - { response: ' stream', done: true }, - ]); - expect(values2).toStrictEqual([ - { response: 'second', done: false }, - { response: ' stream', done: true }, - ]); - expect(responseQueue).toHaveLength(0); - }); - }); -}); diff --git a/packages/kernel-language-model-service/src/test-utils/queue/model.ts b/packages/kernel-language-model-service/src/test-utils/queue/model.ts deleted file mode 100644 index c30ebe017..000000000 --- a/packages/kernel-language-model-service/src/test-utils/queue/model.ts +++ /dev/null @@ -1,80 +0,0 @@ -import { objectResponseFormatter } from './response.ts'; -import type { ResponseFormatter } from './response.ts'; -import type { Tokenizer } from './tokenizer.ts'; -import { whitespaceTokenizer } from './tokenizer.ts'; -import { - makeAbortableAsyncIterable, - makeEmptyStreamWithAbort, - mapAsyncIterable, - normalizeToAsyncIterable, -} from './utils.ts'; -import type { StreamWithAbort } from './utils.ts'; -import type { LanguageModel, ModelInfo } from '../../types.ts'; - -/** - * Queue-based language model with helper methods for configuring responses. - * Responses are queued and consumed by sample() calls. - * - * @template Response - The type of response generated by the model - */ -export type QueueLanguageModel = - // QueueLanguageModel does not support any sample options - LanguageModel & { - /** - * Pushes a streaming response to the queue for the next sample() call. - * The text will be tokenized and streamed token by token. - * - * @param text - The complete text to stream - */ - push: (text: string) => void; - }; - -/** - * Make a queue-based language model instance. - * - * @template Response - The type of response generated by the model - * @param options - Configuration options for the model - * @param options.tokenizer - The tokenizer function to use. Defaults to whitespace splitting. - * @param options.responseFormatter - The function to use to format each yielded token into a response. Defaults to an object with a response and done property. - * @param options.responseQueue - For testing only. The queue to use for responses. Defaults to an empty array. - * @returns A queue-based language model instance. - */ -export const makeQueueModel = < - Response extends object = { response: string; done: boolean }, ->({ - tokenizer = whitespaceTokenizer, - responseFormatter = objectResponseFormatter as ResponseFormatter, - // Available for testing - responseQueue = [], -}: { - tokenizer?: Tokenizer; - responseFormatter?: ResponseFormatter; - responseQueue?: StreamWithAbort[]; -} = {}): QueueLanguageModel => { - const makeStreamWithAbort = (text: string): StreamWithAbort => - makeAbortableAsyncIterable( - mapAsyncIterable( - normalizeToAsyncIterable(tokenizer(text)), - responseFormatter, - ), - ); - - return harden({ - getInfo: async (): Promise>> => ({ - model: 'test', - }), - load: async (): Promise => { - // No-op: queue model doesn't require loading - }, - unload: async (): Promise => { - // No-op: queue model doesn't require unloading - }, - sample: async (): Promise> => { - return responseQueue.shift() ?? makeEmptyStreamWithAbort(); - }, - push: (text: string): void => { - const streamWithAbort = makeStreamWithAbort(text); - responseQueue.push(streamWithAbort); - }, - }); -}; diff --git a/packages/kernel-language-model-service/src/test-utils/queue/response.test.ts b/packages/kernel-language-model-service/src/test-utils/queue/response.test.ts deleted file mode 100644 index 3d4cb7ac2..000000000 --- a/packages/kernel-language-model-service/src/test-utils/queue/response.test.ts +++ /dev/null @@ -1,24 +0,0 @@ -import { describe, it, expect } from 'vitest'; - -import { objectResponseFormatter } from './response.ts'; - -describe('objectResponseFormatter', () => { - it.each([ - { - response: 'hello', - done: false, - expected: { response: 'hello', done: false }, - }, - { - response: 'world', - done: true, - expected: { response: 'world', done: true }, - }, - { response: '', done: false, expected: { response: '', done: false } }, - ])( - 'formats response "$response" with done=$done', - ({ response, done, expected }) => { - expect(objectResponseFormatter(response, done)).toStrictEqual(expected); - }, - ); -}); diff --git a/packages/kernel-language-model-service/src/test-utils/queue/response.ts b/packages/kernel-language-model-service/src/test-utils/queue/response.ts deleted file mode 100644 index d7fee9682..000000000 --- a/packages/kernel-language-model-service/src/test-utils/queue/response.ts +++ /dev/null @@ -1,15 +0,0 @@ -export type ResponseFormatter = ( - response: string, - done: boolean, -) => FormattedResponse; - -// Default response formatter that returns an object with a response and done property -export const objectResponseFormatter: ResponseFormatter<{ - response: string; - done: boolean; -}> = (response, done) => ({ response, done }); - -export type ObjectResponse = { - response: string; - done: boolean; -}; diff --git a/packages/kernel-language-model-service/src/test-utils/queue/service.test.ts b/packages/kernel-language-model-service/src/test-utils/queue/service.test.ts deleted file mode 100644 index 9a107305a..000000000 --- a/packages/kernel-language-model-service/src/test-utils/queue/service.test.ts +++ /dev/null @@ -1,63 +0,0 @@ -import '@ocap/repo-tools/test-utils/mock-endoify'; -import { describe, it, expect, vi, beforeEach } from 'vitest'; - -import type { QueueLanguageModel } from './model.ts'; -import * as model from './model.ts'; -import { makeQueueService } from './service.ts'; - -vi.mock('./model.ts', () => ({ - makeQueueModel: vi.fn(), -})); - -describe('makeQueueService', () => { - let mockMakeQueueModel: ReturnType; - let mockModel: QueueLanguageModel<{ response: string; done: boolean }>; - - beforeEach(() => { - vi.clearAllMocks(); - - mockModel = { - getInfo: vi.fn(), - load: vi.fn(), - unload: vi.fn(), - sample: vi.fn(), - push: vi.fn(), - } as unknown as QueueLanguageModel<{ response: string; done: boolean }>; - - mockMakeQueueModel = vi.mocked(model.makeQueueModel); - mockMakeQueueModel.mockReturnValue(mockModel); - }); - - it('creates service with makeInstance method', () => { - const service = makeQueueService(); - expect(service).toMatchObject({ - makeInstance: expect.any(Function), - }); - }); - - it('makeInstance calls makeQueueModel with options', async () => { - const service = makeQueueService(); - const config = { - model: 'test', - options: { - tokenizer: vi.fn(), - }, - }; - - const result = await service.makeInstance(config); - - expect(mockMakeQueueModel).toHaveBeenCalledWith(config.options); - expect(result).toBe(mockModel); - }); - - it('makeInstance calls makeQueueModel with undefined options', async () => { - const service = makeQueueService(); - const config = { - model: 'test', - }; - - await service.makeInstance(config); - - expect(mockMakeQueueModel).toHaveBeenCalledWith(undefined); - }); -}); diff --git a/packages/kernel-language-model-service/src/test-utils/queue/service.ts b/packages/kernel-language-model-service/src/test-utils/queue/service.ts deleted file mode 100644 index c2ec5269a..000000000 --- a/packages/kernel-language-model-service/src/test-utils/queue/service.ts +++ /dev/null @@ -1,58 +0,0 @@ -import type { QueueLanguageModel } from './model.ts'; -import { makeQueueModel } from './model.ts'; -import type { ObjectResponse, ResponseFormatter } from './response.ts'; -import type { Tokenizer } from './tokenizer.ts'; -import type { InstanceConfig, LanguageModelService } from '../../types.ts'; - -type QueueLanguageModelServiceConfig = { - tokenizer?: Tokenizer; - responseFormatter?: ResponseFormatter; -}; - -/** - * Queue-based language model service that returns QueueLanguageModel instances. - * This is a minimal implementation of LanguageModelService that uses a queue for responses. - * - * @template Config - The type of configuration accepted by the service - * @template Response - The type of response generated by created models - */ -export type QueueLanguageModelService< - Response extends object = ObjectResponse, -> = LanguageModelService< - QueueLanguageModelServiceConfig, - unknown, - Response -> & { - /** - * Creates a new queue-based language model instance. - * The configuration is ignored - all instances use the 'test' model. - * - * @param config - The configuration for the model instance - * @param config.tokenizer - The tokenizer function to use. Defaults to whitespace splitting. - * @param config.responseFormatter - The function to use to format each yielded token into a response. Defaults to an object with a response and done property. - * @returns A promise that resolves to a queue-based language model instance - */ - makeInstance: ( - config: InstanceConfig>, - ) => Promise>; -}; - -/** - * Creates a queue-based language model service. - * This is a minimal implementation of LanguageModelService that uses a queue for responses. - * - * @template Config - The type of configuration accepted by the service - * @template Response - The type of response generated by created models - * @returns A hardened queue-based language model service - */ -export const makeQueueService = < - Response extends object = ObjectResponse, ->(): QueueLanguageModelService => { - const makeInstance = async ( - config: InstanceConfig>, - ): Promise> => { - return makeQueueModel(config.options); - }; - - return harden({ makeInstance }); -}; diff --git a/packages/kernel-language-model-service/src/test-utils/queue/tokenizer.test.ts b/packages/kernel-language-model-service/src/test-utils/queue/tokenizer.test.ts deleted file mode 100644 index 89512dd22..000000000 --- a/packages/kernel-language-model-service/src/test-utils/queue/tokenizer.test.ts +++ /dev/null @@ -1,23 +0,0 @@ -import { describe, it, expect } from 'vitest'; - -import { whitespaceTokenizer } from './tokenizer.ts'; - -describe('whitespaceTokenizer', () => { - it.each([ - { text: 'hello world', expected: ['hello', ' world'] }, - { text: 'hello', expected: ['hello'] }, - { text: 'hello world test', expected: ['hello', ' world', ' test'] }, - { text: ' hello world ', expected: [' ', ' hello', ' ', ' world', ' '] }, - { text: 'hello world', expected: ['hello', ' ', ' ', ' world'] }, - { text: 'hello\tworld', expected: ['hello', '\tworld'] }, - { text: 'hello\nworld', expected: ['hello', '\nworld'] }, - { text: 'hello\n\nworld', expected: ['hello', '\n', '\nworld'] }, - { text: ' hello ', expected: [' hello', ' '] }, - { text: '\t\nhello', expected: ['\t', '\nhello'] }, - { text: ' ', expected: [' ', ' '] }, - { text: '', expected: [] }, - { text: 'a b c d', expected: ['a', ' b', ' c', ' d'] }, - ])('tokenizes "$text" to $expected', ({ text, expected }) => { - expect(whitespaceTokenizer(text)).toStrictEqual(expected); - }); -}); diff --git a/packages/kernel-language-model-service/src/test-utils/queue/tokenizer.ts b/packages/kernel-language-model-service/src/test-utils/queue/tokenizer.ts deleted file mode 100644 index f24213e11..000000000 --- a/packages/kernel-language-model-service/src/test-utils/queue/tokenizer.ts +++ /dev/null @@ -1,58 +0,0 @@ -/** - * Tokenizer function that converts a string into tokens. - * Can return either a synchronous array or an async iterable. - * - * @param text - The text to tokenize - * @returns Either an array of tokens or an async iterable of tokens - */ -export type Tokenizer = (text: string) => string[] | AsyncIterable; - -/** - * Split text by whitespace. - * For each word, attach at most one whitespace character from the whitespace - * immediately preceding it. Any extra whitespace becomes separate tokens. - * - * @param text - The text to tokenize - * @returns An array of tokens - */ -export const whitespaceTokenizer = (text: string): string[] => { - const tokens: string[] = []; - // Match words with optional preceding whitespace (captured in group 1) - const regex = /(\s*)(\S+)/gu; - let match: RegExpExecArray | null; - let lastIndex = 0; - - while ((match = regex.exec(text)) !== null) { - const [, whitespace, word] = match; - const matchIndex = match.index; - if (!word) { - continue; - } - const whitespaceStr = whitespace ?? ''; - const whitespaceLength = whitespaceStr.length; - - // Process whitespace before the word - if (whitespaceLength > 0) { - // Add all but one whitespace character as separate tokens (before the word) - for (const char of whitespaceStr.slice(0, whitespaceLength - 1)) { - tokens.push(char); - } - // Attach the last whitespace character to the word - tokens.push(whitespaceStr[whitespaceLength - 1] + word); - } else { - tokens.push(word); - } - - lastIndex = matchIndex + whitespaceLength + word.length; - } - - // Add any trailing whitespace as separate tokens - if (lastIndex < text.length) { - const trailing = text.slice(lastIndex); - for (const char of trailing) { - tokens.push(char); - } - } - - return tokens; -}; diff --git a/packages/kernel-language-model-service/src/test-utils/queue/utils.test.ts b/packages/kernel-language-model-service/src/test-utils/queue/utils.test.ts deleted file mode 100644 index 8ea4a4bbe..000000000 --- a/packages/kernel-language-model-service/src/test-utils/queue/utils.test.ts +++ /dev/null @@ -1,145 +0,0 @@ -import { describe, it, expect } from 'vitest'; - -import { - makeAbortableAsyncIterable, - makeEmptyStreamWithAbort, - mapAsyncIterable, - normalizeToAsyncIterable, -} from './utils.ts'; - -describe('normalizeToAsyncIterable', () => { - it.each([ - { input: [1, 2, 3], expected: [1, 2, 3] }, - { input: [], expected: [] }, - { input: ['a', 'b'], expected: ['a', 'b'] }, - ])( - 'normalizes array $input to async iterable', - async ({ input, expected }) => { - const result = normalizeToAsyncIterable<(typeof input)[number]>(input); - const values: (typeof input)[number][] = []; - for await (const value of result) { - values.push(value); - } - expect(values).toStrictEqual(expected); - }, - ); - - it('returns async iterable unchanged', async () => { - const asyncIter = (async function* () { - yield 1; - yield 2; - })(); - const result = normalizeToAsyncIterable(asyncIter); - const values: number[] = []; - for await (const value of result) { - values.push(value); - } - expect(values).toStrictEqual([1, 2]); - }); -}); - -describe('mapAsyncIterable', () => { - it.each([ - { input: [1, 2, 3], expected: [false, false, true] }, - { input: [1], expected: [true] }, - { input: ['a', 'b', 'c'], expected: [false, false, true] }, - ])('maps $input with done flag', async ({ input, expected }) => { - const iterable = (async function* () { - yield* input; - })(); - const result = mapAsyncIterable(iterable, (_value, done) => done); - const values: boolean[] = []; - for await (const value of result) { - values.push(value); - } - expect(values).toStrictEqual(expected); - }); - - it('maps values correctly', async () => { - const iterable = (async function* () { - yield 1; - yield 2; - })(); - const result = mapAsyncIterable(iterable, (value, _done) => value * 2); - const values: number[] = []; - for await (const value of result) { - values.push(value); - } - expect(values).toStrictEqual([2, 4]); - }); - - it('handles empty iterable', async () => { - const iterable = (async function* () { - // Empty iterable for testing - })(); - const result = mapAsyncIterable(iterable, (_value, done) => done); - const values: boolean[] = []; - for await (const value of result) { - values.push(value); - } - expect(values).toStrictEqual([]); - }); -}); - -describe('makeAbortableAsyncIterable', () => { - it('yields values until abort', async () => { - const iterable = (async function* () { - yield 1; - yield 2; - yield 3; - })(); - const { stream, abort } = makeAbortableAsyncIterable(iterable); - const values: number[] = []; - for await (const value of stream) { - values.push(value); - if (value === 2) { - await abort(); - } - } - expect(values).toStrictEqual([1, 2]); - }); - - it('stops yielding after abort', async () => { - const iterable = (async function* () { - yield 1; - yield 2; - yield 3; - })(); - const { stream, abort } = makeAbortableAsyncIterable(iterable); - await abort(); - const values: number[] = []; - for await (const value of stream) { - values.push(value); - } - expect(values).toStrictEqual([]); - }); - - it('completes normally when not aborted', async () => { - const iterable = (async function* () { - yield 1; - yield 2; - })(); - const { stream } = makeAbortableAsyncIterable(iterable); - const values: number[] = []; - for await (const value of stream) { - values.push(value); - } - expect(values).toStrictEqual([1, 2]); - }); -}); - -describe('makeEmptyStreamWithAbort', () => { - it('returns empty stream', async () => { - const { stream } = makeEmptyStreamWithAbort(); - const values: number[] = []; - for await (const value of stream) { - values.push(value); - } - expect(values).toStrictEqual([]); - }); - - it('provides no-op abort function', async () => { - const { abort } = makeEmptyStreamWithAbort(); - expect(await abort()).toBeUndefined(); - }); -}); diff --git a/packages/kernel-language-model-service/src/test-utils/queue/utils.ts b/packages/kernel-language-model-service/src/test-utils/queue/utils.ts deleted file mode 100644 index 5d72694c0..000000000 --- a/packages/kernel-language-model-service/src/test-utils/queue/utils.ts +++ /dev/null @@ -1,98 +0,0 @@ -/** - * Normalize an array or async iterable to an async iterable. - * - * @param value - The value to normalize. - * @returns The normalized value. - */ -export const normalizeToAsyncIterable = ( - value: Type[] | AsyncIterable, -): AsyncIterable => - Array.isArray(value) - ? (async function* () { - yield* value; - })() - : value; - -/** - * Map an async iterable to a new async iterable. - * The mapper receives both the value and whether it's the last item. - * - * @param iterable - The iterable to map. - * @param mapper - The mapper function that receives (value, done). - * @returns The mapped iterable. - */ -export const mapAsyncIterable = ( - iterable: AsyncIterable, - mapper: (value: Type, done: boolean) => Result, -): AsyncIterable => - (async function* () { - const iterator = iterable[Symbol.asyncIterator](); - let current = await iterator.next(); - - if (current.done) { - return; - } - - let next = await iterator.next(); - while (!next.done) { - yield mapper(current.value, false); - current = next; - next = await iterator.next(); - } - - yield mapper(current.value, true); - })(); - -/** - * Creates a queue-based language model instance. - * This is a minimal implementation of LanguageModel that uses a queue for responses. - * - * @template Options - The type of options supported by the model - * @template Response - The type of response generated by the model - * @returns A hardened queue-based language model instance with helper methods - */ -export type StreamWithAbort = { - stream: AsyncIterable; - abort: () => Promise; -}; - -/** - * Make an async iterable abortable. - * - * @param iterable - The iterable to make abortable. - * @returns A tuple containing the abortable iterable and the abort function. - */ -export const makeAbortableAsyncIterable = ( - iterable: AsyncIterable, -): StreamWithAbort => { - let didAbort = false; - return { - stream: (async function* () { - for await (const value of iterable) { - if (didAbort) { - break; - } - yield value; - } - })(), - abort: async () => { - didAbort = true; - }, - }; -}; - -/** - * Make an empty stream with abort. - * - * @returns A stream with abort. - */ -export const makeEmptyStreamWithAbort = < - Response, ->(): StreamWithAbort => ({ - stream: (async function* () { - // Empty stream - })() as AsyncIterable, - abort: async () => { - // No-op abort - }, -}); diff --git a/packages/kernel-language-model-service/src/types.ts b/packages/kernel-language-model-service/src/types.ts index 5895986c4..aafaabddd 100644 --- a/packages/kernel-language-model-service/src/types.ts +++ b/packages/kernel-language-model-service/src/types.ts @@ -1,3 +1,121 @@ +/** + * A role in a chat conversation. + */ +export type ChatRole = 'system' | 'user' | 'assistant'; + +/** + * A single message in a chat conversation. + */ +export type ChatMessage = { + role: ChatRole; + content: string; +}; + +/** + * Token usage statistics for a completion request. + */ +export type Usage = { + prompt_tokens: number; + completion_tokens: number; + total_tokens: number; +}; + +/** + * Options shared by raw sampling requests. + */ +export type SampleOptions = { + temperature?: number; + top_p?: number; + seed?: number; + max_tokens?: number; + stop?: string | string[]; +}; + +/** + * Parameters for a raw token-prediction request (bypasses chat template). + */ +export type SampleParams = { + model: string; + prompt: string; +} & SampleOptions; + +/** + * Result of a raw token-prediction request. + */ +export type SampleResult = { + text: string; +}; + +/** + * Parameters for a chat completion request. + */ +export type ChatParams = { + model: string; + messages: ChatMessage[]; + max_tokens?: number; + temperature?: number; + top_p?: number; + stop?: string | string[]; + seed?: number; + n?: number; + /** When `true`, the response is an SSE stream of {@link ChatStreamChunk}s. */ + stream?: boolean; +}; + +/** + * A partial message delta from a streaming chat completion response. + */ +export type ChatStreamDelta = { + role?: ChatRole; + content?: string; +}; + +/** + * A single chunk from a streaming chat completion response. + */ +export type ChatStreamChunk = { + id: string; + model: string; + choices: { + delta: ChatStreamDelta; + index: number; + finish_reason: string | null; + }[]; +}; + +/** + * A single choice in a chat completion response. + */ +export type ChatChoice = { + message: ChatMessage; + index: number; + finish_reason: string | null; +}; + +/** + * Result of a chat completion request. + */ +export type ChatResult = { + id: string; + model: string; + choices: ChatChoice[]; + usage: Usage; +}; + +/** + * Minimal service interface required by `makeChatClient`. + */ +export type ChatService = { + chat: (params: ChatParams) => Promise; +}; + +/** + * Minimal service interface required by `makeSampleClient`. + */ +export type SampleService = { + sample: (params: SampleParams) => Promise; +}; + /** * Configuration information for a language model. * Contains the model identifier and optional configuration parameters. diff --git a/packages/kernel-test-local/package.json b/packages/kernel-test-local/package.json index 70d2cd939..085cabbcb 100644 --- a/packages/kernel-test-local/package.json +++ b/packages/kernel-test-local/package.json @@ -13,7 +13,8 @@ }, "type": "module", "scripts": { - "clean": "rimraf --glob './*.tsbuildinfo' ./.eslintcache ./coverage ./.turbo ./logs", + "build": "ocap bundle src/vats", + "clean": "rimraf --glob './*.tsbuildinfo' ./.eslintcache ./coverage ./.turbo ./logs './src/vats/*.bundle'", "lint": "yarn lint:eslint && yarn lint:misc --check && yarn constraints && yarn lint:dependencies", "lint:dependencies": "depcheck --quiet", "lint:eslint": "eslint . --cache", @@ -29,7 +30,12 @@ "test:dev:quiet": "yarn test:dev --reporter @ocap/repo-tools/vitest-reporters/silent" }, "dependencies": { + "@endo/eventual-send": "^1.3.4", + "@metamask/kernel-node-runtime": "workspace:^", + "@metamask/kernel-store": "workspace:^", + "@metamask/kernel-utils": "workspace:^", "@metamask/logger": "workspace:^", + "@metamask/ocap-kernel": "workspace:^", "@ocap/kernel-agents": "workspace:^", "@ocap/kernel-language-model-service": "workspace:^", "@ocap/repo-tools": "workspace:^" @@ -39,6 +45,8 @@ "@metamask/eslint-config": "^15.0.0", "@metamask/eslint-config-nodejs": "^15.0.0", "@metamask/eslint-config-typescript": "^15.0.0", + "@metamask/kernel-cli": "workspace:^", + "@metamask/kernel-shims": "workspace:^", "@ocap/kernel-agents-repl": "workspace:^", "@types/node": "^22.13.1", "@typescript-eslint/eslint-plugin": "^8.29.0", diff --git a/packages/kernel-test-local/test/e2e/agents.test.ts b/packages/kernel-test-local/src/agents.e2e.test.ts similarity index 98% rename from packages/kernel-test-local/test/e2e/agents.test.ts rename to packages/kernel-test-local/src/agents.e2e.test.ts index 7ab684286..9297915f9 100644 --- a/packages/kernel-test-local/test/e2e/agents.test.ts +++ b/packages/kernel-test-local/src/agents.e2e.test.ts @@ -19,8 +19,8 @@ import { vi, } from 'vitest'; -import { DEFAULT_MODEL } from '../../src/constants.ts'; -import { filterTransports, randomLetter } from '../../src/utils.ts'; +import { DEFAULT_MODEL } from './constants.ts'; +import { filterTransports, randomLetter } from './utils.ts'; const logger = new Logger({ tags: ['test'], diff --git a/packages/kernel-test-local/src/constants.ts b/packages/kernel-test-local/src/constants.ts index 86b329e2e..a096d674a 100644 --- a/packages/kernel-test-local/src/constants.ts +++ b/packages/kernel-test-local/src/constants.ts @@ -10,10 +10,8 @@ export const TEST_MODELS = ['llama3.1:latest', 'gpt-oss:20b']; export const OLLAMA_API_BASE = 'http://localhost:11434'; export const OLLAMA_TAGS_ENDPOINT = `${OLLAMA_API_BASE}/api/tags`; -// extract ignored logger tags from environment variable - /** - * The tags to ignore for the local tests. + * Logger tags to ignore, parsed from the LOGGER_IGNORE environment variable. */ export const IGNORE_TAGS = // eslint-disable-next-line n/no-process-env diff --git a/packages/kernel-test-local/src/lms-chat.e2e.test.ts b/packages/kernel-test-local/src/lms-chat.e2e.test.ts new file mode 100644 index 000000000..631e54190 --- /dev/null +++ b/packages/kernel-test-local/src/lms-chat.e2e.test.ts @@ -0,0 +1,31 @@ +import '@metamask/kernel-shims/endoify-node'; + +import { makeOpenV1NodejsService } from '@ocap/kernel-language-model-service'; +import { fetchMock } from '@ocap/repo-tools/test-utils/fetch-mock'; +import { afterAll, beforeAll, describe, it } from 'vitest'; + +import { runLmsChatKernelTest } from './lms-chat.ts'; + +describe.sequential('lms-kernel (e2e)', () => { + beforeAll(() => { + fetchMock.disableMocks(); + }); + + afterAll(() => { + fetchMock.enableMocks(); + }); + + // eslint-disable-next-line vitest/expect-expect + it( + 'sends a chat message through the kernel to Ollama and receives a response', + { timeout: 60_000 }, + async () => { + const { chat } = makeOpenV1NodejsService({ + endowments: { fetch }, + baseUrl: 'http://localhost:11434', + apiKey: 'test-api-key', + }); + await runLmsChatKernelTest(chat); + }, + ); +}); diff --git a/packages/kernel-test-local/src/lms-chat.test.ts b/packages/kernel-test-local/src/lms-chat.test.ts new file mode 100644 index 000000000..cc484268b --- /dev/null +++ b/packages/kernel-test-local/src/lms-chat.test.ts @@ -0,0 +1,18 @@ +import '@metamask/kernel-shims/endoify-node'; + +import { makeOpenV1NodejsService } from '@ocap/kernel-language-model-service'; +import { makeMockOpenV1Fetch } from '@ocap/kernel-language-model-service/test-utils'; +import { describe, it } from 'vitest'; + +import { runLmsChatKernelTest } from './lms-chat.ts'; + +describe.sequential('lms-kernel', () => { + // eslint-disable-next-line vitest/expect-expect + it('sends a chat message through the kernel and receives a response', async () => { + const { chat } = makeOpenV1NodejsService({ + endowments: { fetch: makeMockOpenV1Fetch(['Hello.']) }, + baseUrl: 'http://localhost:11434', + }); + await runLmsChatKernelTest(chat); + }); +}); diff --git a/packages/kernel-test-local/src/lms-chat.ts b/packages/kernel-test-local/src/lms-chat.ts new file mode 100644 index 000000000..b5b27f23b --- /dev/null +++ b/packages/kernel-test-local/src/lms-chat.ts @@ -0,0 +1,73 @@ +import { NodejsPlatformServices } from '@metamask/kernel-node-runtime'; +import { makeSQLKernelDatabase } from '@metamask/kernel-store/sqlite/nodejs'; +import { waitUntilQuiescent } from '@metamask/kernel-utils'; +import { + Logger, + makeArrayTransport, + makeConsoleTransport, +} from '@metamask/logger'; +import type { LogEntry } from '@metamask/logger'; +import { Kernel } from '@metamask/ocap-kernel'; +import type { + ChatParams, + ChatResult, +} from '@ocap/kernel-language-model-service'; +import { + LANGUAGE_MODEL_SERVICE_NAME, + makeKernelLanguageModelService, +} from '@ocap/kernel-language-model-service'; +import { expect } from 'vitest'; + +import { DEFAULT_MODEL } from './constants.ts'; +import { filterTransports } from './utils.ts'; + +const getBundleSpec = (name: string): string => + new URL(`./vats/${name}.bundle`, import.meta.url).toString(); + +export const runLmsChatKernelTest = async ( + chat: (params: ChatParams & { stream?: true & false }) => Promise, +): Promise => { + const kernelDatabase = await makeSQLKernelDatabase({ + dbFilename: ':memory:', + }); + + const entries: LogEntry[] = []; + const logger = new Logger({ + transports: [ + filterTransports(makeConsoleTransport(), makeArrayTransport(entries)), + ], + }); + + const platformServices = new NodejsPlatformServices({ + logger: logger.subLogger({ tags: ['vat-worker-manager'] }), + }); + + const kernel = await Kernel.make(platformServices, kernelDatabase, { + resetStorage: true, + logger, + }); + + const { name, service } = makeKernelLanguageModelService(chat); + kernel.registerKernelServiceObject(name, service); + + await kernel.launchSubcluster({ + bootstrap: 'main', + services: [LANGUAGE_MODEL_SERVICE_NAME], + vats: { + main: { + bundleSpec: getBundleSpec('lms-chat-vat'), + parameters: { model: DEFAULT_MODEL }, + }, + }, + }); + await waitUntilQuiescent(100); + + const responseEntry = entries.find((entry) => + entry.message?.startsWith('lms-chat response:'), + ); + expect(responseEntry).toBeDefined(); + expect(responseEntry?.message?.length).toBeGreaterThan( + 'lms-chat response: '.length, + ); + expect(responseEntry?.message).toMatch(/^lms-chat response: [hH]ello[.!]?$/u); +}; diff --git a/packages/kernel-test-local/src/utils.ts b/packages/kernel-test-local/src/utils.ts index 219200830..74584cf5c 100644 --- a/packages/kernel-test-local/src/utils.ts +++ b/packages/kernel-test-local/src/utils.ts @@ -23,7 +23,7 @@ export const filterTransports = ( /** * Generate a random letter. * - * @returns a random letter. + * @returns A random letter. */ export function randomLetter(): string { return String.fromCharCode(Math.floor(Math.random() * 26) + 97); diff --git a/packages/kernel-test-local/src/utils.test.ts b/packages/kernel-test-local/src/utils.unit.test.ts similarity index 100% rename from packages/kernel-test-local/src/utils.test.ts rename to packages/kernel-test-local/src/utils.unit.test.ts diff --git a/packages/kernel-test-local/src/vats/lms-chat-vat.ts b/packages/kernel-test-local/src/vats/lms-chat-vat.ts new file mode 100644 index 000000000..2c01356e8 --- /dev/null +++ b/packages/kernel-test-local/src/vats/lms-chat-vat.ts @@ -0,0 +1,41 @@ +import type { ERef } from '@endo/eventual-send'; +import { makeDefaultExo } from '@metamask/kernel-utils/exo'; +import type { Logger } from '@metamask/logger'; +import { makeChatClient } from '@ocap/kernel-language-model-service'; +import type { ChatService } from '@ocap/kernel-language-model-service'; + +/** + * A vat that uses a kernel language model service to perform a chat completion + * and logs the response. Used by lms-chat.test.ts and lms-chat.e2e.test.ts to verify the full + * kernel → LMS service → Ollama round-trip. + * + * @param vatPowers - Vat powers, expected to include a logger. + * @param parameters - Vat parameters. + * @param parameters.model - The model to use for chat completion. + * @returns A default Exo instance. + */ +// eslint-disable-next-line @typescript-eslint/explicit-function-return-type +export function buildRootObject( + vatPowers: Record, + { model }: { model: string }, +) { + const logger = vatPowers.logger as Logger; + const tlog = (message: string): void => { + logger.subLogger({ tags: ['test', 'lms-chat'] }).log(message); + }; + + return makeDefaultExo('root', { + async bootstrap( + _roots: unknown, + { languageModelService }: { languageModelService: ERef }, + ) { + const client = makeChatClient(languageModelService, model); + const result = await client.chat.completions.create({ + messages: [ + { role: 'user', content: 'Reply with exactly one word: hello.' }, + ], + }); + tlog(`lms-chat response: ${result.choices[0]?.message.content ?? ''}`); + }, + }); +} diff --git a/packages/kernel-test-local/test/e2e/suite.test.ts b/packages/kernel-test-local/test/suite.test.ts similarity index 97% rename from packages/kernel-test-local/test/e2e/suite.test.ts rename to packages/kernel-test-local/test/suite.test.ts index c6c7cee81..835ebeb3e 100644 --- a/packages/kernel-test-local/test/e2e/suite.test.ts +++ b/packages/kernel-test-local/test/suite.test.ts @@ -13,7 +13,7 @@ import { DEFAULT_MODEL, OLLAMA_API_BASE, OLLAMA_TAGS_ENDPOINT, -} from '../../src/constants.ts'; +} from '../src/constants.ts'; describe.sequential('test suite', () => { beforeAll(() => { diff --git a/packages/kernel-test-local/tsconfig.json b/packages/kernel-test-local/tsconfig.json index 8539255af..cb233df63 100644 --- a/packages/kernel-test-local/tsconfig.json +++ b/packages/kernel-test-local/tsconfig.json @@ -8,9 +8,14 @@ }, "references": [ { "path": "../kernel-agents" }, + { "path": "../kernel-shims" }, { "path": "../kernel-agents-repl" }, { "path": "../kernel-language-model-service" }, + { "path": "../kernel-store" }, + { "path": "../kernel-utils" }, { "path": "../logger" }, + { "path": "../ocap-kernel" }, + { "path": "../kernel-node-runtime" }, { "path": "../repo-tools" } ], "include": [ @@ -18,6 +23,6 @@ "./src", "./vitest.config.ts", "./vitest.config.e2e.ts", - "./test/e2e" + "./test" ] } diff --git a/packages/kernel-test-local/vitest.config.e2e.ts b/packages/kernel-test-local/vitest.config.e2e.ts index 97c0ce5ca..ef76f427c 100644 --- a/packages/kernel-test-local/vitest.config.e2e.ts +++ b/packages/kernel-test-local/vitest.config.e2e.ts @@ -10,9 +10,9 @@ export default defineConfig((args) => { defineProject({ test: { name: 'kernel-test-local-e2e', - testTimeout: 30_000, + testTimeout: 60_000, hookTimeout: 10_000, - include: ['./test/e2e/**/*.test.ts'], + include: ['./src/**/*.e2e.test.ts', './test/**/*.test.ts'], }, }), ); diff --git a/packages/kernel-test-local/vitest.config.ts b/packages/kernel-test-local/vitest.config.ts index 6eee6669c..0586515a0 100644 --- a/packages/kernel-test-local/vitest.config.ts +++ b/packages/kernel-test-local/vitest.config.ts @@ -13,6 +13,7 @@ export default defineConfig((args) => { testTimeout: 30_000, hookTimeout: 10_000, include: ['./src/**/*.test.ts'], + exclude: ['**/*.e2e.test.ts'], }, }), ); diff --git a/packages/kernel-test/src/lms-chat.test.ts b/packages/kernel-test/src/lms-chat.test.ts new file mode 100644 index 000000000..c655b61b9 --- /dev/null +++ b/packages/kernel-test/src/lms-chat.test.ts @@ -0,0 +1,51 @@ +import { makeSQLKernelDatabase } from '@metamask/kernel-store/sqlite/nodejs'; +import { waitUntilQuiescent } from '@metamask/kernel-utils'; +import { + LANGUAGE_MODEL_SERVICE_NAME, + makeKernelLanguageModelService, + makeOpenV1NodejsService, +} from '@ocap/kernel-language-model-service'; +import { makeMockOpenV1Fetch } from '@ocap/kernel-language-model-service/test-utils'; +import { describe, expect, it } from 'vitest'; + +import { + getBundleSpec, + makeKernel, + makeTestLogger, + runTestVats, +} from './utils.ts'; + +describe('lms-chat vat', () => { + it('receives chat response via makeChatClient', async () => { + const kernelDatabase = await makeSQLKernelDatabase({ + dbFilename: ':memory:', + }); + const { logger, entries } = makeTestLogger(); + const kernel = await makeKernel(kernelDatabase, true, logger); + + const { chat } = makeOpenV1NodejsService({ + endowments: { fetch: makeMockOpenV1Fetch(['My name is Alice.']) }, + baseUrl: 'http://localhost:11434', + }); + const { name, service } = makeKernelLanguageModelService(chat); + kernel.registerKernelServiceObject(name, service); + + await runTestVats(kernel, { + bootstrap: 'main', + services: [LANGUAGE_MODEL_SERVICE_NAME], + vats: { + main: { + bundleSpec: getBundleSpec('lms-chat-vat'), + parameters: { name: 'Alice' }, + }, + }, + }); + await waitUntilQuiescent(100); + + expect( + entries.some((entry) => + entry.message.includes('response: My name is Alice.'), + ), + ).toBe(true); + }); +}); diff --git a/packages/kernel-test/src/lms-sample.test.ts b/packages/kernel-test/src/lms-sample.test.ts new file mode 100644 index 000000000..999f152a9 --- /dev/null +++ b/packages/kernel-test/src/lms-sample.test.ts @@ -0,0 +1,52 @@ +import { makeSQLKernelDatabase } from '@metamask/kernel-store/sqlite/nodejs'; +import { waitUntilQuiescent } from '@metamask/kernel-utils'; +import type { LogEntry } from '@metamask/logger'; +import { + LANGUAGE_MODEL_SERVICE_NAME, + makeKernelLanguageModelService, +} from '@ocap/kernel-language-model-service'; +import { makeMockSample } from '@ocap/kernel-language-model-service/test-utils'; +import { describe, expect, it, vi } from 'vitest'; + +import { + getBundleSpec, + makeKernel, + makeTestLogger, + runTestVats, +} from './utils.ts'; + +describe('lms-sample vat', () => { + it('receives sample response via makeSampleClient', async () => { + const kernelDatabase = await makeSQLKernelDatabase({ + dbFilename: ':memory:', + }); + const { logger, entries } = makeTestLogger(); + const kernel = await makeKernel(kernelDatabase, true, logger); + + const chat = vi.fn(); + const { name, service } = makeKernelLanguageModelService( + chat, + makeMockSample(['The sky is blue.']), + ); + kernel.registerKernelServiceObject(name, service); + + await runTestVats(kernel, { + bootstrap: 'main', + services: [LANGUAGE_MODEL_SERVICE_NAME], + vats: { + main: { + bundleSpec: getBundleSpec('lms-sample-vat'), + parameters: { prompt: 'What color is the sky?' }, + }, + }, + }); + await waitUntilQuiescent(100); + + expect( + entries.some( + (entry: LogEntry) => + entry.message?.includes('response: The sky is blue.') ?? false, + ), + ).toBe(true); + }); +}); diff --git a/packages/kernel-test/src/lms-user.test.ts b/packages/kernel-test/src/lms-user.test.ts deleted file mode 100644 index 00200ae1d..000000000 --- a/packages/kernel-test/src/lms-user.test.ts +++ /dev/null @@ -1,43 +0,0 @@ -import { makeSQLKernelDatabase } from '@metamask/kernel-store/sqlite/nodejs'; -import { waitUntilQuiescent } from '@metamask/kernel-utils'; -import { describe, expect, it } from 'vitest'; - -import { - extractTestLogs, - getBundleSpec, - makeKernel, - makeTestLogger, - runTestVats, -} from './utils.ts'; - -const testSubcluster = { - bootstrap: 'main', - forceReset: true, - vats: { - main: { - bundleSpec: getBundleSpec('lms-user-vat'), - parameters: { - name: 'Alice', - }, - }, - languageModelService: { - bundleSpec: getBundleSpec('lms-queue-vat'), - }, - }, -}; - -describe('lms-user vat', () => { - it('logs response from language model', async () => { - const kernelDatabase = await makeSQLKernelDatabase({ - dbFilename: ':memory:', - }); - const { logger, entries } = makeTestLogger(); - const kernel = await makeKernel(kernelDatabase, true, logger); - - await runTestVats(kernel, testSubcluster); - await waitUntilQuiescent(100); - - const testLogs = extractTestLogs(entries); - expect(testLogs).toContain('response: My name is Alice.'); - }); -}); diff --git a/packages/kernel-test/src/vats/lms-chat-vat.ts b/packages/kernel-test/src/vats/lms-chat-vat.ts new file mode 100644 index 000000000..5831bcb86 --- /dev/null +++ b/packages/kernel-test/src/vats/lms-chat-vat.ts @@ -0,0 +1,41 @@ +import type { ERef } from '@endo/eventual-send'; +import { makeDefaultExo } from '@metamask/kernel-utils/exo'; +import { makeChatClient } from '@ocap/kernel-language-model-service'; +import type { ChatService } from '@ocap/kernel-language-model-service'; + +import { unwrapTestLogger } from '../test-powers.ts'; +import type { TestPowers } from '../test-powers.ts'; + +/** + * A vat that uses a language model service to generate text. + * + * @param vatPowers - The powers of the vat. + * @param parameters - The parameters of the vat. + * @param parameters.name - The name of the vat. + * @returns A default Exo instance. + */ +// eslint-disable-next-line @typescript-eslint/explicit-function-return-type +export function buildRootObject( + vatPowers: TestPowers, + { name = 'anonymous' }: { name?: string } = {}, +) { + const tlog = unwrapTestLogger(vatPowers, name); + const root = makeDefaultExo('root', { + async bootstrap( + _roots: unknown, + { languageModelService }: { languageModelService: ERef }, + ) { + const client = makeChatClient(languageModelService, 'test'); + const result = await client.chat.completions.create({ + messages: [ + { + role: 'user', + content: `Hello, my name is ${name}. What is your name?`, + }, + ], + }); + tlog(`response: ${result.choices[0]?.message.content ?? ''}`); + }, + }); + return root; +} diff --git a/packages/kernel-test/src/vats/lms-queue-vat.ts b/packages/kernel-test/src/vats/lms-queue-vat.ts deleted file mode 100644 index e7872f1a6..000000000 --- a/packages/kernel-test/src/vats/lms-queue-vat.ts +++ /dev/null @@ -1,64 +0,0 @@ -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; - unload: () => Promise; - sample: (prompt: string) => Promise<{ - stream: AsyncIterable; - 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 { - 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); - }, - }); - }, - }); -} diff --git a/packages/kernel-test/src/vats/lms-sample-vat.ts b/packages/kernel-test/src/vats/lms-sample-vat.ts new file mode 100644 index 000000000..e95ed0901 --- /dev/null +++ b/packages/kernel-test/src/vats/lms-sample-vat.ts @@ -0,0 +1,35 @@ +import type { ERef } from '@endo/eventual-send'; +import { makeDefaultExo } from '@metamask/kernel-utils/exo'; +import { makeSampleClient } from '@ocap/kernel-language-model-service'; +import type { SampleService } from '@ocap/kernel-language-model-service'; + +import { unwrapTestLogger } from '../test-powers.ts'; +import type { TestPowers } from '../test-powers.ts'; + +/** + * A vat that uses a kernel language model service to perform a raw sample + * completion and logs the response. Used to verify the full kernel → LMS + * service round-trip for the sample path. + * + * @param vatPowers - The powers of the vat. + * @param parameters - The parameters of the vat. + * @param parameters.prompt - The prompt to sample from. + * @returns A default Exo instance. + */ +// eslint-disable-next-line @typescript-eslint/explicit-function-return-type +export function buildRootObject( + vatPowers: TestPowers, + { prompt = 'Hello' }: { prompt?: string } = {}, +) { + const tlog = unwrapTestLogger(vatPowers, 'lms-sample'); + return makeDefaultExo('root', { + async bootstrap( + _roots: unknown, + { languageModelService }: { languageModelService: ERef }, + ) { + const client = makeSampleClient(languageModelService, 'test'); + const result = await client.sample({ prompt }); + tlog(`response: ${result.text}`); + }, + }); +} diff --git a/packages/kernel-test/src/vats/lms-user-vat.ts b/packages/kernel-test/src/vats/lms-user-vat.ts deleted file mode 100644 index 4a0c1b0e5..000000000 --- a/packages/kernel-test/src/vats/lms-user-vat.ts +++ /dev/null @@ -1,49 +0,0 @@ -import { E } from '@endo/eventual-send'; -import { makeDefaultExo } from '@metamask/kernel-utils/exo'; -import { makeEventualIterator } from '@ocap/remote-iterables'; - -import { unwrapTestLogger } from '../test-powers.ts'; -import type { TestPowers } from '../test-powers.ts'; - -/** - * A vat that uses a language model service to generate text. - * - * @param vatPowers - The powers of the vat. - * @param vatPowers.logger - The logger of the vat. - * @param parameters - The parameters of the vat. - * @param parameters.name - The name of the vat. - * @returns A default Exo instance. - */ -// eslint-disable-next-line @typescript-eslint/explicit-function-return-type -export function buildRootObject( - vatPowers: TestPowers, - { name = 'anonymous' }: { name?: string } = {}, -) { - const tlog = unwrapTestLogger(vatPowers, name); - let languageModel: unknown; - const root = makeDefaultExo('root', { - async bootstrap( - { languageModelService }: { languageModelService: unknown }, - _kernelServices: unknown, - ) { - languageModel = await E(languageModelService).makeInstance({ - model: 'test', - }); - await E(languageModel).push(`My name is ${name}.`); - const response = await E(root).ask('Hello, what is your name?'); - tlog(`response: ${response}`); - }, - async ask(prompt: string) { - let response = ''; - const sampleResult = await E(languageModel).sample(prompt); - const stream = await E(sampleResult).getStream(); - const iterator = makeEventualIterator(stream); - for await (const chunk of iterator) { - response += (chunk as { response: string }).response; - } - return response; - }, - }); - - return root; -} diff --git a/yarn.lock b/yarn.lock index bd2208b3e..76034a9f0 100644 --- a/yarn.lock +++ b/yarn.lock @@ -3780,6 +3780,7 @@ __metadata: resolution: "@ocap/kernel-language-model-service@workspace:packages/kernel-language-model-service" dependencies: "@arethetypeswrong/cli": "npm:^0.17.4" + "@endo/eventual-send": "npm:^1.3.4" "@metamask/auto-changelog": "npm:^5.3.0" "@metamask/eslint-config": "npm:^15.0.0" "@metamask/eslint-config-nodejs": "npm:^15.0.0" @@ -3821,10 +3822,17 @@ __metadata: resolution: "@ocap/kernel-test-local@workspace:packages/kernel-test-local" dependencies: "@arethetypeswrong/cli": "npm:^0.17.4" + "@endo/eventual-send": "npm:^1.3.4" "@metamask/eslint-config": "npm:^15.0.0" "@metamask/eslint-config-nodejs": "npm:^15.0.0" "@metamask/eslint-config-typescript": "npm:^15.0.0" + "@metamask/kernel-cli": "workspace:^" + "@metamask/kernel-node-runtime": "workspace:^" + "@metamask/kernel-shims": "workspace:^" + "@metamask/kernel-store": "workspace:^" + "@metamask/kernel-utils": "workspace:^" "@metamask/logger": "workspace:^" + "@metamask/ocap-kernel": "workspace:^" "@ocap/kernel-agents": "workspace:^" "@ocap/kernel-agents-repl": "workspace:^" "@ocap/kernel-language-model-service": "workspace:^"