feat: refactor API key handling to shuffle keys and attempt each once
This commit is contained in:
+5
-4
@@ -10,19 +10,20 @@ export async function chat(systemPrompt, userContent) {
|
||||
const headers = { 'Content-Type': 'application/json' };
|
||||
if (provider === 'claude') headers['anthropic-version'] = '2023-06-01';
|
||||
|
||||
const shuffled = [...apiKeys].sort(() => Math.random() - 0.5);
|
||||
let lastError;
|
||||
for (let i = 0; i < apiKeys.length; i++) {
|
||||
if (provider !== 'ollama') headers['Authorization'] = `Bearer ${apiKeys[i]}`;
|
||||
for (let i = 0; i < shuffled.length; i++) {
|
||||
if (provider !== 'ollama') headers['Authorization'] = `Bearer ${shuffled[i]}`;
|
||||
try {
|
||||
const resp = await axios.post(
|
||||
`${baseURL.replace(/\/$/, '')}/chat/completions`,
|
||||
{ model, messages: [{ role: 'system', content: systemPrompt }, { role: 'user', content: userContent }], temperature: 0.2 },
|
||||
{ headers, timeout: 30000 }
|
||||
{ headers }
|
||||
);
|
||||
return resp.data.choices[0].message.content;
|
||||
} catch (e) {
|
||||
lastError = e;
|
||||
console.log(` [LLM] key[${i + 1}/${apiKeys.length}] 失敗: ${e.message}`);
|
||||
console.log(` [LLM] key[${i + 1}/${shuffled.length}] 失敗: ${e.message}`);
|
||||
}
|
||||
}
|
||||
console.error(' [LLM] 所有 API Key 均失敗,終止流程');
|
||||
|
||||
+14
-14
@@ -48,18 +48,18 @@ describe('chat - key rotation', async () => {
|
||||
assert.equal(result, 'hello');
|
||||
});
|
||||
|
||||
it('rotates to second key when first fails', async () => {
|
||||
process.env.OPENAI_API_KEY = 'key1,key2';
|
||||
mockAxiosPost([new Error('rate limit'), makeOkResponse('from key2')]);
|
||||
const result = await chat('sys', 'user');
|
||||
assert.equal(result, 'from key2');
|
||||
});
|
||||
|
||||
it('rotates through all keys and succeeds on last', async () => {
|
||||
process.env.OPENAI_API_KEY = 'k1,k2,k3';
|
||||
mockAxiosPost([new Error('fail'), new Error('fail'), makeOkResponse('from k3')]);
|
||||
const result = await chat('sys', 'user');
|
||||
assert.equal(result, 'from k3');
|
||||
it('shuffles keys and tries each exactly once', async () => {
|
||||
process.env.OPENAI_API_KEY = 'key1,key2,key3';
|
||||
const usedKeys = [];
|
||||
mock.method(axios, 'post', async (_url, _body, opts) => {
|
||||
usedKeys.push(opts.headers['Authorization'].replace('Bearer ', ''));
|
||||
throw new Error('fail');
|
||||
});
|
||||
const exitMock = mock.method(process, 'exit', () => { throw new Error('exit:1'); });
|
||||
await assert.rejects(() => chat('sys', 'user'), /exit:1/);
|
||||
assert.equal(exitMock.mock.calls[0].arguments[0], 1);
|
||||
assert.equal(usedKeys.length, 3);
|
||||
assert.deepEqual([...usedKeys].sort(), ['key1', 'key2', 'key3']);
|
||||
});
|
||||
|
||||
it('calls process.exit(1) when all keys fail', async () => {
|
||||
@@ -93,7 +93,7 @@ describe('chat - key rotation', async () => {
|
||||
assert.equal(capturedHeaders['Authorization'], 'Bearer sk-test');
|
||||
});
|
||||
|
||||
it('uses 30s timeout', async () => {
|
||||
it('does not set timeout', async () => {
|
||||
process.env.OPENAI_API_KEY = 'sk-test';
|
||||
let capturedOpts;
|
||||
mock.method(axios, 'post', async (_url, _body, opts) => {
|
||||
@@ -101,7 +101,7 @@ describe('chat - key rotation', async () => {
|
||||
return makeOkResponse();
|
||||
});
|
||||
await chat('sys', 'user');
|
||||
assert.equal(capturedOpts.timeout, 30000);
|
||||
assert.equal(capturedOpts.timeout, undefined);
|
||||
});
|
||||
|
||||
it('does not pass httpsAgent to axios', async () => {
|
||||
|
||||
Reference in New Issue
Block a user