74 lines
1.8 KiB
TypeScript
74 lines
1.8 KiB
TypeScript
import { AIStream, type AIStreamCallbacksAndOptions } from './ai-stream';
|
|
import { createStreamDataTransformer } from './stream-data';
|
|
|
|
// from replicate SDK
|
|
interface Prediction {
|
|
id: string;
|
|
status: 'starting' | 'processing' | 'succeeded' | 'failed' | 'canceled';
|
|
version: string;
|
|
input: object;
|
|
output?: any;
|
|
source: 'api' | 'web';
|
|
error?: any;
|
|
logs?: string;
|
|
metrics?: {
|
|
predict_time?: number;
|
|
};
|
|
webhook?: string;
|
|
webhook_events_filter?: ('start' | 'output' | 'logs' | 'completed')[];
|
|
created_at: string;
|
|
updated_at?: string;
|
|
completed_at?: string;
|
|
urls: {
|
|
get: string;
|
|
cancel: string;
|
|
stream?: string;
|
|
};
|
|
}
|
|
|
|
/**
|
|
* Stream predictions from Replicate.
|
|
* Only certain models are supported and you must pass `stream: true` to
|
|
* replicate.predictions.create().
|
|
* @see https://github.com/replicate/replicate-javascript#streaming
|
|
*
|
|
* @example
|
|
* const response = await replicate.predictions.create({
|
|
* stream: true,
|
|
* input: {
|
|
* prompt: messages.join('\n')
|
|
* },
|
|
* version: '2c1608e18606fad2812020dc541930f2d0495ce32eee50074220b87300bc16e1'
|
|
* })
|
|
*
|
|
* const stream = await ReplicateStream(response)
|
|
* return new StreamingTextResponse(stream)
|
|
*
|
|
*/
|
|
export async function ReplicateStream(
|
|
res: Prediction,
|
|
cb?: AIStreamCallbacksAndOptions,
|
|
options?: {
|
|
headers?: Record<string, string>;
|
|
},
|
|
): Promise<ReadableStream> {
|
|
const url = res.urls?.stream;
|
|
|
|
if (!url) {
|
|
if (res.error) throw new Error(res.error);
|
|
else throw new Error('Missing stream URL in Replicate response');
|
|
}
|
|
|
|
const eventStream = await fetch(url, {
|
|
method: 'GET',
|
|
headers: {
|
|
Accept: 'text/event-stream',
|
|
...options?.headers,
|
|
},
|
|
});
|
|
|
|
return AIStream(eventStream, undefined, cb).pipeThrough(
|
|
createStreamDataTransformer(),
|
|
);
|
|
}
|