From de242b7f7b4f6938410cbe0f0377510969f5f296 Mon Sep 17 00:00:00 2001 From: grypez <143971198+grypez@users.noreply.github.com> Date: Tue, 17 Mar 2026 15:23:28 -0400 Subject: [PATCH 01/20] feat(klms): language model service MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add @ocap/kernel-language-model-service — a new package that bridges the kernel with OpenAI-compatible language model backends. Core types and kernel service: - ChatParams / ChatResult / ChatService types for request/response shapes - makeKernelLanguageModelService: registers a kernel service object that dispatches chat calls to a provided chat function - LANGUAGE_MODEL_SERVICE_NAME constant Backends: - Open /v1 backend with SSE streaming (makeOpenV1NodejsService) - Ollama Node.js backend (OllamaNodejsService) Client API: - makeChatClient: returns an OpenAI-SDK-compatible client object that routes chat.completions.create() calls through a CapTP ChatService - makeSampleClient: convenience wrapper for single-token sampling Test utilities: - makeMockOpenV1Fetch: deterministic SSE mock for unit/integration tests Co-Authored-By: Claude Sonnet 4.6 --- .../package.json | 1 + .../src/client.test.ts | 107 +++++ .../src/client.ts | 85 ++++ .../src/index.ts | 6 + .../src/kernel-service.test.ts | 100 +++++ .../src/kernel-service.ts | 41 ++ .../src/ollama/base.ts | 93 ++++- .../src/ollama/nodejs.ts | 26 ++ .../src/ollama/types.ts | 22 +- .../src/open-v1/base.test.ts | 246 ++++++++++++ .../src/open-v1/base.ts | 136 +++++++ .../src/open-v1/nodejs.ts | 39 ++ .../src/open-v1/types.ts | 50 +++ .../src/test-utils/index.ts | 6 +- .../src/test-utils/mock-fetch.ts | 31 ++ .../src/test-utils/mock-sample.ts | 18 + .../src/test-utils/queue/README.md | 49 --- .../src/test-utils/queue/model.test.ts | 371 ------------------ .../src/test-utils/queue/model.ts | 80 ---- .../src/test-utils/queue/response.test.ts | 24 -- .../src/test-utils/queue/response.ts | 15 - .../src/test-utils/queue/service.test.ts | 63 --- .../src/test-utils/queue/service.ts | 58 --- .../src/test-utils/queue/tokenizer.test.ts | 23 -- .../src/test-utils/queue/tokenizer.ts | 58 --- .../src/test-utils/queue/utils.test.ts | 145 ------- .../src/test-utils/queue/utils.ts | 98 ----- .../src/types.ts | 118 ++++++ packages/kernel-test/src/lms-chat.test.ts | 51 +++ packages/kernel-test/src/lms-sample.test.ts | 52 +++ packages/kernel-test/src/lms-user.test.ts | 43 -- packages/kernel-test/src/vats/lms-chat-vat.ts | 41 ++ .../kernel-test/src/vats/lms-queue-vat.ts | 64 --- .../kernel-test/src/vats/lms-sample-vat.ts | 35 ++ packages/kernel-test/src/vats/lms-user-vat.ts | 49 --- yarn.lock | 8 + 36 files changed, 1301 insertions(+), 1151 deletions(-) create mode 100644 packages/kernel-language-model-service/src/client.test.ts create mode 100644 packages/kernel-language-model-service/src/client.ts create mode 100644 packages/kernel-language-model-service/src/kernel-service.test.ts create mode 100644 packages/kernel-language-model-service/src/kernel-service.ts create mode 100644 packages/kernel-language-model-service/src/open-v1/base.test.ts create mode 100644 packages/kernel-language-model-service/src/open-v1/base.ts create mode 100644 packages/kernel-language-model-service/src/open-v1/nodejs.ts create mode 100644 packages/kernel-language-model-service/src/open-v1/types.ts create mode 100644 packages/kernel-language-model-service/src/test-utils/mock-fetch.ts create mode 100644 packages/kernel-language-model-service/src/test-utils/mock-sample.ts delete mode 100644 packages/kernel-language-model-service/src/test-utils/queue/README.md delete mode 100644 packages/kernel-language-model-service/src/test-utils/queue/model.test.ts delete mode 100644 packages/kernel-language-model-service/src/test-utils/queue/model.ts delete mode 100644 packages/kernel-language-model-service/src/test-utils/queue/response.test.ts delete mode 100644 packages/kernel-language-model-service/src/test-utils/queue/response.ts delete mode 100644 packages/kernel-language-model-service/src/test-utils/queue/service.test.ts delete mode 100644 packages/kernel-language-model-service/src/test-utils/queue/service.ts delete mode 100644 packages/kernel-language-model-service/src/test-utils/queue/tokenizer.test.ts delete mode 100644 packages/kernel-language-model-service/src/test-utils/queue/tokenizer.ts delete mode 100644 packages/kernel-language-model-service/src/test-utils/queue/utils.test.ts delete mode 100644 packages/kernel-language-model-service/src/test-utils/queue/utils.ts create mode 100644 packages/kernel-test/src/lms-chat.test.ts create mode 100644 packages/kernel-test/src/lms-sample.test.ts delete mode 100644 packages/kernel-test/src/lms-user.test.ts create mode 100644 packages/kernel-test/src/vats/lms-chat-vat.ts delete mode 100644 packages/kernel-test/src/vats/lms-queue-vat.ts create mode 100644 packages/kernel-test/src/vats/lms-sample-vat.ts delete mode 100644 packages/kernel-test/src/vats/lms-user-vat.ts diff --git a/packages/kernel-language-model-service/package.json b/packages/kernel-language-model-service/package.json index 30e3c4feae..48a23ea2a9 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 0000000000..22dd6febee --- /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 0000000000..24973d6daa --- /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 c20e32944a..1b31cdf675 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 0000000000..a6ff8712f1 --- /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 0000000000..9b6e2b661b --- /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 37b9234224..064c631079 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 f0bdc2e14d..ed7d3c37a1 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 65f7e8289d..3630932d21 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 0000000000..6ecb36c8b5 --- /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 0000000000..54a396843f --- /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 0000000000..a1eb20ac05 --- /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 0000000000..997dd8c991 --- /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 9fc946e73c..8fcf6100e8 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 0000000000..a122eb8a42 --- /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 0000000000..454b33a0cf --- /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 9b16967a87..0000000000 --- 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 479fd54ba7..0000000000 --- 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 c30ebe0176..0000000000 --- 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 3d4cb7ac25..0000000000 --- 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 d7fee96829..0000000000 --- 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 9a107305a3..0000000000 --- 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 c2ec5269a3..0000000000 --- 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 89512dd22c..0000000000 --- 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 f24213e113..0000000000 --- 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 8ea4a4bbed..0000000000 --- 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 5d72694c03..0000000000 --- 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 5895986c41..aafaabddd5 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/src/lms-chat.test.ts b/packages/kernel-test/src/lms-chat.test.ts new file mode 100644 index 0000000000..c655b61b9d --- /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 0000000000..999f152a96 --- /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 00200ae1d2..0000000000 --- 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 0000000000..5831bcb86b --- /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 e7872f1a68..0000000000 --- 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 0000000000..e95ed0901a --- /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 4a0c1b0e53..0000000000 --- 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 9ba23ce663..e5a93d6b52 100644 --- a/yarn.lock +++ b/yarn.lock @@ -3950,6 +3950,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" @@ -3991,10 +3992,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:^" From 3827d5398261fff6c837f45b050d9abe34fbf137 Mon Sep 17 00:00:00 2001 From: grypez <143971198+grypez@users.noreply.github.com> Date: Mon, 30 Mar 2026 14:50:42 -0400 Subject: [PATCH 02/20] feat(klms): add listModels to OpenV1BaseService --- .../package.json | 10 +++++ .../src/open-v1/base.test.ts | 38 +++++++++++++++++++ .../src/open-v1/base.ts | 13 +++++++ .../src/open-v1/nodejs.ts | 2 + tsconfig.packages.json | 3 ++ 5 files changed, 66 insertions(+) diff --git a/packages/kernel-language-model-service/package.json b/packages/kernel-language-model-service/package.json index 48a23ea2a9..e32b821cfe 100644 --- a/packages/kernel-language-model-service/package.json +++ b/packages/kernel-language-model-service/package.json @@ -23,6 +23,16 @@ "default": "./dist/index.cjs" } }, + "./open-v1/nodejs": { + "import": { + "types": "./dist/open-v1/nodejs.d.mts", + "default": "./dist/open-v1/nodejs.mjs" + }, + "require": { + "types": "./dist/open-v1/nodejs.d.cts", + "default": "./dist/open-v1/nodejs.cjs" + } + }, "./ollama/nodejs": { "import": { "types": "./dist/ollama/nodejs.d.mts", 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 index 6ecb36c8b5..f1d0c6ff12 100644 --- a/packages/kernel-language-model-service/src/open-v1/base.test.ts +++ b/packages/kernel-language-model-service/src/open-v1/base.test.ts @@ -173,6 +173,44 @@ describe('OpenV1BaseService', () => { }); }); + describe('listModels', () => { + it('gETs /v1/models and returns model IDs', async () => { + const modelsFetch = makeMockFetch({ + data: [{ id: 'model-a' }, { id: 'model-b' }], + }); + const modelsService = new OpenV1BaseService( + modelsFetch, + 'http://localhost:8080', + 'sk-test', + ); + + const result = await modelsService.listModels(); + + expect(modelsFetch).toHaveBeenCalledWith( + 'http://localhost:8080/v1/models', + expect.any(Object), + ); + expect(result).toStrictEqual(['model-a', 'model-b']); + }); + + it('includes Authorization header when apiKey is set', async () => { + const modelsFetch = makeMockFetch({ data: [{ id: 'model-a' }] }); + const modelsService = new OpenV1BaseService( + modelsFetch, + 'http://localhost:8080', + 'sk-key', + ); + + await modelsService.listModels(); + + const [, init] = (modelsFetch as ReturnType).mock + .calls[0] as [string, RequestInit]; + expect((init.headers as Record).Authorization).toBe( + 'Bearer sk-key', + ); + }); + }); + describe('chat with stream: true', () => { it('pOSTs to /v1/chat/completions with stream: true in body', async () => { const streamFetch = makeMockStreamFetch([makeStreamChunk('hi')]); diff --git a/packages/kernel-language-model-service/src/open-v1/base.ts b/packages/kernel-language-model-service/src/open-v1/base.ts index 54a396843f..91ccc0e22e 100644 --- a/packages/kernel-language-model-service/src/open-v1/base.ts +++ b/packages/kernel-language-model-service/src/open-v1/base.ts @@ -124,6 +124,19 @@ export class OpenV1BaseService { } } + /** + * Lists models available at the `/v1/models` endpoint. + * + * @returns A promise resolving to an array of model ID strings. + */ + async listModels(): Promise { + const response = await this.#fetch(`${this.#baseUrl}/v1/models`, { + headers: this.#makeHeaders(), + }); + const data = (await response.json()) as { data: { id: string }[] }; + return harden(data.data.map((model) => model.id)); + } + /** * @returns Headers for the request, including Authorization if an API key is set. */ diff --git a/packages/kernel-language-model-service/src/open-v1/nodejs.ts b/packages/kernel-language-model-service/src/open-v1/nodejs.ts index a1eb20ac05..4b2e290beb 100644 --- a/packages/kernel-language-model-service/src/open-v1/nodejs.ts +++ b/packages/kernel-language-model-service/src/open-v1/nodejs.ts @@ -24,6 +24,7 @@ export const makeOpenV1NodejsService = (config: { (params: ChatParams & { stream: true }): AsyncIterable; (params: ChatParams & { stream?: false }): Promise; }; + listModels: () => Promise; } => { const { endowments, baseUrl, apiKey } = config; if (!endowments?.fetch) { @@ -35,5 +36,6 @@ export const makeOpenV1NodejsService = (config: { (params: ChatParams & { stream: true }): AsyncIterable; (params: ChatParams & { stream?: false }): Promise; }, + listModels: service.listModels.bind(service), }); }; diff --git a/tsconfig.packages.json b/tsconfig.packages.json index bc9a147e46..78e532c3bb 100644 --- a/tsconfig.packages.json +++ b/tsconfig.packages.json @@ -9,6 +9,9 @@ */ "paths": { "@metamask/*": ["../*/src"], + "@ocap/kernel-language-model-service/open-v1/nodejs": [ + "../kernel-language-model-service/src/open-v1/nodejs.ts" + ], "@metamask/kernel-store/sqlite/nodejs": [ "../kernel-store/src/sqlite/nodejs.ts" ], From bf86032f67547e7fe012e504f68cbbdf0a7397e4 Mon Sep 17 00:00:00 2001 From: grypez <143971198+grypez@users.noreply.github.com> Date: Mon, 30 Mar 2026 15:27:07 -0400 Subject: [PATCH 03/20] feat(klms): add streaming overload to OllamaBaseService.sample --- .../package.json | 1 + .../src/ollama/base.test.ts | 65 +++++++++ .../src/ollama/base.ts | 129 ++++++++++++++---- .../src/ollama/nodejs.ts | 11 +- .../tsconfig.build.json | 5 +- .../tsconfig.json | 6 +- .../sample-agent.e2e.test.ts} | 0 yarn.lock | 1 + 8 files changed, 186 insertions(+), 32 deletions(-) rename packages/kernel-test-local/{test/e2e/agents.test.ts => src/sample-agent.e2e.test.ts} (100%) diff --git a/packages/kernel-language-model-service/package.json b/packages/kernel-language-model-service/package.json index e32b821cfe..764f7d7017 100644 --- a/packages/kernel-language-model-service/package.json +++ b/packages/kernel-language-model-service/package.json @@ -115,6 +115,7 @@ }, "dependencies": { "@endo/eventual-send": "^1.3.4", + "@metamask/kernel-utils": "workspace:^", "@metamask/superstruct": "^3.2.1", "ollama": "^0.5.16", "ses": "^1.14.0" diff --git a/packages/kernel-language-model-service/src/ollama/base.test.ts b/packages/kernel-language-model-service/src/ollama/base.test.ts index e1819b84b1..e2262a4fa4 100644 --- a/packages/kernel-language-model-service/src/ollama/base.test.ts +++ b/packages/kernel-language-model-service/src/ollama/base.test.ts @@ -69,6 +69,71 @@ describe('OllamaBaseService', () => { }); }); + describe('sample', () => { + const mockChunk = { + response: 'hello', + done: false, + } as GenerateResponse; + + it('gETs non-streaming response when stream is omitted', async () => { + vi.mocked(mockClient.generate).mockResolvedValue({ + response: 'hello world', + } as GenerateResponse); + + const result = await service.sample({ + model: 'llama2:7b', + prompt: 'Hi', + }); + + expect(mockClient.generate).toHaveBeenCalledWith( + expect.objectContaining({ stream: false, raw: true }), + ); + expect(result).toStrictEqual({ text: 'hello world' }); + }); + + it('passes options sub-object to generate', async () => { + vi.mocked(mockClient.generate).mockResolvedValue({ + response: 'hi', + } as GenerateResponse); + + await service.sample({ + model: 'llama2:7b', + prompt: 'Hi', + temperature: 0.5, + seed: 42, + }); + + expect(mockClient.generate).toHaveBeenCalledWith( + expect.objectContaining({ + options: expect.objectContaining({ temperature: 0.5, seed: 42 }), + }), + ); + }); + + it('returns streaming result when stream is true', async () => { + vi.mocked(mockClient.generate).mockResolvedValue( + makeMockAbortableAsyncIterator([mockChunk]), + ); + + const { stream, abort } = await service.sample({ + model: 'llama2:7b', + prompt: 'Hi', + stream: true, + }); + + const chunks: GenerateResponse[] = []; + for await (const chunk of stream) { + chunks.push(chunk); + } + await abort(); + + expect(mockClient.generate).toHaveBeenCalledWith( + expect.objectContaining({ stream: true, raw: true }), + ); + expect(chunks).toStrictEqual([mockChunk]); + }); + }); + describe('makeInstance', () => { it('should create instance with model', async () => { const config: InstanceConfig = { diff --git a/packages/kernel-language-model-service/src/ollama/base.ts b/packages/kernel-language-model-service/src/ollama/base.ts index 064c631079..9a36046eb3 100644 --- a/packages/kernel-language-model-service/src/ollama/base.ts +++ b/packages/kernel-language-model-service/src/ollama/base.ts @@ -1,4 +1,9 @@ -import type { GenerateResponse, ListResponse } from 'ollama'; +import { ifDefined } from '@metamask/kernel-utils'; +import type { + AbortableAsyncIterator, + GenerateResponse, + ListResponse, +} from 'ollama'; import type { ChatParams, @@ -16,6 +21,14 @@ import type { OllamaModelOptions, } from './types.ts'; +/** + * Result of a streaming raw token-prediction request. + */ +export type SampleStreamResult = { + stream: AsyncIterable; + abort: () => Promise; +}; + /** * Base service for interacting with Ollama language models. * Provides a generic interface for creating and managing Ollama model instances. @@ -70,15 +83,13 @@ export class OllamaBaseService 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 }), - }, + options: ifDefined({ + temperature, + top_p: params.top_p, + seed, + num_predict: params.max_tokens, + stop: stopArr, + }), }); const promptTokens = response.prompt_eval_count ?? 0; const completionTokens = response.eval_count ?? 0; @@ -107,34 +118,96 @@ export class OllamaBaseService * Performs a raw token-prediction request via Ollama's generate API with raw=true, * bypassing the model's chat template. * + * When `params.stream` is `true`, returns a streaming result with an async + * iterable of {@link GenerateResponse} chunks and an abort handle. + * When `params.stream` is `false` or omitted, awaits and returns the full + * {@link SampleResult}. + * * @param params - The raw sample parameters. - * @returns A hardened raw sample result. + * @returns A streaming result when `stream: true`, or the full result otherwise. */ - 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]; + sample(params: SampleParams & { stream: true }): Promise; + + /** + * @param params - The raw sample parameters. + * @returns A promise resolving to the full sample result. + */ + sample(params: SampleParams & { stream?: false }): Promise; + + /** + * @param params - The raw sample parameters. + * @returns A streaming result or full result depending on `params.stream`. + */ + async sample( + params: SampleParams, + ): Promise | Promise { + if (params.stream === true) { + return this.#streamingSample(params); } + return this.#nonStreamingSample(params); + } + + /** + * @param params - The raw sample parameters. + * @returns A promise resolving to the full sample result. + */ + async #nonStreamingSample(params: SampleParams): Promise { + const ollama = await this.#makeClient(); const response = await ollama.generate({ - model, - prompt, + model: params.model, + prompt: params.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 }), - }, + options: this.#buildSampleOptions(params), }); return harden({ text: response.response }); } + /** + * @param params - The raw sample parameters. + * @returns A promise resolving to a streaming result. + */ + async #streamingSample(params: SampleParams): Promise { + const ollama = await this.#makeClient(); + const response: AbortableAsyncIterator = + await ollama.generate({ + model: params.model, + prompt: params.prompt, + raw: true, + stream: true, + options: this.#buildSampleOptions(params), + }); + return harden({ + stream: (async function* () { + for await (const chunk of response) { + yield harden(chunk); + } + })(), + abort: async () => response.abort(), + }); + } + + /** + * @param params - The raw sample parameters. + * @returns The options sub-object for an Ollama generate request. + */ + #buildSampleOptions(params: SampleParams): Record { + const { temperature, seed } = params; + let stopArr: string[] | undefined; + if (params.stop !== undefined) { + stopArr = Array.isArray(params.stop) ? params.stop : [params.stop]; + } + return harden( + ifDefined({ + temperature, + top_p: params.top_p, + seed, + num_predict: params.max_tokens, + stop: stopArr, + }), + ); + } + /** * Creates a new language model instance with the specified configuration. * The returned instance is hardened for object capability security. diff --git a/packages/kernel-language-model-service/src/ollama/nodejs.ts b/packages/kernel-language-model-service/src/ollama/nodejs.ts index ed7d3c37a1..62f514f5c4 100644 --- a/packages/kernel-language-model-service/src/ollama/nodejs.ts +++ b/packages/kernel-language-model-service/src/ollama/nodejs.ts @@ -7,6 +7,7 @@ import type { SampleResult, } from '../types.ts'; import { OllamaBaseService } from './base.ts'; +import type { SampleStreamResult } from './base.ts'; import { defaultClientConfig } from './constants.ts'; import type { OllamaClient, OllamaNodejsConfig } from './types.ts'; @@ -52,11 +53,17 @@ export const makeOllamaNodejsKernelService = ( config: OllamaNodejsConfig, ): { chat: (params: ChatParams) => Promise; - sample: (params: SampleParams) => Promise; + sample: { + (params: SampleParams & { stream: true }): Promise; + (params: SampleParams & { stream?: false }): Promise; + }; } => { const service = new OllamaNodejsService(config); return harden({ chat: async (params: ChatParams) => service.chat(params), - sample: async (params: SampleParams) => service.sample(params), + sample: service.sample.bind(service) as { + (params: SampleParams & { stream: true }): Promise; + (params: SampleParams & { stream?: false }): Promise; + }, }); }; diff --git a/packages/kernel-language-model-service/tsconfig.build.json b/packages/kernel-language-model-service/tsconfig.build.json index 960e113284..ee6e5d2743 100644 --- a/packages/kernel-language-model-service/tsconfig.build.json +++ b/packages/kernel-language-model-service/tsconfig.build.json @@ -7,7 +7,10 @@ "rootDir": "./src", "types": ["chrome", "ses"] }, - "references": [{ "path": "../streams/tsconfig.build.json" }], + "references": [ + { "path": "../kernel-utils/tsconfig.build.json" }, + { "path": "../streams/tsconfig.build.json" } + ], "files": [], "include": ["./src"] } diff --git a/packages/kernel-language-model-service/tsconfig.json b/packages/kernel-language-model-service/tsconfig.json index 430f58857d..dfdb69c630 100644 --- a/packages/kernel-language-model-service/tsconfig.json +++ b/packages/kernel-language-model-service/tsconfig.json @@ -5,7 +5,11 @@ "lib": ["ES2022", "DOM"], "types": ["vitest", "chrome", "ses"] }, - "references": [{ "path": "../repo-tools" }, { "path": "../streams" }], + "references": [ + { "path": "../kernel-utils" }, + { "path": "../repo-tools" }, + { "path": "../streams" } + ], "include": [ "../../vitest.config.ts", "./src", diff --git a/packages/kernel-test-local/test/e2e/agents.test.ts b/packages/kernel-test-local/src/sample-agent.e2e.test.ts similarity index 100% rename from packages/kernel-test-local/test/e2e/agents.test.ts rename to packages/kernel-test-local/src/sample-agent.e2e.test.ts diff --git a/yarn.lock b/yarn.lock index e5a93d6b52..af1601d16b 100644 --- a/yarn.lock +++ b/yarn.lock @@ -3955,6 +3955,7 @@ __metadata: "@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-utils": "workspace:^" "@metamask/streams": "workspace:^" "@metamask/superstruct": "npm:^3.2.1" "@ocap/repo-tools": "workspace:^" From 67faee0c2ac0c82299a81111c1517cadc2e76f96 Mon Sep 17 00:00:00 2001 From: grypez <143971198+grypez@users.noreply.github.com> Date: Tue, 17 Mar 2026 15:23:42 -0400 Subject: [PATCH 04/20] test(kernel-test-local): kernel-LMS integration tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add an integration test and e2e test that exercise the full kernel → LMS service → Ollama round-trip via a bundled vat. - lms-chat-vat: bundled vat that sends a single chat message and logs the response, used to verify the round-trip through the kernel - lms-chat.ts: shared test helper (runLmsChatKernelTest) - lms-chat.test.ts: uses a mock fetch — CI-safe, no network - lms-chat.e2e.test.ts: uses real fetch against local Ollama - agents.e2e.test.ts: json/repl agent e2e tests against local Ollama - test/suite.test.ts: pre-flight check that Ollama is running Lay out the package to match kernel-test: all vats, helpers, and test files live under src/; test/ holds only the Ollama pre-flight suite. Co-Authored-By: Claude Sonnet 4.6 --- packages/kernel-test-local/package.json | 10 ++- packages/kernel-test-local/src/constants.ts | 4 +- .../src/lms-chat.e2e.test.ts | 31 ++++++++ .../kernel-test-local/src/lms-chat.test.ts | 18 +++++ packages/kernel-test-local/src/lms-chat.ts | 73 +++++++++++++++++++ .../src/sample-agent.e2e.test.ts | 38 +++++----- packages/kernel-test-local/src/utils.ts | 2 +- .../src/{utils.test.ts => utils.unit.test.ts} | 0 .../src/vats/lms-chat-vat.ts | 41 +++++++++++ .../kernel-test-local/test/e2e/suite.test.ts | 48 ------------ packages/kernel-test-local/test/suite.test.ts | 40 ++++++++++ packages/kernel-test-local/tsconfig.json | 7 +- .../kernel-test-local/vitest.config.e2e.ts | 4 +- packages/kernel-test-local/vitest.config.ts | 1 + 14 files changed, 241 insertions(+), 76 deletions(-) create mode 100644 packages/kernel-test-local/src/lms-chat.e2e.test.ts create mode 100644 packages/kernel-test-local/src/lms-chat.test.ts create mode 100644 packages/kernel-test-local/src/lms-chat.ts rename packages/kernel-test-local/src/{utils.test.ts => utils.unit.test.ts} (100%) create mode 100644 packages/kernel-test-local/src/vats/lms-chat-vat.ts delete mode 100644 packages/kernel-test-local/test/e2e/suite.test.ts create mode 100644 packages/kernel-test-local/test/suite.test.ts diff --git a/packages/kernel-test-local/package.json b/packages/kernel-test-local/package.json index 70d2cd9391..085cabbcbc 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/src/constants.ts b/packages/kernel-test-local/src/constants.ts index 86b329e2e0..a096d674a1 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 0000000000..9493e44083 --- /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 { LMS_BASE_URL } from './constants.ts'; +import { runLmsChatKernelTest } from './lms-chat.ts'; + +describe('lms-kernel (e2e)', () => { + beforeAll(() => { + fetchMock.disableMocks(); + }); + + afterAll(() => { + fetchMock.enableMocks(); + }); + + // eslint-disable-next-line vitest/expect-expect + it( + 'sends a chat message through the kernel and receives a response', + { timeout: 60_000 }, + async () => { + const { chat } = makeOpenV1NodejsService({ + endowments: { fetch }, + baseUrl: LMS_BASE_URL, + }); + 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 0000000000..cc484268bb --- /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 0000000000..b5b27f23ba --- /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/sample-agent.e2e.test.ts b/packages/kernel-test-local/src/sample-agent.e2e.test.ts index 7ab6842861..40dff73b7d 100644 --- a/packages/kernel-test-local/src/sample-agent.e2e.test.ts +++ b/packages/kernel-test-local/src/sample-agent.e2e.test.ts @@ -6,7 +6,7 @@ import { getMoonPhase } from '@ocap/kernel-agents/capabilities/examples'; import { count, add, multiply } from '@ocap/kernel-agents/capabilities/math'; import { makeJsonAgent } from '@ocap/kernel-agents/json'; import { makeReplAgent } from '@ocap/kernel-agents-repl'; -import { OllamaNodejsService } from '@ocap/kernel-language-model-service/ollama/nodejs'; +import { makeOllamaNodejsKernelService } from '@ocap/kernel-language-model-service/ollama/nodejs'; import { fetchMock } from '@ocap/repo-tools/test-utils/fetch-mock'; import { afterAll, @@ -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'], @@ -72,10 +72,10 @@ describe.each([ fetchMock.enableMocks(); }); - let languageModelService: OllamaNodejsService; + let service: ReturnType; beforeEach(() => { result = undefined; - languageModelService = new OllamaNodejsService({ endowments: { fetch } }); + service = makeOllamaNodejsKernelService({ endowments: { fetch } }); printLogger.log(`\n<== New ${strategy.toUpperCase()} ===`); }); @@ -85,13 +85,19 @@ describe.each([ printLogger.log(`=== End ${strategy.toUpperCase()} ==>`); }); + const makeLanguageModel = (svc: typeof service) => ({ + getInfo: async () => ({ model: DEFAULT_MODEL, options: {} }), + load: async () => undefined, + unload: async () => undefined, + sample: async (prompt: string) => + svc.sample({ model: DEFAULT_MODEL, prompt, stream: true }), + }); + it( 'processes a semantic request', { retry, timeout }, catchErrorAsResult(async () => { - const languageModel = await languageModelService.makeInstance({ - model: DEFAULT_MODEL, - }); + const languageModel = makeLanguageModel(service); const agent = makeAgent({ languageModel, capabilities: {}, logger }); expect(agent).toBeDefined(); @@ -121,9 +127,7 @@ describe.each([ 'uses tools', { retry, timeout }, catchErrorAsResult(async () => { - const languageModel = await languageModelService.makeInstance({ - model: DEFAULT_MODEL, - }); + const languageModel = makeLanguageModel(service); const getMoonPhaseSpy = vi.spyOn(getMoonPhase, 'func'); const agent = makeAgent({ languageModel, @@ -143,9 +147,7 @@ describe.each([ 'performs multi-step calculations', { retry, timeout }, catchErrorAsResult(async () => { - const languageModel = await languageModelService.makeInstance({ - model: DEFAULT_MODEL, - }); + const languageModel = makeLanguageModel(service); const capabilities = {}; const agent = makeAgent({ languageModel, capabilities, logger }); expect(agent).toBeDefined(); @@ -162,9 +164,7 @@ describe.each([ // Caveat: We don't expect the solution to be correct. { retry, timeout: 120_000 }, catchErrorAsResult(async () => { - const languageModel = await languageModelService.makeInstance({ - model: DEFAULT_MODEL, - }); + const languageModel = makeLanguageModel(service); const capabilities = {}; const agent = makeAgent({ languageModel, capabilities, logger }); expect(agent).toBeDefined(); @@ -188,9 +188,7 @@ describe.each([ { retry, timeout }, // TODO: This functionality is not yet implemented. catchErrorAsResult(async () => { - const languageModel = await languageModelService.makeInstance({ - model: DEFAULT_MODEL, - }); + const languageModel = makeLanguageModel(service); const capabilities = {}; const agent = makeAgent({ languageModel, capabilities, logger }); expect(agent).toBeDefined(); diff --git a/packages/kernel-test-local/src/utils.ts b/packages/kernel-test-local/src/utils.ts index 2192008301..74584cf5cb 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 0000000000..2c01356e82 --- /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/e2e/suite.test.ts deleted file mode 100644 index c6c7cee816..0000000000 --- a/packages/kernel-test-local/test/e2e/suite.test.ts +++ /dev/null @@ -1,48 +0,0 @@ -/** - * Pre-test verification suite that checks: - * - * - Ollama service is running and accessible - * - Required models are available - * - * These tests run sequentially and must pass before the main test suite. - */ -import { fetchMock } from '@ocap/repo-tools/test-utils/fetch-mock'; -import { afterAll, beforeAll, describe, expect, it } from 'vitest'; - -import { - DEFAULT_MODEL, - OLLAMA_API_BASE, - OLLAMA_TAGS_ENDPOINT, -} from '../../src/constants.ts'; - -describe.sequential('test suite', () => { - beforeAll(() => { - fetchMock.disableMocks(); - }); - - afterAll(() => { - fetchMock.enableMocks(); - }); - - it(`connects to Ollama instance`, async () => { - const response = await fetch(OLLAMA_API_BASE); - expect(response.ok).toBe(true); - }); - - it(`can access ${DEFAULT_MODEL} model`, async () => { - const response = await fetch(OLLAMA_TAGS_ENDPOINT); - expect(response.ok).toBe(true); - - const data = (await response.json()) as { - models: { name: string }[]; - }; - expect(data?.models).toBeDefined(); - expect(Array.isArray(data.models)).toBe(true); - - const llamaModel = data.models.find( - (foundModel: { name: string }) => foundModel.name === DEFAULT_MODEL, - ); - expect(llamaModel).toBeDefined(); - expect(llamaModel?.name).toBe(DEFAULT_MODEL); - }); -}); diff --git a/packages/kernel-test-local/test/suite.test.ts b/packages/kernel-test-local/test/suite.test.ts new file mode 100644 index 0000000000..3938565f0c --- /dev/null +++ b/packages/kernel-test-local/test/suite.test.ts @@ -0,0 +1,40 @@ +/** + * Pre-test verification suite that checks: + * + * - Configured LMS service is running and accessible + * - Required model is available + * + * These tests run sequentially and must pass before the main test suite. + */ +import '@ocap/repo-tools/test-utils/mock-endoify'; + +import { makeOpenV1NodejsService } from '@ocap/kernel-language-model-service/open-v1/nodejs'; +import { fetchMock } from '@ocap/repo-tools/test-utils/fetch-mock'; +import { afterAll, beforeAll, describe, expect, it } from 'vitest'; + +import { LMS_BASE_URL, LMS_CHAT_MODEL } from '../src/constants.ts'; + +describe.sequential('test suite', () => { + let service: ReturnType; + + beforeAll(() => { + fetchMock.disableMocks(); + service = makeOpenV1NodejsService({ + endowments: { fetch }, + baseUrl: LMS_BASE_URL, + }); + }); + + afterAll(() => { + fetchMock.enableMocks(); + }); + + it('connects to LMS service', async () => { + expect(await service.listModels()).toBeDefined(); + }); + + it(`can access ${LMS_CHAT_MODEL} model`, async () => { + const models = await service.listModels(); + expect(models).toContain(LMS_CHAT_MODEL); + }); +}); diff --git a/packages/kernel-test-local/tsconfig.json b/packages/kernel-test-local/tsconfig.json index 8539255afc..cb233df630 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 97c0ce5cac..ef76f427c1 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 6eee6669cf..0586515a02 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'], }, }), ); From dbbe997d28701e55e851b3611b0149830c85c6a5 Mon Sep 17 00:00:00 2001 From: grypez <143971198+grypez@users.noreply.github.com> Date: Mon, 30 Mar 2026 08:46:04 -0400 Subject: [PATCH 05/20] feat(kernel-agents): add makeChatAgent for chat-completion-based agents Adds a chat-based capability-augmented agent alongside the existing text-completion JSON agent. Uses ChatService.chat() with a messages[] array instead of LanguageModel.sample() with a raw prompt string. - makeChatAgent({ chat, capabilities }) implements the Agent interface - Capabilities are described to the model via a JSON system prompt - Model signals completion by invoking the auto-injected 'end' capability - Capability results are injected as 'user' turns for model-provider compatibility - Exported as @ocap/kernel-agents/chat, mirroring @ocap/kernel-agents/json Usage: const client = makeChatClient(lmsServiceRef, 'qwen2.5:0.5b'); const agent = makeChatAgent({ chat: (messages) => client.chat.completions.create({ messages }), capabilities: { walletBalance, walletSend }, }); const result = await agent.task('what is my balance?'); Co-Authored-By: Claude Sonnet 4.6 --- packages/kernel-agents/package.json | 10 + .../src/strategies/chat-agent.test.ts | 171 +++++++++++++ .../src/strategies/chat-agent.ts | 235 ++++++++++++++++++ 3 files changed, 416 insertions(+) create mode 100644 packages/kernel-agents/src/strategies/chat-agent.test.ts create mode 100644 packages/kernel-agents/src/strategies/chat-agent.ts diff --git a/packages/kernel-agents/package.json b/packages/kernel-agents/package.json index 0b7d4bc9a6..2f8c290725 100644 --- a/packages/kernel-agents/package.json +++ b/packages/kernel-agents/package.json @@ -33,6 +33,16 @@ "default": "./dist/strategies/json-agent.cjs" } }, + "./chat": { + "import": { + "types": "./dist/strategies/chat-agent.d.mts", + "default": "./dist/strategies/chat-agent.mjs" + }, + "require": { + "types": "./dist/strategies/chat-agent.d.cts", + "default": "./dist/strategies/chat-agent.cjs" + } + }, "./agent": { "import": { "types": "./dist/agent.d.mts", diff --git a/packages/kernel-agents/src/strategies/chat-agent.test.ts b/packages/kernel-agents/src/strategies/chat-agent.test.ts new file mode 100644 index 0000000000..9fac5cd6de --- /dev/null +++ b/packages/kernel-agents/src/strategies/chat-agent.test.ts @@ -0,0 +1,171 @@ +import '@ocap/repo-tools/test-utils/mock-endoify'; + +import { describe, expect, it, vi } from 'vitest'; + +import { makeChatAgent } from './chat-agent.ts'; +import type { BoundChat } from './chat-agent.ts'; +import { capability } from '../capabilities/capability.ts'; + +const makeChat = (responses: string[]): BoundChat => { + let call = 0; + return async () => { + const index = call; + call += 1; + return { + id: String(index), + model: 'test', + choices: [ + { + message: { + role: 'assistant' as const, + content: responses[index] ?? '', + }, + index: 0, + finish_reason: 'stop', + }, + ], + usage: { prompt_tokens: 0, completion_tokens: 0, total_tokens: 0 }, + }; + }; +}; + +const noCapabilities = {}; + +describe('makeChatAgent', () => { + it('returns plain text response when model does not invoke a capability', async () => { + const chat = makeChat(['Hello, world!']); + const agent = makeChatAgent({ chat, capabilities: noCapabilities }); + + const result = await agent.task('say hello'); + expect(result).toBe('Hello, world!'); + }); + + it('returns result when model invokes end capability', async () => { + const chat = makeChat([ + '{"name": "end", "args": {"final": "the answer is 42"}}', + ]); + const agent = makeChatAgent({ chat, capabilities: noCapabilities }); + + const result = await agent.task('what is the answer?'); + expect(result).toBe('the answer is 42'); + }); + + it('dispatches a user capability and continues to end', async () => { + const add = vi.fn(async ({ a, b }: { a: number; b: number }) => a + b); + const addCap = capability(add, { + description: 'Add two numbers', + args: { + a: { type: 'number' }, + b: { type: 'number' }, + }, + returns: { type: 'number' }, + }); + + const chat = makeChat([ + '{"name": "add", "args": {"a": 3, "b": 4}}', + '{"name": "end", "args": {"final": "7"}}', + ]); + const agent = makeChatAgent({ + chat, + capabilities: { add: addCap }, + }); + + const result = await agent.task('add 3 and 4'); + expect(add).toHaveBeenCalledWith({ a: 3, b: 4 }); + expect(result).toBe('7'); + }); + + it('injects tool result into messages before next turn', async () => { + const messages: string[][] = []; + const chat: BoundChat = async (chatMsgs) => { + messages.push(chatMsgs.map((chatMsg) => chatMsg.content)); + const turn = messages.length - 1; + return { + id: String(turn), + model: 'test', + choices: [ + { + message: { + role: 'assistant' as const, + content: + turn === 0 + ? '{"name": "ping", "args": {}}' + : '{"name": "end", "args": {"final": "done"}}', + }, + index: 0, + finish_reason: 'stop', + }, + ], + usage: { prompt_tokens: 0, completion_tokens: 0, total_tokens: 0 }, + }; + }; + + const ping = capability(async () => 'pong', { + description: 'Ping', + args: {}, + returns: { type: 'string' }, + }); + const agent = makeChatAgent({ chat, capabilities: { ping } }); + await agent.task('ping'); + + // Second turn messages should include the tool result + expect( + messages[1]?.some((content) => content.includes('[Result of ping]')), + ).toBe(true); + }); + + it('appends error message for unknown capability and continues', async () => { + const chat = makeChat([ + '{"name": "nonexistent", "args": {}}', + '{"name": "end", "args": {"final": "recovered"}}', + ]); + const agent = makeChatAgent({ chat, capabilities: noCapabilities }); + + const result = await agent.task('do something'); + expect(result).toBe('recovered'); + }); + + it('throws when invocation budget is exceeded', async () => { + // Always invokes a capability but never ends + const ping = capability(async () => 'pong', { + description: 'Ping', + args: {}, + }); + const neverEnd = makeChat( + Array.from({ length: 20 }, () => '{"name": "ping", "args": {}}'), + ); + const agent = makeChatAgent({ + chat: neverEnd, + capabilities: { ping }, + }); + + await expect( + agent.task('go', undefined, { invocationBudget: 3 }), + ).rejects.toThrow('Invocation budget exceeded'); + }); + + it('applies judgment to end result', async () => { + const chat = makeChat(['{"name": "end", "args": {"final": 99}}']); + const agent = makeChatAgent({ chat, capabilities: noCapabilities }); + + const isString = (result: unknown): result is string => + typeof result === 'string'; + await expect(agent.task('go', isString)).rejects.toThrow('Invalid result'); + }); + + it('accumulates experiences across tasks', async () => { + const chat = makeChat(['hello', 'world']); + const agent = makeChatAgent({ chat, capabilities: noCapabilities }); + + await agent.task('first'); + await agent.task('second'); + + const exps = []; + for await (const exp of agent.experiences) { + exps.push(exp); + } + expect(exps).toHaveLength(2); + expect(exps[0]?.objective.intent).toBe('first'); + expect(exps[1]?.objective.intent).toBe('second'); + }); +}); diff --git a/packages/kernel-agents/src/strategies/chat-agent.ts b/packages/kernel-agents/src/strategies/chat-agent.ts new file mode 100644 index 0000000000..6b10c39026 --- /dev/null +++ b/packages/kernel-agents/src/strategies/chat-agent.ts @@ -0,0 +1,235 @@ +import { mergeDisjointRecords } from '@metamask/kernel-utils'; +import type { Logger } from '@metamask/logger'; +import type { + ChatMessage, + ChatResult, +} from '@ocap/kernel-language-model-service'; + +import { + extractCapabilitySchemas, + extractCapabilities, +} from '../capabilities/capability.ts'; +import { makeEnd } from '../capabilities/end.ts'; +import type { Agent } from '../types/agent.ts'; +import { Message } from '../types/messages.ts'; +import type { CapabilityRecord, Experience } from '../types.ts'; + +/** + * Adapts a raw {@link ChatMessage} into the typed {@link Message} hierarchy + * so that chat turns can be recorded in {@link Experience.history}. + */ +class ChatTurn extends Message { + /** + * @param chatMessage - The raw chat message to wrap. + * @param chatMessage.role - The sender role of the message. + * @param chatMessage.content - The text content of the message. + */ + constructor({ role, content }: ChatMessage) { + super(role, { content }); + } +} + +/** + * A bound chat function with the model already configured. + * Construct one from a {@link ChatService} using `makeChatClient`: + * + * ```ts + * const client = makeChatClient(serviceRef, model); + * const chat = (messages) => client.chat.completions.create({ messages }); + * ``` + */ +export type BoundChat = (messages: ChatMessage[]) => Promise; + +export type MakeChatAgentArgs = { + /** + * Bound chat function — model is pre-configured by the caller. + * + * @see {@link BoundChat} + */ + chat: BoundChat; + /** + * Capabilities the agent may invoke, expressed as a {@link CapabilityRecord}. + * An `end` capability is automatically added to signal task completion. + */ + capabilities: CapabilityRecord; +}; + +/** + * Build the system prompt that instructs the model to invoke capabilities + * by responding with JSON objects. + * + * @param capSchemas - Serialized capability schemas. + * @returns The system prompt string. + */ +function buildSystemPrompt(capSchemas: Record): string { + return [ + 'You are a capability-augmented assistant.', + 'To invoke a capability, respond with ONLY a JSON object:', + ' {"name": "", "args": {}}', + 'Do not include any other text when invoking a capability.', + '', + 'Available capabilities:', + JSON.stringify(capSchemas, null, 2), + '', + 'When you have a final answer, invoke the "end" capability:', + ' {"name": "end", "args": {"final": ""}}', + ].join('\n'); +} + +/** + * Extract the first JSON object from the model's response and validate that + * it looks like a capability invocation (`{name, args}`). + * + * @param content - Raw assistant message content. + * @returns Parsed invocation, or `null` if none found. + */ +function parseInvocation( + content: string, +): { name: string; args: Record } | null { + const match = /\{[\s\S]*\}/u.exec(content); + if (!match) { + return null; + } + try { + const parsed = JSON.parse(match[0]) as unknown; + if ( + parsed !== null && + typeof parsed === 'object' && + 'name' in parsed && + typeof (parsed as Record).name === 'string' + ) { + const { name, args } = parsed as { + name: string; + args?: Record; + }; + return { name, args: args ?? {} }; + } + } catch { + // not a valid capability invocation + } + return null; +} + +/** + * Make a chat-based capability-augmented agent. + * + * Unlike {@link makeJsonAgent} which uses raw text completion, this agent + * drives the loop via a chat messages array, making it compatible with any + * OpenAI-compatible or Ollama chat endpoint. + * + * Capabilities are described to the model via a JSON system prompt. + * The model signals completion by invoking the auto-injected `end` capability. + * + * @param args - Construction arguments. + * @param args.chat - Bound chat function (model already configured). + * @param args.capabilities - Capabilities the agent may invoke. + * @returns A kernel agent implementing the {@link Agent} interface. + */ +export const makeChatAgent = ({ + chat, + capabilities: agentCapabilities, +}: MakeChatAgentArgs): Agent => { + const experienceLog: Experience[] = []; + + return { + task: async ( + intent: string, + judgment?: (result: unknown) => result is Result, + { + invocationBudget = 10, + logger, + }: { invocationBudget?: number; logger?: Logger } = {}, + ): Promise => { + const [end, didEnd, getEnd] = makeEnd(); + const capabilities = mergeDisjointRecords(agentCapabilities, { + end, + }) as CapabilityRecord; + + const effectiveJudgment = + judgment ?? ((result: unknown): result is Result => true); + const objective = { intent, judgment: effectiveJudgment }; + const context = { capabilities }; + + const capSchemas = extractCapabilitySchemas(capabilities); + const capFunctions = extractCapabilities(capabilities); + + const chatHistory: ChatMessage[] = [ + { role: 'system', content: buildSystemPrompt(capSchemas) }, + { role: 'user', content: intent }, + ]; + + const history = chatHistory.map((chatMsg) => new ChatTurn(chatMsg)); + const experience: Experience = { objective, context, history }; + experienceLog.push(experience); + + try { + for (let step = 0; step < invocationBudget; step++) { + logger?.info(`Step ${step + 1} of ${invocationBudget}`); + + const chatResult = await chat(chatHistory); + const assistantMessage = chatResult.choices[0]?.message; + if (!assistantMessage) { + throw new Error('No response from model'); + } + + chatHistory.push(assistantMessage); + history.push(new ChatTurn(assistantMessage)); + + const invocation = parseInvocation(assistantMessage.content); + if (!invocation) { + // Plain text — treat as final answer without capability invocation + const result = assistantMessage.content as unknown as Result; + Object.assign(experience, { result }); + return result; + } + + const { name, args } = invocation; + logger?.info(`Invoking capability: ${name}`, args); + + const cap = capFunctions[name]; + if (!cap) { + const errorContent = `[Error]: Unknown capability "${name}"`; + chatHistory.push({ role: 'user', content: errorContent }); + history.push(new ChatTurn({ role: 'user', content: errorContent })); + continue; + } + + let toolResult: unknown; + try { + toolResult = await cap(args as never); + } catch (error) { + const errorContent = `[Error calling ${name}]: ${(error as Error).message}`; + chatHistory.push({ role: 'user', content: errorContent }); + history.push(new ChatTurn({ role: 'user', content: errorContent })); + continue; + } + + const resultContent = `[Result of ${name}]: ${JSON.stringify(toolResult)}`; + chatHistory.push({ role: 'user', content: resultContent }); + history.push(new ChatTurn({ role: 'user', content: resultContent })); + + if (didEnd()) { + const result = getEnd(); + if (!effectiveJudgment(result)) { + throw new Error(`Invalid result: ${JSON.stringify(result)}`); + } + Object.assign(experience, { result }); + return result; + } + } + throw new Error('Invocation budget exceeded'); + } catch (error) { + if (error instanceof Error) { + Object.assign(experience, { error }); + } + throw error; + } + }, + + get experiences() { + return (async function* () { + yield* experienceLog; + })(); + }, + }; +}; From 8ab31ab86cff02521a443372d46343206a06c9bf Mon Sep 17 00:00:00 2001 From: grypez <143971198+grypez@users.noreply.github.com> Date: Mon, 30 Mar 2026 12:51:28 -0400 Subject: [PATCH 06/20] feat(kernel-agents): use standard tool-calling interface in makeChatAgent Replace the custom JSON-in-system-prompt approach with the standard chat completions tool-calling interface: capabilities are passed as `tools`, responses arrive via `tool_calls`, and results are returned as `role: "tool"` messages. Add a `glm-4.7-flash` / llama.cpp e2e test and a `LMS_PROVIDER` constant so the suite can be aimed at either Ollama or llama.cpp without OOM. Co-Authored-By: Claude Sonnet 4.6 --- .../src/strategies/chat-agent.test.ts | 249 +++++++++++------- .../src/strategies/chat-agent.ts | 188 ++++++------- .../src/ollama/base.ts | 2 +- .../src/open-v1/types.ts | 23 ++ .../src/types.ts | 31 ++- packages/kernel-test-local/package.json | 1 + .../src/chat-agent.e2e.test.ts | 45 ++++ packages/kernel-test-local/src/constants.ts | 19 ++ packages/kernel-test-local/src/lms-chat.ts | 4 +- .../src/sample-agent.e2e.test.ts | 4 +- 10 files changed, 363 insertions(+), 203 deletions(-) create mode 100644 packages/kernel-test-local/src/chat-agent.e2e.test.ts diff --git a/packages/kernel-agents/src/strategies/chat-agent.test.ts b/packages/kernel-agents/src/strategies/chat-agent.test.ts index 9fac5cd6de..f9f068efdc 100644 --- a/packages/kernel-agents/src/strategies/chat-agent.test.ts +++ b/packages/kernel-agents/src/strategies/chat-agent.test.ts @@ -1,56 +1,67 @@ import '@ocap/repo-tools/test-utils/mock-endoify'; +import type { + ChatMessage, + ChatResult, + ToolCall, +} from '@ocap/kernel-language-model-service'; import { describe, expect, it, vi } from 'vitest'; import { makeChatAgent } from './chat-agent.ts'; import type { BoundChat } from './chat-agent.ts'; import { capability } from '../capabilities/capability.ts'; -const makeChat = (responses: string[]): BoundChat => { - let call = 0; - return async () => { - const index = call; - call += 1; - return { - id: String(index), - model: 'test', - choices: [ - { - message: { - role: 'assistant' as const, - content: responses[index] ?? '', - }, - index: 0, - finish_reason: 'stop', - }, - ], - usage: { prompt_tokens: 0, completion_tokens: 0, total_tokens: 0 }, - }; - }; -}; +const makeToolCall = ( + id: string, + name: string, + args: Record, +): ToolCall => ({ + id, + type: 'function', + function: { name, arguments: JSON.stringify(args) }, +}); + +const makeTextResponse = (content: string): ChatResult => ({ + id: '0', + model: 'test', + choices: [ + { + message: { role: 'assistant', content }, + index: 0, + finish_reason: 'stop', + }, + ], + usage: { prompt_tokens: 0, completion_tokens: 0, total_tokens: 0 }, +}); + +const makeToolCallResponse = ( + id: string, + toolCalls: ToolCall[], +): ChatResult => ({ + id, + model: 'test', + choices: [ + { + message: { role: 'assistant', content: '', tool_calls: toolCalls }, + index: 0, + finish_reason: 'tool_calls', + }, + ], + usage: { prompt_tokens: 0, completion_tokens: 0, total_tokens: 0 }, +}); const noCapabilities = {}; describe('makeChatAgent', () => { - it('returns plain text response when model does not invoke a capability', async () => { - const chat = makeChat(['Hello, world!']); + it('returns plain text response when model does not invoke a tool', async () => { + const chat: BoundChat = async () => makeTextResponse('Hello, world!'); const agent = makeChatAgent({ chat, capabilities: noCapabilities }); const result = await agent.task('say hello'); expect(result).toBe('Hello, world!'); }); - it('returns result when model invokes end capability', async () => { - const chat = makeChat([ - '{"name": "end", "args": {"final": "the answer is 42"}}', - ]); - const agent = makeChatAgent({ chat, capabilities: noCapabilities }); - - const result = await agent.task('what is the answer?'); - expect(result).toBe('the answer is 42'); - }); - - it('dispatches a user capability and continues to end', async () => { + it('dispatches a tool call and returns final text answer', async () => { const add = vi.fn(async ({ a, b }: { a: number; b: number }) => a + b); const addCap = capability(add, { description: 'Add two numbers', @@ -61,100 +72,158 @@ describe('makeChatAgent', () => { returns: { type: 'number' }, }); - const chat = makeChat([ - '{"name": "add", "args": {"a": 3, "b": 4}}', - '{"name": "end", "args": {"final": "7"}}', - ]); - const agent = makeChatAgent({ - chat, - capabilities: { add: addCap }, - }); + let call = 0; + const chat: BoundChat = async () => { + call += 1; + if (call === 1) { + return makeToolCallResponse('0', [ + makeToolCall('c1', 'add', { a: 3, b: 4 }), + ]); + } + return makeTextResponse('7'); + }; + + const agent = makeChatAgent({ chat, capabilities: { add: addCap } }); const result = await agent.task('add 3 and 4'); expect(add).toHaveBeenCalledWith({ a: 3, b: 4 }); expect(result).toBe('7'); }); - it('injects tool result into messages before next turn', async () => { - const messages: string[][] = []; - const chat: BoundChat = async (chatMsgs) => { - messages.push(chatMsgs.map((chatMsg) => chatMsg.content)); - const turn = messages.length - 1; - return { - id: String(turn), - model: 'test', - choices: [ - { - message: { - role: 'assistant' as const, - content: - turn === 0 - ? '{"name": "ping", "args": {}}' - : '{"name": "end", "args": {"final": "done"}}', - }, - index: 0, - finish_reason: 'stop', - }, - ], - usage: { prompt_tokens: 0, completion_tokens: 0, total_tokens: 0 }, - }; - }; - + it('injects tool result message before next turn', async () => { + const recorded: ChatMessage[][] = []; const ping = capability(async () => 'pong', { description: 'Ping', args: {}, returns: { type: 'string' }, }); + + let call = 0; + const chat: BoundChat = async ({ messages }) => { + recorded.push([...messages]); + call += 1; + if (call === 1) { + return makeToolCallResponse('0', [makeToolCall('c1', 'ping', {})]); + } + return makeTextResponse('done'); + }; + const agent = makeChatAgent({ chat, capabilities: { ping } }); await agent.task('ping'); - // Second turn messages should include the tool result + // Second turn must include the tool result message + const secondTurn = recorded[1] ?? []; expect( - messages[1]?.some((content) => content.includes('[Result of ping]')), + secondTurn.some( + (message) => message.role === 'tool' && message.tool_call_id === 'c1', + ), ).toBe(true); + expect(secondTurn.some((message) => message.content === '"pong"')).toBe( + true, + ); }); - it('appends error message for unknown capability and continues', async () => { - const chat = makeChat([ - '{"name": "nonexistent", "args": {}}', - '{"name": "end", "args": {"final": "recovered"}}', - ]); - const agent = makeChatAgent({ chat, capabilities: noCapabilities }); + it('injects error message for unknown tool and continues', async () => { + const recorded: ChatMessage[][] = []; + let call = 0; + const chat: BoundChat = async ({ messages }) => { + recorded.push([...messages]); + call += 1; + if (call === 1) { + return makeToolCallResponse('0', [ + makeToolCall('c1', 'nonexistent', {}), + ]); + } + return makeTextResponse('recovered'); + }; + const agent = makeChatAgent({ chat, capabilities: noCapabilities }); const result = await agent.task('do something'); + expect(result).toBe('recovered'); + const secondTurn = recorded[1] ?? []; + expect( + secondTurn.some( + (message) => + message.role === 'tool' && + message.content.includes('Unknown capability'), + ), + ).toBe(true); }); it('throws when invocation budget is exceeded', async () => { - // Always invokes a capability but never ends const ping = capability(async () => 'pong', { description: 'Ping', args: {}, }); - const neverEnd = makeChat( - Array.from({ length: 20 }, () => '{"name": "ping", "args": {}}'), - ); - const agent = makeChatAgent({ - chat: neverEnd, - capabilities: { ping }, - }); + const chat: BoundChat = async () => + makeToolCallResponse('0', [makeToolCall('c1', 'ping', {})]); + + const agent = makeChatAgent({ chat, capabilities: { ping } }); await expect( agent.task('go', undefined, { invocationBudget: 3 }), ).rejects.toThrow('Invocation budget exceeded'); }); - it('applies judgment to end result', async () => { - const chat = makeChat(['{"name": "end", "args": {"final": 99}}']); + it('applies judgment to final answer', async () => { + const chat: BoundChat = async () => makeTextResponse('hello'); + const agent = makeChatAgent({ chat, capabilities: noCapabilities }); + + const isNumber = (result: unknown): result is number => + typeof result === 'number'; + await expect(agent.task('go', isNumber)).rejects.toThrow('Invalid result'); + }); + + it('passes tools to the chat function', async () => { + const recordedTools: unknown[] = []; + const ping = capability(async () => 'pong', { + description: 'Ping the server', + args: {}, + returns: { type: 'string' }, + }); + + const chat: BoundChat = async ({ tools }) => { + recordedTools.push(tools); + return makeTextResponse('done'); + }; + + const agent = makeChatAgent({ chat, capabilities: { ping } }); + await agent.task('go'); + + expect(recordedTools[0]).toStrictEqual([ + { + type: 'function', + function: { + name: 'ping', + description: 'Ping the server', + parameters: { type: 'object', properties: {}, required: [] }, + }, + }, + ]); + }); + + it('passes undefined tools when there are no capabilities', async () => { + let recordedTools: unknown = 'not-set'; + const chat: BoundChat = async ({ tools }) => { + recordedTools = tools; + return makeTextResponse('done'); + }; + const agent = makeChatAgent({ chat, capabilities: noCapabilities }); + await agent.task('go'); - const isString = (result: unknown): result is string => - typeof result === 'string'; - await expect(agent.task('go', isString)).rejects.toThrow('Invalid result'); + expect(recordedTools).toBeUndefined(); }); it('accumulates experiences across tasks', async () => { - const chat = makeChat(['hello', 'world']); + let call = 0; + const responses = ['hello', 'world']; + const chat: BoundChat = async () => { + const response = makeTextResponse(responses[call] ?? ''); + call += 1; + return response; + }; const agent = makeChatAgent({ chat, capabilities: noCapabilities }); await agent.task('first'); diff --git a/packages/kernel-agents/src/strategies/chat-agent.ts b/packages/kernel-agents/src/strategies/chat-agent.ts index 6b10c39026..2bfd4a8ee3 100644 --- a/packages/kernel-agents/src/strategies/chat-agent.ts +++ b/packages/kernel-agents/src/strategies/chat-agent.ts @@ -1,15 +1,14 @@ -import { mergeDisjointRecords } from '@metamask/kernel-utils'; import type { Logger } from '@metamask/logger'; import type { ChatMessage, ChatResult, + Tool, } from '@ocap/kernel-language-model-service'; import { extractCapabilitySchemas, extractCapabilities, } from '../capabilities/capability.ts'; -import { makeEnd } from '../capabilities/end.ts'; import type { Agent } from '../types/agent.ts'; import { Message } from '../types/messages.ts'; import type { CapabilityRecord, Experience } from '../types.ts'; @@ -25,7 +24,7 @@ class ChatTurn extends Message { * @param chatMessage.content - The text content of the message. */ constructor({ role, content }: ChatMessage) { - super(role, { content }); + super(role, { content: content ?? '' }); } } @@ -35,10 +34,14 @@ class ChatTurn extends Message { * * ```ts * const client = makeChatClient(serviceRef, model); - * const chat = (messages) => client.chat.completions.create({ messages }); + * const chat = ({ messages, tools }) => + * client.chat.completions.create({ messages, tools }); * ``` */ -export type BoundChat = (messages: ChatMessage[]) => Promise; +export type BoundChat = (params: { + messages: ChatMessage[]; + tools?: Tool[]; +}) => Promise; export type MakeChatAgentArgs = { /** @@ -49,76 +52,42 @@ export type MakeChatAgentArgs = { chat: BoundChat; /** * Capabilities the agent may invoke, expressed as a {@link CapabilityRecord}. - * An `end` capability is automatically added to signal task completion. */ capabilities: CapabilityRecord; }; /** - * Build the system prompt that instructs the model to invoke capabilities - * by responding with JSON objects. + * Convert a {@link CapabilityRecord} to the {@link Tool} array expected by + * the chat completions API. * - * @param capSchemas - Serialized capability schemas. - * @returns The system prompt string. + * @param capabilities - The capabilities to convert. + * @returns An array of tool definitions. */ -function buildSystemPrompt(capSchemas: Record): string { - return [ - 'You are a capability-augmented assistant.', - 'To invoke a capability, respond with ONLY a JSON object:', - ' {"name": "", "args": {}}', - 'Do not include any other text when invoking a capability.', - '', - 'Available capabilities:', - JSON.stringify(capSchemas, null, 2), - '', - 'When you have a final answer, invoke the "end" capability:', - ' {"name": "end", "args": {"final": ""}}', - ].join('\n'); -} - -/** - * Extract the first JSON object from the model's response and validate that - * it looks like a capability invocation (`{name, args}`). - * - * @param content - Raw assistant message content. - * @returns Parsed invocation, or `null` if none found. - */ -function parseInvocation( - content: string, -): { name: string; args: Record } | null { - const match = /\{[\s\S]*\}/u.exec(content); - if (!match) { - return null; - } - try { - const parsed = JSON.parse(match[0]) as unknown; - if ( - parsed !== null && - typeof parsed === 'object' && - 'name' in parsed && - typeof (parsed as Record).name === 'string' - ) { - const { name, args } = parsed as { - name: string; - args?: Record; - }; - return { name, args: args ?? {} }; - } - } catch { - // not a valid capability invocation - } - return null; +function buildTools(capabilities: CapabilityRecord): Tool[] { + const schemas = extractCapabilitySchemas(capabilities); + return Object.entries(schemas).map(([name, schema]) => ({ + type: 'function' as const, + function: { + name, + description: schema.description, + parameters: { + type: 'object' as const, + properties: schema.args, + required: Object.keys(schema.args), + }, + }, + })); } /** * Make a chat-based capability-augmented agent. * * Unlike {@link makeJsonAgent} which uses raw text completion, this agent - * drives the loop via a chat messages array, making it compatible with any - * OpenAI-compatible or Ollama chat endpoint. + * drives the loop via a chat messages array and the standard tool-calling + * interface, making it compatible with any OpenAI-compatible chat endpoint. * - * Capabilities are described to the model via a JSON system prompt. - * The model signals completion by invoking the auto-injected `end` capability. + * Capabilities are exposed to the model as tools via the `tools` parameter. + * The model signals completion by returning a message without tool calls. * * @param args - Construction arguments. * @param args.chat - Bound chat function (model already configured). @@ -140,23 +109,15 @@ export const makeChatAgent = ({ logger, }: { invocationBudget?: number; logger?: Logger } = {}, ): Promise => { - const [end, didEnd, getEnd] = makeEnd(); - const capabilities = mergeDisjointRecords(agentCapabilities, { - end, - }) as CapabilityRecord; - const effectiveJudgment = judgment ?? ((result: unknown): result is Result => true); const objective = { intent, judgment: effectiveJudgment }; - const context = { capabilities }; + const context = { capabilities: agentCapabilities }; - const capSchemas = extractCapabilitySchemas(capabilities); - const capFunctions = extractCapabilities(capabilities); + const capFunctions = extractCapabilities(agentCapabilities); + const tools = buildTools(agentCapabilities); - const chatHistory: ChatMessage[] = [ - { role: 'system', content: buildSystemPrompt(capSchemas) }, - { role: 'user', content: intent }, - ]; + const chatHistory: ChatMessage[] = [{ role: 'user', content: intent }]; const history = chatHistory.map((chatMsg) => new ChatTurn(chatMsg)); const experience: Experience = { objective, context, history }; @@ -166,7 +127,10 @@ export const makeChatAgent = ({ for (let step = 0; step < invocationBudget; step++) { logger?.info(`Step ${step + 1} of ${invocationBudget}`); - const chatResult = await chat(chatHistory); + const chatResult = await chat({ + messages: chatHistory, + ...(tools.length > 0 && { tools }), + }); const assistantMessage = chatResult.choices[0]?.message; if (!assistantMessage) { throw new Error('No response from model'); @@ -175,46 +139,56 @@ export const makeChatAgent = ({ chatHistory.push(assistantMessage); history.push(new ChatTurn(assistantMessage)); - const invocation = parseInvocation(assistantMessage.content); - if (!invocation) { - // Plain text — treat as final answer without capability invocation + const { tool_calls: toolCalls } = assistantMessage; + if (!toolCalls?.length) { + // No tool calls — model has a final answer const result = assistantMessage.content as unknown as Result; + if (!effectiveJudgment(result)) { + throw new Error(`Invalid result: ${JSON.stringify(result)}`); + } Object.assign(experience, { result }); return result; } - const { name, args } = invocation; - logger?.info(`Invoking capability: ${name}`, args); - - const cap = capFunctions[name]; - if (!cap) { - const errorContent = `[Error]: Unknown capability "${name}"`; - chatHistory.push({ role: 'user', content: errorContent }); - history.push(new ChatTurn({ role: 'user', content: errorContent })); - continue; - } - - let toolResult: unknown; - try { - toolResult = await cap(args as never); - } catch (error) { - const errorContent = `[Error calling ${name}]: ${(error as Error).message}`; - chatHistory.push({ role: 'user', content: errorContent }); - history.push(new ChatTurn({ role: 'user', content: errorContent })); - continue; - } - - const resultContent = `[Result of ${name}]: ${JSON.stringify(toolResult)}`; - chatHistory.push({ role: 'user', content: resultContent }); - history.push(new ChatTurn({ role: 'user', content: resultContent })); + for (const toolCall of toolCalls) { + const { name, arguments: argsJson } = toolCall.function; + logger?.info(`Invoking capability: ${name}`); + + const cap = capFunctions[name]; + if (!cap) { + const errorContent = `Unknown capability "${name}"`; + const toolMsg: ChatMessage = { + role: 'tool', + tool_call_id: toolCall.id, + content: errorContent, + }; + chatHistory.push(toolMsg); + history.push(new ChatTurn(toolMsg)); + continue; + } - if (didEnd()) { - const result = getEnd(); - if (!effectiveJudgment(result)) { - throw new Error(`Invalid result: ${JSON.stringify(result)}`); + let toolResult: unknown; + try { + toolResult = await cap(JSON.parse(argsJson) as never); + } catch (error) { + const errorContent = `Error calling ${name}: ${(error as Error).message}`; + const toolMsg: ChatMessage = { + role: 'tool', + tool_call_id: toolCall.id, + content: errorContent, + }; + chatHistory.push(toolMsg); + history.push(new ChatTurn(toolMsg)); + continue; } - Object.assign(experience, { result }); - return result; + + const toolMsg: ChatMessage = { + role: 'tool', + tool_call_id: toolCall.id, + content: JSON.stringify(toolResult), + }; + chatHistory.push(toolMsg); + history.push(new ChatTurn(toolMsg)); } } throw new Error('Invocation budget exceeded'); diff --git a/packages/kernel-language-model-service/src/ollama/base.ts b/packages/kernel-language-model-service/src/ollama/base.ts index 9a36046eb3..a0872d7c87 100644 --- a/packages/kernel-language-model-service/src/ollama/base.ts +++ b/packages/kernel-language-model-service/src/ollama/base.ts @@ -81,7 +81,7 @@ export class OllamaBaseService } const response = await ollama.chat({ model, - messages, + messages: messages.map(({ role, content }) => ({ role, content })), stream: false, options: ifDefined({ temperature, diff --git a/packages/kernel-language-model-service/src/open-v1/types.ts b/packages/kernel-language-model-service/src/open-v1/types.ts index 997dd8c991..51e6ed85f6 100644 --- a/packages/kernel-language-model-service/src/open-v1/types.ts +++ b/packages/kernel-language-model-service/src/open-v1/types.ts @@ -1,4 +1,5 @@ import { + any, array, boolean, literal, @@ -18,6 +19,8 @@ export type { ChatRole, ChatStreamChunk, ChatStreamDelta, + Tool, + ToolCall, Usage, } from '../types.ts'; @@ -25,11 +28,30 @@ const ChatRoleStruct = union([ literal('system'), literal('user'), literal('assistant'), + literal('tool'), ]); +const ToolCallStruct = object({ + id: string(), + type: literal('function'), + index: optional(number()), + function: object({ name: string(), arguments: string() }), +}); + const ChatMessageStruct = object({ role: ChatRoleStruct, content: string(), + tool_calls: optional(array(ToolCallStruct)), + tool_call_id: optional(string()), +}); + +const ToolStruct = object({ + type: literal('function'), + function: object({ + name: string(), + description: optional(string()), + parameters: optional(any()), + }), }); const StopStruct = optional(union([string(), array(string())])); @@ -40,6 +62,7 @@ const StopStruct = optional(union([string(), array(string())])); export const ChatParamsStruct = object({ model: size(string(), 1, Infinity), messages: array(ChatMessageStruct), + tools: optional(array(ToolStruct)), max_tokens: optional(number()), temperature: optional(number()), top_p: optional(number()), diff --git a/packages/kernel-language-model-service/src/types.ts b/packages/kernel-language-model-service/src/types.ts index aafaabddd5..3eb606483f 100644 --- a/packages/kernel-language-model-service/src/types.ts +++ b/packages/kernel-language-model-service/src/types.ts @@ -1,7 +1,33 @@ /** * A role in a chat conversation. */ -export type ChatRole = 'system' | 'user' | 'assistant'; +export type ChatRole = 'system' | 'user' | 'assistant' | 'tool'; + +/** + * A single tool call made by the model. + */ +export type ToolCall = { + id: string; + type: 'function'; + index?: number; + function: { name: string; arguments: string }; +}; + +/** + * A tool definition passed to the model. + */ +export type Tool = { + type: 'function'; + function: { + name: string; + description?: string; + parameters?: { + type: 'object'; + properties?: Record; + required?: string[]; + }; + }; +}; /** * A single message in a chat conversation. @@ -9,6 +35,8 @@ export type ChatRole = 'system' | 'user' | 'assistant'; export type ChatMessage = { role: ChatRole; content: string; + tool_calls?: ToolCall[]; + tool_call_id?: string; }; /** @@ -52,6 +80,7 @@ export type SampleResult = { export type ChatParams = { model: string; messages: ChatMessage[]; + tools?: Tool[]; max_tokens?: number; temperature?: number; top_p?: number; diff --git a/packages/kernel-test-local/package.json b/packages/kernel-test-local/package.json index 085cabbcbc..cfe3184650 100644 --- a/packages/kernel-test-local/package.json +++ b/packages/kernel-test-local/package.json @@ -23,6 +23,7 @@ "build:docs": "typedoc", "test": "vitest run --config vitest.config.ts", "test:e2e:local": "vitest run --config vitest.config.e2e.ts", + "test:e2e:local:llama-cpp": "LMS_PROVIDER=llama-cpp yarn test:e2e:local", "test:clean": "yarn test --no-cache --coverage.clean", "test:dev": "yarn test --mode development", "test:verbose": "yarn test --reporter verbose", diff --git a/packages/kernel-test-local/src/chat-agent.e2e.test.ts b/packages/kernel-test-local/src/chat-agent.e2e.test.ts new file mode 100644 index 0000000000..9a55f4fe8d --- /dev/null +++ b/packages/kernel-test-local/src/chat-agent.e2e.test.ts @@ -0,0 +1,45 @@ +import '@ocap/repo-tools/test-utils/mock-endoify'; + +import { add } from '@ocap/kernel-agents/capabilities/math'; +import { makeChatAgent } from '@ocap/kernel-agents/chat'; +import type { BoundChat } from '@ocap/kernel-agents/chat'; +import { makeOpenV1NodejsService } from '@ocap/kernel-language-model-service'; +import { fetchMock } from '@ocap/repo-tools/test-utils/fetch-mock'; +import { afterAll, beforeAll, describe, expect, it, vi } from 'vitest'; + +import { LMS_BASE_URL, LMS_CHAT_MODEL } from './constants.ts'; + +describe.sequential('makeChatAgent tool calling (e2e)', () => { + beforeAll(() => { + fetchMock.disableMocks(); + }); + + afterAll(() => { + fetchMock.enableMocks(); + }); + + it( + 'invokes the add tool and returns the computed sum', + { timeout: 60_000 }, + async () => { + const service = makeOpenV1NodejsService({ + endowments: { fetch }, + baseUrl: LMS_BASE_URL, + }); + + const chat: BoundChat = async ({ messages, tools }) => + service.chat({ model: LMS_CHAT_MODEL, messages, tools }); + + const addSpy = vi.spyOn(add, 'func'); + const agent = makeChatAgent({ chat, capabilities: { add } }); + + const result = await agent.task( + 'What is 123 plus 456? Use the add tool to compute it.', + ); + + expect(addSpy).toHaveBeenCalled(); + expect(typeof result).toBe('string'); + expect(result as string).toContain('579'); + }, + ); +}); diff --git a/packages/kernel-test-local/src/constants.ts b/packages/kernel-test-local/src/constants.ts index a096d674a1..ad399e25e4 100644 --- a/packages/kernel-test-local/src/constants.ts +++ b/packages/kernel-test-local/src/constants.ts @@ -10,6 +10,25 @@ 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`; +/** + * Supported LMS providers for E2E tests. + * Select with: LMS_PROVIDER=llama-cpp yarn test:e2e:local + */ +const LMS_PROVIDERS = { + ollama: { baseUrl: OLLAMA_API_BASE, model: DEFAULT_MODEL }, + 'llama-cpp': { baseUrl: 'http://localhost:8080', model: 'glm-4.7-flash' }, +} as const; + +export type LmsProvider = keyof typeof LMS_PROVIDERS; + +// eslint-disable-next-line n/no-process-env +const rawProvider = process?.env?.LMS_PROVIDER ?? 'ollama'; +export const LMS_PROVIDER: LmsProvider = + rawProvider in LMS_PROVIDERS ? (rawProvider as LmsProvider) : 'ollama'; + +export const LMS_BASE_URL = LMS_PROVIDERS[LMS_PROVIDER].baseUrl; +export const LMS_CHAT_MODEL = LMS_PROVIDERS[LMS_PROVIDER].model; + /** * Logger tags to ignore, parsed from the LOGGER_IGNORE environment variable. */ diff --git a/packages/kernel-test-local/src/lms-chat.ts b/packages/kernel-test-local/src/lms-chat.ts index b5b27f23ba..bf76cf78e4 100644 --- a/packages/kernel-test-local/src/lms-chat.ts +++ b/packages/kernel-test-local/src/lms-chat.ts @@ -18,7 +18,7 @@ import { } from '@ocap/kernel-language-model-service'; import { expect } from 'vitest'; -import { DEFAULT_MODEL } from './constants.ts'; +import { LMS_CHAT_MODEL } from './constants.ts'; import { filterTransports } from './utils.ts'; const getBundleSpec = (name: string): string => @@ -56,7 +56,7 @@ export const runLmsChatKernelTest = async ( vats: { main: { bundleSpec: getBundleSpec('lms-chat-vat'), - parameters: { model: DEFAULT_MODEL }, + parameters: { model: LMS_CHAT_MODEL }, }, }, }); diff --git a/packages/kernel-test-local/src/sample-agent.e2e.test.ts b/packages/kernel-test-local/src/sample-agent.e2e.test.ts index 40dff73b7d..0402e78b57 100644 --- a/packages/kernel-test-local/src/sample-agent.e2e.test.ts +++ b/packages/kernel-test-local/src/sample-agent.e2e.test.ts @@ -19,7 +19,7 @@ import { vi, } from 'vitest'; -import { DEFAULT_MODEL } from './constants.ts'; +import { DEFAULT_MODEL, LMS_PROVIDER } from './constants.ts'; import { filterTransports, randomLetter } from './utils.ts'; const logger = new Logger({ @@ -33,7 +33,7 @@ const makeJsonAgentWithMathCapabilities = (args: MakeAgentArgs) => capabilities: { count, add, multiply, ...args.capabilities }, }); -describe.each([ +describe.skipIf(LMS_PROVIDER !== 'ollama').each([ ['json', makeJsonAgentWithMathCapabilities], ['repl', makeReplAgent], ])( From 70ae122a6e42fbb9dda0f422ac7d4cffb57522d3 Mon Sep 17 00:00:00 2001 From: grypez <143971198+grypez@users.noreply.github.com> Date: Mon, 30 Mar 2026 15:43:48 -0400 Subject: [PATCH 07/20] types(klms): fixup --- packages/kernel-language-model-service/src/ollama/base.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/packages/kernel-language-model-service/src/ollama/base.ts b/packages/kernel-language-model-service/src/ollama/base.ts index a0872d7c87..bd47f9ed57 100644 --- a/packages/kernel-language-model-service/src/ollama/base.ts +++ b/packages/kernel-language-model-service/src/ollama/base.ts @@ -139,8 +139,8 @@ export class OllamaBaseService * @returns A streaming result or full result depending on `params.stream`. */ async sample( - params: SampleParams, - ): Promise | Promise { + params: SampleParams & { stream?: boolean }, + ): Promise { if (params.stream === true) { return this.#streamingSample(params); } From f6fcb2b165dedbe474c9988c5920b141551a5cbf Mon Sep 17 00:00:00 2001 From: grypez <143971198+grypez@users.noreply.github.com> Date: Mon, 30 Mar 2026 18:50:21 -0400 Subject: [PATCH 08/20] fix superstruct and type issues --- .../src/client.test.ts | 27 +++++++- .../src/client.ts | 62 +++++++++++++------ .../src/ollama/base.ts | 5 +- .../src/open-v1/types.ts | 3 +- .../src/types.ts | 5 +- 5 files changed, 78 insertions(+), 24 deletions(-) diff --git a/packages/kernel-language-model-service/src/client.test.ts b/packages/kernel-language-model-service/src/client.test.ts index 22dd6febee..f493c56b37 100644 --- a/packages/kernel-language-model-service/src/client.test.ts +++ b/packages/kernel-language-model-service/src/client.test.ts @@ -1,7 +1,7 @@ import { describe, it, expect, vi } from 'vitest'; import { makeChatClient, makeSampleClient } from './client.ts'; -import type { ChatResult, SampleResult } from './types.ts'; +import type { ChatResult, ChatStreamChunk, SampleResult } from './types.ts'; const MODEL = 'glm-4.7-flash'; @@ -67,6 +67,31 @@ describe('makeChatClient', () => { }), ).rejects.toThrow('model is required'); }); + + it('returns async iterable when stream: true', async () => { + async function* makeStream(): AsyncGenerator { + yield { + id: 'chunk-1', + model: MODEL, + choices: [ + { delta: { content: 'hello' }, index: 0, finish_reason: null }, + ], + }; + } + const stream = makeStream(); + const lmsRef = { chat: vi.fn().mockReturnValue(stream) }; + + const client = makeChatClient(lmsRef, MODEL); + const result = await client.chat.completions.create({ + messages: [{ role: 'user', content: 'hi' }], + stream: true, + }); + + expect(lmsRef.chat).toHaveBeenCalledWith( + expect.objectContaining({ stream: true }), + ); + expect(result).toBe(stream); + }); }); describe('makeSampleClient', () => { diff --git a/packages/kernel-language-model-service/src/client.ts b/packages/kernel-language-model-service/src/client.ts index 24973d6daa..30acdc70d7 100644 --- a/packages/kernel-language-model-service/src/client.ts +++ b/packages/kernel-language-model-service/src/client.ts @@ -5,6 +5,7 @@ import type { ChatParams, ChatResult, ChatService, + ChatStreamChunk, SampleParams, SampleResult, SampleService, @@ -17,6 +18,7 @@ import type { * ```ts * const client = makeChatClient(lmsRef, 'gpt-4o'); * const result = await client.chat.completions.create({ messages }); + * const stream = await client.chat.completions.create({ messages, stream: true }); * ``` * * @param lmsRef - Reference to a service with a `chat` method. @@ -29,27 +31,49 @@ export const makeChatClient = ( ): { chat: { completions: { - create: ( - params: Omit & { model?: string }, - ) => Promise; + create( + params: Omit & { model?: string; stream: true }, + ): Promise>; + create( + params: Omit & { model?: string; stream?: false }, + ): 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 })); - }, - }), - }), - }); +} => { + type BaseParams = Omit & { model?: string }; + + /** + * @param params - Chat completion parameters with `stream: true`. + * @returns A promise resolving to an async iterable of stream chunks. + */ + function create( + params: BaseParams & { stream: true }, + ): Promise>; + /** + * @param params - Chat completion parameters. + * @returns A promise resolving to the full chat result. + */ + function create(params: BaseParams & { stream?: false }): Promise; + /** + * @param params - Chat completion parameters. + * @returns A promise resolving to a stream or full result depending on `stream`. + */ + async function create( + params: BaseParams, + ): Promise | ChatResult> { + const model = params.model ?? defaultModel; + if (!model) { + throw new Error('model is required'); + } + const fullParams = harden({ ...params, model }); + if (fullParams.stream === true) { + return E(lmsRef).chat(fullParams as ChatParams & { stream: true }); + } + return E(lmsRef).chat(fullParams as ChatParams & { stream?: false }); + } + + return harden({ chat: { completions: { create } } }); +}; /** * Wraps a remote service reference with raw token-prediction ergonomics. diff --git a/packages/kernel-language-model-service/src/ollama/base.ts b/packages/kernel-language-model-service/src/ollama/base.ts index bd47f9ed57..edd304ff90 100644 --- a/packages/kernel-language-model-service/src/ollama/base.ts +++ b/packages/kernel-language-model-service/src/ollama/base.ts @@ -81,7 +81,10 @@ export class OllamaBaseService } const response = await ollama.chat({ model, - messages: messages.map(({ role, content }) => ({ role, content })), + messages: messages.map(({ role, content }) => ({ + role, + content: content ?? '', + })), stream: false, options: ifDefined({ temperature, diff --git a/packages/kernel-language-model-service/src/open-v1/types.ts b/packages/kernel-language-model-service/src/open-v1/types.ts index 51e6ed85f6..b63ace1374 100644 --- a/packages/kernel-language-model-service/src/open-v1/types.ts +++ b/packages/kernel-language-model-service/src/open-v1/types.ts @@ -3,6 +3,7 @@ import { array, boolean, literal, + nullable, number, object, optional, @@ -40,7 +41,7 @@ const ToolCallStruct = object({ const ChatMessageStruct = object({ role: ChatRoleStruct, - content: string(), + content: nullable(string()), tool_calls: optional(array(ToolCallStruct)), tool_call_id: optional(string()), }); diff --git a/packages/kernel-language-model-service/src/types.ts b/packages/kernel-language-model-service/src/types.ts index 3eb606483f..52fa41d5be 100644 --- a/packages/kernel-language-model-service/src/types.ts +++ b/packages/kernel-language-model-service/src/types.ts @@ -34,7 +34,7 @@ export type Tool = { */ export type ChatMessage = { role: ChatRole; - content: string; + content: string | null; tool_calls?: ToolCall[]; tool_call_id?: string; }; @@ -135,7 +135,8 @@ export type ChatResult = { * Minimal service interface required by `makeChatClient`. */ export type ChatService = { - chat: (params: ChatParams) => Promise; + chat(params: ChatParams & { stream: true }): AsyncIterable; + chat(params: ChatParams & { stream?: false }): Promise; }; /** From 16cce23df349e2e25faa3f1d84a0a29584ab505b Mon Sep 17 00:00:00 2001 From: grypez <143971198+grypez@users.noreply.github.com> Date: Mon, 30 Mar 2026 20:59:52 -0400 Subject: [PATCH 09/20] fix ollama chat dropping tool data --- .../src/ollama/base.test.ts | 108 +++++++++++++++++- .../src/ollama/base.ts | 28 ++++- .../src/ollama/types.ts | 2 + 3 files changed, 135 insertions(+), 3 deletions(-) diff --git a/packages/kernel-language-model-service/src/ollama/base.test.ts b/packages/kernel-language-model-service/src/ollama/base.test.ts index e2262a4fa4..e8372fb387 100644 --- a/packages/kernel-language-model-service/src/ollama/base.test.ts +++ b/packages/kernel-language-model-service/src/ollama/base.test.ts @@ -1,4 +1,4 @@ -import type { GenerateResponse, ListResponse } from 'ollama'; +import type { ChatResponse, GenerateResponse, ListResponse } from 'ollama'; import { describe, it, expect, vi, beforeEach } from 'vitest'; import type { OllamaClient, OllamaModelOptions } from './types.ts'; @@ -15,6 +15,7 @@ describe('OllamaBaseService', () => { mockClient = { list: vi.fn(), generate: vi.fn(), + chat: vi.fn(), }; mockMakeClient = vi.fn().mockResolvedValue(mockClient); @@ -134,6 +135,111 @@ describe('OllamaBaseService', () => { }); }); + describe('chat', () => { + const makeChatResponse = ( + overrides: Partial = {}, + ): ChatResponse => + ({ + model: 'llama3', + message: { role: 'assistant', content: 'hello' }, + done: true, + done_reason: 'stop', + prompt_eval_count: 5, + eval_count: 3, + ...overrides, + }) as ChatResponse; + + it('forwards tools to the Ollama client', async () => { + vi.mocked(mockClient.chat).mockResolvedValue(makeChatResponse()); + + await service.chat({ + model: 'llama3', + messages: [{ role: 'user', content: 'hi' }], + tools: [ + { + type: 'function', + function: { name: 'get_time', description: 'Returns the time' }, + }, + ], + }); + + expect(mockClient.chat).toHaveBeenCalledWith( + expect.objectContaining({ + tools: [ + { + type: 'function', + function: { name: 'get_time', description: 'Returns the time' }, + }, + ], + }), + ); + }); + + it('forwards message tool_calls with parsed arguments', async () => { + vi.mocked(mockClient.chat).mockResolvedValue(makeChatResponse()); + + await service.chat({ + model: 'llama3', + messages: [ + { + role: 'assistant', + content: null, + tool_calls: [ + { + id: 'tc-1', + type: 'function', + function: { name: 'get_time', arguments: '{"tz":"UTC"}' }, + }, + ], + }, + ], + }); + + expect(mockClient.chat).toHaveBeenCalledWith( + expect.objectContaining({ + messages: [ + expect.objectContaining({ + tool_calls: [ + { function: { name: 'get_time', arguments: { tz: 'UTC' } } }, + ], + }), + ], + }), + ); + }); + + it('maps response tool_calls back with stringified arguments', async () => { + vi.mocked(mockClient.chat).mockResolvedValue( + makeChatResponse({ + message: { + role: 'assistant', + content: '', + tool_calls: [ + { function: { name: 'get_time', arguments: { tz: 'UTC' } } }, + ], + }, + }), + ); + + const result = await service.chat({ + model: 'llama3', + messages: [{ role: 'user', content: 'what time is it?' }], + }); + + expect(result.choices[0]?.message).toStrictEqual({ + role: 'assistant', + content: '', + tool_calls: [ + { + id: 'tool-0', + type: 'function', + function: { name: 'get_time', arguments: '{"tz":"UTC"}' }, + }, + ], + }); + }); + }); + describe('makeInstance', () => { it('should create instance with model', async () => { const config: InstanceConfig = { diff --git a/packages/kernel-language-model-service/src/ollama/base.ts b/packages/kernel-language-model-service/src/ollama/base.ts index edd304ff90..0e04b351c2 100644 --- a/packages/kernel-language-model-service/src/ollama/base.ts +++ b/packages/kernel-language-model-service/src/ollama/base.ts @@ -19,6 +19,7 @@ import type { OllamaModel, OllamaClient, OllamaModelOptions, + OllamaTool, } from './types.ts'; /** @@ -73,7 +74,7 @@ export class OllamaBaseService * @returns A hardened chat result. */ async chat(params: ChatParams): Promise { - const { model, messages, temperature, seed, stop } = params; + const { model, messages, tools, temperature, seed, stop } = params; const ollama = await this.#makeClient(); let stopArr: string[] | undefined; if (stop !== undefined) { @@ -81,10 +82,22 @@ export class OllamaBaseService } const response = await ollama.chat({ model, - messages: messages.map(({ role, content }) => ({ + // eslint-disable-next-line camelcase + messages: messages.map(({ role, content, tool_calls }) => ({ role, content: content ?? '', + // eslint-disable-next-line camelcase + ...(tool_calls && { + // eslint-disable-next-line camelcase + tool_calls: tool_calls.map(({ function: fn }) => ({ + function: { + name: fn.name, + arguments: JSON.parse(fn.arguments) as Record, + }, + })), + }), })), + ...(tools && { tools: tools as OllamaTool[] }), stream: false, options: ifDefined({ temperature, @@ -96,6 +109,7 @@ export class OllamaBaseService }); const promptTokens = response.prompt_eval_count ?? 0; const completionTokens = response.eval_count ?? 0; + const { tool_calls: responseToolCalls } = response.message; return harden({ id: 'ollama-chat', model: response.model, @@ -104,6 +118,16 @@ export class OllamaBaseService message: { role: response.message.role as ChatRole, content: response.message.content, + ...(responseToolCalls && { + tool_calls: responseToolCalls.map((tc, index) => ({ + id: `tool-${index}`, + type: 'function' as const, + function: { + name: tc.function.name, + arguments: JSON.stringify(tc.function.arguments), + }, + })), + }), }, index: 0, finish_reason: response.done_reason ?? 'stop', diff --git a/packages/kernel-language-model-service/src/ollama/types.ts b/packages/kernel-language-model-service/src/ollama/types.ts index 3630932d21..99e4561966 100644 --- a/packages/kernel-language-model-service/src/ollama/types.ts +++ b/packages/kernel-language-model-service/src/ollama/types.ts @@ -8,6 +8,7 @@ import type { Config, ChatRequest, ChatResponse, + Tool as OllamaTool, } from 'ollama'; import type { LanguageModel } from '../types.ts'; @@ -30,6 +31,7 @@ export type { GenerateRequest, GenerateResponse, OllamaClient, + OllamaTool, ChatRequest, ChatResponse, }; From 8a08a1362edeba4dae09ea7b4923a167fc04c549 Mon Sep 17 00:00:00 2001 From: grypez <143971198+grypez@users.noreply.github.com> Date: Tue, 31 Mar 2026 19:50:38 -0400 Subject: [PATCH 10/20] address review Made-with: Cursor --- .../package.json | 10 ++ .../src/ollama/base.test.ts | 24 +++ .../src/ollama/base.ts | 3 +- .../src/open-v1/base.test.ts | 142 +++++++++++++++++- .../src/open-v1/base.ts | 25 ++- .../src/open-v1/response-json.test.ts | 48 ++++++ .../src/open-v1/response-json.ts | 33 ++++ .../src/open-v1/types.ts | 54 ++++++- .../src/utils/parse-tool-arguments.test.ts | 48 ++++++ .../src/utils/parse-tool-arguments.ts | 51 +++++++ 10 files changed, 422 insertions(+), 16 deletions(-) create mode 100644 packages/kernel-language-model-service/src/open-v1/response-json.test.ts create mode 100644 packages/kernel-language-model-service/src/open-v1/response-json.ts create mode 100644 packages/kernel-language-model-service/src/utils/parse-tool-arguments.test.ts create mode 100644 packages/kernel-language-model-service/src/utils/parse-tool-arguments.ts diff --git a/packages/kernel-language-model-service/package.json b/packages/kernel-language-model-service/package.json index 764f7d7017..4767775ce6 100644 --- a/packages/kernel-language-model-service/package.json +++ b/packages/kernel-language-model-service/package.json @@ -53,6 +53,16 @@ "default": "./dist/test-utils/index.cjs" } }, + "./utils/parse-tool-arguments": { + "import": { + "types": "./dist/utils/parse-tool-arguments.d.mts", + "default": "./dist/utils/parse-tool-arguments.mjs" + }, + "require": { + "types": "./dist/utils/parse-tool-arguments.d.cts", + "default": "./dist/utils/parse-tool-arguments.cjs" + } + }, "./package.json": "./package.json" }, "files": [ diff --git a/packages/kernel-language-model-service/src/ollama/base.test.ts b/packages/kernel-language-model-service/src/ollama/base.test.ts index e8372fb387..bc6c6a2d41 100644 --- a/packages/kernel-language-model-service/src/ollama/base.test.ts +++ b/packages/kernel-language-model-service/src/ollama/base.test.ts @@ -238,6 +238,30 @@ describe('OllamaBaseService', () => { ], }); }); + + it('rejects when message tool_calls contain invalid JSON arguments', async () => { + await expect( + service.chat({ + model: 'llama3', + messages: [ + { + role: 'assistant', + content: null, + tool_calls: [ + { + id: 'tc-1', + type: 'function', + function: { + name: 'get_time', + arguments: '{not valid json', + }, + }, + ], + }, + ], + }), + ).rejects.toThrow('Invalid tool arguments JSON'); + }); }); describe('makeInstance', () => { diff --git a/packages/kernel-language-model-service/src/ollama/base.ts b/packages/kernel-language-model-service/src/ollama/base.ts index 0e04b351c2..c567cb1651 100644 --- a/packages/kernel-language-model-service/src/ollama/base.ts +++ b/packages/kernel-language-model-service/src/ollama/base.ts @@ -21,6 +21,7 @@ import type { OllamaModelOptions, OllamaTool, } from './types.ts'; +import { parseToolArguments } from '../utils/parse-tool-arguments.ts'; /** * Result of a streaming raw token-prediction request. @@ -92,7 +93,7 @@ export class OllamaBaseService tool_calls: tool_calls.map(({ function: fn }) => ({ function: { name: fn.name, - arguments: JSON.parse(fn.arguments) as Record, + arguments: parseToolArguments(fn.arguments), }, })), }), 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 index f1d0c6ff12..75535d8548 100644 --- a/packages/kernel-language-model-service/src/open-v1/base.test.ts +++ b/packages/kernel-language-model-service/src/open-v1/base.test.ts @@ -18,8 +18,24 @@ const makeChatResult = (): ChatResult => ({ 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 makeOkResponse = (init: { + text?: () => Promise; + // eslint-disable-next-line n/no-unsupported-features/node-builtins -- Response.body type in tests + body?: ReadableStream | null; +}): Response => + ({ + ok: true, + status: 200, + statusText: 'OK', + ...init, + }) as Response; + +const makeMockFetch = (payload: unknown): typeof globalThis.fetch => + vi.fn().mockResolvedValue( + makeOkResponse({ + text: vi.fn().mockResolvedValue(JSON.stringify(payload)), + }), + ); const makeSSEStream = ( chunks: ChatStreamChunk[], @@ -48,7 +64,7 @@ const makeStreamChunk = (content: string): ChatStreamChunk => ({ const makeMockStreamFetch = ( chunks: ChatStreamChunk[], ): typeof globalThis.fetch => - vi.fn().mockResolvedValue({ body: makeSSEStream(chunks) }); + vi.fn().mockResolvedValue(makeOkResponse({ body: makeSSEStream(chunks) })); describe('OpenV1BaseService', () => { let service: OpenV1BaseService; @@ -64,7 +80,7 @@ describe('OpenV1BaseService', () => { }); describe('chat', () => { - it('pOSTs to /v1/chat/completions with serialized params', async () => { + it('posts to /v1/chat/completions with serialized params', async () => { const params = { model: MODEL, messages: [{ role: 'user' as const, content: 'hello' }], @@ -171,10 +187,51 @@ describe('OpenV1BaseService', () => { expect.any(Object), ); }); + + it('throws a clear error when the HTTP status is not ok', async () => { + const errFetch = vi.fn().mockResolvedValue({ + ok: false, + status: 429, + statusText: 'Too Many Requests', + text: vi.fn().mockResolvedValue('{"error":"rate_limited"}'), + } as unknown as Response); + const errService = new OpenV1BaseService( + errFetch, + 'http://localhost:8080', + ); + + await expect( + errService.chat({ + model: MODEL, + messages: [{ role: 'user', content: 'hi' }], + }), + ).rejects.toThrow( + 'HTTP 429 Too Many Requests — {"error":"rate_limited"}', + ); + }); + + it('throws when the response body is not valid JSON', async () => { + const badJsonFetch = vi.fn().mockResolvedValue( + makeOkResponse({ + text: vi.fn().mockResolvedValue('not-json'), + }), + ); + const badService = new OpenV1BaseService( + badJsonFetch, + 'http://localhost:8080', + ); + + await expect( + badService.chat({ + model: MODEL, + messages: [{ role: 'user', content: 'hi' }], + }), + ).rejects.toThrow(SyntaxError); + }); }); describe('listModels', () => { - it('gETs /v1/models and returns model IDs', async () => { + it('gets /v1/models and returns model IDs', async () => { const modelsFetch = makeMockFetch({ data: [{ id: 'model-a' }, { id: 'model-b' }], }); @@ -209,10 +266,27 @@ describe('OpenV1BaseService', () => { 'Bearer sk-key', ); }); + + it('throws a clear error when the HTTP status is not ok', async () => { + const errFetch = vi.fn().mockResolvedValue({ + ok: false, + status: 401, + statusText: 'Unauthorized', + text: vi.fn().mockResolvedValue('invalid key'), + } as unknown as Response); + const errService = new OpenV1BaseService( + errFetch, + 'http://localhost:8080', + ); + + await expect(errService.listModels()).rejects.toThrow( + 'HTTP 401 Unauthorized — invalid key', + ); + }); }); describe('chat with stream: true', () => { - it('pOSTs to /v1/chat/completions with stream: true in body', async () => { + it('posts to /v1/chat/completions with stream: true in body', async () => { const streamFetch = makeMockStreamFetch([makeStreamChunk('hi')]); const streamService = new OpenV1BaseService( streamFetch, @@ -263,7 +337,7 @@ describe('OpenV1BaseService', () => { it('throws when response body is null', async () => { const nullBodyFetch: typeof globalThis.fetch = vi .fn() - .mockResolvedValue({ body: null }); + .mockResolvedValue(makeOkResponse({ body: null })); const streamService = new OpenV1BaseService( nullBodyFetch, 'http://localhost:11434', @@ -280,5 +354,59 @@ describe('OpenV1BaseService', () => { } }).rejects.toThrow('No response body for streaming'); }); + + it('throws a clear error when the HTTP status is not ok before reading the stream', async () => { + const errFetch = vi.fn().mockResolvedValue({ + ok: false, + status: 500, + statusText: 'Internal Server Error', + text: vi.fn().mockResolvedValue('upstream failure'), + body: makeSSEStream([makeStreamChunk('x')]), + } as unknown as Response); + const errService = new OpenV1BaseService( + errFetch, + 'http://localhost:11434', + ); + + await expect(async () => { + // eslint-disable-next-line @typescript-eslint/no-unused-vars + for await (const _ of errService.chat({ + model: MODEL, + messages: [{ role: 'user', content: 'hi' }], + stream: true, + })) { + // drain + } + }).rejects.toThrow('HTTP 500 Internal Server Error — upstream failure'); + }); + + it('throws a descriptive error when an SSE data line is not valid JSON', async () => { + const encoder = new TextEncoder(); + // eslint-disable-next-line n/no-unsupported-features/node-builtins + const badBody = new ReadableStream({ + start(controller) { + controller.enqueue(encoder.encode('data: not-json\n\n')); + controller.close(); + }, + }); + const badFetch = vi + .fn() + .mockResolvedValue(makeOkResponse({ body: badBody })); + const badService = new OpenV1BaseService( + badFetch, + 'http://localhost:11434', + ); + + await expect(async () => { + // eslint-disable-next-line @typescript-eslint/no-unused-vars + for await (const _ of badService.chat({ + model: MODEL, + messages: [{ role: 'user', content: 'hi' }], + stream: true, + })) { + // drain + } + }).rejects.toThrow(SyntaxError); + }); }); }); diff --git a/packages/kernel-language-model-service/src/open-v1/base.ts b/packages/kernel-language-model-service/src/open-v1/base.ts index 91ccc0e22e..d784c10208 100644 --- a/packages/kernel-language-model-service/src/open-v1/base.ts +++ b/packages/kernel-language-model-service/src/open-v1/base.ts @@ -1,7 +1,13 @@ import { assert } from '@metamask/superstruct'; import type { ChatParams, ChatResult, ChatStreamChunk } from '../types.ts'; -import { ChatParamsStruct } from './types.ts'; +import { checkResponseOk, readAndCheckResponse } from './response-json.ts'; +import { + ChatParamsStruct, + ChatResultStruct, + ChatStreamChunkStruct, + ListModelsResponseStruct, +} from './types.ts'; /** * Base service for any Open /v1-compatible HTTP endpoint. @@ -75,8 +81,10 @@ export class OpenV1BaseService { headers: this.#makeHeaders(), body: JSON.stringify({ ...params, stream: false }), }); - const result = (await response.json()) as ChatResult; - return harden(result); + const body = await readAndCheckResponse(response); + const json: unknown = JSON.parse(body); + assert(json, ChatResultStruct); + return harden(json as ChatResult); } /** @@ -89,6 +97,7 @@ export class OpenV1BaseService { headers: this.#makeHeaders(), body: JSON.stringify({ ...params, stream: true }), }); + await checkResponseOk(response); if (!response.body) { throw new Error('No response body for streaming'); } @@ -111,7 +120,9 @@ export class OpenV1BaseService { return; } if (data) { - yield harden(JSON.parse(data) as ChatStreamChunk); + const json: unknown = JSON.parse(data); + assert(json, ChatStreamChunkStruct); + yield harden(json as ChatStreamChunk); } } } @@ -133,8 +144,10 @@ export class OpenV1BaseService { const response = await this.#fetch(`${this.#baseUrl}/v1/models`, { headers: this.#makeHeaders(), }); - const data = (await response.json()) as { data: { id: string }[] }; - return harden(data.data.map((model) => model.id)); + const body = await readAndCheckResponse(response); + const json: unknown = JSON.parse(body); + assert(json, ListModelsResponseStruct); + return harden(json.data.map((model) => model.id)); } /** diff --git a/packages/kernel-language-model-service/src/open-v1/response-json.test.ts b/packages/kernel-language-model-service/src/open-v1/response-json.test.ts new file mode 100644 index 0000000000..05e4196cec --- /dev/null +++ b/packages/kernel-language-model-service/src/open-v1/response-json.test.ts @@ -0,0 +1,48 @@ +import { describe, expect, it, vi } from 'vitest'; + +import { checkResponseOk, readAndCheckResponse } from './response-json.ts'; + +describe('readAndCheckResponse', () => { + it('returns body text when response is ok', async () => { + const response = { + ok: true, + status: 200, + statusText: 'OK', + text: async () => Promise.resolve('{"x":1}'), + } as Response; + expect(await readAndCheckResponse(response)).toBe('{"x":1}'); + }); + + it('throws with status and body snippet when response is not ok', async () => { + const response = { + ok: false, + status: 503, + statusText: 'Service Unavailable', + text: async () => Promise.resolve('upstream down'), + } as Response; + await expect(readAndCheckResponse(response)).rejects.toThrow( + 'HTTP 503 Service Unavailable — upstream down', + ); + }); +}); + +describe('checkResponseOk', () => { + it('resolves without reading the body when ok', async () => { + const text = vi.fn(); + const response = { ok: true, text } as unknown as Response; + await checkResponseOk(response); + expect(text).not.toHaveBeenCalled(); + }); + + it('reads body and throws when not ok', async () => { + const response = { + ok: false, + status: 500, + statusText: 'Error', + text: async () => Promise.resolve('fail'), + } as Response; + await expect(checkResponseOk(response)).rejects.toThrow( + 'HTTP 500 Error — fail', + ); + }); +}); diff --git a/packages/kernel-language-model-service/src/open-v1/response-json.ts b/packages/kernel-language-model-service/src/open-v1/response-json.ts new file mode 100644 index 0000000000..a22af0b034 --- /dev/null +++ b/packages/kernel-language-model-service/src/open-v1/response-json.ts @@ -0,0 +1,33 @@ +/** + * Read the response body as UTF-8 text. Throws if {@link Response.ok} is false. + * + * @param response - The fetch response. + * @returns The full response body text. + */ +export async function readAndCheckResponse( + response: Response, +): Promise { + const body = await response.text(); + if (!response.ok) { + throw new Error( + `HTTP ${response.status} ${response.statusText} — ${body.slice(0, 500)}`, + ); + } + return body; +} + +/** + * If {@link Response.ok} is false, read the body as text and throw. + * When OK, returns without reading the body so the stream remains available. + * + * @param response - The fetch response. + */ +export async function checkResponseOk(response: Response): Promise { + if (response.ok) { + return; + } + const body = await response.text(); + throw new Error( + `HTTP ${response.status} ${response.statusText} — ${body.slice(0, 500)}`, + ); +} diff --git a/packages/kernel-language-model-service/src/open-v1/types.ts b/packages/kernel-language-model-service/src/open-v1/types.ts index b63ace1374..38a2a6da92 100644 --- a/packages/kernel-language-model-service/src/open-v1/types.ts +++ b/packages/kernel-language-model-service/src/open-v1/types.ts @@ -1,5 +1,4 @@ import { - any, array, boolean, literal, @@ -10,6 +9,7 @@ import { size, string, union, + unknown, } from '@metamask/superstruct'; export type { @@ -51,7 +51,7 @@ const ToolStruct = object({ function: object({ name: string(), description: optional(string()), - parameters: optional(any()), + parameters: optional(unknown()), }), }); @@ -72,3 +72,53 @@ export const ChatParamsStruct = object({ n: optional(number()), stream: optional(boolean()), }); + +const UsageStruct = object({ + prompt_tokens: number(), + completion_tokens: number(), + total_tokens: number(), +}); + +const ChatChoiceStruct = object({ + message: ChatMessageStruct, + index: number(), + finish_reason: nullable(string()), +}); + +/** + * Superstruct schema for a non-streaming `/v1/chat/completions` response body. + */ +export const ChatResultStruct = object({ + id: string(), + model: string(), + choices: array(ChatChoiceStruct), + usage: UsageStruct, +}); + +const ChatStreamDeltaStruct = object({ + role: optional(ChatRoleStruct), + content: optional(string()), + tool_calls: optional(array(unknown())), +}); + +/** + * Superstruct schema for one SSE `data:` JSON object from `/v1/chat/completions` when `stream: true`. + */ +export const ChatStreamChunkStruct = object({ + id: string(), + model: string(), + choices: array( + object({ + delta: ChatStreamDeltaStruct, + index: number(), + finish_reason: nullable(string()), + }), + ), +}); + +/** + * Superstruct schema for a `/v1/models` response body (only fields we read). + */ +export const ListModelsResponseStruct = object({ + data: array(object({ id: string() })), +}); diff --git a/packages/kernel-language-model-service/src/utils/parse-tool-arguments.test.ts b/packages/kernel-language-model-service/src/utils/parse-tool-arguments.test.ts new file mode 100644 index 0000000000..60aae1cd6a --- /dev/null +++ b/packages/kernel-language-model-service/src/utils/parse-tool-arguments.test.ts @@ -0,0 +1,48 @@ +import { describe, expect, it } from 'vitest'; + +import { parseJsonObject, parseToolArguments } from './parse-tool-arguments.ts'; + +const chatLabels = { + invalidJson: 'Chat completion: invalid JSON in response body', + notObject: 'Chat completion: response must be a JSON object', +} as const; + +describe('parseJsonObject', () => { + it('uses custom labels for invalid JSON', () => { + expect(() => parseJsonObject('{x', chatLabels)).toThrow(SyntaxError); + expect(() => parseJsonObject('{x', chatLabels)).toThrow( + /Chat completion: invalid JSON in response body/u, + ); + }); + + it('uses custom labels when the value is not an object', () => { + expect(() => parseJsonObject('[]', chatLabels)).toThrow( + 'Chat completion: response must be a JSON object', + ); + }); +}); + +describe('parseToolArguments', () => { + it('returns a plain object for valid JSON object text', () => { + expect(parseToolArguments('{"a":1}')).toStrictEqual({ a: 1 }); + }); + + it('throws SyntaxError with context when JSON is invalid', () => { + expect(() => parseToolArguments('{not json')).toThrow(SyntaxError); + expect(() => parseToolArguments('{not json')).toThrow( + /Invalid tool arguments JSON/u, + ); + }); + + it('throws when the value is not a plain object', () => { + expect(() => parseToolArguments('[]')).toThrow( + 'Tool arguments must be a JSON object', + ); + expect(() => parseToolArguments('null')).toThrow( + 'Tool arguments must be a JSON object', + ); + expect(() => parseToolArguments('"x"')).toThrow( + 'Tool arguments must be a JSON object', + ); + }); +}); diff --git a/packages/kernel-language-model-service/src/utils/parse-tool-arguments.ts b/packages/kernel-language-model-service/src/utils/parse-tool-arguments.ts new file mode 100644 index 0000000000..2891292b08 --- /dev/null +++ b/packages/kernel-language-model-service/src/utils/parse-tool-arguments.ts @@ -0,0 +1,51 @@ +export type ParseJsonObjectLabels = { + /** Prefix for invalid JSON (message continues with preview and parse error). */ + invalidJson: string; + /** Message when JSON parses but the top-level value is not a plain object. */ + notObject: string; +}; + +/** + * Parse JSON text and ensure the top-level value is a plain object. + * + * @param json - Raw JSON text. + * @param labels - Human-readable labels for thrown {@link SyntaxError} messages. + * @returns The parsed object. + * @throws {SyntaxError} When JSON is invalid or the value is not a plain object. + */ +export function parseJsonObject( + json: string, + labels: ParseJsonObjectLabels, +): Record { + let parsed: unknown; + try { + parsed = JSON.parse(json); + } catch (cause) { + const preview = json.length > 200 ? `${json.slice(0, 200)}…` : json; + throw new SyntaxError( + `${labels.invalidJson} (${preview}): ${ + cause instanceof Error ? cause.message : String(cause) + }`, + { cause }, + ); + } + if (typeof parsed !== 'object' || parsed === null || Array.isArray(parsed)) { + throw new SyntaxError(labels.notObject); + } + return parsed as Record; +} + +/** + * Parse a tool call `function.arguments` string (JSON) into a plain object. + * Used for OpenAI-style chat `tool_calls` and when adapting messages for Ollama. + * + * @param json - Raw JSON object text from the model. + * @returns Parsed object for APIs that expect a record. + * @throws {SyntaxError} When JSON is invalid or the value is not a plain object. + */ +export function parseToolArguments(json: string): Record { + return parseJsonObject(json, { + invalidJson: 'Invalid tool arguments JSON', + notObject: 'Tool arguments must be a JSON object', + }); +} From 7c8311ad63f5bf1d9a4d3a800c790ba6208ad303 Mon Sep 17 00:00:00 2001 From: grypez <143971198+grypez@users.noreply.github.com> Date: Tue, 31 Mar 2026 19:55:13 -0400 Subject: [PATCH 11/20] validate tool args against spec Made-with: Cursor --- packages/kernel-agents/package.json | 1 + .../validate-capability-args.test.ts | 59 +++++++++ .../capabilities/validate-capability-args.ts | 17 +++ .../src/strategies/chat-agent.ts | 21 +-- packages/kernel-utils/package.json | 10 ++ packages/kernel-utils/src/index.test.ts | 2 + packages/kernel-utils/src/index.ts | 4 + .../src/json-schema-to-struct.test.ts | 73 +++++++++++ .../kernel-utils/src/json-schema-to-struct.ts | 120 ++++++++++++++++++ yarn.lock | 1 + 10 files changed, 300 insertions(+), 8 deletions(-) create mode 100644 packages/kernel-agents/src/capabilities/validate-capability-args.test.ts create mode 100644 packages/kernel-agents/src/capabilities/validate-capability-args.ts create mode 100644 packages/kernel-utils/src/json-schema-to-struct.test.ts create mode 100644 packages/kernel-utils/src/json-schema-to-struct.ts diff --git a/packages/kernel-agents/package.json b/packages/kernel-agents/package.json index 2f8c290725..d8d54ad835 100644 --- a/packages/kernel-agents/package.json +++ b/packages/kernel-agents/package.json @@ -196,6 +196,7 @@ "@metamask/kernel-errors": "workspace:^", "@metamask/kernel-utils": "workspace:^", "@metamask/logger": "workspace:^", + "@metamask/superstruct": "^3.2.1", "@ocap/kernel-language-model-service": "workspace:^", "partial-json": "^0.1.7", "ses": "^1.14.0" diff --git a/packages/kernel-agents/src/capabilities/validate-capability-args.test.ts b/packages/kernel-agents/src/capabilities/validate-capability-args.test.ts new file mode 100644 index 0000000000..21bdd4b0ca --- /dev/null +++ b/packages/kernel-agents/src/capabilities/validate-capability-args.test.ts @@ -0,0 +1,59 @@ +import { describe, expect, it } from 'vitest'; + +import { validateCapabilityArgs } from './validate-capability-args.ts'; + +describe('validateCapabilityArgs', () => { + it('accepts values matching primitive arg schemas', () => { + expect(() => + validateCapabilityArgs( + { a: 1, b: 2 }, + { + description: 'add', + args: { + a: { type: 'number' }, + b: { type: 'number' }, + }, + }, + ), + ).not.toThrow(); + }); + + it('throws when a required argument is missing', () => { + expect(() => + validateCapabilityArgs( + { a: 1 }, + { + description: 'add', + args: { + a: { type: 'number' }, + b: { type: 'number' }, + }, + }, + ), + ).toThrow(/At path: b -- Expected a number/u); + }); + + it('throws when a value does not match the schema', () => { + expect(() => + validateCapabilityArgs( + { a: 'not-a-number' }, + { + description: 'x', + args: { a: { type: 'number' } }, + }, + ), + ).toThrow(/Expected a number/u); + }); + + it('does nothing when there are no declared arguments', () => { + expect(() => + validateCapabilityArgs( + { extra: 1 }, + { + description: 'ping', + args: {}, + }, + ), + ).not.toThrow(); + }); +}); diff --git a/packages/kernel-agents/src/capabilities/validate-capability-args.ts b/packages/kernel-agents/src/capabilities/validate-capability-args.ts new file mode 100644 index 0000000000..7b4668f26e --- /dev/null +++ b/packages/kernel-agents/src/capabilities/validate-capability-args.ts @@ -0,0 +1,17 @@ +import { methodArgsToStruct } from '@metamask/kernel-utils/json-schema-to-struct'; +import { assert } from '@metamask/superstruct'; + +import type { CapabilitySchema } from '../types.ts'; + +/** + * Assert `values` match the capability's declared argument schemas using Superstruct. + * + * @param values - Parsed tool arguments (a plain object). + * @param capabilitySchema - {@link CapabilitySchema} for this capability. + */ +export function validateCapabilityArgs( + values: Record, + capabilitySchema: CapabilitySchema, +): void { + assert(values, methodArgsToStruct(capabilitySchema.args)); +} diff --git a/packages/kernel-agents/src/strategies/chat-agent.ts b/packages/kernel-agents/src/strategies/chat-agent.ts index 2bfd4a8ee3..3885e60966 100644 --- a/packages/kernel-agents/src/strategies/chat-agent.ts +++ b/packages/kernel-agents/src/strategies/chat-agent.ts @@ -4,11 +4,10 @@ import type { ChatResult, Tool, } from '@ocap/kernel-language-model-service'; +import { parseToolArguments } from '@ocap/kernel-language-model-service/utils/parse-tool-arguments'; -import { - extractCapabilitySchemas, - extractCapabilities, -} from '../capabilities/capability.ts'; +import { extractCapabilitySchemas } from '../capabilities/capability.ts'; +import { validateCapabilityArgs } from '../capabilities/validate-capability-args.ts'; import type { Agent } from '../types/agent.ts'; import { Message } from '../types/messages.ts'; import type { CapabilityRecord, Experience } from '../types.ts'; @@ -25,9 +24,12 @@ class ChatTurn extends Message { */ constructor({ role, content }: ChatMessage) { super(role, { content: content ?? '' }); + harden(this); } } +harden(ChatTurn); + /** * A bound chat function with the model already configured. * Construct one from a {@link ChatService} using `makeChatClient`: @@ -114,7 +116,6 @@ export const makeChatAgent = ({ const objective = { intent, judgment: effectiveJudgment }; const context = { capabilities: agentCapabilities }; - const capFunctions = extractCapabilities(agentCapabilities); const tools = buildTools(agentCapabilities); const chatHistory: ChatMessage[] = [{ role: 'user', content: intent }]; @@ -154,8 +155,10 @@ export const makeChatAgent = ({ const { name, arguments: argsJson } = toolCall.function; logger?.info(`Invoking capability: ${name}`); - const cap = capFunctions[name]; - if (!cap) { + const spec = Object.hasOwn(agentCapabilities, name) + ? agentCapabilities[name] + : undefined; + if (spec === undefined) { const errorContent = `Unknown capability "${name}"`; const toolMsg: ChatMessage = { role: 'tool', @@ -169,7 +172,9 @@ export const makeChatAgent = ({ let toolResult: unknown; try { - toolResult = await cap(JSON.parse(argsJson) as never); + const args = parseToolArguments(argsJson); + validateCapabilityArgs(args, spec.schema); + toolResult = await spec.func(args as never); } catch (error) { const errorContent = `Error calling ${name}: ${(error as Error).message}`; const toolMsg: ChatMessage = { diff --git a/packages/kernel-utils/package.json b/packages/kernel-utils/package.json index 0f6d0ebff0..4a5178b13f 100644 --- a/packages/kernel-utils/package.json +++ b/packages/kernel-utils/package.json @@ -39,6 +39,16 @@ "default": "./dist/exo.cjs" } }, + "./json-schema-to-struct": { + "import": { + "types": "./dist/json-schema-to-struct.d.mts", + "default": "./dist/json-schema-to-struct.mjs" + }, + "require": { + "types": "./dist/json-schema-to-struct.d.cts", + "default": "./dist/json-schema-to-struct.cjs" + } + }, "./discoverable": { "import": { "types": "./dist/discoverable.d.mts", diff --git a/packages/kernel-utils/src/index.test.ts b/packages/kernel-utils/src/index.test.ts index e930280adf..1fe47eb9ed 100644 --- a/packages/kernel-utils/src/index.test.ts +++ b/packages/kernel-utils/src/index.test.ts @@ -24,11 +24,13 @@ describe('index', () => { 'isTypedArray', 'isTypedObject', 'isVatBundle', + 'jsonSchemaToStruct', 'makeCounter', 'makeDefaultExo', 'makeDefaultInterface', 'makeDiscoverableExo', 'mergeDisjointRecords', + 'methodArgsToStruct', 'prettifySmallcaps', 'retry', 'retryWithBackoff', diff --git a/packages/kernel-utils/src/index.ts b/packages/kernel-utils/src/index.ts index c50e445926..5857728891 100644 --- a/packages/kernel-utils/src/index.ts +++ b/packages/kernel-utils/src/index.ts @@ -3,6 +3,10 @@ export { makeDefaultInterface, makeDefaultExo } from './exo.ts'; export { GET_DESCRIPTION, makeDiscoverableExo } from './discoverable.ts'; export type { DiscoverableExo } from './discoverable.ts'; export type { JsonSchema, MethodSchema } from './schema.ts'; +export { + jsonSchemaToStruct, + methodArgsToStruct, +} from './json-schema-to-struct.ts'; export { fetchValidatedJson } from './fetchValidatedJson.ts'; export { abortableDelay, delay, ifDefined, makeCounter } from './misc.ts'; export { stringify } from './stringify.ts'; diff --git a/packages/kernel-utils/src/json-schema-to-struct.test.ts b/packages/kernel-utils/src/json-schema-to-struct.test.ts new file mode 100644 index 0000000000..3d25a7d211 --- /dev/null +++ b/packages/kernel-utils/src/json-schema-to-struct.test.ts @@ -0,0 +1,73 @@ +import { assert } from '@metamask/superstruct'; +import { describe, expect, it } from 'vitest'; + +import { + jsonSchemaToStruct, + methodArgsToStruct, +} from './json-schema-to-struct.ts'; + +describe('jsonSchemaToStruct', () => { + it('validates string, number, and boolean', () => { + assert('x', jsonSchemaToStruct({ type: 'string' })); + assert(1, jsonSchemaToStruct({ type: 'number' })); + assert(true, jsonSchemaToStruct({ type: 'boolean' })); + }); + + it('validates arrays recursively', () => { + assert( + [1, 2], + jsonSchemaToStruct({ type: 'array', items: { type: 'number' } }), + ); + }); + + it('validates nested objects and required keys', () => { + const struct = jsonSchemaToStruct({ + type: 'object', + properties: { + a: { type: 'string' }, + b: { type: 'number' }, + }, + }); + assert({ a: 'hi', b: 1 }, struct); + expect(() => assert({ a: 'hi' }, struct)).toThrow( + /Missing required property "b"/u, + ); + }); + + it('rejects unknown keys when additionalProperties is false', () => { + const struct = jsonSchemaToStruct({ + type: 'object', + properties: { a: { type: 'string' } }, + additionalProperties: false, + }); + assert({ a: 'x' }, struct); + expect(() => assert({ a: 'x', b: 1 }, struct)).toThrow(/path: b/u); + }); + + it('allows unknown keys on objects when additionalProperties is not false', () => { + const struct = jsonSchemaToStruct({ + type: 'object', + properties: { a: { type: 'number' } }, + }); + assert({ a: 1, extra: 'ignored' }, struct); + }); +}); + +describe('methodArgsToStruct', () => { + it('builds an object struct for method args', () => { + const struct = methodArgsToStruct({ + a: { type: 'number' }, + b: { type: 'number' }, + }); + assert({ a: 1, b: 2 }, struct); + expect(() => assert({ a: 1 }, struct)).toThrow(/path: b/u); + }); + + it('accepts an empty args map', () => { + assert({}, methodArgsToStruct({})); + }); + + it('allows extra keys when the args map is empty', () => { + assert({ extra: 1 }, methodArgsToStruct({})); + }); +}); diff --git a/packages/kernel-utils/src/json-schema-to-struct.ts b/packages/kernel-utils/src/json-schema-to-struct.ts new file mode 100644 index 0000000000..449b50e895 --- /dev/null +++ b/packages/kernel-utils/src/json-schema-to-struct.ts @@ -0,0 +1,120 @@ +import { + array, + assert, + boolean, + define, + number, + object, + optional, + string, + StructError, +} from '@metamask/superstruct'; +import type { Struct } from '@metamask/superstruct'; + +import type { JsonSchema } from './schema.ts'; + +type ObjectJsonSchema = Extract; + +/** + * Object schema where unknown property names are allowed (when `additionalProperties` is not `false`). + * Known properties are validated; required keys must be present. + * + * @param schema - JSON Schema with `type: 'object'`, `properties`, and optional `required`. + * @returns A Superstruct validator for plain objects matching the loose object rules. + */ +function looseObjectStruct(schema: ObjectJsonSchema): Struct { + const { properties } = schema; + const required = new Set(schema.required ?? Object.keys(properties)); + return define('JsonSchemaObject', (value) => { + if (typeof value !== 'object' || value === null || Array.isArray(value)) { + return 'Expected a plain object'; + } + const obj = value as Record; + for (const key of required) { + if (!(key in obj)) { + return `Missing required property "${key}"`; + } + } + for (const [key, subSchema] of Object.entries(properties)) { + if (!(key in obj)) { + continue; + } + try { + assert(obj[key], jsonSchemaToStruct(subSchema)); + } catch (caught) { + if (caught instanceof StructError) { + return `At ${key}: ${caught.message}`; + } + throw caught; + } + } + return true; + }) as Struct; +} + +/** + * Build a Superstruct {@link Struct} from our {@link JsonSchema} subset (primitives, + * arrays, objects). Used to reuse Superstruct validation for values described by + * discoverable exo / capability {@link MethodSchema} argument shapes. + * + * @param schema - JSON Schema value (must include `type`). + * @returns A struct that validates the same shapes as the prior hand-rolled checks. + */ +export function jsonSchemaToStruct(schema: JsonSchema): Struct { + switch (schema.type) { + case 'string': + return string() as Struct; + case 'number': + return number() as Struct; + case 'boolean': + return boolean() as Struct; + case 'array': + return array(jsonSchemaToStruct(schema.items)) as Struct; + case 'object': { + const { properties } = schema; + const required = new Set(schema.required ?? Object.keys(properties)); + const shape = Object.fromEntries( + Object.entries(properties).map(([key, subSchema]) => { + const fieldStruct = jsonSchemaToStruct(subSchema); + return [key, required.has(key) ? fieldStruct : optional(fieldStruct)]; + }), + ); + if (schema.additionalProperties === false) { + return object(shape) as Struct; + } + if (Object.keys(properties).length === 0) { + return looseObjectStruct(schema); + } + return looseObjectStruct(schema); + } + default: { + const _never: never = schema; + throw new TypeError(`Unsupported JSON schema: ${String(_never)}`); + } + } +} + +/** + * Build a Superstruct object struct for a method/capability `args` map + * (name → per-argument {@link JsonSchema}). All listed arguments are required. + * + * @param args - Same shape as {@link MethodSchema.args}. + * @returns A struct that validates a plain object with one field per declared argument. + */ +export function methodArgsToStruct( + args: Record, +): Struct> { + const entries = Object.entries(args); + if (entries.length === 0) { + return define('EmptyCapabilityArgs', (value) => { + if (typeof value !== 'object' || value === null || Array.isArray(value)) { + return 'Expected a plain object'; + } + return true; + }) as Struct>; + } + const shape = Object.fromEntries( + entries.map(([name, jsonSchema]) => [name, jsonSchemaToStruct(jsonSchema)]), + ); + return object(shape) as Struct>; +} diff --git a/yarn.lock b/yarn.lock index af1601d16b..b8c87be4fa 100644 --- a/yarn.lock +++ b/yarn.lock @@ -3914,6 +3914,7 @@ __metadata: "@metamask/kernel-errors": "workspace:^" "@metamask/kernel-utils": "workspace:^" "@metamask/logger": "workspace:^" + "@metamask/superstruct": "npm:^3.2.1" "@ocap/kernel-language-model-service": "workspace:^" "@ocap/repo-tools": "workspace:^" "@ts-bridge/cli": "npm:^0.6.3" From 20aa2ebaa7dcfd7720458af0914a931b1549df11 Mon Sep 17 00:00:00 2001 From: grypez <143971198+grypez@users.noreply.github.com> Date: Tue, 31 Mar 2026 20:41:48 -0400 Subject: [PATCH 12/20] differentiate chat message types Made-with: Cursor --- .../src/strategies/chat-agent.ts | 13 ++-- .../src/ollama/base.ts | 70 +++++++++++-------- .../src/open-v1/types.ts | 28 +++++++- .../src/types.ts | 36 ++++++++-- 4 files changed, 106 insertions(+), 41 deletions(-) diff --git a/packages/kernel-agents/src/strategies/chat-agent.ts b/packages/kernel-agents/src/strategies/chat-agent.ts index 3885e60966..2ee5474f1f 100644 --- a/packages/kernel-agents/src/strategies/chat-agent.ts +++ b/packages/kernel-agents/src/strategies/chat-agent.ts @@ -22,8 +22,12 @@ class ChatTurn extends Message { * @param chatMessage.role - The sender role of the message. * @param chatMessage.content - The text content of the message. */ - constructor({ role, content }: ChatMessage) { - super(role, { content: content ?? '' }); + constructor(chatMessage: ChatMessage) { + const content = + chatMessage.role === 'assistant' + ? (chatMessage.content ?? '') + : chatMessage.content; + super(chatMessage.role, { content }); harden(this); } } @@ -132,10 +136,11 @@ export const makeChatAgent = ({ messages: chatHistory, ...(tools.length > 0 && { tools }), }); - const assistantMessage = chatResult.choices[0]?.message; - if (!assistantMessage) { + const choiceMessage = chatResult.choices[0]?.message; + if (!choiceMessage || choiceMessage.role !== 'assistant') { throw new Error('No response from model'); } + const assistantMessage = choiceMessage; chatHistory.push(assistantMessage); history.push(new ChatTurn(assistantMessage)); diff --git a/packages/kernel-language-model-service/src/ollama/base.ts b/packages/kernel-language-model-service/src/ollama/base.ts index c567cb1651..4eb75393e2 100644 --- a/packages/kernel-language-model-service/src/ollama/base.ts +++ b/packages/kernel-language-model-service/src/ollama/base.ts @@ -6,9 +6,9 @@ import type { } from 'ollama'; import type { + AssistantMessage, ChatParams, ChatResult, - ChatRole, LanguageModelService, SampleParams, SampleResult, @@ -83,21 +83,30 @@ export class OllamaBaseService } const response = await ollama.chat({ model, - // eslint-disable-next-line camelcase - messages: messages.map(({ role, content, tool_calls }) => ({ - role, - content: content ?? '', - // eslint-disable-next-line camelcase - ...(tool_calls && { - // eslint-disable-next-line camelcase - tool_calls: tool_calls.map(({ function: fn }) => ({ - function: { - name: fn.name, - arguments: parseToolArguments(fn.arguments), - }, - })), - }), - })), + messages: messages.map((chatMessage) => { + if (chatMessage.role === 'tool') { + return { + role: 'tool' as const, + content: chatMessage.content, + tool_call_id: chatMessage.tool_call_id, + }; + } + if (chatMessage.role === 'assistant') { + return { + role: 'assistant' as const, + content: chatMessage.content ?? '', + ...(chatMessage.tool_calls && { + tool_calls: chatMessage.tool_calls.map(({ function: fn }) => ({ + function: { + name: fn.name, + arguments: parseToolArguments(fn.arguments), + }, + })), + }), + }; + } + return { role: chatMessage.role, content: chatMessage.content }; + }), ...(tools && { tools: tools as OllamaTool[] }), stream: false, options: ifDefined({ @@ -111,25 +120,26 @@ export class OllamaBaseService const promptTokens = response.prompt_eval_count ?? 0; const completionTokens = response.eval_count ?? 0; const { tool_calls: responseToolCalls } = response.message; + const assistantMessage: AssistantMessage = { + role: 'assistant', + content: response.message.content, + ...(responseToolCalls && { + tool_calls: responseToolCalls.map((tc, index) => ({ + id: `tool-${index}`, + type: 'function' as const, + function: { + name: tc.function.name, + arguments: JSON.stringify(tc.function.arguments), + }, + })), + }), + }; return harden({ id: 'ollama-chat', model: response.model, choices: [ { - message: { - role: response.message.role as ChatRole, - content: response.message.content, - ...(responseToolCalls && { - tool_calls: responseToolCalls.map((tc, index) => ({ - id: `tool-${index}`, - type: 'function' as const, - function: { - name: tc.function.name, - arguments: JSON.stringify(tc.function.arguments), - }, - })), - }), - }, + message: assistantMessage, index: 0, finish_reason: response.done_reason ?? 'stop', }, diff --git a/packages/kernel-language-model-service/src/open-v1/types.ts b/packages/kernel-language-model-service/src/open-v1/types.ts index 38a2a6da92..8ca42abdfe 100644 --- a/packages/kernel-language-model-service/src/open-v1/types.ts +++ b/packages/kernel-language-model-service/src/open-v1/types.ts @@ -39,13 +39,35 @@ const ToolCallStruct = object({ function: object({ name: string(), arguments: string() }), }); -const ChatMessageStruct = object({ - role: ChatRoleStruct, +const SystemMessageStruct = object({ + role: literal('system'), + content: string(), +}); + +const UserMessageStruct = object({ + role: literal('user'), + content: string(), +}); + +const AssistantMessageStruct = object({ + role: literal('assistant'), content: nullable(string()), tool_calls: optional(array(ToolCallStruct)), - tool_call_id: optional(string()), }); +const ToolMessageStruct = object({ + role: literal('tool'), + content: string(), + tool_call_id: string(), +}); + +const ChatMessageStruct = union([ + SystemMessageStruct, + UserMessageStruct, + AssistantMessageStruct, + ToolMessageStruct, +]); + const ToolStruct = object({ type: literal('function'), function: object({ diff --git a/packages/kernel-language-model-service/src/types.ts b/packages/kernel-language-model-service/src/types.ts index 52fa41d5be..d91ce64454 100644 --- a/packages/kernel-language-model-service/src/types.ts +++ b/packages/kernel-language-model-service/src/types.ts @@ -30,15 +30,43 @@ export type Tool = { }; /** - * A single message in a chat conversation. + * Chat message from the model or system (no tool metadata). */ -export type ChatMessage = { - role: ChatRole; +export type SystemMessage = { role: 'system'; content: string }; + +/** + * End-user turn (no tool metadata). + */ +export type UserMessage = { role: 'user'; content: string }; + +/** + * Assistant turn; may include tool calls. `content` may be null when the model + * only returns tool calls. + */ +export type AssistantMessage = { + role: 'assistant'; content: string | null; tool_calls?: ToolCall[]; - tool_call_id?: string; }; +/** + * Result of a tool invocation, correlated by {@link ToolCall.id}. + */ +export type ToolMessage = { + role: 'tool'; + content: string; + tool_call_id: string; +}; + +/** + * A single message in a chat conversation (discriminated by `role`). + */ +export type ChatMessage = + | SystemMessage + | UserMessage + | AssistantMessage + | ToolMessage; + /** * Token usage statistics for a completion request. */ From d9ead80d6a3563c408054fa9539e67dc4406a9d9 Mon Sep 17 00:00:00 2001 From: grypez <143971198+grypez@users.noreply.github.com> Date: Tue, 31 Mar 2026 20:43:09 -0400 Subject: [PATCH 13/20] tighten streaming chat delta types Made-with: Cursor --- .../src/open-v1/base.test.ts | 73 +++++++++++++++++++ .../src/open-v1/types.ts | 25 ++++--- .../src/types.ts | 17 ++++- 3 files changed, 103 insertions(+), 12 deletions(-) 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 index 75535d8548..82a2ab5a9e 100644 --- a/packages/kernel-language-model-service/src/open-v1/base.test.ts +++ b/packages/kernel-language-model-service/src/open-v1/base.test.ts @@ -1,3 +1,4 @@ +import { StructError } from '@metamask/superstruct'; import { describe, it, expect, vi, beforeEach } from 'vitest'; import type { ChatResult, ChatStreamChunk } from '../types.ts'; @@ -334,6 +335,78 @@ describe('OpenV1BaseService', () => { expect(received).toStrictEqual(expected); }); + it('accepts streaming delta with assistant role and tool_calls fragments', async () => { + const toolChunk: ChatStreamChunk = { + id: 'chat-1', + model: MODEL, + choices: [ + { + delta: { + role: 'assistant', + tool_calls: [ + { + index: 0, + id: 'call_1', + type: 'function', + function: { name: 'fn' }, + }, + ], + }, + index: 0, + finish_reason: null, + }, + ], + }; + const streamFetch = makeMockStreamFetch([toolChunk]); + 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([toolChunk]); + }); + + it('rejects streaming chunk when delta role is not assistant', async () => { + const badChunk = { + id: 'chat-1', + model: MODEL, + choices: [ + { + delta: { role: 'user', content: 'x' }, + index: 0, + finish_reason: null, + }, + ], + }; + const streamFetch = makeMockStreamFetch([ + badChunk as unknown as ChatStreamChunk, + ]); + const streamService = new OpenV1BaseService( + streamFetch, + '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(StructError); + }); + it('throws when response body is null', async () => { const nullBodyFetch: typeof globalThis.fetch = vi .fn() diff --git a/packages/kernel-language-model-service/src/open-v1/types.ts b/packages/kernel-language-model-service/src/open-v1/types.ts index 8ca42abdfe..e97f70c26f 100644 --- a/packages/kernel-language-model-service/src/open-v1/types.ts +++ b/packages/kernel-language-model-service/src/open-v1/types.ts @@ -17,21 +17,14 @@ export type { ChatMessage, ChatParams, ChatResult, - ChatRole, ChatStreamChunk, ChatStreamDelta, + ChatStreamToolCallDelta, Tool, ToolCall, Usage, } from '../types.ts'; -const ChatRoleStruct = union([ - literal('system'), - literal('user'), - literal('assistant'), - literal('tool'), -]); - const ToolCallStruct = object({ id: string(), type: literal('function'), @@ -117,10 +110,22 @@ export const ChatResultStruct = object({ usage: UsageStruct, }); +const ChatStreamToolCallDeltaStruct = object({ + index: optional(number()), + id: optional(string()), + type: optional(literal('function')), + function: optional( + object({ + name: optional(string()), + arguments: optional(string()), + }), + ), +}); + const ChatStreamDeltaStruct = object({ - role: optional(ChatRoleStruct), + role: optional(literal('assistant')), content: optional(string()), - tool_calls: optional(array(unknown())), + tool_calls: optional(array(ChatStreamToolCallDeltaStruct)), }); /** diff --git a/packages/kernel-language-model-service/src/types.ts b/packages/kernel-language-model-service/src/types.ts index d91ce64454..1cc8000dbd 100644 --- a/packages/kernel-language-model-service/src/types.ts +++ b/packages/kernel-language-model-service/src/types.ts @@ -120,11 +120,24 @@ export type ChatParams = { }; /** - * A partial message delta from a streaming chat completion response. + * Partial tool-call fragment in a streaming completion delta (OpenAI-style SSE). + * Fields arrive incrementally across multiple chunks. + */ +export type ChatStreamToolCallDelta = { + index?: number; + id?: string; + type?: 'function'; + function?: { name?: string; arguments?: string }; +}; + +/** + * Delta for `/v1/chat/completions` streaming: assistant message fragments only. + * Each SSE event may include any subset of fields; the full message is accumulated client-side. */ export type ChatStreamDelta = { - role?: ChatRole; + role?: 'assistant'; content?: string; + tool_calls?: ChatStreamToolCallDelta[]; }; /** From c6c05a50039de527da0f6aef63a7c79401b2e9d4 Mon Sep 17 00:00:00 2001 From: grypez <143971198+grypez@users.noreply.github.com> Date: Tue, 31 Mar 2026 21:12:55 -0400 Subject: [PATCH 14/20] feat(klms): normalize streaming chunks to assistant deltas Made-with: Cursor --- .../src/client.test.ts | 6 +- .../src/open-v1/base.test.ts | 61 ++++++++-------- .../src/open-v1/base.ts | 4 +- .../src/open-v1/normalize-stream-chunk.ts | 72 +++++++++++++++++++ .../src/open-v1/types.ts | 6 +- .../src/types.ts | 14 ++-- 6 files changed, 124 insertions(+), 39 deletions(-) create mode 100644 packages/kernel-language-model-service/src/open-v1/normalize-stream-chunk.ts diff --git a/packages/kernel-language-model-service/src/client.test.ts b/packages/kernel-language-model-service/src/client.test.ts index f493c56b37..62ecdc2c18 100644 --- a/packages/kernel-language-model-service/src/client.test.ts +++ b/packages/kernel-language-model-service/src/client.test.ts @@ -74,7 +74,11 @@ describe('makeChatClient', () => { id: 'chunk-1', model: MODEL, choices: [ - { delta: { content: 'hello' }, index: 0, finish_reason: null }, + { + delta: { role: 'assistant', content: 'hello' }, + index: 0, + finish_reason: null, + }, ], }; } 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 index 82a2ab5a9e..ef6c391725 100644 --- a/packages/kernel-language-model-service/src/open-v1/base.test.ts +++ b/packages/kernel-language-model-service/src/open-v1/base.test.ts @@ -3,6 +3,11 @@ import { describe, it, expect, vi, beforeEach } from 'vitest'; import type { ChatResult, ChatStreamChunk } from '../types.ts'; import { OpenV1BaseService } from './base.ts'; +import { normalizeStreamChunk } from './normalize-stream-chunk.ts'; +import type { + ChatStreamChunkWire, + ChatStreamDeltaWire, +} from './normalize-stream-chunk.ts'; const MODEL = 'glm-4.7-flash'; @@ -39,7 +44,7 @@ const makeMockFetch = (payload: unknown): typeof globalThis.fetch => ); const makeSSEStream = ( - chunks: ChatStreamChunk[], + chunks: ChatStreamChunkWire[], // eslint-disable-next-line n/no-unsupported-features/node-builtins ): ReadableStream => { const encoder = new TextEncoder(); @@ -56,14 +61,19 @@ const makeSSEStream = ( }); }; -const makeStreamChunk = (content: string): ChatStreamChunk => ({ +const makeWireStreamChunk = ( + delta: ChatStreamDeltaWire, +): ChatStreamChunkWire => ({ id: 'chat-1', model: MODEL, - choices: [{ delta: { content }, index: 0, finish_reason: null }], + choices: [{ delta, index: 0, finish_reason: null }], }); +const makeStreamChunk = (content: string): ChatStreamChunkWire => + makeWireStreamChunk({ content }); + const makeMockStreamFetch = ( - chunks: ChatStreamChunk[], + chunks: ChatStreamChunkWire[], ): typeof globalThis.fetch => vi.fn().mockResolvedValue(makeOkResponse({ body: makeSSEStream(chunks) })); @@ -316,8 +326,9 @@ describe('OpenV1BaseService', () => { }); it('yields parsed chunks and stops at [DONE]', async () => { - const expected = [makeStreamChunk('Hello'), makeStreamChunk(', world!')]; - const streamFetch = makeMockStreamFetch(expected); + const wire = [makeStreamChunk('Hello'), makeStreamChunk(', world!')]; + const expected = wire.map((chunk) => normalizeStreamChunk(chunk)); + const streamFetch = makeMockStreamFetch(wire); const streamService = new OpenV1BaseService( streamFetch, 'http://localhost:11434', @@ -336,28 +347,17 @@ describe('OpenV1BaseService', () => { }); it('accepts streaming delta with assistant role and tool_calls fragments', async () => { - const toolChunk: ChatStreamChunk = { - id: 'chat-1', - model: MODEL, - choices: [ + const wireTool = makeWireStreamChunk({ + tool_calls: [ { - delta: { - role: 'assistant', - tool_calls: [ - { - index: 0, - id: 'call_1', - type: 'function', - function: { name: 'fn' }, - }, - ], - }, index: 0, - finish_reason: null, + id: 'call_1', + type: 'function', + function: { name: 'fn' }, }, ], - }; - const streamFetch = makeMockStreamFetch([toolChunk]); + }); + const streamFetch = makeMockStreamFetch([wireTool]); const streamService = new OpenV1BaseService( streamFetch, 'http://localhost:11434', @@ -372,24 +372,25 @@ describe('OpenV1BaseService', () => { received.push(chunk); } - expect(received).toStrictEqual([toolChunk]); + expect(received).toStrictEqual([normalizeStreamChunk(wireTool)]); }); it('rejects streaming chunk when delta role is not assistant', async () => { - const badChunk = { + const badChunk: ChatStreamChunkWire = { id: 'chat-1', model: MODEL, choices: [ { - delta: { role: 'user', content: 'x' }, + delta: { + role: 'user', + content: 'x', + } as unknown as ChatStreamDeltaWire, index: 0, finish_reason: null, }, ], }; - const streamFetch = makeMockStreamFetch([ - badChunk as unknown as ChatStreamChunk, - ]); + const streamFetch = makeMockStreamFetch([badChunk]); const streamService = new OpenV1BaseService( streamFetch, 'http://localhost:11434', diff --git a/packages/kernel-language-model-service/src/open-v1/base.ts b/packages/kernel-language-model-service/src/open-v1/base.ts index d784c10208..cd9c75ecfe 100644 --- a/packages/kernel-language-model-service/src/open-v1/base.ts +++ b/packages/kernel-language-model-service/src/open-v1/base.ts @@ -1,6 +1,8 @@ import { assert } from '@metamask/superstruct'; import type { ChatParams, ChatResult, ChatStreamChunk } from '../types.ts'; +import { normalizeStreamChunk } from './normalize-stream-chunk.ts'; +import type { ChatStreamChunkWire } from './normalize-stream-chunk.ts'; import { checkResponseOk, readAndCheckResponse } from './response-json.ts'; import { ChatParamsStruct, @@ -122,7 +124,7 @@ export class OpenV1BaseService { if (data) { const json: unknown = JSON.parse(data); assert(json, ChatStreamChunkStruct); - yield harden(json as ChatStreamChunk); + yield harden(normalizeStreamChunk(json as ChatStreamChunkWire)); } } } diff --git a/packages/kernel-language-model-service/src/open-v1/normalize-stream-chunk.ts b/packages/kernel-language-model-service/src/open-v1/normalize-stream-chunk.ts new file mode 100644 index 0000000000..7693524456 --- /dev/null +++ b/packages/kernel-language-model-service/src/open-v1/normalize-stream-chunk.ts @@ -0,0 +1,72 @@ +import type { + AssistantStreamDelta, + ChatStreamChunk, + ChatStreamToolCallDelta, +} from '../types.ts'; + +/** + * `delta` object as emitted on the wire by OpenAI-style SSE (role often omitted + * after the first chunk). + */ +export type ChatStreamDeltaWire = { + role?: 'assistant'; + content?: string; + tool_calls?: ChatStreamToolCallDelta[]; +}; + +/** + * One stream event before normalization to {@link ChatStreamChunk}. + */ +export type ChatStreamChunkWire = { + id: string; + model: string; + choices: { + delta: ChatStreamDeltaWire; + index: number; + finish_reason: string | null; + }[]; +}; + +/** + * Coerce a wire delta into an assistant-only delta with required `role`. + * + * @param delta - Parsed `delta` from one SSE `choices[]` entry. + * @returns Delta with `role: 'assistant'` and any wire fields carried over. + */ +export function normalizeAssistantStreamDelta( + delta: ChatStreamDeltaWire, +): AssistantStreamDelta { + if (delta.role === undefined || delta.role === 'assistant') { + const out: AssistantStreamDelta = { role: 'assistant' }; + if ('content' in delta) { + out.content = delta.content; + } + if ('tool_calls' in delta) { + out.tool_calls = delta.tool_calls; + } + return out; + } + throw new TypeError( + `Expected stream delta role to be "assistant" or omitted, received: ${String(delta.role)}`, + ); +} + +/** + * Normalize every choice delta in a parsed SSE JSON object. + * + * @param chunk - One parsed `data:` JSON object from the stream. + * @returns The same chunk with each `delta` normalized to {@link AssistantStreamDelta}. + */ +export function normalizeStreamChunk( + chunk: ChatStreamChunkWire, +): ChatStreamChunk { + return { + id: chunk.id, + model: chunk.model, + choices: chunk.choices.map((choice) => ({ + index: choice.index, + finish_reason: choice.finish_reason, + delta: normalizeAssistantStreamDelta(choice.delta), + })), + }; +} diff --git a/packages/kernel-language-model-service/src/open-v1/types.ts b/packages/kernel-language-model-service/src/open-v1/types.ts index e97f70c26f..112de6c827 100644 --- a/packages/kernel-language-model-service/src/open-v1/types.ts +++ b/packages/kernel-language-model-service/src/open-v1/types.ts @@ -13,6 +13,7 @@ import { } from '@metamask/superstruct'; export type { + AssistantStreamDelta, ChatChoice, ChatMessage, ChatParams, @@ -122,7 +123,8 @@ const ChatStreamToolCallDeltaStruct = object({ ), }); -const ChatStreamDeltaStruct = object({ +/** Wire `delta` before streaming normalization adds `role: 'assistant'`. */ +const ChatStreamDeltaWireStruct = object({ role: optional(literal('assistant')), content: optional(string()), tool_calls: optional(array(ChatStreamToolCallDeltaStruct)), @@ -136,7 +138,7 @@ export const ChatStreamChunkStruct = object({ model: string(), choices: array( object({ - delta: ChatStreamDeltaStruct, + delta: ChatStreamDeltaWireStruct, index: number(), finish_reason: nullable(string()), }), diff --git a/packages/kernel-language-model-service/src/types.ts b/packages/kernel-language-model-service/src/types.ts index 1cc8000dbd..f514ad02b5 100644 --- a/packages/kernel-language-model-service/src/types.ts +++ b/packages/kernel-language-model-service/src/types.ts @@ -131,15 +131,19 @@ export type ChatStreamToolCallDelta = { }; /** - * Delta for `/v1/chat/completions` streaming: assistant message fragments only. - * Each SSE event may include any subset of fields; the full message is accumulated client-side. + * Assistant fragment from a streaming `/v1/chat/completions` response after + * the Open /v1 client normalizes each SSE event. `role` is always `'assistant'`; + * `content` / `tool_calls` are present only when the wire event included them. */ -export type ChatStreamDelta = { - role?: 'assistant'; +export type AssistantStreamDelta = { + role: 'assistant'; content?: string; tool_calls?: ChatStreamToolCallDelta[]; }; +/** @alias {@link AssistantStreamDelta} */ +export type ChatStreamDelta = AssistantStreamDelta; + /** * A single chunk from a streaming chat completion response. */ @@ -147,7 +151,7 @@ export type ChatStreamChunk = { id: string; model: string; choices: { - delta: ChatStreamDelta; + delta: AssistantStreamDelta; index: number; finish_reason: string | null; }[]; From 8c2430b651c7a1109444f3f22e8da03eab6abffa Mon Sep 17 00:00:00 2001 From: grypez <143971198+grypez@users.noreply.github.com> Date: Tue, 31 Mar 2026 21:13:02 -0400 Subject: [PATCH 15/20] fix(klms): mock OpenV1 fetch implements Response.text Made-with: Cursor --- .../src/test-utils/mock-fetch.ts | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) 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 index a122eb8a42..cecaa286da 100644 --- a/packages/kernel-language-model-service/src/test-utils/mock-fetch.ts +++ b/packages/kernel-language-model-service/src/test-utils/mock-fetch.ts @@ -14,7 +14,7 @@ export const makeMockOpenV1Fetch = ( return async (_url, _init) => { const content = responses[idx] ?? ''; idx += 1; - const result = harden({ + const result = { id: `chat-${idx}`, model, choices: [ @@ -25,7 +25,14 @@ export const makeMockOpenV1Fetch = ( }, ], usage: { prompt_tokens: 0, completion_tokens: 0, total_tokens: 0 }, - }); - return { json: async () => result } as unknown as globalThis.Response; + }; + const bodyText = JSON.stringify(result); + return { + ok: true, + status: 200, + statusText: 'OK', + text: async () => bodyText, + json: async () => JSON.parse(bodyText), + } as unknown as globalThis.Response; }; }; From 8f0704a24fff5d7e4756aa2c885dcddc39e8d025 Mon Sep 17 00:00:00 2001 From: grypez <143971198+grypez@users.noreply.github.com> Date: Tue, 31 Mar 2026 21:15:03 -0400 Subject: [PATCH 16/20] fix(klms): strip Open /v1 JSON before Superstruct assert Made-with: Cursor --- .../src/open-v1/base.ts | 23 +- .../src/open-v1/strip-open-v1-json.test.ts | 98 ++++++++ .../src/open-v1/strip-open-v1-json.ts | 233 ++++++++++++++++++ 3 files changed, 348 insertions(+), 6 deletions(-) create mode 100644 packages/kernel-language-model-service/src/open-v1/strip-open-v1-json.test.ts create mode 100644 packages/kernel-language-model-service/src/open-v1/strip-open-v1-json.ts diff --git a/packages/kernel-language-model-service/src/open-v1/base.ts b/packages/kernel-language-model-service/src/open-v1/base.ts index cd9c75ecfe..e450de111b 100644 --- a/packages/kernel-language-model-service/src/open-v1/base.ts +++ b/packages/kernel-language-model-service/src/open-v1/base.ts @@ -4,6 +4,11 @@ import type { ChatParams, ChatResult, ChatStreamChunk } from '../types.ts'; import { normalizeStreamChunk } from './normalize-stream-chunk.ts'; import type { ChatStreamChunkWire } from './normalize-stream-chunk.ts'; import { checkResponseOk, readAndCheckResponse } from './response-json.ts'; +import { + stripChatResultJson, + stripChatStreamChunkJson, + stripListModelsResponseJson, +} from './strip-open-v1-json.ts'; import { ChatParamsStruct, ChatResultStruct, @@ -85,8 +90,9 @@ export class OpenV1BaseService { }); const body = await readAndCheckResponse(response); const json: unknown = JSON.parse(body); - assert(json, ChatResultStruct); - return harden(json as ChatResult); + const stripped = stripChatResultJson(json); + assert(stripped, ChatResultStruct); + return harden(stripped as ChatResult); } /** @@ -123,8 +129,11 @@ export class OpenV1BaseService { } if (data) { const json: unknown = JSON.parse(data); - assert(json, ChatStreamChunkStruct); - yield harden(normalizeStreamChunk(json as ChatStreamChunkWire)); + const stripped = stripChatStreamChunkJson(json); + assert(stripped, ChatStreamChunkStruct); + yield harden( + normalizeStreamChunk(stripped as ChatStreamChunkWire), + ); } } } @@ -148,8 +157,10 @@ export class OpenV1BaseService { }); const body = await readAndCheckResponse(response); const json: unknown = JSON.parse(body); - assert(json, ListModelsResponseStruct); - return harden(json.data.map((model) => model.id)); + const stripped = stripListModelsResponseJson(json); + assert(stripped, ListModelsResponseStruct); + const { data } = stripped as { data: { id: string }[] }; + return harden(data.map((model) => model.id)); } /** diff --git a/packages/kernel-language-model-service/src/open-v1/strip-open-v1-json.test.ts b/packages/kernel-language-model-service/src/open-v1/strip-open-v1-json.test.ts new file mode 100644 index 0000000000..f47eb804c3 --- /dev/null +++ b/packages/kernel-language-model-service/src/open-v1/strip-open-v1-json.test.ts @@ -0,0 +1,98 @@ +import { describe, expect, it } from 'vitest'; + +import { + stripChatResultJson, + stripChatStreamChunkJson, + stripListModelsResponseJson, +} from './strip-open-v1-json.ts'; + +describe('stripChatResultJson', () => { + it('drops top-level provider fields before strict validation', () => { + const stripped = stripChatResultJson({ + object: 'chat.completion', + id: 'chat-1', + model: 'm', + choices: [ + { + index: 0, + finish_reason: 'stop', + logprobs: null, + message: { + role: 'assistant', + content: 'hi', + refusal: null, + }, + }, + ], + usage: { + prompt_tokens: 1, + completion_tokens: 2, + total_tokens: 3, + }, + }); + expect(stripped).toStrictEqual({ + id: 'chat-1', + model: 'm', + choices: [ + { + index: 0, + finish_reason: 'stop', + message: { role: 'assistant', content: 'hi' }, + }, + ], + usage: { + prompt_tokens: 1, + completion_tokens: 2, + total_tokens: 3, + }, + }); + }); +}); + +describe('stripListModelsResponseJson', () => { + it('keeps only model id entries', () => { + const stripped = stripListModelsResponseJson({ + object: 'list', + data: [ + { + id: 'llama3.1:latest', + object: 'model', + created: 123, + owned_by: 'library', + }, + ], + }); + expect(stripped).toStrictEqual({ + data: [{ id: 'llama3.1:latest' }], + }); + }); +}); + +describe('stripChatStreamChunkJson', () => { + it('drops extra keys on stream events', () => { + const stripped = stripChatStreamChunkJson({ + object: 'chat.completion.chunk', + id: 'c1', + model: 'm', + choices: [ + { + index: 0, + finish_reason: null, + logprobs: null, + delta: { content: 'x', role: 'assistant' }, + }, + ], + }); + expect(stripped).toStrictEqual({ + id: 'c1', + model: 'm', + choices: [ + { + index: 0, + finish_reason: null, + delta: { content: 'x', role: 'assistant' }, + }, + ], + }); + }); +}); diff --git a/packages/kernel-language-model-service/src/open-v1/strip-open-v1-json.ts b/packages/kernel-language-model-service/src/open-v1/strip-open-v1-json.ts new file mode 100644 index 0000000000..c3e8622261 --- /dev/null +++ b/packages/kernel-language-model-service/src/open-v1/strip-open-v1-json.ts @@ -0,0 +1,233 @@ +/** + * Reduce permissive Open /v1 JSON bodies to the shapes validated by our + * Superstruct schemas, dropping provider-specific keys (e.g. `object`). + */ + +/** + * @param value - Unknown JSON value. + * @returns Plain object record, or `null` if not a non-null object. + */ +function asRecord(value: unknown): Record | null { + if (typeof value !== 'object' || value === null || Array.isArray(value)) { + return null; + } + return value as Record; +} + +/** + * @param call - Raw tool call from JSON. + * @returns Plain object matching the non-streaming tool-call struct. + */ +function stripToolCall(call: unknown): unknown { + const row = asRecord(call); + if (!row) { + return call; + } + const fnRow = asRecord(row.function); + return { + id: row.id, + type: 'function', + ...(typeof row.index === 'number' ? { index: row.index } : {}), + function: { + name: typeof fnRow?.name === 'string' ? fnRow.name : '', + arguments: typeof fnRow?.arguments === 'string' ? fnRow.arguments : '', + }, + }; +} + +/** + * @param raw - Raw `message` object from JSON. + * @returns Plain object matching a `ChatMessage` union variant. + */ +function stripMessage(raw: unknown): unknown { + const row = asRecord(raw); + if (!row || typeof row.role !== 'string') { + return raw; + } + switch (row.role) { + case 'system': + return { role: 'system', content: row.content }; + case 'user': + return { role: 'user', content: row.content }; + case 'assistant': { + const out: Record = { + role: 'assistant', + content: row.content === undefined ? null : row.content, + }; + if (Array.isArray(row.tool_calls)) { + out.tool_calls = row.tool_calls.map(stripToolCall); + } + return out; + } + case 'tool': + return { + role: 'tool', + content: row.content, + tool_call_id: row.tool_call_id, + }; + default: + return raw; + } +} + +/** + * @param raw - Raw choice from JSON. + * @returns Plain object matching one chat choice in a result. + */ +function stripChoice(raw: unknown): unknown { + const row = asRecord(raw); + if (!row) { + return raw; + } + return { + message: stripMessage(row.message), + index: row.index, + finish_reason: row.finish_reason === undefined ? null : row.finish_reason, + }; +} + +/** + * @param raw - Raw `usage` object from JSON. + * @returns Plain object matching token `usage`. + */ +function stripUsage(raw: unknown): unknown { + const row = asRecord(raw); + if (!row) { + return raw; + } + return { + prompt_tokens: row.prompt_tokens, + completion_tokens: row.completion_tokens, + total_tokens: row.total_tokens, + }; +} + +/** + * Keep only fields used for non-streaming chat result validation. + * + * @param json - Parsed `/v1/chat/completions` response body. + * @returns Plain object with only `id`, `model`, `choices`, `usage`. + */ +export function stripChatResultJson(json: unknown): unknown { + const row = asRecord(json); + if (!row) { + return json; + } + return { + id: row.id, + model: row.model, + choices: Array.isArray(row.choices) ? row.choices.map(stripChoice) : [], + usage: stripUsage(row.usage), + }; +} + +/** + * @param item - One element of `data` from `/v1/models`. + * @returns `{ id }` only. + */ +function stripModelEntry(item: unknown): unknown { + const row = asRecord(item); + if (!row) { + return item; + } + return { id: row.id }; +} + +/** + * Keep only fields used for `/v1/models` response validation. + * + * @param json - Parsed `/v1/models` response body. + * @returns Plain object with only `data: { id }[]`. + */ +export function stripListModelsResponseJson(json: unknown): unknown { + const row = asRecord(json); + if (!row || !Array.isArray(row.data)) { + return json; + } + return { + data: row.data.map(stripModelEntry), + }; +} + +/** + * @param raw - One streaming tool-call fragment. + * @returns Plain object matching one streaming tool-call delta fragment. + */ +function stripStreamToolCallDelta(raw: unknown): unknown { + const part = asRecord(raw); + if (!part) { + return raw; + } + const fnRow = asRecord(part.function); + const fnOut: Record = {}; + if (fnRow) { + if ('name' in fnRow) { + fnOut.name = fnRow.name; + } + if ('arguments' in fnRow) { + fnOut.arguments = fnRow.arguments; + } + } + return { + ...(typeof part.index === 'number' ? { index: part.index } : {}), + ...(typeof part.id === 'string' ? { id: part.id } : {}), + ...(part.type === 'function' ? { type: 'function' } : {}), + ...(Object.keys(fnOut).length > 0 ? { function: fnOut } : {}), + }; +} + +/** + * @param raw - Raw `delta` from a stream chunk. + * @returns Plain object matching a streaming `delta` (wire shape). + */ +function stripStreamDelta(raw: unknown): unknown { + const deltaRow = asRecord(raw); + if (!deltaRow) { + return raw; + } + const out: Record = {}; + if ('role' in deltaRow) { + out.role = deltaRow.role; + } + if ('content' in deltaRow) { + out.content = deltaRow.content; + } + if (Array.isArray(deltaRow.tool_calls)) { + out.tool_calls = deltaRow.tool_calls.map(stripStreamToolCallDelta); + } + return out; +} + +/** + * @param raw - Raw choice in a stream chunk. + * @returns Stripped choice for one streaming chunk. + */ +function stripStreamChoice(raw: unknown): unknown { + const row = asRecord(raw); + if (!row) { + return raw; + } + return { + delta: stripStreamDelta(row.delta), + index: row.index, + finish_reason: row.finish_reason === undefined ? null : row.finish_reason, + }; +} + +/** + * Keep only fields used for streaming chunk validation (wire shape). + * + * @param json - Parsed one SSE `data:` JSON object. + * @returns Plain object with `id`, `model`, `choices` only. + */ +export function stripChatStreamChunkJson(json: unknown): unknown { + const row = asRecord(json); + if (!row || !Array.isArray(row.choices)) { + return json; + } + return { + id: row.id, + model: row.model, + choices: row.choices.map(stripStreamChoice), + }; +} From 1d2349ac924d158a52df381079b5e3b4bf139bd4 Mon Sep 17 00:00:00 2001 From: grypez <143971198+grypez@users.noreply.github.com> Date: Tue, 31 Mar 2026 22:06:36 -0400 Subject: [PATCH 17/20] fix(klms): wrap Open /v1 SSE chunk errors with Error and cause Made-with: Cursor --- .../src/open-v1/base.test.ts | 20 +++++++++++++++++-- .../src/open-v1/base.ts | 7 ++++++- 2 files changed, 24 insertions(+), 3 deletions(-) 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 index ef6c391725..298f43a538 100644 --- a/packages/kernel-language-model-service/src/open-v1/base.test.ts +++ b/packages/kernel-language-model-service/src/open-v1/base.test.ts @@ -405,7 +405,15 @@ describe('OpenV1BaseService', () => { })) { // drain } - }).rejects.toThrow(StructError); + }).rejects.toSatisfy((rejection: unknown) => { + if (!(rejection instanceof Error)) { + return false; + } + if (!rejection.message.startsWith('Error parsing JSON: ')) { + return false; + } + return rejection.cause instanceof StructError; + }); }); it('throws when response body is null', async () => { @@ -480,7 +488,15 @@ describe('OpenV1BaseService', () => { })) { // drain } - }).rejects.toThrow(SyntaxError); + }).rejects.toSatisfy((rejection: unknown) => { + if (!(rejection instanceof Error)) { + return false; + } + if (rejection.message !== 'Error parsing JSON: not-json') { + return false; + } + return rejection.cause instanceof SyntaxError; + }); }); }); }); diff --git a/packages/kernel-language-model-service/src/open-v1/base.ts b/packages/kernel-language-model-service/src/open-v1/base.ts index e450de111b..a0b8a7f505 100644 --- a/packages/kernel-language-model-service/src/open-v1/base.ts +++ b/packages/kernel-language-model-service/src/open-v1/base.ts @@ -127,13 +127,18 @@ export class OpenV1BaseService { if (data === '[DONE]') { return; } - if (data) { + if (!data) { + continue; + } + try { const json: unknown = JSON.parse(data); const stripped = stripChatStreamChunkJson(json); assert(stripped, ChatStreamChunkStruct); yield harden( normalizeStreamChunk(stripped as ChatStreamChunkWire), ); + } catch (cause: unknown) { + throw new Error(`Error parsing JSON: ${data}`, { cause }); } } } From 3d6e7135e1710b0d9dc592cc6c41f58e3e084396 Mon Sep 17 00:00:00 2001 From: grypez <143971198+grypez@users.noreply.github.com> Date: Tue, 31 Mar 2026 22:10:49 -0400 Subject: [PATCH 18/20] refactor(kernel-utils): build object shape only for strict objects Made-with: Cursor --- .../kernel-utils/src/json-schema-to-struct.ts | 20 +++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/packages/kernel-utils/src/json-schema-to-struct.ts b/packages/kernel-utils/src/json-schema-to-struct.ts index 449b50e895..fb591a10ea 100644 --- a/packages/kernel-utils/src/json-schema-to-struct.ts +++ b/packages/kernel-utils/src/json-schema-to-struct.ts @@ -72,19 +72,19 @@ export function jsonSchemaToStruct(schema: JsonSchema): Struct { return array(jsonSchemaToStruct(schema.items)) as Struct; case 'object': { const { properties } = schema; - const required = new Set(schema.required ?? Object.keys(properties)); - const shape = Object.fromEntries( - Object.entries(properties).map(([key, subSchema]) => { - const fieldStruct = jsonSchemaToStruct(subSchema); - return [key, required.has(key) ? fieldStruct : optional(fieldStruct)]; - }), - ); if (schema.additionalProperties === false) { + const required = new Set(schema.required ?? Object.keys(properties)); + const shape = Object.fromEntries( + Object.entries(properties).map(([key, subSchema]) => { + const fieldStruct = jsonSchemaToStruct(subSchema); + return [ + key, + required.has(key) ? fieldStruct : optional(fieldStruct), + ]; + }), + ); return object(shape) as Struct; } - if (Object.keys(properties).length === 0) { - return looseObjectStruct(schema); - } return looseObjectStruct(schema); } default: { From 9da54329f8a426e8e9eadc801b37b0ff84024257 Mon Sep 17 00:00:00 2001 From: grypez <143971198+grypez@users.noreply.github.com> Date: Wed, 1 Apr 2026 09:33:27 -0400 Subject: [PATCH 19/20] Harden makeChatAgent return Co-authored-by: Dimitris Marlagkoutsos --- packages/kernel-agents/src/strategies/chat-agent.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/kernel-agents/src/strategies/chat-agent.ts b/packages/kernel-agents/src/strategies/chat-agent.ts index 2ee5474f1f..e3d2f65ab1 100644 --- a/packages/kernel-agents/src/strategies/chat-agent.ts +++ b/packages/kernel-agents/src/strategies/chat-agent.ts @@ -106,7 +106,7 @@ export const makeChatAgent = ({ }: MakeChatAgentArgs): Agent => { const experienceLog: Experience[] = []; - return { + return harden({ task: async ( intent: string, judgment?: (result: unknown) => result is Result, From 6b2324477440fea8f231ccdfaa32d510cddfb1f7 Mon Sep 17 00:00:00 2001 From: grypez <143971198+grypez@users.noreply.github.com> Date: Wed, 1 Apr 2026 09:45:07 -0400 Subject: [PATCH 20/20] And the paren --- packages/kernel-agents/src/strategies/chat-agent.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/kernel-agents/src/strategies/chat-agent.ts b/packages/kernel-agents/src/strategies/chat-agent.ts index e3d2f65ab1..99c6a187c7 100644 --- a/packages/kernel-agents/src/strategies/chat-agent.ts +++ b/packages/kernel-agents/src/strategies/chat-agent.ts @@ -215,5 +215,5 @@ export const makeChatAgent = ({ yield* experienceLog; })(); }, - }; + }); };