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<T, any, never, never> 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.
This commit is contained in:
Kit Langton
@@ -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<T extends Tools> {
|
||||
* 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<T extends Tools> {
|
||||
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 = <T extends Tools>(
|
||||
client: LLMClient,
|
||||
options: RunOptions<T>,
|
||||
): Stream.Stream<LLMEvent, LLMError, RequestExecutor.Service> => {
|
||||
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<LLMEvent, LLMError, RequestExecutor.Service> =>
|
||||
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 = <T extends Tools>(
|
||||
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<Dispatched> => {
|
||||
const dispatch = (tools: Tools, call: ToolCallPart): Effect.Effect<ToolResultValue> => {
|
||||
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<Dispatched> =>
|
||||
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<Schema.Top, Schema.Top>,
|
||||
input: unknown,
|
||||
): Effect.Effect<ToolResultValue, ToolFailure> => {
|
||||
const decode = Schema.decodeUnknownEffect(tool.parameters) as unknown as (
|
||||
input: unknown,
|
||||
) => Effect.Effect<unknown, { readonly message?: string }>
|
||||
const encode = Schema.encodeEffect(tool.success) as unknown as (
|
||||
value: unknown,
|
||||
) => Effect.Effect<unknown, { readonly message?: string }>
|
||||
|
||||
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<ToolResultValue, ToolFailure> =>
|
||||
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<LLMEvent> =>
|
||||
result.type === "error"
|
||||
|
||||
+55
-10
@@ -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<T> = Schema.Codec<T, any, never, never>
|
||||
|
||||
/**
|
||||
* 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<Parameters extends Schema.Top, Success extends Schema.Top> {
|
||||
export interface Tool<Parameters extends ToolSchema<any>, Success extends ToolSchema<any>> {
|
||||
readonly description: string
|
||||
readonly parameters: Parameters
|
||||
readonly success: Success
|
||||
readonly execute: (
|
||||
params: Schema.Schema.Type<Parameters>,
|
||||
) => Effect.Effect<Schema.Schema.Type<Success>, ToolFailure>
|
||||
/** @internal */
|
||||
readonly _decode: (input: unknown) => Effect.Effect<Schema.Schema.Type<Parameters>, Schema.SchemaError>
|
||||
/** @internal */
|
||||
readonly _encode: (value: Schema.Schema.Type<Success>) => Effect.Effect<unknown, Schema.SchemaError>
|
||||
/** @internal */
|
||||
readonly _definition: ToolDefinitionClass
|
||||
}
|
||||
|
||||
export type AnyTool = Tool<ToolSchema<any>, ToolSchema<any>>
|
||||
|
||||
/**
|
||||
* 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<Parameters extends Schema.Top, Success extends Schema.Top>
|
||||
* })
|
||||
* ```
|
||||
*/
|
||||
export const tool = <Parameters extends Schema.Top, Success extends Schema.Top>(
|
||||
config: Tool<Parameters, Success>,
|
||||
): Tool<Parameters, Success> => config
|
||||
export const tool = <Parameters extends ToolSchema<any>, Success extends ToolSchema<any>>(config: {
|
||||
readonly description: string
|
||||
readonly parameters: Parameters
|
||||
readonly success: Success
|
||||
readonly execute: (
|
||||
params: Schema.Schema.Type<Parameters>,
|
||||
) => Effect.Effect<Schema.Schema.Type<Success>, ToolFailure>
|
||||
}): Tool<Parameters, Success> => ({
|
||||
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<string, Tool<any, any>>
|
||||
export type Tools = Record<string, AnyTool>
|
||||
|
||||
/**
|
||||
* 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<ToolDefinitionClass> =>
|
||||
Object.entries(tools).map(([name, item]) =>
|
||||
new ToolDefinition({
|
||||
name,
|
||||
description: item.description,
|
||||
inputSchema: Schema.toJsonSchemaDocument(item.parameters).schema as Record<string, unknown>,
|
||||
description: item._definition.description,
|
||||
inputSchema: item._definition.inputSchema,
|
||||
}),
|
||||
)
|
||||
|
||||
const toJsonSchema = (schema: Schema.Top): Record<string, unknown> => {
|
||||
const document = Schema.toJsonSchemaDocument(schema)
|
||||
if (Object.keys(document.definitions).length === 0) return document.schema as Record<string, unknown>
|
||||
return { ...document.schema, $defs: document.definitions } as Record<string, unknown>
|
||||
}
|
||||
|
||||
export { ToolFailure }
|
||||
|
||||
export * as Tool from "./tool"
|
||||
|
||||
@@ -80,9 +80,9 @@ const gemini = Adapter.define<FakeDraft, FakeDraft>({
|
||||
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" },
|
||||
|
||||
@@ -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<typeof Response>[0], init?: ResponseInit) => HttpClientResponse.HttpClientResponse
|
||||
}
|
||||
|
||||
export type Handler = (input: HandlerInput) => Effect.Effect<Response>
|
||||
export type Handler = (input: HandlerInput) => Effect.Effect<HttpClientResponse.HttpClientResponse>
|
||||
|
||||
const handlerLayer = (handler: Handler): Layer.Layer<HttpClient.HttpClient> =>
|
||||
Layer.succeed(
|
||||
@@ -16,8 +17,11 @@ const handlerLayer = (handler: Handler): Layer.Layer<HttpClient.HttpClient> =>
|
||||
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<string>) =>
|
||||
dynamicResponse(() =>
|
||||
dynamicResponse((input) =>
|
||||
Effect.sync(() => {
|
||||
const encoder = new TextEncoder()
|
||||
const stream = new ReadableStream({
|
||||
@@ -53,6 +57,26 @@ export const truncatedStream = (chunks: ReadonlyArray<string>) =>
|
||||
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<string>, 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)
|
||||
}),
|
||||
)
|
||||
}),
|
||||
)
|
||||
}
|
||||
|
||||
@@ -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 } }],
|
||||
})
|
||||
@@ -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: "!" }),
|
||||
|
||||
@@ -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<string>) =>
|
||||
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"])
|
||||
}),
|
||||
)
|
||||
})
|
||||
|
||||
Reference in New Issue
Block a user