From 48fc9e3cc3a1626a7cb77e2541c62c612ddbcf7b Mon Sep 17 00:00:00 2001 From: Shoubhit Dash Date: Mon, 29 Jun 2026 16:26:33 +0530 Subject: [PATCH] feat(llm): add response reducer (#34417) --- packages/llm/src/route/client.ts | 17 +- packages/llm/src/schema/events.ts | 245 +++++++++++++++++- packages/llm/test/adapter.test.ts | 1 + ...nttokencount-on-identical-second-call.json | 4 +- .../llm/test/provider/openai-chat.test.ts | 19 +- packages/llm/test/response.test.ts | 60 +++++ 6 files changed, 324 insertions(+), 22 deletions(-) create mode 100644 packages/llm/test/response.test.ts diff --git a/packages/llm/src/route/client.ts b/packages/llm/src/route/client.ts index 5b5bc5ab2d5..b261cd7f231 100644 --- a/packages/llm/src/route/client.ts +++ b/packages/llm/src/route/client.ts @@ -374,17 +374,12 @@ const streamRequestWith = (runtime: TransportRuntime) => (request: LLMRequest) = const generateWith = (stream: Interface["stream"]) => Effect.fn("LLM.generate")(function* (request: LLMRequest) { - return new LLMResponse( - yield* stream(request).pipe( - Stream.runFold( - () => ({ events: [] as LLMEvent[], usage: undefined as LLMResponse["usage"] }), - (acc, event) => { - acc.events.push(event) - if ("usage" in event && event.usage !== undefined) acc.usage = event.usage - return acc - }, - ), - ), + const state = yield* stream(request).pipe(Stream.runFold(LLMResponse.empty, LLMResponse.reduce)) + const response = LLMResponse.complete(state) + if (response) return response + return yield* ProviderShared.eventError( + `${request.model.provider}/${request.model.route.id}`, + "Provider stream ended without a terminal finish event", ) }) diff --git a/packages/llm/src/schema/events.ts b/packages/llm/src/schema/events.ts index 3e46013521e..e0e09aa703f 100644 --- a/packages/llm/src/schema/events.ts +++ b/packages/llm/src/schema/events.ts @@ -1,7 +1,7 @@ import { Schema } from "effect" import { ContentBlockID, FinishReason, ProtocolID, ProviderMetadata, RouteID, ToolCallID } from "./ids" import { ModelSchema } from "./options" -import { ToolOutput, ToolResultValue } from "./messages" +import { Message, ToolCallPart, ToolOutput, ToolResultPart, ToolResultValue, type ContentPart } from "./messages" import { ProviderFailureClassification } from "./errors" /** @@ -335,9 +335,231 @@ const responseUsage = (events: ReadonlyArray) => undefined, ) +interface ContentAssembly { + readonly contentIndex: number + readonly text: string + readonly providerMetadata?: ProviderMetadata +} + +interface ToolInputAssembly { + readonly name: string + readonly text: string + readonly providerMetadata?: ProviderMetadata +} + +interface ResponseState { + readonly events: ReadonlyArray + readonly message: Message + readonly usage?: Usage + readonly finishReason?: FinishReason + readonly textParts: Readonly> + readonly reasoningParts: Readonly> + readonly toolInputs: Readonly> +} + +const emptyResponseState = (): ResponseState => ({ + events: [], + message: Message.assistant([]), + textParts: {}, + reasoningParts: {}, + toolInputs: {}, +}) + +const appendEvent = (state: ResponseState, event: LLMEvent): ResponseState => { + const events = [...state.events, event] + if (LLMEvent.is.finish(event)) { + return { + ...state, + events, + usage: event.usage ?? state.usage, + finishReason: event.reason, + } + } + if (LLMEvent.is.providerError(event)) { + return { + ...state, + events, + finishReason: state.finishReason ?? "error", + } + } + return { + ...state, + events, + usage: "usage" in event && event.usage !== undefined ? event.usage : state.usage, + } +} + +const textContent = (text: string, providerMetadata: ProviderMetadata | undefined): ContentPart => + providerMetadata === undefined ? { type: "text", text } : { type: "text", text, providerMetadata } + +const reasoningContent = (text: string, providerMetadata: ProviderMetadata | undefined): ContentPart => + providerMetadata === undefined ? { type: "reasoning", text } : { type: "reasoning", text, providerMetadata } + +const contentWith = (state: ResponseState, content: ReadonlyArray): ResponseState => ({ + ...state, + message: Message.assistant(content), +}) + +const appendContent = (state: ResponseState, part: ContentPart) => + contentWith(state, [...state.message.content, part]) + +const replaceContent = (state: ResponseState, index: number, part: ContentPart) => + contentWith( + state, + state.message.content.map((item, itemIndex) => (itemIndex === index ? part : item)), + ) + +const ensureText = (state: ResponseState, id: string, providerMetadata?: ProviderMetadata): ResponseState => { + if (state.textParts[id]) return state + return { + ...appendContent(state, textContent("", providerMetadata)), + textParts: { + ...state.textParts, + [id]: { contentIndex: state.message.content.length, text: "", providerMetadata }, + }, + } +} + +const reduceTextDelta = (state: ResponseState, event: TextDelta): ResponseState => { + const started = ensureText(state, event.id, event.providerMetadata) + const current = started.textParts[event.id] + if (!current) return started + const text = current.text + event.text + const providerMetadata = event.providerMetadata ?? current.providerMetadata + return { + ...replaceContent(started, current.contentIndex, textContent(text, providerMetadata)), + textParts: { ...started.textParts, [event.id]: { ...current, text, providerMetadata } }, + } +} + +const reduceTextEnd = (state: ResponseState, event: TextEnd): ResponseState => { + const current = state.textParts[event.id] + if (!current) return state + const providerMetadata = event.providerMetadata ?? current.providerMetadata + return { + ...replaceContent(state, current.contentIndex, textContent(current.text, providerMetadata)), + textParts: { ...state.textParts, [event.id]: { ...current, providerMetadata } }, + } +} + +const ensureReasoning = (state: ResponseState, id: string, providerMetadata?: ProviderMetadata): ResponseState => { + if (state.reasoningParts[id]) return state + return { + ...appendContent(state, reasoningContent("", providerMetadata)), + reasoningParts: { + ...state.reasoningParts, + [id]: { contentIndex: state.message.content.length, text: "", providerMetadata }, + }, + } +} + +const reduceReasoningDelta = (state: ResponseState, event: ReasoningDelta): ResponseState => { + const started = ensureReasoning(state, event.id, event.providerMetadata) + const current = started.reasoningParts[event.id] + if (!current) return started + const text = current.text + event.text + const providerMetadata = event.providerMetadata ?? current.providerMetadata + return { + ...replaceContent(started, current.contentIndex, reasoningContent(text, providerMetadata)), + reasoningParts: { ...started.reasoningParts, [event.id]: { ...current, text, providerMetadata } }, + } +} + +const reduceReasoningEnd = (state: ResponseState, event: ReasoningEnd): ResponseState => { + const current = state.reasoningParts[event.id] + if (!current) return state + const providerMetadata = event.providerMetadata ?? current.providerMetadata + return { + ...replaceContent(state, current.contentIndex, reasoningContent(current.text, providerMetadata)), + reasoningParts: { ...state.reasoningParts, [event.id]: { ...current, providerMetadata } }, + } +} + +const reduceToolInputStart = (state: ResponseState, event: ToolInputStart): ResponseState => ({ + ...state, + toolInputs: { + ...state.toolInputs, + [event.id]: { name: event.name, text: "", providerMetadata: event.providerMetadata }, + }, +}) + +const reduceToolInputDelta = (state: ResponseState, event: ToolInputDelta): ResponseState => { + const current = state.toolInputs[event.id] ?? { name: event.name, text: "" } + return { + ...state, + toolInputs: { ...state.toolInputs, [event.id]: { ...current, text: current.text + event.text } }, + } +} + +const reduceToolInputEnd = (state: ResponseState, event: ToolInputEnd): ResponseState => { + const current = state.toolInputs[event.id] ?? { name: event.name, text: "" } + return { + ...state, + toolInputs: { + ...state.toolInputs, + [event.id]: { ...current, name: event.name, providerMetadata: event.providerMetadata ?? current.providerMetadata }, + }, + } +} + +const toolCallContent = (event: ToolCall): ContentPart => + ToolCallPart.make({ + id: event.id, + name: event.name, + input: event.input, + ...(event.providerExecuted === undefined ? {} : { providerExecuted: event.providerExecuted }), + ...(event.providerMetadata === undefined ? {} : { providerMetadata: event.providerMetadata }), + }) + +const toolResultContent = (event: ToolResult): ContentPart => + ToolResultPart.make({ + id: event.id, + name: event.name, + result: event.result, + ...(event.providerExecuted === undefined ? {} : { providerExecuted: event.providerExecuted }), + ...(event.providerMetadata === undefined ? {} : { providerMetadata: event.providerMetadata }), + }) + +const reduceToolCall = (state: ResponseState, event: ToolCall): ResponseState => { + const { [event.id]: _finished, ...toolInputs } = state.toolInputs + return { ...appendContent(state, toolCallContent(event)), toolInputs } +} + +const reduceResponseState = (state: ResponseState, event: LLMEvent): ResponseState => { + const next = appendEvent(state, event) + switch (event.type) { + case "text-start": + return ensureText(next, event.id, event.providerMetadata) + case "text-delta": + return reduceTextDelta(next, event) + case "text-end": + return reduceTextEnd(next, event) + case "reasoning-start": + return ensureReasoning(next, event.id, event.providerMetadata) + case "reasoning-delta": + return reduceReasoningDelta(next, event) + case "reasoning-end": + return reduceReasoningEnd(next, event) + case "tool-input-start": + return reduceToolInputStart(next, event) + case "tool-input-delta": + return reduceToolInputDelta(next, event) + case "tool-input-end": + return reduceToolInputEnd(next, event) + case "tool-call": + return reduceToolCall(next, event) + case "tool-result": + return appendContent(next, toolResultContent(event)) + default: + return next + } +} + export class LLMResponse extends Schema.Class("LLM.Response")({ + message: Message, events: Schema.Array(LLMEvent), usage: Schema.optional(Usage), + finishReason: FinishReason, }) { /** Concatenated assistant text assembled from streamed `text-delta` events. */ get text() { @@ -356,8 +578,29 @@ export class LLMResponse extends Schema.Class("LLM.Response")({ } export namespace LLMResponse { + export type State = ResponseState export type Output = LLMResponse | { readonly events: ReadonlyArray; readonly usage?: Usage } + /** Initial reducer state for assembling one provider attempt. */ + export const empty = emptyResponseState + + /** Purely fold one provider-neutral event into the attempt assembly state. */ + export const reduce = reduceResponseState + + /** Return a completed response only after a terminal finish or provider error. */ + export const complete = (state: State): LLMResponse | undefined => + state.finishReason === undefined + ? undefined + : new LLMResponse({ + message: state.message, + events: [...state.events], + usage: state.usage, + finishReason: state.finishReason, + }) + + /** Convenience reducer for callers that already have a collected event list. */ + export const fromEvents = (events: ReadonlyArray) => complete(events.reduce(reduce, empty())) + /** Concatenate assistant text from a response or collected event list. */ export const text = (response: Output) => responseText(response.events) diff --git a/packages/llm/test/adapter.test.ts b/packages/llm/test/adapter.test.ts index 8e182948fa2..f71c6ef9be5 100644 --- a/packages/llm/test/adapter.test.ts +++ b/packages/llm/test/adapter.test.ts @@ -115,6 +115,7 @@ describe("llm route", () => { expect(events.map((event) => event.type)).toEqual(["text-delta", "finish"]) expect(response.events.map((event) => event.type)).toEqual(["text-delta", "finish"]) + expect(response.message.content).toEqual([{ type: "text", text: 'echo:{"body":"hello"}' }]) }), ) diff --git a/packages/llm/test/fixtures/recordings/gemini-cache/reports-cachedcontenttokencount-on-identical-second-call.json b/packages/llm/test/fixtures/recordings/gemini-cache/reports-cachedcontenttokencount-on-identical-second-call.json index 0145756887e..209aadfc1a4 100644 --- a/packages/llm/test/fixtures/recordings/gemini-cache/reports-cachedcontenttokencount-on-identical-second-call.json +++ b/packages/llm/test/fixtures/recordings/gemini-cache/reports-cachedcontenttokencount-on-identical-second-call.json @@ -21,7 +21,7 @@ "headers": { "content-type": "text/event-stream" }, - "body": "" + "body": "data: {\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"text\":\"Hi.\"}]},\"finishReason\":\"STOP\"}],\"usageMetadata\":{\"promptTokenCount\":1200,\"candidatesTokenCount\":2,\"totalTokenCount\":1202}}\n\n" } }, { @@ -39,7 +39,7 @@ "headers": { "content-type": "text/event-stream" }, - "body": "" + "body": "data: {\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"text\":\"Hi.\"}]},\"finishReason\":\"STOP\"}],\"usageMetadata\":{\"cachedContentTokenCount\":1100,\"promptTokenCount\":1200,\"candidatesTokenCount\":2,\"totalTokenCount\":1202}}\n\n" } } ] diff --git a/packages/llm/test/provider/openai-chat.test.ts b/packages/llm/test/provider/openai-chat.test.ts index 02995c00ce6..b736dc9dd33 100644 --- a/packages/llm/test/provider/openai-chat.test.ts +++ b/packages/llm/test/provider/openai-chat.test.ts @@ -1,7 +1,7 @@ import { describe, expect } from "bun:test" import { Effect, Schema, Stream } from "effect" import { HttpClientRequest } from "effect/unstable/http" -import { LLM, LLMError, Message, Model, ToolCallPart, Usage } from "../../src" +import { LLM, LLMError, LLMEvent, Message, Model, ToolCallPart, Usage } from "../../src" import * as Azure from "../../src/providers/azure" import * as OpenAI from "../../src/providers/openai" import * as OpenAIChat from "../../src/protocols/openai-chat" @@ -597,19 +597,22 @@ describe("OpenAI Chat route", () => { }), deltaChunk({ tool_calls: [{ index: 0, function: { arguments: ':"weather"}' } }] }), ) - const response = yield* LLMClient.generate( - LLM.updateRequest(request, { - tools: [{ name: "lookup", description: "Lookup data", inputSchema: { type: "object" } }], - }), - ).pipe(Effect.provide(fixedResponse(body))) + const input = LLM.updateRequest(request, { + tools: [{ name: "lookup", description: "Lookup data", inputSchema: { type: "object" } }], + }) + const events = Array.from( + yield* LLMClient.stream(input).pipe(Stream.runCollect, Effect.provide(fixedResponse(body))), + ) + const error = yield* LLMClient.generate(input).pipe(Effect.provide(fixedResponse(body)), Effect.flip) - expect(response.events).toEqual([ + expect(events).toEqual([ { type: "step-start", index: 0 }, { type: "tool-input-start", id: "call_1", name: "lookup", providerMetadata: undefined }, { type: "tool-input-delta", id: "call_1", name: "lookup", text: '{"query"' }, { type: "tool-input-delta", id: "call_1", name: "lookup", text: ':"weather"}' }, ]) - expect(response.toolCalls).toEqual([]) + expect(events.filter(LLMEvent.is.toolCall)).toEqual([]) + expect(error.message).toContain("Provider stream ended without a terminal finish event") }), ) diff --git a/packages/llm/test/response.test.ts b/packages/llm/test/response.test.ts new file mode 100644 index 00000000000..8c8de191511 --- /dev/null +++ b/packages/llm/test/response.test.ts @@ -0,0 +1,60 @@ +import { describe, expect, test } from "bun:test" +import { LLMEvent, LLMResponse } from "../src" + +const reduce = (events: ReadonlyArray) => events.reduce(LLMResponse.reduce, LLMResponse.empty()) + +describe("LLMResponse reducer", () => { + test("assembles interleaved reasoning and text with end metadata", () => { + const response = LLMResponse.fromEvents([ + LLMEvent.reasoningStart({ id: "r1" }), + LLMEvent.reasoningDelta({ id: "r1", text: "I should " }), + LLMEvent.textStart({ id: "t1" }), + LLMEvent.reasoningDelta({ id: "r1", text: "compare..." }), + LLMEvent.reasoningEnd({ id: "r1", providerMetadata: { anthropic: { signature: "sig" } } }), + LLMEvent.textDelta({ id: "t1", text: "Answer" }), + LLMEvent.textEnd({ id: "t1" }), + LLMEvent.finish({ reason: "stop", usage: { outputTokens: 5 } }), + ]) + + expect(response?.finishReason).toBe("stop") + expect(response?.usage).toMatchObject({ outputTokens: 5 }) + expect(response?.message.content).toEqual([ + { + type: "reasoning", + text: "I should compare...", + providerMetadata: { anthropic: { signature: "sig" } }, + }, + { type: "text", text: "Answer" }, + ]) + expect(response?.events).toHaveLength(8) + }) + + test("preserves partial content without completing a failed stream", () => { + const state = reduce([LLMEvent.textStart({ id: "t1" }), LLMEvent.textDelta({ id: "t1", text: "partial" })]) + + expect(LLMResponse.complete(state)).toBeUndefined() + expect(state.message.content).toEqual([{ type: "text", text: "partial" }]) + }) + + test("assembles tool-call content only after the completed tool call event", () => { + const pending = reduce([ + LLMEvent.toolInputStart({ id: "call_1", name: "lookup" }), + LLMEvent.toolInputDelta({ id: "call_1", name: "lookup", text: '{"query"' }), + ]) + + expect(pending.message.content).toEqual([]) + expect(pending.toolInputs.call_1?.text).toBe('{"query"') + + const response = LLMResponse.fromEvents([ + ...pending.events, + LLMEvent.toolInputDelta({ id: "call_1", name: "lookup", text: ':"weather"}' }), + LLMEvent.toolInputEnd({ id: "call_1", name: "lookup" }), + LLMEvent.toolCall({ id: "call_1", name: "lookup", input: { query: "weather" } }), + LLMEvent.finish({ reason: "tool-calls" }), + ]) + + expect(response?.message.content).toEqual([ + { type: "tool-call", id: "call_1", name: "lookup", input: { query: "weather" } }, + ]) + }) +})