Preserve native LLM tool context (#27116)

This commit is contained in:
Kit Langton
2026-05-12 16:16:58 -04:00
committed by GitHub
parent b9e7cbf13c
commit dd14413a64
18 changed files with 244 additions and 75 deletions

View File

@@ -51,7 +51,7 @@ const request = LLM.request({
const raiseEvent = (event: FakeEvent): import("../src/schema").LLMEvent =>
event.type === "finish"
? { type: "request-finish", reason: event.reason }
? { type: "finish", reason: event.reason }
: { type: "text-delta", id: "text-0", text: event.text }
const fakeProtocol = Protocol.make<FakeBody, FakeEvent, FakeEvent, void>({
@@ -112,8 +112,8 @@ describe("llm route", () => {
const events = Array.from(yield* llm.stream(request).pipe(Stream.runCollect))
const response = yield* llm.generate(request)
expect(events.map((event) => event.type)).toEqual(["text-delta", "request-finish"])
expect(response.events.map((event) => event.type)).toEqual(["text-delta", "request-finish"])
expect(events.map((event) => event.type)).toEqual(["text-delta", "finish"])
expect(response.events.map((event) => event.type)).toEqual(["text-delta", "finish"])
}),
)

View File

@@ -127,7 +127,7 @@ describe("llm constructors", () => {
LLMResponse.text({
events: [
{ type: "text-delta", id: "text-0", text: "hi" },
{ type: "request-finish", reason: "stop" },
{ type: "finish", reason: "stop" },
],
}),
).toBe("hi")

View File

@@ -124,7 +124,7 @@ describe("Anthropic Messages route", () => {
providerMetadata: { anthropic: { signature: "sig_1" } },
})
expect(response.events.at(-1)).toMatchObject({
type: "request-finish",
type: "finish",
reason: "stop",
providerMetadata: { anthropic: { stopSequence: "\n\nHuman:" } },
})
@@ -182,7 +182,7 @@ describe("Anthropic Messages route", () => {
},
{ type: "step-finish", index: 0, reason: "tool-calls", usage, providerMetadata: undefined },
{
type: "request-finish",
type: "finish",
reason: "tool-calls",
providerMetadata: undefined,
usage,
@@ -275,7 +275,7 @@ describe("Anthropic Messages route", () => {
providerMetadata: { anthropic: { blockType: "web_search_tool_result" } },
})
expect(response.text).toBe("Found it.")
expect(response.events.at(-1)).toMatchObject({ type: "request-finish", reason: "stop" })
expect(response.events.at(-1)).toMatchObject({ type: "finish", reason: "stop" })
}),
)

View File

@@ -169,12 +169,12 @@ describe("Bedrock Converse route", () => {
const response = yield* LLMClient.generate(baseRequest).pipe(Effect.provide(fixedBytes(body)))
expect(response.text).toBe("Hello!")
const finishes = response.events.filter((event) => event.type === "request-finish")
const finishes = response.events.filter((event) => event.type === "finish")
// Bedrock splits the finish across `messageStop` (carries reason) and
// `metadata` (carries usage). We consolidate them into a single
// terminal `request-finish` event with both.
// terminal `finish` event with both.
expect(finishes).toHaveLength(1)
expect(finishes[0]).toMatchObject({ type: "request-finish", reason: "stop" })
expect(finishes[0]).toMatchObject({ type: "finish", reason: "stop" })
expect(response.usage).toMatchObject({
inputTokens: 5,
outputTokens: 2,
@@ -213,7 +213,7 @@ describe("Bedrock Converse route", () => {
{ type: "tool-input-delta", id: "tool_1", name: "lookup", text: '{"query"' },
{ type: "tool-input-delta", id: "tool_1", name: "lookup", text: ':"weather"}' },
])
expect(response.events.at(-1)).toMatchObject({ type: "request-finish", reason: "tool-calls" })
expect(response.events.at(-1)).toMatchObject({ type: "finish", reason: "tool-calls" })
}),
)

View File

@@ -232,7 +232,7 @@ describe("Gemini route", () => {
{ type: "text-end", id: "text-0" },
{ type: "step-finish", index: 0, reason: "stop", usage, providerMetadata: undefined },
{
type: "request-finish",
type: "finish",
reason: "stop",
usage,
},
@@ -291,7 +291,7 @@ describe("Gemini route", () => {
},
{ type: "step-finish", index: 0, reason: "tool-calls", usage, providerMetadata: undefined },
{
type: "request-finish",
type: "finish",
reason: "tool-calls",
usage,
},
@@ -325,7 +325,7 @@ describe("Gemini route", () => {
{ type: "tool-call", id: "tool_0", name: "lookup", input: { query: "weather" } },
{ type: "tool-call", id: "tool_1", name: "lookup", input: { query: "news" } },
])
expect(response.events.at(-1)).toMatchObject({ type: "request-finish", reason: "tool-calls" })
expect(response.events.at(-1)).toMatchObject({ type: "finish", reason: "tool-calls" })
}),
)
@@ -344,10 +344,10 @@ describe("Gemini route", () => {
),
)
expect(length.events.map((event) => event.type)).toEqual(["step-start", "step-finish", "request-finish"])
expect(length.events.at(-1)).toMatchObject({ type: "request-finish", reason: "length" })
expect(filtered.events.map((event) => event.type)).toEqual(["step-start", "step-finish", "request-finish"])
expect(filtered.events.at(-1)).toMatchObject({ type: "request-finish", reason: "content-filter" })
expect(length.events.map((event) => event.type)).toEqual(["step-start", "step-finish", "finish"])
expect(length.events.at(-1)).toMatchObject({ type: "finish", reason: "length" })
expect(filtered.events.map((event) => event.type)).toEqual(["step-start", "step-finish", "finish"])
expect(filtered.events.at(-1)).toMatchObject({ type: "finish", reason: "content-filter" })
}),
)

View File

@@ -249,7 +249,7 @@ describe("OpenAI Chat route", () => {
{ type: "text-end", id: "text-0" },
{ type: "step-finish", index: 0, reason: "stop", usage, providerMetadata: undefined },
{
type: "request-finish",
type: "finish",
reason: "stop",
usage,
},
@@ -288,7 +288,7 @@ describe("OpenAI Chat route", () => {
providerMetadata: undefined,
},
{ type: "step-finish", index: 0, reason: "tool-calls", usage: undefined, providerMetadata: undefined },
{ type: "request-finish", reason: "tool-calls", usage: undefined },
{ type: "finish", reason: "tool-calls", usage: undefined },
])
}),
)

