73 lines
2.0 KiB
TypeScript
73 lines
2.0 KiB
TypeScript
import {
|
|
AIStreamCallbacksAndOptions,
|
|
createCallbacksTransformer,
|
|
readableFromAsyncIterable,
|
|
} from './ai-stream';
|
|
import { createStreamDataTransformer } from './stream-data';
|
|
|
|
interface AWSBedrockResponse {
|
|
body?: AsyncIterable<{
|
|
chunk?: { bytes?: Uint8Array };
|
|
}>;
|
|
}
|
|
|
|
async function* asDeltaIterable(
|
|
response: AWSBedrockResponse,
|
|
extractTextDeltaFromChunk: (chunk: any) => string,
|
|
) {
|
|
const decoder = new TextDecoder();
|
|
for await (const chunk of response.body ?? []) {
|
|
const bytes = chunk.chunk?.bytes;
|
|
|
|
if (bytes != null) {
|
|
const chunkText = decoder.decode(bytes);
|
|
const chunkJSON = JSON.parse(chunkText);
|
|
const delta = extractTextDeltaFromChunk(chunkJSON);
|
|
|
|
if (delta != null) {
|
|
yield delta;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
export function AWSBedrockAnthropicMessagesStream(
|
|
response: AWSBedrockResponse,
|
|
callbacks?: AIStreamCallbacksAndOptions,
|
|
): ReadableStream {
|
|
return AWSBedrockStream(response, callbacks, chunk => chunk.delta?.text);
|
|
}
|
|
|
|
export function AWSBedrockAnthropicStream(
|
|
response: AWSBedrockResponse,
|
|
callbacks?: AIStreamCallbacksAndOptions,
|
|
): ReadableStream {
|
|
return AWSBedrockStream(response, callbacks, chunk => chunk.completion);
|
|
}
|
|
|
|
export function AWSBedrockCohereStream(
|
|
response: AWSBedrockResponse,
|
|
callbacks?: AIStreamCallbacksAndOptions,
|
|
): ReadableStream {
|
|
return AWSBedrockStream(response, callbacks, chunk => chunk?.text);
|
|
}
|
|
|
|
export function AWSBedrockLlama2Stream(
|
|
response: AWSBedrockResponse,
|
|
callbacks?: AIStreamCallbacksAndOptions,
|
|
): ReadableStream {
|
|
return AWSBedrockStream(response, callbacks, chunk => chunk.generation);
|
|
}
|
|
|
|
export function AWSBedrockStream(
|
|
response: AWSBedrockResponse,
|
|
callbacks: AIStreamCallbacksAndOptions | undefined,
|
|
extractTextDeltaFromChunk: (chunk: any) => string,
|
|
) {
|
|
return readableFromAsyncIterable(
|
|
asDeltaIterable(response, extractTextDeltaFromChunk),
|
|
)
|
|
.pipeThrough(createCallbacksTransformer(callbacks))
|
|
.pipeThrough(createStreamDataTransformer());
|
|
}
|