From 13f8297b06f9fa6801f15008ffafff74dbc7aa1a Mon Sep 17 00:00:00 2001 From: Eric Allam Date: Wed, 11 Dec 2024 22:22:04 +0000 Subject: [PATCH] Improve the way we stream from tasks to the server --- apps/webapp/server.ts | 13 ++++ packages/core/src/v3/runMetadata/manager.ts | 3 +- .../core/src/v3/runMetadata/metadataStream.ts | 76 +++++++++---------- references/nextjs-realtime/package.json | 3 +- .../components/TriggerButtonWithStreaming.tsx | 1 - references/nextjs-realtime/src/trigger/ai.ts | 1 - 6 files changed, 53 insertions(+), 44 deletions(-) diff --git a/apps/webapp/server.ts b/apps/webapp/server.ts index aec725c80..26e30343a 100644 --- a/apps/webapp/server.ts +++ b/apps/webapp/server.ts @@ -81,6 +81,19 @@ if (process.env.HTTP_SERVER_DISABLED !== "true") { }); if (process.env.DASHBOARD_AND_API_DISABLED !== "true") { + if (process.env.ALLOW_ONLY_REALTIME_API === "true") { + // Block all requests that do not start with /realtime + app.use((req, res, next) => { + // Make sure /healthcheck is still accessible + if (!req.url.startsWith("/realtime") && req.url !== "/healthcheck") { + res.status(404).send("Not Found"); + return; + } + + next(); + }); + } + app.use(apiRateLimiter); app.all( diff --git a/packages/core/src/v3/runMetadata/manager.ts b/packages/core/src/v3/runMetadata/manager.ts index ea04d3692..83d8e7488 100644 --- a/packages/core/src/v3/runMetadata/manager.ts +++ b/packages/core/src/v3/runMetadata/manager.ts @@ -233,13 +233,14 @@ export class StandardMetadataManager implements RunMetadataManager { // Add the key to the special stream metadata object this.appendKey(`$$streams`, key); this.setKey("$$streamsVersion", this.streamsVersion); + this.setKey("$$streamsBaseUrl", this.streamsBaseUrl); await this.flush(); const streamInstance = new MetadataStream({ key, runId: this.runId, - iterator: $value[Symbol.asyncIterator](), + source: $value, baseUrl: this.streamsBaseUrl, headers: this.apiClient.getHeaders(), signal, diff --git a/packages/core/src/v3/runMetadata/metadataStream.ts b/packages/core/src/v3/runMetadata/metadataStream.ts index dfd9965c0..1d6143f5c 100644 --- a/packages/core/src/v3/runMetadata/metadataStream.ts +++ b/packages/core/src/v3/runMetadata/metadataStream.ts @@ -2,7 +2,7 @@ export type MetadataOptions = { baseUrl: string; runId: string; key: string; - iterator: AsyncIterator; + source: AsyncIterable; headers?: Record; signal?: AbortSignal; version?: "v1" | "v2"; @@ -10,57 +10,40 @@ export type MetadataOptions = { export class MetadataStream { private controller = new AbortController(); - private serverQueue: Array>> = []; - private consumerQueue: Array>> = []; - private serverIterator: AsyncIterator; - private consumerIterator: AsyncIterator; + private serverStream: ReadableStream; + private consumerStream: ReadableStream; private streamPromise: Promise; constructor(private options: MetadataOptions) { - const { serverIterator, consumerIterator } = this.createTeeIterators(); - this.serverIterator = serverIterator; - this.consumerIterator = consumerIterator; + const [serverStream, consumerStream] = this.createTeeStreams(); + this.serverStream = serverStream; + this.consumerStream = consumerStream; this.streamPromise = this.initializeServerStream(); } - private createTeeIterators() { - const teeIterator = (queue: Array>>): AsyncIterator => ({ - next: () => { - if (queue.length === 0) { - const result = this.options.iterator.next(); - this.serverQueue.push(result); - this.consumerQueue.push(result); + private createTeeStreams() { + const readableSource = new ReadableStream({ + start: async (controller) => { + for await (const value of this.options.source) { + controller.enqueue(value); } - return queue.shift()!; + + controller.close(); }, }); - return { - serverIterator: teeIterator(this.serverQueue), - consumerIterator: teeIterator(this.consumerQueue), - }; + return readableSource.tee(); } - private initializeServerStream(): Promise { - const serverIterator = this.serverIterator; - - const serverStream = new ReadableStream({ - async pull(controller) { - try { - const { value, done } = await serverIterator.next(); - if (done) { - controller.close(); - return; - } - - controller.enqueue(JSON.stringify(value) + "\n"); - } catch (err) { - controller.error(err); - } - }, - cancel: () => this.controller.abort(), - }); + private initializeServerStream(): Promise { + const serverStream = this.serverStream.pipeThrough( + new TransformStream({ + async transform(chunk, controller) { + controller.enqueue(JSON.stringify(chunk) + "\n"); + }, + }) + ); return fetch( `${this.options.baseUrl}/realtime/${this.options.version ?? "v1"}/streams/${ @@ -82,6 +65,19 @@ export class MetadataStream { } public [Symbol.asyncIterator]() { - return this.consumerIterator; + return streamToAsyncIterator(this.consumerStream); + } +} + +async function* streamToAsyncIterator(stream: ReadableStream): AsyncIterableIterator { + const reader = stream.getReader(); + try { + while (true) { + const { done, value } = await reader.read(); + if (done) return; + yield value; + } + } finally { + reader.releaseLock(); } } diff --git a/references/nextjs-realtime/package.json b/references/nextjs-realtime/package.json index 7a859d9dc..75df02777 100644 --- a/references/nextjs-realtime/package.json +++ b/references/nextjs-realtime/package.json @@ -7,7 +7,8 @@ "build": "next build", "start": "next start", "lint": "next lint", - "dev:trigger": "trigger dev" + "dev:trigger": "trigger dev", + "deploy": "trigger deploy" }, "dependencies": { "@ai-sdk/openai": "^1.0.1", diff --git a/references/nextjs-realtime/src/components/TriggerButtonWithStreaming.tsx b/references/nextjs-realtime/src/components/TriggerButtonWithStreaming.tsx index 07895daf9..38ee63128 100644 --- a/references/nextjs-realtime/src/components/TriggerButtonWithStreaming.tsx +++ b/references/nextjs-realtime/src/components/TriggerButtonWithStreaming.tsx @@ -23,7 +23,6 @@ export default function TriggerButton({ accessToken }: { accessToken: string }) >("openai-streaming", { accessToken, baseURL: process.env.NEXT_PUBLIC_TRIGGER_API_URL, - experimental_throttleInMs: 100, }); const openWeatherReport = useCallback(() => { diff --git a/references/nextjs-realtime/src/trigger/ai.ts b/references/nextjs-realtime/src/trigger/ai.ts index 70cd498fe..26d4c63a9 100644 --- a/references/nextjs-realtime/src/trigger/ai.ts +++ b/references/nextjs-realtime/src/trigger/ai.ts @@ -108,7 +108,6 @@ export const openaiStreaming = schemaTask({ }); const stream = await metadata.stream("openai", result.fullStream); - await metadata.stream("openaiText", result.textStream); }, });