View File

@@ -231,7 +231,7 @@ describe("OpenAI-compatible Chat route", () => {
expect(response.text).toBe("Hello!")
expect(response.usage).toMatchObject({ inputTokens: 5, outputTokens: 2, totalTokens: 7 })
expect(response.events.at(-1)).toMatchObject({ type: "request-finish", reason: "stop" })
expect(response.events.at(-1)).toMatchObject({ type: "finish", reason: "stop" })
}),
)
})

View File

@@ -366,7 +366,7 @@ describe("OpenAI Responses route", () => {
usage,
},
{
type: "request-finish",
type: "finish",
reason: "stop",
providerMetadata: { openai: { responseId: "resp_1", serviceTier: "default" } },
usage,
@@ -447,7 +447,7 @@ describe("OpenAI Responses route", () => {
},
{ type: "step-finish", index: 0, reason: "tool-calls", usage, providerMetadata: undefined },
{
type: "request-finish",
type: "finish",
reason: "tool-calls",
providerMetadata: undefined,
usage,

View File

@@ -120,8 +120,8 @@ export const runWeatherToolLoop = (request: LLMRequest) =>
export const expectFinish = (
events: ReadonlyArray<LLMEvent>,
reason: Extract<LLMEvent, { readonly type: "request-finish" }>["reason"],
) => expect(events.at(-1)).toMatchObject({ type: "request-finish", reason })
reason: Extract<LLMEvent, { readonly type: "finish" }>["reason"],
) => expect(events.at(-1)).toMatchObject({ type: "finish", reason })
export const expectWeatherToolCall = (response: LLMResponse) =>
expect(response.toolCalls).toMatchObject([
@@ -129,10 +129,12 @@ export const expectWeatherToolCall = (response: LLMResponse) =>
])
export const expectWeatherToolLoop = (events: ReadonlyArray<LLMEvent>) => {
const finishes = events.filter(LLMEvent.is.requestFinish)
expect(finishes).toHaveLength(2)
expect(finishes[0]?.reason).toBe("tool-calls")
expect(finishes.at(-1)?.reason).toBe("stop")
const finishes = events.filter(LLMEvent.is.finish)
expect(finishes).toHaveLength(1)
expect(finishes[0]?.reason).toBe("stop")
const stepFinishes = events.filter(LLMEvent.is.stepFinish)
expect(stepFinishes.map((event) => event.reason)).toEqual(["tool-calls", "stop"])
const toolCalls = events.filter(LLMEvent.is.toolCall)
expect(toolCalls).toHaveLength(1)
@@ -272,7 +274,7 @@ export const eventSummary = (events: ReadonlyArray<LLMEvent>) => {
summary.push({ type: "tool-error", name: event.name, message: event.message })
continue
}
if (event.type === "request-finish") {
if (event.type === "finish") {
summary.push({ type: "finish", reason: event.reason, usage: usageSummary(event.usage) })
}
}

View File

@@ -44,6 +44,11 @@ describe("llm schema", () => {
expect(() => Schema.decodeUnknownSync(LLMEvent)({ type: "bogus" })).toThrow()
})
test("finish constructors accept usage input", () => {
expect(LLMEvent.stepFinish({ index: 0, reason: "stop", usage: { inputTokens: 1 } }).usage).toBeInstanceOf(Usage)
expect(LLMEvent.finish({ reason: "stop", usage: { outputTokens: 2 } }).usage).toBeInstanceOf(Usage)
})
test("content part tagged union exposes guards", () => {
expect(ContentPart.guards.text({ type: "text", text: "hi" })).toBe(true)
expect(ContentPart.guards.media({ type: "text", text: "hi" })).toBe(false)

View File

@@ -4,7 +4,8 @@ import { GenerationOptions, LLM, LLMEvent, LLMRequest, LLMResponse, ToolChoice }
import { LLMClient } from "../src/route"
import * as AnthropicMessages from "../src/protocols/anthropic-messages"
import * as OpenAIChat from "../src/protocols/openai-chat"
import { tool, ToolFailure } from "../src/tool"
import { tool, ToolFailure, type ToolExecuteContext } from "../src/tool"
import { ToolRuntime } from "../src/tool-runtime"
import { it } from "./lib/effect"
import * as TestToolRuntime from "./lib/tool-runtime"
import { dynamicResponse, scriptedResponses } from "./lib/http"
@@ -129,7 +130,7 @@ describe("LLMClient tools", () => {
name: "get_weather",
result: { type: "json", value: { temperature: 22, condition: "sunny" } },
})
expect(events.at(-1)?.type).toBe("request-finish")
expect(events.at(-1)?.type).toBe("finish")
expect(LLMResponse.text({ events })).toBe("It's sunny in Paris.")
}),
)
@@ -148,11 +149,40 @@ describe("LLMClient tools", () => {
),
)
expect(events.filter(LLMEvent.is.requestFinish)).toHaveLength(1)
expect(events.filter(LLMEvent.is.finish)).toHaveLength(1)
expect(events.find(LLMEvent.is.toolResult)).toMatchObject({ type: "tool-result", id: "call_1" })
}),
)
it.effect("passes tool call context to execute", () =>
Effect.gen(function* () {
let context: ToolExecuteContext | undefined
const contextual = tool({
description: "Capture tool context.",
parameters: Schema.Struct({ value: Schema.String }),
success: Schema.Struct({ ok: Schema.Boolean }),
execute: (_params, ctx) =>
Effect.sync(() => {
context = ctx
return { ok: true }
}),
})
const events = Array.from(
yield* TestToolRuntime.runTools({ request: baseRequest, tools: { contextual } }).pipe(
Stream.runCollect,
Effect.provide(
scriptedResponses([
sseEvents(toolCallChunk("call_ctx", "contextual", '{"value":"x"}'), finishChunk("tool_calls")),
]),
),
),
)
expect(events.some(LLMEvent.is.toolResult)).toBe(true)
expect(context).toEqual({ id: "call_ctx", name: "contextual" })
}),
)
it.effect("can expose tool schemas without executing tool calls", () =>
Effect.gen(function* () {
const layer = scriptedResponses([
@@ -319,7 +349,7 @@ describe("LLMClient tools", () => {
"text-delta",
"text-end",
"step-finish",
"request-finish",
"finish",
])
expect(LLMResponse.text({ events })).toBe("Done.")
}),
@@ -343,7 +373,57 @@ describe("LLMClient tools", () => {
),
)
expect(events.filter(LLMEvent.is.requestFinish)).toHaveLength(2)
expect(events.filter(LLMEvent.is.finish)).toHaveLength(1)
expect(events.filter(LLMEvent.is.stepStart).map((event) => event.index)).toEqual([0, 1])
expect(events.filter(LLMEvent.is.stepFinish).map((event) => event.index)).toEqual([0, 1])
}),
)
it.effect("emits one final finish with aggregate usage", () =>
Effect.gen(function* () {
let calls = 0
const events = Array.from(
yield* ToolRuntime.stream({
request: baseRequest,
tools: { get_weather },
stopWhen: ToolRuntime.stepCountIs(2),
stream: () =>
Stream.fromIterable<LLMEvent>(
calls++ === 0
? [
LLMEvent.stepStart({ index: 0 }),
LLMEvent.toolCall({ id: "call_1", name: "get_weather", input: { city: "Paris" } }),
LLMEvent.stepFinish({
index: 0,
reason: "tool-calls",
usage: { inputTokens: 1, outputTokens: 2, totalTokens: 3 },
}),
LLMEvent.finish({
reason: "tool-calls",
usage: { inputTokens: 1, outputTokens: 2, totalTokens: 3 },
}),
]
: [
LLMEvent.stepStart({ index: 0 }),
LLMEvent.textDelta({ id: "text_1", text: "Done." }),
LLMEvent.stepFinish({
index: 0,
reason: "stop",
usage: { inputTokens: 4, outputTokens: 5, totalTokens: 9 },
}),
LLMEvent.finish({ reason: "stop", usage: { inputTokens: 4, outputTokens: 5, totalTokens: 9 } }),
],
),
}).pipe(Stream.runCollect),
)
expect(events.filter(LLMEvent.is.stepFinish).map((event) => event.index)).toEqual([0, 1])
expect(events.filter(LLMEvent.is.finish)).toHaveLength(1)
expect(events.find(LLMEvent.is.finish)?.usage).toMatchObject({
inputTokens: 5,
outputTokens: 7,
totalTokens: 12,
})
}),
)
@@ -362,7 +442,7 @@ describe("LLMClient tools", () => {
}).pipe(Stream.runCollect, Effect.provide(layer)),
)
expect(events.filter(LLMEvent.is.requestFinish)).toHaveLength(1)
expect(events.filter(LLMEvent.is.finish)).toHaveLength(1)
expect(events.find(LLMEvent.is.toolResult)).toMatchObject({ type: "tool-result", id: "call_1" })
}),
)