From ca198f739e8fc968ed42856a66d66ff5692172d7 Mon Sep 17 00:00:00 2001 From: Kit Langton Date: Sun, 26 Apr 2026 11:53:58 -0400 Subject: [PATCH] refactor(llm): cache tool codecs and tighten ToolRuntime types Simplify pass after the typed ToolRuntime initial drop. Findings from a parallel review (code reuse + quality + perf): src/tool.ts - Tool now carries memoized decode/encode codecs and a precomputed ToolDefinition, derived once at tool() construction time. The runtime no longer rebuilds Schema closures or JSON Schema docs per call/per run. - Constrains parameters/success to Schema.Codec so the codecs have no service requirements. Drops the 'as unknown as' casts the runtime needed previously. - Fixes a latent bug: schemas with $ref now correctly emit $defs on ToolDefinition.inputSchema (toJsonSchemaDocument's definitions were silently dropped before). src/tool-runtime.ts - Uses LLMRequest constructor instead of 'as LLMRequest' casts. - Default tool dispatch concurrency is 10 (was 'unbounded'); exposed via RunOptions.concurrency. Unbounded is still available for handlers that do not share a saturable resource. - Drops dead 'usage' state, the single-use Dispatched interface, and the DEFAULT_MAX_STEPS constant per the inline-when-used style rule. - accumulate() now factors text-delta and reasoning-delta into one helper. test/lib/openai-chunks.ts (new) - Shared deltaChunk / usageChunk / toolCallChunk / finishChunk helpers. test/lib/http.ts - scriptedResponses moved here from tool-runtime.test.ts so future multi-step adapter tests can reuse it. Also picks up parallel work that swapped HandlerInput to a 'respond' callback for cleaner Response construction. test/tool-runtime.test.ts - Uses LLMEvent.guards for typed event filtering instead of cast-and-check. - Concurrent test now uses sseEvents + deltaChunk instead of a hand-rolled body string. Includes parallel callsite updates in test/adapter.test.ts and test/provider/openai-compatible-chat.test.ts that adopt the 'respond' API in lib/http.ts. --- packages/llm/src/tool-runtime.ts | 138 +++++------- packages/llm/src/tool.ts | 65 +++++- packages/llm/test/adapter.test.ts | 4 +- packages/llm/test/lib/http.ts | 38 +++- packages/llm/test/lib/openai-chunks.ts | 27 +++ .../provider/openai-compatible-chat.test.ts | 2 +- packages/llm/test/tool-runtime.test.ts | 207 ++++++------------ 7 files changed, 234 insertions(+), 247 deletions(-) create mode 100644 packages/llm/test/lib/openai-chunks.ts diff --git a/packages/llm/src/tool-runtime.ts b/packages/llm/src/tool-runtime.ts index 69b8211f68..e3f36bd568 100644 --- a/packages/llm/src/tool-runtime.ts +++ b/packages/llm/src/tool-runtime.ts @@ -1,19 +1,19 @@ -import { Effect, Schema, Stream } from "effect" +import { Effect, Stream } from "effect" +import type { Concurrency } from "effect/Types" import type { LLMClient } from "./adapter" import type { RequestExecutor } from "./executor" import * as LLM from "./llm" -import type { - ContentPart, - FinishReason, - LLMError, - LLMEvent, +import { + type ContentPart, + type FinishReason, + type LLMError, + type LLMEvent, LLMRequest, - ToolCallPart, - ToolResultValue, - Usage, + type ToolCallPart, + type ToolResultValue, } from "./schema" import { ToolFailure } from "./schema" -import { type Tool, type Tools, toDefinitions } from "./tool" +import { type AnyTool, type Tools, toDefinitions } from "./tool" export interface RuntimeState { readonly step: number @@ -29,6 +29,13 @@ export interface RunOptions { * simply stops and the last `request-finish` event is the terminal signal. */ readonly maxSteps?: number + /** + * How many tool handlers to dispatch in parallel within a single step. + * Defaults to 10. Use `"unbounded"` only when handlers do not share an + * external dependency that can be saturated (rate-limited APIs, single + * connections, etc). + */ + readonly concurrency?: Concurrency /** * Optional predicate evaluated after each step's `request-finish` event. If * it returns `true`, the loop stops even if the model wanted to continue. @@ -36,8 +43,6 @@ export interface RunOptions { readonly stopWhen?: (state: RuntimeState) => boolean } -const DEFAULT_MAX_STEPS = 10 - /** * Run a model with a typed tool record. The runtime streams the model, on * each `tool-call` event decodes the input against the tool's `parameters` @@ -54,23 +59,18 @@ export const run = ( client: LLMClient, options: RunOptions, ): Stream.Stream => { - const maxSteps = options.maxSteps ?? DEFAULT_MAX_STEPS + const maxSteps = options.maxSteps ?? 10 + const concurrency = options.concurrency ?? 10 const tools = options.tools as Tools - const definitions = toDefinitions(tools) - const initialRequest: LLMRequest = { + const initialRequest = new LLMRequest({ ...options.request, - tools: [...options.request.tools, ...definitions], - } as LLMRequest + tools: [...options.request.tools, ...toDefinitions(tools)], + }) const loop = (request: LLMRequest, step: number): Stream.Stream => Stream.unwrap( Effect.gen(function* () { - const state: StepState = { - assistantContent: [], - toolCalls: [], - finishReason: undefined, - usage: undefined, - } + const state: StepState = { assistantContent: [], toolCalls: [], finishReason: undefined } const modelStream = client.stream(request).pipe( Stream.tap((event) => Effect.sync(() => accumulate(state, event))), @@ -82,24 +82,25 @@ export const run = ( if (options.stopWhen?.({ step, request })) return Stream.empty if (step + 1 >= maxSteps) return Stream.empty - const dispatched = yield* Effect.forEach(state.toolCalls, (call) => dispatch(tools, call), { - concurrency: "unbounded", - }) - const followUp: LLMRequest = { + const dispatched = yield* Effect.forEach( + state.toolCalls, + (call) => dispatch(tools, call).pipe(Effect.map((result) => [call, result] as const)), + { concurrency }, + ) + const followUp = new LLMRequest({ ...request, messages: [ ...request.messages, LLM.assistant(state.assistantContent), - ...dispatched.map(({ call, result }) => + ...dispatched.map(([call, result]) => LLM.toolMessage({ id: call.id, name: call.name, result }), ), ], - } as LLMRequest + }) - const dispatchEvents = Stream.fromIterable( - dispatched.flatMap(({ call, result }) => emitEvents(call, result)), + return Stream.fromIterable(dispatched.flatMap(([call, result]) => emitEvents(call, result))).pipe( + Stream.concat(loop(followUp, step + 1)), ) - return dispatchEvents.pipe(Stream.concat(loop(followUp, step + 1))) }), ) @@ -114,26 +115,15 @@ interface StepState { assistantContent: ContentPart[] toolCalls: ToolCallPart[] finishReason: FinishReason | undefined - usage: Usage | undefined } const accumulate = (state: StepState, event: LLMEvent) => { if (event.type === "text-delta") { - const last = state.assistantContent.at(-1) - if (last?.type === "text") { - state.assistantContent[state.assistantContent.length - 1] = { ...last, text: `${last.text}${event.text}` } - } else { - state.assistantContent.push({ type: "text", text: event.text }) - } + appendStreamingText(state, "text", event.text) return } if (event.type === "reasoning-delta") { - const last = state.assistantContent.at(-1) - if (last?.type === "reasoning") { - state.assistantContent[state.assistantContent.length - 1] = { ...last, text: `${last.text}${event.text}` } - } else { - state.assistantContent.push({ type: "reasoning", text: event.text }) - } + appendStreamingText(state, "reasoning", event.text) return } if (event.type === "tool-call") { @@ -144,67 +134,45 @@ const accumulate = (state: StepState, event: LLMEvent) => { } if (event.type === "request-finish") { state.finishReason = event.reason - if (event.usage !== undefined) state.usage = event.usage + } +} + +const appendStreamingText = (state: StepState, type: "text" | "reasoning", text: string) => { + const last = state.assistantContent.at(-1) + if (last?.type === type) { + state.assistantContent[state.assistantContent.length - 1] = { ...last, text: `${last.text}${text}` } return } - if (event.type === "step-finish" && event.usage !== undefined) { - state.usage = event.usage - } + state.assistantContent.push({ type, text }) } -interface Dispatched { - readonly call: ToolCallPart - readonly result: ToolResultValue -} - -const dispatch = (tools: Tools, call: ToolCallPart): Effect.Effect => { +const dispatch = (tools: Tools, call: ToolCallPart): Effect.Effect => { const tool = tools[call.name] - if (!tool) { - return Effect.succeed({ - call, - result: { type: "error" as const, value: `Unknown tool: ${call.name}` }, - }) - } + if (!tool) return Effect.succeed({ type: "error" as const, value: `Unknown tool: ${call.name}` }) return decodeAndExecute(tool, call.input).pipe( - Effect.map((result): Dispatched => ({ call, result })), - Effect.catchTag( - "LLM.ToolFailure", - (failure): Effect.Effect => - Effect.succeed({ call, result: { type: "error" as const, value: failure.message } }), + Effect.catchTag("LLM.ToolFailure", (failure) => + Effect.succeed({ type: "error" as const, value: failure.message } satisfies ToolResultValue), ), ) } -const decodeAndExecute = ( - tool: Tool, - input: unknown, -): Effect.Effect => { - const decode = Schema.decodeUnknownEffect(tool.parameters) as unknown as ( - input: unknown, - ) => Effect.Effect - const encode = Schema.encodeEffect(tool.success) as unknown as ( - value: unknown, - ) => Effect.Effect - - return decode(input).pipe( - Effect.mapError( - (error) => new ToolFailure({ message: `Invalid tool input: ${error.message ?? String(error)}` }), - ), - Effect.flatMap((decoded) => tool.execute(decoded as never)), +const decodeAndExecute = (tool: AnyTool, input: unknown): Effect.Effect => + tool._decode(input).pipe( + Effect.mapError((error) => new ToolFailure({ message: `Invalid tool input: ${error.message}` })), + Effect.flatMap((decoded) => tool.execute(decoded)), Effect.flatMap((value) => - encode(value).pipe( + tool._encode(value).pipe( Effect.mapError( (error) => new ToolFailure({ - message: `Tool returned an invalid value for its success schema: ${error.message ?? String(error)}`, + message: `Tool returned an invalid value for its success schema: ${error.message}`, }), ), ), ), Effect.map((encoded): ToolResultValue => ({ type: "json", value: encoded })), ) -} const emitEvents = (call: ToolCallPart, result: ToolResultValue): ReadonlyArray => result.type === "error" diff --git a/packages/llm/src/tool.ts b/packages/llm/src/tool.ts index 7c3e64743c..e3f9791a50 100644 --- a/packages/llm/src/tool.ts +++ b/packages/llm/src/tool.ts @@ -2,6 +2,13 @@ import { Effect, Schema } from "effect" import type { ToolDefinition as ToolDefinitionClass } from "./schema" import { ToolDefinition, ToolFailure } from "./schema" +/** + * Schema constraint for tool parameters / success values: no decoding or + * encoding services are allowed. Tools should be self-contained — anything + * beyond pure data transformation belongs in the handler closure. + */ +export type ToolSchema = Schema.Codec + /** * A type-safe LLM tool. Each tool bundles its own description, parameter * Schema, success Schema, and execute handler. The handler closes over any @@ -10,20 +17,31 @@ import { ToolDefinition, ToolFailure } from "./schema" * * Errors must be expressed as `ToolFailure`. Unmapped errors and defects fail * the stream. + * + * Internally each tool also carries memoized codecs and a precomputed + * `ToolDefinition` so the runtime doesn't rebuild them per invocation. */ -export interface Tool { +export interface Tool, Success extends ToolSchema> { readonly description: string readonly parameters: Parameters readonly success: Success readonly execute: ( params: Schema.Schema.Type, ) => Effect.Effect, ToolFailure> + /** @internal */ + readonly _decode: (input: unknown) => Effect.Effect, Schema.SchemaError> + /** @internal */ + readonly _encode: (value: Schema.Schema.Type) => Effect.Effect + /** @internal */ + readonly _definition: ToolDefinitionClass } +export type AnyTool = Tool, ToolSchema> + /** - * Helper that returns its argument unchanged. Its only purpose is to give - * TypeScript the inference points for `parameters` / `success` / `execute` at - * the call site so consumers don't have to spell out the type parameters. + * Constructs a typed tool. The Schema codecs and JSON-schema-shaped + * `ToolDefinition` are derived once at this call site so the runtime can + * reuse them across every invocation without recomputing. * * ```ts * const getWeather = tool({ @@ -34,29 +52,56 @@ export interface Tool * }) * ``` */ -export const tool = ( - config: Tool, -): Tool => config +export const tool = , Success extends ToolSchema>(config: { + readonly description: string + readonly parameters: Parameters + readonly success: Success + readonly execute: ( + params: Schema.Schema.Type, + ) => Effect.Effect, ToolFailure> +}): Tool => ({ + description: config.description, + parameters: config.parameters, + success: config.success, + execute: config.execute, + _decode: Schema.decodeUnknownEffect(config.parameters), + _encode: Schema.encodeEffect(config.success), + _definition: new ToolDefinition({ + name: "", + description: config.description, + inputSchema: toJsonSchema(config.parameters), + }), +}) /** * A record of named tools. The record key becomes the tool name on the wire. */ -export type Tools = Record> +export type Tools = Record /** * Convert a tools record into the `ToolDefinition[]` shape that * `LLMRequest.tools` expects. The runtime calls this internally; consumers * that build `LLMRequest` themselves can use it too. + * + * Tool names come from the record keys, so the per-tool cached + * `_definition` is rebuilt with the correct name here. The JSON Schema body + * is reused. */ export const toDefinitions = (tools: Tools): ReadonlyArray => Object.entries(tools).map(([name, item]) => new ToolDefinition({ name, - description: item.description, - inputSchema: Schema.toJsonSchemaDocument(item.parameters).schema as Record, + description: item._definition.description, + inputSchema: item._definition.inputSchema, }), ) +const toJsonSchema = (schema: Schema.Top): Record => { + const document = Schema.toJsonSchemaDocument(schema) + if (Object.keys(document.definitions).length === 0) return document.schema as Record + return { ...document.schema, $defs: document.definitions } as Record +} + export { ToolFailure } export * as Tool from "./tool" diff --git a/packages/llm/test/adapter.test.ts b/packages/llm/test/adapter.test.ts index a4f0b4461e..f5785bbbef 100644 --- a/packages/llm/test/adapter.test.ts +++ b/packages/llm/test/adapter.test.ts @@ -80,9 +80,9 @@ const gemini = Adapter.define({ protocol: "gemini", }) -const echoLayer = dynamicResponse(({ text }) => +const echoLayer = dynamicResponse(({ text, respond }) => Effect.succeed( - new Response( + respond( encodeJson([ { type: "text", text: `echo:${text}` }, { type: "finish", reason: "stop" }, diff --git a/packages/llm/test/lib/http.ts b/packages/llm/test/lib/http.ts index 6ae8bb00ad..f14de847e9 100644 --- a/packages/llm/test/lib/http.ts +++ b/packages/llm/test/lib/http.ts @@ -1,13 +1,14 @@ -import { Effect, Layer } from "effect" +import { Effect, Layer, Ref } from "effect" import { HttpClient, HttpClientRequest, HttpClientResponse } from "effect/unstable/http" import { RequestExecutor } from "../../src/executor" export type HandlerInput = { readonly request: HttpClientRequest.HttpClientRequest readonly text: string + readonly respond: (body: ConstructorParameters[0], init?: ResponseInit) => HttpClientResponse.HttpClientResponse } -export type Handler = (input: HandlerInput) => Effect.Effect +export type Handler = (input: HandlerInput) => Effect.Effect const handlerLayer = (handler: Handler): Layer.Layer => Layer.succeed( @@ -16,8 +17,11 @@ const handlerLayer = (handler: Handler): Layer.Layer => Effect.gen(function* () { const web = yield* HttpClientRequest.toWeb(request).pipe(Effect.orDie) const text = yield* Effect.promise(() => web.text()) - const response = yield* handler({ request, text }) - return HttpClientResponse.fromWeb(request, response) + return yield* handler({ + request, + text, + respond: (body, init) => HttpClientResponse.fromWeb(request, new Response(body, init)), + }) }), ), ) @@ -32,7 +36,7 @@ const SSE_HEADERS = { "content-type": "text/event-stream" } as const * fixture tests where the request shape is irrelevant. */ export const fixedResponse = (body: string, init: ResponseInit = { headers: SSE_HEADERS }) => - executorWith(handlerLayer(() => Effect.succeed(new Response(body, init)))) + executorWith(handlerLayer((input) => Effect.succeed(input.respond(body, init)))) /** * Layer that builds a response per request. Useful for echo servers. @@ -44,7 +48,7 @@ export const dynamicResponse = (handler: Handler) => executorWith(handlerLayer(h * exercise transport errors that surface during parsing. */ export const truncatedStream = (chunks: ReadonlyArray) => - dynamicResponse(() => + dynamicResponse((input) => Effect.sync(() => { const encoder = new TextEncoder() const stream = new ReadableStream({ @@ -53,6 +57,26 @@ export const truncatedStream = (chunks: ReadonlyArray) => controller.error(new Error("connection reset")) }, }) - return new Response(stream, { headers: SSE_HEADERS }) + return input.respond(stream, { headers: SSE_HEADERS }) }), ) + +/** + * Layer that returns successive bodies on each request. Useful for scripting + * multi-step model exchanges (e.g. tool-call loops). The last body in the + * array is reused if the test makes more requests than scripted. + */ +export const scriptedResponses = (bodies: ReadonlyArray, init: ResponseInit = { headers: SSE_HEADERS }) => { + if (bodies.length === 0) throw new Error("scriptedResponses requires at least one body") + return Layer.unwrap( + Effect.gen(function* () { + const cursor = yield* Ref.make(0) + return dynamicResponse((input) => + Effect.gen(function* () { + const index = yield* Ref.getAndUpdate(cursor, (n) => n + 1) + return input.respond(bodies[index] ?? bodies[bodies.length - 1], init) + }), + ) + }), + ) +} diff --git a/packages/llm/test/lib/openai-chunks.ts b/packages/llm/test/lib/openai-chunks.ts new file mode 100644 index 0000000000..77a7c919e1 --- /dev/null +++ b/packages/llm/test/lib/openai-chunks.ts @@ -0,0 +1,27 @@ +/** + * Shared chunk shapes for OpenAI Chat / OpenAI-compatible Chat fixture tests. + * Multiple test files build the same `{ id, choices: [{ delta, finish_reason }], usage }` + * envelope; consolidating here keeps tool-call event shapes consistent. + */ + +const FIXTURE_ID = "chatcmpl_fixture" + +export const deltaChunk = (delta: object, finishReason: string | null = null) => ({ + id: FIXTURE_ID, + choices: [{ delta, finish_reason: finishReason }], + usage: null, +}) + +export const usageChunk = (usage: object) => ({ + id: FIXTURE_ID, + choices: [], + usage, +}) + +export const finishChunk = (reason: string) => deltaChunk({}, reason) + +export const toolCallChunk = (id: string, name: string, args: string, index = 0) => + deltaChunk({ + role: "assistant", + tool_calls: [{ index, id, function: { name, arguments: args } }], + }) diff --git a/packages/llm/test/provider/openai-compatible-chat.test.ts b/packages/llm/test/provider/openai-compatible-chat.test.ts index 2394e3c492..8e802bbc68 100644 --- a/packages/llm/test/provider/openai-compatible-chat.test.ts +++ b/packages/llm/test/provider/openai-compatible-chat.test.ts @@ -212,7 +212,7 @@ describe("OpenAI-compatible Chat adapter", () => { { role: "user", content: "Say hello." }, ], }) - return new Response( + return input.respond( sseEvents( deltaChunk({ role: "assistant", content: "Hello" }), deltaChunk({ content: "!" }), diff --git a/packages/llm/test/tool-runtime.test.ts b/packages/llm/test/tool-runtime.test.ts index 16e371717c..922bc9c3a4 100644 --- a/packages/llm/test/tool-runtime.test.ts +++ b/packages/llm/test/tool-runtime.test.ts @@ -1,12 +1,13 @@ import { describe, expect } from "bun:test" -import { Effect, Layer, Ref, Schema, Stream } from "effect" -import { LLM } from "../src" +import { Effect, Layer, Schema, Stream } from "effect" +import { LLM, LLMEvent } from "../src" import { client } from "../src/adapter" import { OpenAIChat } from "../src/provider/openai-chat" import { tool, ToolFailure } from "../src/tool" import { ToolRuntime } from "../src/tool-runtime" import { testEffect } from "./lib/effect" -import { dynamicResponse } from "./lib/http" +import { scriptedResponses } from "./lib/http" +import { deltaChunk, finishChunk, toolCallChunk } from "./lib/openai-chunks" import { sseEvents } from "./lib/sse" const model = OpenAIChat.model({ @@ -23,38 +24,6 @@ const baseRequest = LLM.request({ const it = testEffect(Layer.empty) -const deltaChunk = (delta: object, finishReason: string | null = null) => ({ - id: "chatcmpl_x", - choices: [{ delta, finish_reason: finishReason }], - usage: null, -}) - -const toolCallChunk = (id: string, name: string, args: string) => - deltaChunk({ - role: "assistant", - tool_calls: [{ index: 0, id, function: { name, arguments: args } }], - }) - -const finishChunk = (reason: string) => deltaChunk({}, reason) - -/** - * Builds an HTTP layer where successive requests return successive bodies. - * Used to script multi-step model exchanges. - */ -const scriptedResponses = (bodies: ReadonlyArray) => - Layer.unwrap( - Effect.gen(function* () { - const cursor = yield* Ref.make(0) - return dynamicResponse(() => - Effect.gen(function* () { - const index = yield* Ref.getAndUpdate(cursor, (n) => n + 1) - const body = bodies[index] ?? bodies.at(-1)! - return new Response(body, { headers: { "content-type": "text/event-stream" } }) - }), - ) - }), - ) - const get_weather = tool({ description: "Get current weather for a city.", parameters: Schema.Struct({ city: Schema.String }), @@ -71,33 +40,25 @@ describe("ToolRuntime", () => { Effect.gen(function* () { const llm = client({ adapters: [OpenAIChat.adapter] }) const layer = scriptedResponses([ - sseEvents( - toolCallChunk("call_1", "get_weather", '{"city":"Paris"}'), - finishChunk("tool_calls"), - ), - sseEvents( - deltaChunk({ role: "assistant", content: "It's sunny in Paris." }), - finishChunk("stop"), - ), + sseEvents(toolCallChunk("call_1", "get_weather", '{"city":"Paris"}'), finishChunk("tool_calls")), + sseEvents(deltaChunk({ role: "assistant", content: "It's sunny in Paris." }), finishChunk("stop")), ]) const events = Array.from( - yield* ToolRuntime.run(llm, { - request: baseRequest, - tools: { get_weather }, - }).pipe(Stream.runCollect, Effect.provide(layer)), + yield* ToolRuntime.run(llm, { request: baseRequest, tools: { get_weather } }).pipe( + Stream.runCollect, + Effect.provide(layer), + ), ) - const types = events.map((event) => event.type) - expect(types).toContain("tool-call") - expect(types).toContain("tool-result") - expect(events.find((event) => event.type === "tool-result")).toMatchObject({ + const result = events.find(LLMEvent.guards["tool-result"]) + expect(result).toMatchObject({ type: "tool-result", id: "call_1", name: "get_weather", result: { type: "json", value: { temperature: 22, condition: "sunny" } }, }) - expect(types.at(-1)).toBe("request-finish") + expect(events.at(-1)?.type).toBe("request-finish") expect(LLM.outputText({ events })).toBe("It's sunny in Paris.") }), ) @@ -106,27 +67,20 @@ describe("ToolRuntime", () => { Effect.gen(function* () { const llm = client({ adapters: [OpenAIChat.adapter] }) const layer = scriptedResponses([ - sseEvents( - toolCallChunk("call_1", "missing_tool", "{}"), - finishChunk("tool_calls"), - ), + sseEvents(toolCallChunk("call_1", "missing_tool", "{}"), finishChunk("tool_calls")), sseEvents(deltaChunk({ role: "assistant", content: "Sorry." }), finishChunk("stop")), ]) const events = Array.from( - yield* ToolRuntime.run(llm, { - request: baseRequest, - tools: { get_weather }, - }).pipe(Stream.runCollect, Effect.provide(layer)), + yield* ToolRuntime.run(llm, { request: baseRequest, tools: { get_weather } }).pipe( + Stream.runCollect, + Effect.provide(layer), + ), ) - const toolError = events.find((event) => event.type === "tool-error") - expect(toolError).toMatchObject({ - type: "tool-error", - id: "call_1", - name: "missing_tool", - }) - expect((toolError as { message: string }).message).toContain("Unknown tool") + const toolError = events.find(LLMEvent.guards["tool-error"]) + expect(toolError).toMatchObject({ type: "tool-error", id: "call_1", name: "missing_tool" }) + expect(toolError?.message).toContain("Unknown tool") }), ) @@ -134,23 +88,20 @@ describe("ToolRuntime", () => { Effect.gen(function* () { const llm = client({ adapters: [OpenAIChat.adapter] }) const layer = scriptedResponses([ - sseEvents( - toolCallChunk("call_1", "get_weather", '{"city":42}'), - finishChunk("tool_calls"), - ), + sseEvents(toolCallChunk("call_1", "get_weather", '{"city":42}'), finishChunk("tool_calls")), sseEvents(deltaChunk({ role: "assistant", content: "Done." }), finishChunk("stop")), ]) const events = Array.from( - yield* ToolRuntime.run(llm, { - request: baseRequest, - tools: { get_weather }, - }).pipe(Stream.runCollect, Effect.provide(layer)), + yield* ToolRuntime.run(llm, { request: baseRequest, tools: { get_weather } }).pipe( + Stream.runCollect, + Effect.provide(layer), + ), ) - const toolError = events.find((event) => event.type === "tool-error") + const toolError = events.find(LLMEvent.guards["tool-error"]) expect(toolError).toMatchObject({ type: "tool-error", id: "call_1", name: "get_weather" }) - expect((toolError as { message: string }).message).toContain("Invalid tool input") + expect(toolError?.message).toContain("Invalid tool input") }), ) @@ -158,38 +109,33 @@ describe("ToolRuntime", () => { Effect.gen(function* () { const llm = client({ adapters: [OpenAIChat.adapter] }) const layer = scriptedResponses([ - sseEvents( - toolCallChunk("call_1", "get_weather", '{"city":"FAIL"}'), - finishChunk("tool_calls"), - ), + sseEvents(toolCallChunk("call_1", "get_weather", '{"city":"FAIL"}'), finishChunk("tool_calls")), sseEvents(deltaChunk({ role: "assistant", content: "Sorry." }), finishChunk("stop")), ]) const events = Array.from( - yield* ToolRuntime.run(llm, { - request: baseRequest, - tools: { get_weather }, - }).pipe(Stream.runCollect, Effect.provide(layer)), + yield* ToolRuntime.run(llm, { request: baseRequest, tools: { get_weather } }).pipe( + Stream.runCollect, + Effect.provide(layer), + ), ) - const toolError = events.find((event) => event.type === "tool-error") + const toolError = events.find(LLMEvent.guards["tool-error"]) expect(toolError).toMatchObject({ type: "tool-error", id: "call_1", name: "get_weather" }) - expect((toolError as { message: string }).message).toBe("Weather lookup failed for FAIL") + expect(toolError?.message).toBe("Weather lookup failed for FAIL") }), ) it.effect("stops when the model finishes without requesting more tools", () => Effect.gen(function* () { const llm = client({ adapters: [OpenAIChat.adapter] }) - const layer = scriptedResponses([ - sseEvents(deltaChunk({ role: "assistant", content: "Done." }), finishChunk("stop")), - ]) + const layer = scriptedResponses([sseEvents(deltaChunk({ role: "assistant", content: "Done." }), finishChunk("stop"))]) const events = Array.from( - yield* ToolRuntime.run(llm, { - request: baseRequest, - tools: { get_weather }, - }).pipe(Stream.runCollect, Effect.provide(layer)), + yield* ToolRuntime.run(llm, { request: baseRequest, tools: { get_weather } }).pipe( + Stream.runCollect, + Effect.provide(layer), + ), ) expect(events.map((event) => event.type)).toEqual(["text-delta", "request-finish"]) @@ -203,22 +149,17 @@ describe("ToolRuntime", () => { // Every script entry asks for another tool call. With maxSteps: 2 the // runtime should run at most two model rounds and then exit even though // the model still wants to keep going. - const toolCallStep = sseEvents( - toolCallChunk("call_x", "get_weather", '{"city":"Paris"}'), - finishChunk("tool_calls"), - ) + const toolCallStep = sseEvents(toolCallChunk("call_x", "get_weather", '{"city":"Paris"}'), finishChunk("tool_calls")) const layer = scriptedResponses([toolCallStep, toolCallStep, toolCallStep]) const events = Array.from( - yield* ToolRuntime.run(llm, { - request: baseRequest, - tools: { get_weather }, - maxSteps: 2, - }).pipe(Stream.runCollect, Effect.provide(layer)), + yield* ToolRuntime.run(llm, { request: baseRequest, tools: { get_weather }, maxSteps: 2 }).pipe( + Stream.runCollect, + Effect.provide(layer), + ), ) - const finishEvents = events.filter((event) => event.type === "request-finish") - expect(finishEvents).toHaveLength(2) + expect(events.filter(LLMEvent.guards["request-finish"])).toHaveLength(2) }), ) @@ -226,10 +167,7 @@ describe("ToolRuntime", () => { Effect.gen(function* () { const llm = client({ adapters: [OpenAIChat.adapter] }) const layer = scriptedResponses([ - sseEvents( - toolCallChunk("call_1", "get_weather", '{"city":"Paris"}'), - finishChunk("tool_calls"), - ), + sseEvents(toolCallChunk("call_1", "get_weather", '{"city":"Paris"}'), finishChunk("tool_calls")), sseEvents(deltaChunk({ role: "assistant", content: "Should not run." }), finishChunk("stop")), ]) @@ -241,53 +179,38 @@ describe("ToolRuntime", () => { }).pipe(Stream.runCollect, Effect.provide(layer)), ) - const finishEvents = events.filter((event) => event.type === "request-finish") - expect(finishEvents).toHaveLength(1) - // No tool-result was emitted because stopWhen fired before dispatch - expect(events.some((event) => event.type === "tool-result")).toBe(false) + expect(events.filter(LLMEvent.guards["request-finish"])).toHaveLength(1) + expect(events.find(LLMEvent.guards["tool-result"])).toBeUndefined() }), ) it.effect("dispatches multiple tool calls in one step concurrently", () => Effect.gen(function* () { const llm = client({ adapters: [OpenAIChat.adapter] }) - // Two tool calls in the same step; each accumulates in its own index. - const body = `data: ${JSON.stringify({ - id: "x", - choices: [ - { - delta: { - role: "assistant", - tool_calls: [ - { index: 0, id: "c1", function: { name: "get_weather", arguments: '{"city":"Paris"}' } }, - { index: 1, id: "c2", function: { name: "get_weather", arguments: '{"city":"Tokyo"}' } }, - ], - }, - finish_reason: null, - }, - ], - usage: null, - })}\n\ndata: ${JSON.stringify({ - id: "x", - choices: [{ delta: {}, finish_reason: "tool_calls" }], - usage: null, - })}\n\ndata: [DONE]\n\n` - const layer = scriptedResponses([ - body, + sseEvents( + deltaChunk({ + role: "assistant", + tool_calls: [ + { index: 0, id: "c1", function: { name: "get_weather", arguments: '{"city":"Paris"}' } }, + { index: 1, id: "c2", function: { name: "get_weather", arguments: '{"city":"Tokyo"}' } }, + ], + }), + finishChunk("tool_calls"), + ), sseEvents(deltaChunk({ role: "assistant", content: "Both done." }), finishChunk("stop")), ]) const events = Array.from( - yield* ToolRuntime.run(llm, { - request: baseRequest, - tools: { get_weather }, - }).pipe(Stream.runCollect, Effect.provide(layer)), + yield* ToolRuntime.run(llm, { request: baseRequest, tools: { get_weather } }).pipe( + Stream.runCollect, + Effect.provide(layer), + ), ) - const results = events.filter((event) => event.type === "tool-result") + const results = events.filter(LLMEvent.guards["tool-result"]) expect(results).toHaveLength(2) - expect(results.map((event) => (event as { id: string }).id).sort()).toEqual(["c1", "c2"]) + expect(results.map((event) => event.id).toSorted()).toEqual(["c1", "c2"]) }), ) })