|
| 1 | +import { SPAN_STATUS_ERROR } from '../../tracing'; |
| 2 | +import type { Span } from '../../types/span'; |
| 3 | +import { endStreamSpan, type StreamResponseState } from '../ai/utils'; |
| 4 | +import type { WorkersAiUsage } from './types'; |
| 5 | +import { setOutputMessagesAttribute } from './utils'; |
| 6 | + |
| 7 | +interface WorkersAiStreamingToolCall { |
| 8 | + index?: number; |
| 9 | + id?: string; |
| 10 | + type?: string; |
| 11 | + function?: { name?: string; arguments?: string }; |
| 12 | + // Some Workers AI models stream tool calls with the name/arguments at the top |
| 13 | + // level of the tool-call object instead of nested under `function`. |
| 14 | + name?: string; |
| 15 | + arguments?: string; |
| 16 | +} |
| 17 | + |
| 18 | +interface WorkersAiStreamChunk { |
| 19 | + // Native Workers AI streaming shape (`env.AI.run` with `stream: true`). |
| 20 | + response?: unknown; |
| 21 | + tool_calls?: unknown[]; |
| 22 | + // OpenAI-compatible streaming shape emitted for models routed through the |
| 23 | + // OpenAI-compatible endpoint (e.g. via `workers-ai-provider`). |
| 24 | + choices?: Array<{ |
| 25 | + delta?: { content?: unknown; tool_calls?: WorkersAiStreamingToolCall[] }; |
| 26 | + finish_reason?: unknown; |
| 27 | + }>; |
| 28 | + usage?: WorkersAiUsage & { prompt_tokens?: number; completion_tokens?: number; total_tokens?: number }; |
| 29 | +} |
| 30 | + |
| 31 | +/** |
| 32 | + * Accumulate a fragmented OpenAI-compatible tool call (delivered across multiple |
| 33 | + * `choices[].delta.tool_calls` chunks) into the index-keyed accumulator. |
| 34 | + */ |
| 35 | +function accumulateStreamingToolCalls( |
| 36 | + toolCalls: WorkersAiStreamingToolCall[], |
| 37 | + accumulator: Record<number, WorkersAiStreamingToolCall>, |
| 38 | +): void { |
| 39 | + for (const toolCall of toolCalls) { |
| 40 | + // Normalize both shapes: name/arguments nested under `function`, or at the top level. |
| 41 | + const name = toolCall.function?.name ?? toolCall.name; |
| 42 | + const args = toolCall.function?.arguments ?? toolCall.arguments; |
| 43 | + |
| 44 | + // A tool call must carry at least a name or argument fragment to be meaningful. |
| 45 | + if (name == null && args == null) { |
| 46 | + continue; |
| 47 | + } |
| 48 | + |
| 49 | + const index = toolCall.index ?? 0; |
| 50 | + const existing = accumulator[index]; |
| 51 | + |
| 52 | + if (!existing) { |
| 53 | + accumulator[index] = { |
| 54 | + index, |
| 55 | + id: toolCall.id, |
| 56 | + type: toolCall.type, |
| 57 | + function: { |
| 58 | + name, |
| 59 | + arguments: args ?? '', |
| 60 | + }, |
| 61 | + }; |
| 62 | + } else if (existing.function) { |
| 63 | + if (name && !existing.function.name) { |
| 64 | + existing.function.name = name; |
| 65 | + } |
| 66 | + if (args) { |
| 67 | + existing.function.arguments = `${existing.function.arguments ?? ''}${args}`; |
| 68 | + } |
| 69 | + } |
| 70 | + } |
| 71 | +} |
| 72 | + |
| 73 | +/** |
| 74 | + * Parse a single SSE line (`data: {...}`) and accumulate its data into the streaming state. |
| 75 | + * |
| 76 | + * Handles both the native Workers AI shape (top-level `response`/`tool_calls`) and the |
| 77 | + * OpenAI-compatible shape (`choices[].delta.content`/`choices[].delta.tool_calls`), because |
| 78 | + * the same `run()` call transparently yields either format depending on the model. |
| 79 | + */ |
| 80 | +function processLine( |
| 81 | + line: string, |
| 82 | + state: StreamResponseState, |
| 83 | + recordOutputs: boolean, |
| 84 | + toolCallAccumulator: Record<number, WorkersAiStreamingToolCall>, |
| 85 | +): void { |
| 86 | + const trimmed = line.trim(); |
| 87 | + if (!trimmed.startsWith('data:')) { |
| 88 | + return; |
| 89 | + } |
| 90 | + |
| 91 | + const data = trimmed.slice('data:'.length).trim(); |
| 92 | + if (!data || data === '[DONE]') { |
| 93 | + return; |
| 94 | + } |
| 95 | + |
| 96 | + let parsed: WorkersAiStreamChunk; |
| 97 | + try { |
| 98 | + parsed = JSON.parse(data) as WorkersAiStreamChunk; |
| 99 | + } catch { |
| 100 | + return; |
| 101 | + } |
| 102 | + |
| 103 | + if (parsed.usage) { |
| 104 | + if (typeof parsed.usage.prompt_tokens === 'number') { |
| 105 | + state.promptTokens = parsed.usage.prompt_tokens; |
| 106 | + } |
| 107 | + if (typeof parsed.usage.completion_tokens === 'number') { |
| 108 | + state.completionTokens = parsed.usage.completion_tokens; |
| 109 | + } |
| 110 | + if (typeof parsed.usage.total_tokens === 'number') { |
| 111 | + state.totalTokens = parsed.usage.total_tokens; |
| 112 | + } |
| 113 | + } |
| 114 | + |
| 115 | + if (recordOutputs && typeof parsed.response === 'string') { |
| 116 | + state.responseTexts.push(parsed.response); |
| 117 | + } |
| 118 | + |
| 119 | + if (recordOutputs && Array.isArray(parsed.tool_calls) && parsed.tool_calls.length > 0) { |
| 120 | + state.toolCalls.push(...parsed.tool_calls); |
| 121 | + } |
| 122 | + |
| 123 | + if (Array.isArray(parsed.choices)) { |
| 124 | + for (const choice of parsed.choices) { |
| 125 | + if (recordOutputs && typeof choice.delta?.content === 'string' && choice.delta.content) { |
| 126 | + state.responseTexts.push(choice.delta.content); |
| 127 | + } |
| 128 | + if (recordOutputs && Array.isArray(choice.delta?.tool_calls)) { |
| 129 | + accumulateStreamingToolCalls(choice.delta.tool_calls, toolCallAccumulator); |
| 130 | + } |
| 131 | + if (typeof choice.finish_reason === 'string') { |
| 132 | + state.finishReasons.push(choice.finish_reason); |
| 133 | + } |
| 134 | + } |
| 135 | + } |
| 136 | +} |
| 137 | + |
| 138 | +/** |
| 139 | + * Wrap a Workers AI streaming response (a server-sent-events `ReadableStream`) so we can |
| 140 | + * accumulate the response text and token usage while passing the original bytes through untouched. |
| 141 | + * |
| 142 | + * The span is ended once the consumer finishes reading (or cancels) the stream. |
| 143 | + */ |
| 144 | +export function instrumentWorkersAiStream( |
| 145 | + stream: ReadableStream<Uint8Array>, |
| 146 | + span: Span, |
| 147 | + recordOutputs: boolean, |
| 148 | +): ReadableStream<Uint8Array> { |
| 149 | + const reader = stream.getReader(); |
| 150 | + const decoder = new TextDecoder(); |
| 151 | + |
| 152 | + const state: StreamResponseState = { |
| 153 | + responseId: '', |
| 154 | + responseModel: '', |
| 155 | + finishReasons: [], |
| 156 | + responseTexts: [], |
| 157 | + toolCalls: [], |
| 158 | + promptTokens: undefined, |
| 159 | + completionTokens: undefined, |
| 160 | + totalTokens: undefined, |
| 161 | + }; |
| 162 | + |
| 163 | + // OpenAI-compatible tool calls arrive fragmented across chunks and are keyed by index; |
| 164 | + // accumulate them here and flatten into `state.toolCalls` once the stream ends. |
| 165 | + const toolCallAccumulator: Record<number, WorkersAiStreamingToolCall> = {}; |
| 166 | + |
| 167 | + let buffer = ''; |
| 168 | + let spanEnded = false; |
| 169 | + |
| 170 | + const finish = (): void => { |
| 171 | + if (spanEnded) { |
| 172 | + return; |
| 173 | + } |
| 174 | + spanEnded = true; |
| 175 | + |
| 176 | + if (recordOutputs) { |
| 177 | + const accumulatedToolCalls = Object.values(toolCallAccumulator); |
| 178 | + if (accumulatedToolCalls.length > 0) { |
| 179 | + state.toolCalls.push(...accumulatedToolCalls); |
| 180 | + } |
| 181 | + |
| 182 | + // Set the authoritative `gen_ai.output.messages` alongside the deprecated response |
| 183 | + // attributes `endStreamSpan` writes, so tool calls survive Relay's lossy migration. |
| 184 | + setOutputMessagesAttribute(span, { |
| 185 | + responseText: state.responseTexts.join(''), |
| 186 | + toolCalls: state.toolCalls, |
| 187 | + }); |
| 188 | + } |
| 189 | + |
| 190 | + endStreamSpan(span, state, recordOutputs); |
| 191 | + }; |
| 192 | + |
| 193 | + const flushBuffer = (isDone: boolean): void => { |
| 194 | + const lines = buffer.split('\n'); |
| 195 | + // Keep the last (potentially incomplete) line in the buffer unless the stream is done. |
| 196 | + buffer = isDone ? '' : (lines.pop() ?? ''); |
| 197 | + for (const line of lines) { |
| 198 | + processLine(line, state, recordOutputs, toolCallAccumulator); |
| 199 | + } |
| 200 | + }; |
| 201 | + |
| 202 | + return new ReadableStream<Uint8Array>({ |
| 203 | + async pull(controller) { |
| 204 | + try { |
| 205 | + const { done, value } = await reader.read(); |
| 206 | + |
| 207 | + if (done) { |
| 208 | + buffer += decoder.decode(); |
| 209 | + flushBuffer(true); |
| 210 | + finish(); |
| 211 | + controller.close(); |
| 212 | + return; |
| 213 | + } |
| 214 | + |
| 215 | + buffer += decoder.decode(value, { stream: true }); |
| 216 | + flushBuffer(false); |
| 217 | + controller.enqueue(value); |
| 218 | + } catch (error) { |
| 219 | + span.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' }); |
| 220 | + finish(); |
| 221 | + controller.error(error); |
| 222 | + } |
| 223 | + }, |
| 224 | + async cancel(reason) { |
| 225 | + finish(); |
| 226 | + await reader.cancel(reason); |
| 227 | + }, |
| 228 | + }); |
| 229 | +} |
0 commit comments