diff --git a/.changeset/hungry-hounds-yell.md b/.changeset/hungry-hounds-yell.md new file mode 100644 index 000000000..1621c8191 --- /dev/null +++ b/.changeset/hungry-hounds-yell.md @@ -0,0 +1,5 @@ +--- +"braintrust": patch +--- + +fix: Fix TTFT in AI SDK v6 diff --git a/js/src/instrumentation/plugins/ai-sdk-plugin.streaming.test.ts b/js/src/instrumentation/plugins/ai-sdk-plugin.streaming.test.ts new file mode 100644 index 000000000..90497a2dc --- /dev/null +++ b/js/src/instrumentation/plugins/ai-sdk-plugin.streaming.test.ts @@ -0,0 +1,219 @@ +/* eslint-disable @typescript-eslint/no-explicit-any */ +/* eslint-disable @typescript-eslint/consistent-type-assertions */ +import { + afterEach, + beforeAll, + beforeEach, + describe, + expect, + test, +} from "vitest"; +import * as ai from "ai"; +import { configureNode } from "../../node/config"; +import { + _exportsForTestingOnly, + initLogger, + Logger, + TestBackgroundLogger, +} from "../../logger"; +import { wrapAISDK } from "../../wrappers/ai-sdk"; +import { aiSDKChannels } from "./ai-sdk-channels"; + +try { + configureNode(); +} catch {} + +const sleep = (ms: number) => + new Promise((resolve) => setTimeout(resolve, ms)); + +describe("AI SDK streaming instrumentation", () => { + let backgroundLogger: TestBackgroundLogger; + let _logger: Logger; + + beforeAll(async () => { + await _exportsForTestingOnly.simulateLoginForTests(); + }); + + beforeEach(() => { + backgroundLogger = _exportsForTestingOnly.useTestBackgroundLogger(); + _logger = initLogger({ + projectName: "ai-sdk-plugin.streaming.test.ts", + projectId: "test-project-id", + }); + }); + + afterEach(() => { + _exportsForTestingOnly.clearTestBackgroundLogger(); + }); + + test("streamText time_to_first_token ignores AI SDK v6 framing chunks", async () => { + expect(await backgroundLogger.drain()).toHaveLength(0); + + const requestDelayMs = 40; + const contentDelayMs = 80; + let sentContent = false; + const model = { + specificationVersion: "v3", + provider: "mock-delayed-provider", + modelId: "mock-delayed-model", + supportedUrls: {}, + doGenerate: async () => { + throw new Error("doGenerate should not be called"); + }, + doStream: async () => { + await sleep(requestDelayMs); + + return { + stream: new ReadableStream({ + start(controller) { + controller.enqueue({ type: "stream-start", warnings: [] }); + controller.enqueue({ + type: "response-metadata", + id: "response-id", + modelId: "mock-delayed-model", + timestamp: new Date(0), + }); + controller.enqueue({ + type: "raw", + rawValue: { + type: "response.created", + response: { id: "response-id" }, + }, + }); + }, + async pull(controller) { + if (sentContent) { + controller.close(); + return; + } + + sentContent = true; + await sleep(contentDelayMs); + controller.enqueue({ type: "text-start", id: "delayed-text" }); + controller.enqueue({ + type: "text-delta", + id: "delayed-text", + delta: "DELAYED", + }); + controller.enqueue({ type: "text-end", id: "delayed-text" }); + controller.enqueue({ + type: "finish", + finishReason: { unified: "stop", raw: "stop" }, + usage: { + inputTokens: { + total: 1, + noCache: 1, + cacheRead: 0, + cacheWrite: 0, + }, + outputTokens: { + total: 1, + text: 1, + reasoning: 0, + }, + }, + }); + }, + }), + warnings: [], + }; + }, + } as any; + + const wrappedAI = wrapAISDK(ai); + const result = wrappedAI.streamText({ + model, + prompt: "Reply with exactly DELAYED.", + includeRawChunks: true, + maxOutputTokens: 16, + }); + + let fullText = ""; + for await (const chunk of result.textStream) { + fullText += chunk; + } + await result.text; + + expect(fullText).toBe("DELAYED"); + + const spans = (await backgroundLogger.drain()) as any[]; + const streamTextSpan = spans.find( + (s) => s?.span_attributes?.name === "streamText", + ); + const doStreamSpan = spans.find( + (s) => s?.span_attributes?.name === "doStream", + ); + const minimumExpectedTTFT = (requestDelayMs + contentDelayMs) / 1000 / 2; + + expect(streamTextSpan?.metrics?.time_to_first_token).toBeGreaterThanOrEqual( + minimumExpectedTTFT, + ); + expect(doStreamSpan?.metrics?.time_to_first_token).toBeGreaterThanOrEqual( + minimumExpectedTTFT, + ); + expect(streamTextSpan?.output?.text).toBe("DELAYED"); + expect(doStreamSpan?.output?.text).toBe("DELAYED"); + }); + + test("streamText time_to_first_token counts streamed tool input arguments", async () => { + expect(await backgroundLogger.drain()).toHaveLength(0); + + const contentDelayMs = 80; + let sentContent = false; + const result = (await aiSDKChannels.streamText.tracePromise( + async () => ({ + baseStream: new ReadableStream({ + start(controller) { + controller.enqueue({ type: "stream-start", warnings: [] }); + controller.enqueue({ type: "text-start", id: "ignored-text" }); + controller.enqueue({ + type: "tool-input-start", + id: "call-1", + toolName: "lookup", + }); + }, + async pull(controller) { + if (sentContent) { + controller.close(); + return; + } + + sentContent = true; + await sleep(contentDelayMs); + controller.enqueue({ + type: "tool-input-delta", + id: "call-1", + inputTextDelta: '{"query"', + }); + controller.close(); + }, + }), + }), + { + arguments: [ + { + model: "mock-tool-model", + prompt: "Call the lookup tool.", + }, + ], + } as any, + )) as any; + + const reader = result.baseStream.getReader(); + while (true) { + const { done } = await reader.read(); + if (done) { + break; + } + } + + const spans = (await backgroundLogger.drain()) as any[]; + const streamTextSpan = spans.find( + (s) => s?.span_attributes?.name === "streamText", + ); + + expect(streamTextSpan?.metrics?.time_to_first_token).toBeGreaterThanOrEqual( + contentDelayMs / 1000 / 2, + ); + }); +}); diff --git a/js/src/instrumentation/plugins/ai-sdk-plugin.ts b/js/src/instrumentation/plugins/ai-sdk-plugin.ts index cbadfff64..749b6fb70 100644 --- a/js/src/instrumentation/plugins/ai-sdk-plugin.ts +++ b/js/src/instrumentation/plugins/ai-sdk-plugin.ts @@ -1213,10 +1213,12 @@ function prepareAISDKChildTracing( }, }); + const streamStartTime = getCurrentUnixTimestamp(); const result = await withCurrent(span, () => Reflect.apply(originalDoStream, resolvedModel, [options]), ); - const streamStartTime = getCurrentUnixTimestamp(); + // firstChunkTime !== streamStartTime because the first few chunks may be actual bookkeeping stuff for the SDK + // instead of actual LLM output that can be streamed to users. let firstChunkTime: number | undefined; const output: Record = {}; let text = ""; @@ -1226,7 +1228,10 @@ function prepareAISDKChildTracing( const transformStream = new TransformStream({ transform(chunk: AISDKModelStreamChunk, controller) { - if (firstChunkTime === undefined) { + if ( + firstChunkTime === undefined && + isAISDKContentStreamChunk(chunk) + ) { firstChunkTime = getCurrentUnixTimestamp(); } @@ -1475,6 +1480,137 @@ function finalizeAISDKChildTracing(event?: { [key: string]: unknown }): void { } } +function extractAISDKStreamPart(chunk: unknown): unknown { + if (!isObject(chunk) || !isObject(chunk.part)) { + return chunk; + } + + return chunk.part; +} + +function stringContent(value: unknown): string | undefined { + return typeof value === "string" && value.length > 0 ? value : undefined; +} + +function rawValueHasAISDKContent(value: unknown): boolean { + if (value === undefined || value === null) { + return false; + } + + if (typeof value === "string") { + return value.length > 0; + } + + if (!isObject(value)) { + return true; + } + + // Raw streams also include provider lifecycle events; only content shapes count. + const delta = value.delta; + if ( + stringContent(value.content) || + stringContent(value.text) || + stringContent(value.delta) || + stringContent(value.arguments) || + stringContent(value.partial_json) || + (isObject(delta) && + (stringContent(delta.content) || + stringContent(delta.text) || + stringContent(delta.arguments) || + stringContent(delta.partial_json) || + stringContent(delta.thinking))) || + (Array.isArray(value.choices) && + value.choices.some((choice) => { + if (!isObject(choice) || !isObject(choice.delta)) { + return false; + } + + const choiceDelta = choice.delta; + if ( + stringContent(choiceDelta.content) || + stringContent(choiceDelta.text) + ) { + return true; + } + + if ( + isObject(choiceDelta.function_call) && + stringContent(choiceDelta.function_call.arguments) + ) { + return true; + } + + return ( + Array.isArray(choiceDelta.tool_calls) && + choiceDelta.tool_calls.some( + (toolCall) => + isObject(toolCall) && + isObject(toolCall.function) && + stringContent(toolCall.function.arguments), + ) + ); + })) + ) { + return true; + } + + return false; +} + +function isAISDKContentStreamChunk(chunk: unknown): boolean { + const part = extractAISDKStreamPart(chunk); + if (typeof part === "string") { + return part.length > 0; + } + + if (!isObject(part) || typeof part.type !== "string") { + return false; + } + + // TTFT should ignore framing like stream-start, metadata, and text/tool starts. + switch (part.type) { + case "text-delta": + return ( + stringContent(part.textDelta) !== undefined || + stringContent(part.delta) !== undefined || + stringContent(part.text) !== undefined || + stringContent(part.content) !== undefined + ); + case "reasoning-delta": + return ( + stringContent(part.delta) !== undefined || + stringContent(part.text) !== undefined || + stringContent(part.content) !== undefined + ); + case "tool-call": + case "object": + case "file": + return true; + case "tool-input-delta": + case "tool-call-delta": + return ( + stringContent(part.argsTextDelta) !== undefined || + stringContent(part.inputTextDelta) !== undefined || + stringContent(part.delta) !== undefined || + stringContent(part.text) !== undefined || + stringContent(part.content) !== undefined + ); + case "raw": + return rawValueHasAISDKContent(part.rawValue); + default: + return false; + } +} + +function isAISDKContentAsyncIterableChunk(chunk: unknown): boolean { + if (isAISDKContentStreamChunk(chunk)) { + return true; + } + + const part = extractAISDKStreamPart(chunk); + return isObject(part) && typeof part.type !== "string"; +} + function patchAISDKStreamingResult(args: { defaultDenyOutputPaths: string[]; endEvent: { denyOutputPaths?: string[]; [key: string]: unknown }; @@ -1497,7 +1633,10 @@ function patchAISDKStreamingResult(args: { const wrappedBaseStream = resultRecord.baseStream.pipeThrough( new TransformStream({ transform(chunk, controller) { - if (firstChunkTime === undefined) { + if ( + firstChunkTime === undefined && + isAISDKContentStreamChunk(chunk) + ) { firstChunkTime = getCurrentUnixTimestamp(); } controller.enqueue(chunk); @@ -1551,8 +1690,11 @@ function patchAISDKStreamingResult(args: { let firstChunkTime: number | undefined; const wrappedStream = createPatchedAsyncIterable(streamField.stream, { - onChunk: () => { - if (firstChunkTime === undefined) { + onChunk: (chunk) => { + if ( + firstChunkTime === undefined && + isAISDKContentAsyncIterableChunk(chunk) + ) { firstChunkTime = getCurrentUnixTimestamp(); } },