| const TAIL_LIMIT = 4 * 1024 |
| const encoder = new TextEncoder() |
|
|
| export async function prepareRequestBody(body: ReadableStream<Uint8Array>) { |
| const reader = body.getReader() |
| const chunks: Uint8Array[] = [] |
| const decoder = new TextDecoder() |
| let text = "" |
| let done = false |
| let bom = 0 |
| let index = 0 |
| let depth = 0 |
| let stringStart = -1 |
| let escaped = false |
| let phase: "key" | "colon" | "value" | "comma" = "key" |
| let key = "" |
| let found: { model: string; start: number; end: number } | undefined |
|
|
| const scan = () => { |
| while (index < text.length && !found) { |
| const char = text[index] |
| if (stringStart >= 0) { |
| if (escaped) escaped = false |
| else if (char === "\\") escaped = true |
| else if (char === '"') { |
| if (depth === 1 && phase === "key") { |
| key = JSON.parse(text.slice(stringStart, index + 1)) |
| phase = "colon" |
| } else if (depth === 1 && phase === "value") { |
| if (key === "model") { |
| const start = bom + utf8Length(text, stringStart + 1) |
| found = { |
| model: JSON.parse(text.slice(stringStart, index + 1)), |
| start, |
| end: bom + utf8Length(text, index), |
| } |
| } |
| phase = "comma" |
| } |
| stringStart = -1 |
| } |
| index++ |
| continue |
| } |
|
|
| if (char === '"') { |
| stringStart = index++ |
| continue |
| } |
| if (char === "{" || char === "[") { |
| if (depth === 1 && phase === "value") phase = "comma" |
| depth++ |
| index++ |
| continue |
| } |
| if (char === "}" || char === "]") { |
| depth-- |
| index++ |
| continue |
| } |
| if (depth !== 1) { |
| index++ |
| continue |
| } |
| if (char === ":" && phase === "colon") phase = "value" |
| else if (char === "," && phase === "comma") phase = "key" |
| else if (phase === "value" && !/\s/.test(char)) phase = "comma" |
| index++ |
| } |
| } |
|
|
| while (!done && !found) { |
| const next = await reader.read() |
| done = next.done |
| if (!next.value) continue |
| if (!chunks.length && next.value[0] === 0xef && next.value[1] === 0xbb && next.value[2] === 0xbf) bom = 3 |
| chunks.push(next.value) |
| text += decoder.decode(next.value, { stream: true }) |
| scan() |
| } |
| if (done) { |
| text += decoder.decode() |
| scan() |
| } |
|
|
| const preview = text.substring(0, 300) |
| text = "" |
| let used = false |
|
|
| return { |
| model: found?.model ?? "", |
| preview, |
| cancel: () => reader.cancel(), |
| stream(providerModel: string, includeUsage: boolean) { |
| if (used) throw new Error("Request body stream already consumed") |
| if (!found) throw new Error("Missing model field") |
| used = true |
|
|
| const initial = replace(chunks, found.start, found.end, providerModel) |
| chunks.length = 0 |
| const output = passthrough(initial, reader, done) |
| if (!includeUsage) return output |
| return appendUsage(output) |
| }, |
| } |
| } |
|
|
| function utf8Length(value: string, end: number) { |
| let length = 0 |
| for (let i = 0; i < end; i++) { |
| const code = value.charCodeAt(i) |
| if (code <= 0x7f) length++ |
| else if (code <= 0x7ff) length += 2 |
| else if (code >= 0xd800 && code <= 0xdbff && i + 1 < end && value.charCodeAt(i + 1) >= 0xdc00) { |
| length += 4 |
| i++ |
| } else length += 3 |
| } |
| return length |
| } |
|
|
| function replace(chunks: Uint8Array[], start: number, end: number, value: string) { |
| let offset = 0 |
| let inserted = false |
| return chunks.flatMap((chunk) => { |
| const chunkStart = offset |
| const chunkEnd = offset + chunk.length |
| offset = chunkEnd |
| if (chunkEnd <= start || chunkStart >= end) return [chunk] |
|
|
| const parts = [chunk.subarray(0, Math.max(0, start - chunkStart))] |
| if (!inserted) { |
| parts.push(encoder.encode(value)) |
| inserted = true |
| } |
| parts.push(chunk.subarray(Math.min(chunk.length, end - chunkStart))) |
| return parts.filter((part) => part.length) |
| }) |
| } |
|
|
| function passthrough( |
| initial: Array<Uint8Array | undefined>, |
| reader: ReadableStreamDefaultReader<Uint8Array>, |
| sourceDone: boolean, |
| ) { |
| let done = sourceDone |
| let index = 0 |
| return new ReadableStream<Uint8Array>({ |
| async pull(controller) { |
| const chunk = initial[index] |
| if (chunk) { |
| initial[index++] = undefined |
| controller.enqueue(chunk) |
| return |
| } |
| initial.length = 0 |
| if (done) { |
| controller.close() |
| return |
| } |
| const next = await reader.read() |
| done = next.done |
| if (next.value) controller.enqueue(next.value) |
| if (done) controller.close() |
| }, |
| cancel(reason) { |
| initial.length = 0 |
| return reader.cancel(reason) |
| }, |
| }) |
| } |
|
|
| function appendUsage(body: ReadableStream<Uint8Array>) { |
| const reader = body.getReader() |
| const decoder = new TextDecoder() |
| let tail = new Uint8Array() |
| let streamText = "" |
| let isStream = false |
| const inspect = (chunk?: Uint8Array) => { |
| streamText += chunk ? decoder.decode(chunk, { stream: true }) : decoder.decode() |
| for (const match of streamText.matchAll(/"stream"\s*:\s*(true|false)/g)) isStream = match[1] === "true" |
| streamText = streamText.slice(-64) |
| } |
| return new ReadableStream<Uint8Array>({ |
| async pull(controller) { |
| while (true) { |
| const next = await reader.read() |
| if (next.done) { |
| inspect() |
| if (!isStream) { |
| if (tail.length) controller.enqueue(tail) |
| controller.close() |
| return |
| } |
| const close = tail.lastIndexOf(125) |
| if (close < 0) { |
| controller.error(new Error("Invalid JSON request body")) |
| return |
| } |
| if (close) controller.enqueue(tail.subarray(0, close)) |
| controller.enqueue(encoder.encode(',"stream_options":{"include_usage":true}}')) |
| if (close + 1 < tail.length) controller.enqueue(tail.subarray(close + 1)) |
| controller.close() |
| return |
| } |
|
|
| const chunk = next.value |
| inspect(chunk) |
| if (tail.length + chunk.length <= TAIL_LIMIT) { |
| const combined = new Uint8Array(tail.length + chunk.length) |
| combined.set(tail) |
| combined.set(chunk, tail.length) |
| tail = combined |
| continue |
| } |
|
|
| const emit = tail.length + chunk.length - TAIL_LIMIT |
| if (emit <= tail.length) { |
| controller.enqueue(tail.subarray(0, emit)) |
| const combined = new Uint8Array(TAIL_LIMIT) |
| combined.set(tail.subarray(emit)) |
| combined.set(chunk, tail.length - emit) |
| tail = combined |
| return |
| } |
|
|
| if (tail.length) controller.enqueue(tail) |
| controller.enqueue(chunk.subarray(0, emit - tail.length)) |
| tail = chunk.slice(emit - tail.length) |
| return |
| } |
| }, |
| cancel(reason) { |
| return reader.cancel(reason) |
| }, |
| }) |
| } |
|
|
