@genkit-ai/core
Version:
Genkit AI framework core libraries.
1 lines • 8.86 kB
Source Map (JSON)
{"version":3,"sources":["../src/streaming.ts"],"sourcesContent":["/**\n * Copyright 2024 Google LLC\n *\n * Licensed under the Apache License, Version 2.0 (the \"License\");\n * you may not use this file except in compliance with the License.\n * You may obtain a copy of the License at\n *\n * http://www.apache.org/licenses/LICENSE-2.0\n *\n * Unless required by applicable law or agreed to in writing, software\n * distributed under the License is distributed on an \"AS IS\" BASIS,\n * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n * See the License for the specific language governing permissions and\n * limitations under the License.\n */\n\nimport { GenkitError } from './error';\n\n/**\n * Error thrown when a stream cannot be found.\n */\nexport class StreamNotFoundError extends GenkitError {\n constructor(message: string) {\n super({ status: 'NOT_FOUND', message });\n this.name = 'StreamNotFoundError';\n }\n}\n\n/**\n * Interface for writing content to a stream.\n * @template S The type of the stream chunks.\n * @template O The type of the final output.\n */\nexport interface ActionStreamInput<S, O> {\n /**\n * Writes a chunk to the stream.\n * @param chunk The chunk data to write.\n */\n write(chunk: S): Promise<void>;\n /**\n * Closes the stream with a final output.\n * @param output The final output data.\n */\n done(output: O): Promise<void>;\n /**\n * Closes the stream with an error.\n * @param err The error that occurred.\n */\n error(err: any): Promise<void>;\n}\n\n/**\n * Subscriber callbacks for receiving stream events.\n * @template S The type of the stream chunks.\n * @template O The type of the final output.\n */\nexport type ActionStreamSubscriber<S, O> = {\n /** Called when a chunk is received. */\n onChunk: (chunk: S) => void;\n /** Called when the stream completes successfully. */\n onDone: (output: O) => void;\n /** Called when the stream encounters an error. */\n onError: (error: any) => void;\n};\n\n/**\n * Interface for managing streaming actions, allowing creation and subscription to streams.\n * Implementations can provide different storage backends (e.g., in-memory, database, cache).\n */\nexport interface StreamManager {\n /**\n * Opens a stream for writing.\n * @param streamId The unique identifier for the stream.\n * @returns An object to write to the stream.\n */\n open<S, O>(streamId: string): Promise<ActionStreamInput<S, O>>;\n /**\n * Subscribes to a stream to receive its events.\n * @param streamId The unique identifier for the stream.\n * @param options The subscriber callbacks.\n * @returns A promise resolving to an object containing an unsubscribe function.\n */\n subscribe<S, O>(\n streamId: string,\n options: ActionStreamSubscriber<S, O>\n ): Promise<{ unsubscribe: () => void }>;\n}\n\ntype StreamState<S, O> =\n | {\n status: 'open';\n chunks: S[];\n subscribers: ActionStreamSubscriber<S, O>[];\n lastTouched: number;\n }\n | { status: 'done'; chunks: S[]; output: O; lastTouched: number }\n | { status: 'error'; chunks: S[]; error: any; lastTouched: number };\n\n/**\n * An in-memory implementation of StreamManager.\n * Useful for testing or single-instance deployments where persistence is not required.\n */\nexport class InMemoryStreamManager implements StreamManager {\n private streams: Map<string, StreamState<any, any>> = new Map();\n\n /**\n * @param options Configuration options.\n * @param options.ttlSeconds Time-to-live for streams in seconds. Defaults to 5 minutes.\n */\n constructor(private options: { ttlSeconds?: number } = {}) {}\n\n private _cleanup() {\n const ttl = (this.options.ttlSeconds ?? 5 * 60) * 1000;\n const now = Date.now();\n for (const [streamId, stream] of this.streams.entries()) {\n if (stream.status !== 'open' && now - stream.lastTouched > ttl) {\n this.streams.delete(streamId);\n }\n }\n }\n\n async open<S, O>(streamId: string): Promise<ActionStreamInput<S, O>> {\n this._cleanup();\n if (this.streams.has(streamId)) {\n throw new Error(`Stream with id ${streamId} already exists.`);\n }\n this.streams.set(streamId, {\n status: 'open',\n chunks: [],\n subscribers: [],\n lastTouched: Date.now(),\n });\n\n return {\n write: async (chunk: S) => {\n const stream = this.streams.get(streamId);\n if (stream?.status === 'open') {\n stream.chunks.push(chunk);\n stream.subscribers.forEach((s) => s.onChunk(chunk));\n stream.lastTouched = Date.now();\n }\n },\n done: async (output: O) => {\n const stream = this.streams.get(streamId);\n if (stream?.status === 'open') {\n this.streams.set(streamId, {\n status: 'done',\n chunks: stream.chunks,\n output,\n lastTouched: Date.now(),\n });\n stream.subscribers.forEach((s) => s.onDone(output));\n }\n },\n error: async (err: any) => {\n const stream = this.streams.get(streamId);\n if (stream?.status === 'open') {\n stream.subscribers.forEach((s) => s.onError(err));\n this.streams.set(streamId, {\n status: 'error',\n chunks: stream.chunks,\n error: err,\n lastTouched: Date.now(),\n });\n }\n },\n };\n }\n\n async subscribe<S, O>(\n streamId: string,\n subscriber: ActionStreamSubscriber<S, O>\n ): Promise<{ unsubscribe: () => void }> {\n const stream = this.streams.get(streamId);\n if (!stream) {\n throw new StreamNotFoundError(`Stream with id ${streamId} not found.`);\n }\n\n if (stream.status === 'done') {\n for (const chunk of stream.chunks) {\n subscriber.onChunk(chunk);\n }\n subscriber.onDone(stream.output);\n } else if (stream.status === 'error') {\n for (const chunk of stream.chunks) {\n subscriber.onChunk(chunk);\n }\n subscriber.onError(stream.error);\n } else {\n stream.chunks.forEach((chunk) => subscriber.onChunk(chunk));\n stream.subscribers.push(subscriber);\n }\n\n return {\n unsubscribe: () => {\n const currentStream = this.streams.get(streamId);\n if (currentStream?.status === 'open') {\n const index = currentStream.subscribers.indexOf(subscriber);\n if (index > -1) {\n currentStream.subscribers.splice(index, 1);\n }\n }\n },\n };\n }\n}\n"],"mappings":"AAgBA,SAAS,mBAAmB;AAKrB,MAAM,4BAA4B,YAAY;AAAA,EACnD,YAAY,SAAiB;AAC3B,UAAM,EAAE,QAAQ,aAAa,QAAQ,CAAC;AACtC,SAAK,OAAO;AAAA,EACd;AACF;AA4EO,MAAM,sBAA+C;AAAA;AAAA;AAAA;AAAA;AAAA,EAO1D,YAAoB,UAAmC,CAAC,GAAG;AAAvC;AAAA,EAAwC;AAAA,EANpD,UAA8C,oBAAI,IAAI;AAAA,EAQtD,WAAW;AACjB,UAAM,OAAO,KAAK,QAAQ,cAAc,IAAI,MAAM;AAClD,UAAM,MAAM,KAAK,IAAI;AACrB,eAAW,CAAC,UAAU,MAAM,KAAK,KAAK,QAAQ,QAAQ,GAAG;AACvD,UAAI,OAAO,WAAW,UAAU,MAAM,OAAO,cAAc,KAAK;AAC9D,aAAK,QAAQ,OAAO,QAAQ;AAAA,MAC9B;AAAA,IACF;AAAA,EACF;AAAA,EAEA,MAAM,KAAW,UAAoD;AACnE,SAAK,SAAS;AACd,QAAI,KAAK,QAAQ,IAAI,QAAQ,GAAG;AAC9B,YAAM,IAAI,MAAM,kBAAkB,QAAQ,kBAAkB;AAAA,IAC9D;AACA,SAAK,QAAQ,IAAI,UAAU;AAAA,MACzB,QAAQ;AAAA,MACR,QAAQ,CAAC;AAAA,MACT,aAAa,CAAC;AAAA,MACd,aAAa,KAAK,IAAI;AAAA,IACxB,CAAC;AAED,WAAO;AAAA,MACL,OAAO,OAAO,UAAa;AACzB,cAAM,SAAS,KAAK,QAAQ,IAAI,QAAQ;AACxC,YAAI,QAAQ,WAAW,QAAQ;AAC7B,iBAAO,OAAO,KAAK,KAAK;AACxB,iBAAO,YAAY,QAAQ,CAAC,MAAM,EAAE,QAAQ,KAAK,CAAC;AAClD,iBAAO,cAAc,KAAK,IAAI;AAAA,QAChC;AAAA,MACF;AAAA,MACA,MAAM,OAAO,WAAc;AACzB,cAAM,SAAS,KAAK,QAAQ,IAAI,QAAQ;AACxC,YAAI,QAAQ,WAAW,QAAQ;AAC7B,eAAK,QAAQ,IAAI,UAAU;AAAA,YACzB,QAAQ;AAAA,YACR,QAAQ,OAAO;AAAA,YACf;AAAA,YACA,aAAa,KAAK,IAAI;AAAA,UACxB,CAAC;AACD,iBAAO,YAAY,QAAQ,CAAC,MAAM,EAAE,OAAO,MAAM,CAAC;AAAA,QACpD;AAAA,MACF;AAAA,MACA,OAAO,OAAO,QAAa;AACzB,cAAM,SAAS,KAAK,QAAQ,IAAI,QAAQ;AACxC,YAAI,QAAQ,WAAW,QAAQ;AAC7B,iBAAO,YAAY,QAAQ,CAAC,MAAM,EAAE,QAAQ,GAAG,CAAC;AAChD,eAAK,QAAQ,IAAI,UAAU;AAAA,YACzB,QAAQ;AAAA,YACR,QAAQ,OAAO;AAAA,YACf,OAAO;AAAA,YACP,aAAa,KAAK,IAAI;AAAA,UACxB,CAAC;AAAA,QACH;AAAA,MACF;AAAA,IACF;AAAA,EACF;AAAA,EAEA,MAAM,UACJ,UACA,YACsC;AACtC,UAAM,SAAS,KAAK,QAAQ,IAAI,QAAQ;AACxC,QAAI,CAAC,QAAQ;AACX,YAAM,IAAI,oBAAoB,kBAAkB,QAAQ,aAAa;AAAA,IACvE;AAEA,QAAI,OAAO,WAAW,QAAQ;AAC5B,iBAAW,SAAS,OAAO,QAAQ;AACjC,mBAAW,QAAQ,KAAK;AAAA,MAC1B;AACA,iBAAW,OAAO,OAAO,MAAM;AAAA,IACjC,WAAW,OAAO,WAAW,SAAS;AACpC,iBAAW,SAAS,OAAO,QAAQ;AACjC,mBAAW,QAAQ,KAAK;AAAA,MAC1B;AACA,iBAAW,QAAQ,OAAO,KAAK;AAAA,IACjC,OAAO;AACL,aAAO,OAAO,QAAQ,CAAC,UAAU,WAAW,QAAQ,KAAK,CAAC;AAC1D,aAAO,YAAY,KAAK,UAAU;AAAA,IACpC;AAEA,WAAO;AAAA,MACL,aAAa,MAAM;AACjB,cAAM,gBAAgB,KAAK,QAAQ,IAAI,QAAQ;AAC/C,YAAI,eAAe,WAAW,QAAQ;AACpC,gBAAM,QAAQ,cAAc,YAAY,QAAQ,UAAU;AAC1D,cAAI,QAAQ,IAAI;AACd,0BAAc,YAAY,OAAO,OAAO,CAAC;AAAA,UAC3C;AAAA,QACF;AAAA,MACF;AAAA,IACF;AAAA,EACF;AACF;","names":[]}