Push model-switching logging into loadCliConfig (#815)

This commit is contained in:
Tommaso Sciortino
2025-06-07 11:12:30 -07:00
committed by GitHub
parent 680f4cdd61
commit 6ea4479064
5 changed files with 45 additions and 112 deletions

View File

@@ -5,10 +5,7 @@
*/
import { describe, it, expect, vi, beforeEach, afterEach } from 'vitest';
import {
getEffectiveModel,
type EffectiveModelCheckResult,
} from './modelCheck.js';
import { getEffectiveModel } from './modelCheck.js';
import {
DEFAULT_GEMINI_MODEL,
DEFAULT_GEMINI_FLASH_MODEL,
@@ -45,13 +42,10 @@ describe('getEffectiveModel', () => {
});
describe('when currentConfiguredModel is not DEFAULT_GEMINI_MODEL', () => {
it('should return the currentConfiguredModel and switched: false without fetching', async () => {
it('should return the currentConfiguredModel without fetching', async () => {
const customModel = 'custom-model-name';
const result = await getEffectiveModel(apiKey, customModel);
expect(result).toEqual({
effectiveModel: customModel,
switched: false,
});
expect(result).toEqual(customModel);
expect(fetch).not.toHaveBeenCalled();
});
});
@@ -62,15 +56,8 @@ describe('getEffectiveModel', () => {
ok: false,
status: 429,
});
const result: EffectiveModelCheckResult = await getEffectiveModel(
apiKey,
DEFAULT_GEMINI_MODEL,
);
expect(result).toEqual({
effectiveModel: DEFAULT_GEMINI_FLASH_MODEL,
switched: true,
originalModelIfSwitched: DEFAULT_GEMINI_MODEL,
});
const result = await getEffectiveModel(apiKey, DEFAULT_GEMINI_MODEL);
expect(result).toEqual(DEFAULT_GEMINI_FLASH_MODEL);
expect(fetch).toHaveBeenCalledTimes(1);
expect(fetch).toHaveBeenCalledWith(
`https://generativelanguage.googleapis.com/v1beta/models/${DEFAULT_GEMINI_MODEL}:generateContent?key=${apiKey}`,
@@ -84,10 +71,7 @@ describe('getEffectiveModel', () => {
status: 200,
});
const result = await getEffectiveModel(apiKey, DEFAULT_GEMINI_MODEL);
expect(result).toEqual({
effectiveModel: DEFAULT_GEMINI_MODEL,
switched: false,
});
expect(result).toEqual(DEFAULT_GEMINI_MODEL);
expect(fetch).toHaveBeenCalledTimes(1);
});
@@ -97,20 +81,14 @@ describe('getEffectiveModel', () => {
status: 500,
});
const result = await getEffectiveModel(apiKey, DEFAULT_GEMINI_MODEL);
expect(result).toEqual({
effectiveModel: DEFAULT_GEMINI_MODEL,
switched: false,
});
expect(result).toEqual(DEFAULT_GEMINI_MODEL);
expect(fetch).toHaveBeenCalledTimes(1);
});
it('should return DEFAULT_GEMINI_MODEL if fetch throws a network error', async () => {
(fetch as vi.Mock).mockRejectedValueOnce(new Error('Network error'));
const result = await getEffectiveModel(apiKey, DEFAULT_GEMINI_MODEL);
expect(result).toEqual({
effectiveModel: DEFAULT_GEMINI_MODEL,
switched: false,
});
expect(result).toEqual(DEFAULT_GEMINI_MODEL);
expect(fetch).toHaveBeenCalledTimes(1);
});
@@ -146,10 +124,7 @@ describe('getEffectiveModel', () => {
expect(mockAbort).toHaveBeenCalledTimes(0); // setTimeout calls controller.abort(), not our direct mockAbort
expect(abortControllerInstance.abort).toHaveBeenCalledTimes(1);
expect(result).toEqual({
effectiveModel: DEFAULT_GEMINI_MODEL,
switched: false,
});
expect(result).toEqual(DEFAULT_GEMINI_MODEL);
expect(fetch).toHaveBeenCalledTimes(1);
});

View File

@@ -9,12 +9,6 @@ import {
DEFAULT_GEMINI_FLASH_MODEL,
} from '../config/config.js';
export interface EffectiveModelCheckResult {
effectiveModel: string;
switched: boolean;
originalModelIfSwitched?: string;
}
/**
* Checks if the default "pro" model is rate-limited and returns a fallback "flash"
* model if necessary. This function is designed to be silent.
@@ -26,10 +20,10 @@ export interface EffectiveModelCheckResult {
export async function getEffectiveModel(
apiKey: string,
currentConfiguredModel: string,
): Promise<EffectiveModelCheckResult> {
): Promise<string> {
if (currentConfiguredModel !== DEFAULT_GEMINI_MODEL) {
// Only check if the user is trying to use the specific pro model we want to fallback from.
return { effectiveModel: currentConfiguredModel, switched: false };
return currentConfiguredModel;
}
const modelToTest = DEFAULT_GEMINI_MODEL;
@@ -59,17 +53,16 @@ export async function getEffectiveModel(
clearTimeout(timeoutId);
if (response.status === 429) {
return {
effectiveModel: fallbackModel,
switched: true,
originalModelIfSwitched: modelToTest,
};
console.log(
`[INFO] Your configured model (${modelToTest}) was temporarily unavailable. Switched to ${fallbackModel} for this session.`,
);
return fallbackModel;
}
// For any other case (success, other error codes), we stick to the original model.
return { effectiveModel: currentConfiguredModel, switched: false };
return currentConfiguredModel;
} catch (_error) {
clearTimeout(timeoutId);
// On timeout or any other fetch error, stick to the original model.
return { effectiveModel: currentConfiguredModel, switched: false };
return currentConfiguredModel;
}
}