Spaces:
Sleeping
Sleeping
File size: 2,136 Bytes
ad02f33 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 |
/**
* Calls the given Gemini model with the given image and/or text
* parts, streaming output (as a generator function).
*/
export async function* streamGemini({
model = 'gemini-pro', // use 'gemini-pro' for text -> text
contents = [],
} = {}) {
// Send the prompt to the Python backend
// Call API defined in main.py
let response = await fetch("/api/generate", {
method: "POST",
headers: { "content-type": "application/json" },
body: JSON.stringify({ model, contents })
});
yield* streamResponseChunks(response);
}
/**
* A helper that streams text output chunks from a fetch() response.
*/
async function* streamResponseChunks(response) {
let buffer = '';
const CHUNK_SEPARATOR = '\n\n';
let processBuffer = async function* (streamDone = false) {
while (true) {
let flush = false;
let chunkSeparatorIndex = buffer.indexOf(CHUNK_SEPARATOR);
if (streamDone && chunkSeparatorIndex < 0) {
flush = true;
chunkSeparatorIndex = buffer.length;
}
if (chunkSeparatorIndex < 0) {
break;
}
let chunk = buffer.substring(0, chunkSeparatorIndex);
buffer = buffer.substring(chunkSeparatorIndex + CHUNK_SEPARATOR.length);
chunk = chunk.replace(/^data:\s*/, '').trim();
if (!chunk) {
if (flush) break;
continue;
}
let { error, text } = JSON.parse(chunk);
if (error) {
console.error(error);
throw new Error(error?.message || JSON.stringify(error));
}
yield text;
if (flush) break;
}
};
const reader = response.body.getReader();
try {
while (true) {
const { done, value } = await reader.read()
if (done) break;
buffer += new TextDecoder().decode(value);
console.log(new TextDecoder().decode(value));
yield* processBuffer();
}
} finally {
reader.releaseLock();
}
yield* processBuffer(true);
}
|