Skip to content

Commit

Permalink
Auto-retry the LLM chat API
Browse files Browse the repository at this point in the history
  • Loading branch information
@ariya
ariya committed Nov 5, 2024
1 parent 3607a13 commit d6e4ae3
Showing 1 changed file with 108 additions and 83 deletions.
191 changes: 108 additions & 83 deletions query-llm.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,16 @@ const CROSS = '✘';
*/
const pipe = (...fns) => arg => fns.reduce((d, fn) => d.then(fn), Promise.resolve(arg));

const MAX_RETRY_ATTEMPT = 3;

/**
* Suspends the execution for a specified amount of time.
*
* @param {number} ms - The amount of time to suspend execution in milliseconds.
* @returns {Promise<void>} - A promise that resolves after the specified time has elapsed.
*/
const sleep = async (ms) => new Promise((resolve) => setTimeout(resolve, ms));

/**
* Tries to parse a string as JSON, but if that fails, tries adding a
* closing curly brace or double quote to fix the JSON.
Expand Down Expand Up @@ -82,7 +92,8 @@ const unJSON = (text) => {
* @returns {Promise<string>} The completion generated by the LLM.
*/

const chat = async (messages, schema, handler) => {
const chat = async (messages, schema, handler = null, attempt = MAX_RETRY_ATTEMPT) => {
const timeout = 17; // seconds
const gemini = LLM_API_BASE_URL.indexOf('generativelanguage.google') > 0;
const stream = LLM_STREAMING && typeof handler === 'function';
const model = LLM_CHAT_MODEL || 'gpt-4o-mini';
Expand Down Expand Up @@ -122,107 +133,121 @@ const chat = async (messages, schema, handler) => {
console.log(`${MAGENTA}${role}:${NORMAL} ${content}`);
});

const response = await fetch(url, {
method: 'POST',
headers: { 'Content-Type': 'application/json', ...auth },
body: JSON.stringify(body)
});
if (!response.ok) {
throw new Error(`HTTP error with the status: ${response.status} ${response.statusText}`);
}
try {

const extract = (data) => {
const { choices, candidates } = data;
const first = choices ? choices[0] : candidates[0];
if (first?.content || first?.message) {
const content = first?.content ? first.content : first.message.content;
const parts = content?.parts;
const answer = parts ? parts.map(part => part.text).join('') : content;
return answer;
const response = await fetch(url, {
method: 'POST',
headers: { 'Content-Type': 'application/json', ...auth },
body: JSON.stringify(body)
});
if (!response.ok) {
throw new Error(`HTTP error with the status: ${response.status} ${response.statusText}`);
}
return '';
}

if (!stream) {
const data = await response.json();
const answer = extract(data).trim();
if (LLM_DEBUG_CHAT) {
if (LLM_JSON_SCHEMA) {
const parsed = unJSON(answer);
const empty = Object.keys(parsed).length === 0;
const formatted = empty ? answer : JSON.stringify(parsed, null, 2);
console.log(`${YELLOW}${formatted}${NORMAL}`);
} else {
console.log(`${YELLOW}${answer}${NORMAL}`);
const extract = (data) => {
const { choices, candidates } = data;
const first = choices ? choices[0] : candidates[0];
if (first?.content || first?.message) {
const content = first?.content ? first.content : first.message.content;
const parts = content?.parts;
const answer = parts ? parts.map(part => part.text).join('') : content;
return answer;
}
return '';
}
(answer.length > 0) && handler && handler(answer);
return answer;
}

const parse = (line) => {
let partial = null;
const prefix = line.substring(0, 6);
if (prefix === 'data: ') {
const payload = line.substring(6);
try {
const data = JSON.parse(payload);
const { choices, candidates } = data;
if (choices) {
const [choice] = choices;
const { delta } = choice;
partial = delta?.content;
} else if (candidates) {
partial = extract(data);
if (!stream) {
const data = await response.json();
const answer = extract(data).trim();
if (LLM_DEBUG_CHAT) {
if (LLM_JSON_SCHEMA) {
const parsed = unJSON(answer);
const empty = Object.keys(parsed).length === 0;
const formatted = empty ? answer : JSON.stringify(parsed, null, 2);
console.log(`${YELLOW}${formatted}${NORMAL}`);
} else {
console.log(`${YELLOW}${answer}${NORMAL}`);
}
} catch (e) {
// ignore
} finally {
return partial;
}
(answer.length > 0) && handler && handler(answer);
return answer;
}
return partial;
}

const reader = response.body.getReader();
const decoder = new TextDecoder();

let answer = '';
let buffer = '';
while (true) {
const { value, done } = await reader.read();
if (done) {
break;
}
const lines = decoder.decode(value).split('\n');
for (let i = 0; i < lines.length; ++i) {
const line = buffer + lines[i];
if (line[0] === ':') {
buffer = '';
continue;
const parse = (line) => {
let partial = null;
const prefix = line.substring(0, 6);
if (prefix === 'data: ') {
const payload = line.substring(6);
try {
const data = JSON.parse(payload);
const { choices, candidates } = data;
if (choices) {
const [choice] = choices;
const { delta } = choice;
partial = delta?.content;
} else if (candidates) {
partial = extract(data);
}
} catch (e) {
// ignore
} finally {
return partial;
}
}
if (line === 'data: [DONE]') {
return partial;
}

const reader = response.body.getReader();
const decoder = new TextDecoder();

let answer = '';
let buffer = '';
while (true) {
const { value, done } = await reader.read();
if (done) {
break;
}
if (line.length > 0) {
const partial = parse(line.trim());
if (partial === null) {
buffer = line;
} else if (partial && partial.length > 0) {
const lines = decoder.decode(value).split('\n');
for (let i = 0; i < lines.length; ++i) {
const line = buffer + lines[i];
if (line[0] === ':') {
buffer = '';
if (answer.length < 1) {
const leading = partial.trim();
answer = leading;
handler && (leading.length > 0) && handler(leading);
} else {
answer += partial;
handler && handler(partial);
continue;
}
if (line === 'data: [DONE]') {
break;
}
if (line.length > 0) {
const partial = parse(line.trim());
if (partial === null) {
buffer = line;
} else if (partial && partial.length > 0) {
buffer = '';
if (answer.length < 1) {
const leading = partial.trim();
answer = leading;
handler && (leading.length > 0) && handler(leading);
} else {
answer += partial;
handler && handler(partial);
}
}
}
}
}
return answer;
} catch (e) {
if (e.name === 'TimeoutError') {
LLM_DEBUG_CHAT && console.log(`Timeout with LLM chat after ${timeout} seconds`);
}
if (attempt > 1 && (e.name === 'TimeoutError' || e.name === 'EvalError')) {
LLM_DEBUG_CHAT && console.log('Retrying...');
await sleep((MAX_RETRY_ATTEMPT - attempt + 1) * 1500);
return await chat(messages, schema, handler, attempt - 1);
} else {
throw e;
}
}
return answer;
}


Expand Down

0 comments on commit d6e4ae3

Please sign in to comment.