Run model availability check in the background to speed up startup (#4256)
This commit is contained in:
parent
d622e596a1
commit
cba272082d
|
@ -57,6 +57,7 @@ import {
|
|||
EditorType,
|
||||
FlashFallbackEvent,
|
||||
logFlashFallback,
|
||||
AuthType,
|
||||
} from '@google/gemini-cli-core';
|
||||
import { validateAuthMethod } from '../config/auth.js';
|
||||
import { useLogger } from './hooks/useLogger.js';
|
||||
|
@ -294,64 +295,70 @@ const App = ({ config, settings, startupWarnings = [], version }: AppProps) => {
|
|||
): Promise<boolean> => {
|
||||
let message: string;
|
||||
|
||||
// Use actual user tier if available, otherwise default to FREE tier behavior (safe default)
|
||||
const isPaidTier =
|
||||
userTier === UserTierId.LEGACY || userTier === UserTierId.STANDARD;
|
||||
if (
|
||||
config.getContentGeneratorConfig().authType ===
|
||||
AuthType.LOGIN_WITH_GOOGLE
|
||||
) {
|
||||
// Use actual user tier if available, otherwise default to FREE tier behavior (safe default)
|
||||
const isPaidTier =
|
||||
userTier === UserTierId.LEGACY || userTier === UserTierId.STANDARD;
|
||||
|
||||
// Check if this is a Pro quota exceeded error
|
||||
if (error && isProQuotaExceededError(error)) {
|
||||
if (isPaidTier) {
|
||||
message = `⚡ You have reached your daily ${currentModel} quota limit.
|
||||
// Check if this is a Pro quota exceeded error
|
||||
if (error && isProQuotaExceededError(error)) {
|
||||
if (isPaidTier) {
|
||||
message = `⚡ You have reached your daily ${currentModel} quota limit.
|
||||
⚡ Automatically switching from ${currentModel} to ${fallbackModel} for the remainder of this session.
|
||||
⚡ To continue accessing the ${currentModel} model today, consider using /auth to switch to using a paid API key from AI Studio at https://aistudio.google.com/apikey`;
|
||||
} else {
|
||||
message = `⚡ You have reached your daily ${currentModel} quota limit.
|
||||
} else {
|
||||
message = `⚡ You have reached your daily ${currentModel} quota limit.
|
||||
⚡ Automatically switching from ${currentModel} to ${fallbackModel} for the remainder of this session.
|
||||
⚡ To increase your limits, upgrade to a Gemini Code Assist Standard or Enterprise plan with higher limits at https://goo.gle/set-up-gemini-code-assist
|
||||
⚡ Or you can utilize a Gemini API Key. See: https://goo.gle/gemini-cli-docs-auth#gemini-api-key
|
||||
⚡ You can switch authentication methods by typing /auth`;
|
||||
}
|
||||
} else if (error && isGenericQuotaExceededError(error)) {
|
||||
if (isPaidTier) {
|
||||
message = `⚡ You have reached your daily quota limit.
|
||||
}
|
||||
} else if (error && isGenericQuotaExceededError(error)) {
|
||||
if (isPaidTier) {
|
||||
message = `⚡ You have reached your daily quota limit.
|
||||
⚡ Automatically switching from ${currentModel} to ${fallbackModel} for the remainder of this session.
|
||||
⚡ To continue accessing the ${currentModel} model today, consider using /auth to switch to using a paid API key from AI Studio at https://aistudio.google.com/apikey`;
|
||||
} else {
|
||||
message = `⚡ You have reached your daily quota limit.
|
||||
} else {
|
||||
message = `⚡ You have reached your daily quota limit.
|
||||
⚡ Automatically switching from ${currentModel} to ${fallbackModel} for the remainder of this session.
|
||||
⚡ To increase your limits, upgrade to a Gemini Code Assist Standard or Enterprise plan with higher limits at https://goo.gle/set-up-gemini-code-assist
|
||||
⚡ Or you can utilize a Gemini API Key. See: https://goo.gle/gemini-cli-docs-auth#gemini-api-key
|
||||
⚡ You can switch authentication methods by typing /auth`;
|
||||
}
|
||||
} else {
|
||||
if (isPaidTier) {
|
||||
// Default fallback message for other cases (like consecutive 429s)
|
||||
message = `⚡ Automatically switching from ${currentModel} to ${fallbackModel} for faster responses for the remainder of this session.
|
||||
}
|
||||
} else {
|
||||
if (isPaidTier) {
|
||||
// Default fallback message for other cases (like consecutive 429s)
|
||||
message = `⚡ Automatically switching from ${currentModel} to ${fallbackModel} for faster responses for the remainder of this session.
|
||||
⚡ Possible reasons for this are that you have received multiple consecutive capacity errors or you have reached your daily ${currentModel} quota limit
|
||||
⚡ To continue accessing the ${currentModel} model today, consider using /auth to switch to using a paid API key from AI Studio at https://aistudio.google.com/apikey`;
|
||||
} else {
|
||||
// Default fallback message for other cases (like consecutive 429s)
|
||||
message = `⚡ Automatically switching from ${currentModel} to ${fallbackModel} for faster responses for the remainder of this session.
|
||||
} else {
|
||||
// Default fallback message for other cases (like consecutive 429s)
|
||||
message = `⚡ Automatically switching from ${currentModel} to ${fallbackModel} for faster responses for the remainder of this session.
|
||||
⚡ Possible reasons for this are that you have received multiple consecutive capacity errors or you have reached your daily ${currentModel} quota limit
|
||||
⚡ To increase your limits, upgrade to a Gemini Code Assist Standard or Enterprise plan with higher limits at https://goo.gle/set-up-gemini-code-assist
|
||||
⚡ Or you can utilize a Gemini API Key. See: https://goo.gle/gemini-cli-docs-auth#gemini-api-key
|
||||
⚡ You can switch authentication methods by typing /auth`;
|
||||
}
|
||||
}
|
||||
|
||||
// Add message to UI history
|
||||
addItem(
|
||||
{
|
||||
type: MessageType.INFO,
|
||||
text: message,
|
||||
},
|
||||
Date.now(),
|
||||
);
|
||||
|
||||
// Set the flag to prevent tool continuation
|
||||
setModelSwitchedFromQuotaError(true);
|
||||
// Set global quota error flag to prevent Flash model calls
|
||||
config.setQuotaErrorOccurred(true);
|
||||
}
|
||||
|
||||
// Add message to UI history
|
||||
addItem(
|
||||
{
|
||||
type: MessageType.INFO,
|
||||
text: message,
|
||||
},
|
||||
Date.now(),
|
||||
);
|
||||
|
||||
// Set the flag to prevent tool continuation
|
||||
setModelSwitchedFromQuotaError(true);
|
||||
// Set global quota error flag to prevent Flash model calls
|
||||
config.setQuotaErrorOccurred(true);
|
||||
// Switch model for future use but return false to stop current retry
|
||||
config.setModel(fallbackModel);
|
||||
logFlashFallback(
|
||||
|
|
|
@ -151,14 +151,12 @@ describe('Server Config (config.ts)', () => {
|
|||
apiKey: 'test-key',
|
||||
};
|
||||
|
||||
(createContentGeneratorConfig as Mock).mockResolvedValue(
|
||||
mockContentConfig,
|
||||
);
|
||||
(createContentGeneratorConfig as Mock).mockReturnValue(mockContentConfig);
|
||||
|
||||
await config.refreshAuth(authType);
|
||||
|
||||
expect(createContentGeneratorConfig).toHaveBeenCalledWith(
|
||||
MODEL, // Should be called with the original model 'gemini-pro'
|
||||
config,
|
||||
authType,
|
||||
);
|
||||
// Verify that contentGeneratorConfig is updated with the new model
|
||||
|
|
|
@ -274,8 +274,8 @@ export class Config {
|
|||
}
|
||||
|
||||
async refreshAuth(authMethod: AuthType) {
|
||||
this.contentGeneratorConfig = await createContentGeneratorConfig(
|
||||
this.model,
|
||||
this.contentGeneratorConfig = createContentGeneratorConfig(
|
||||
this,
|
||||
authMethod,
|
||||
);
|
||||
|
||||
|
|
|
@ -64,12 +64,18 @@ describe('createContentGenerator', () => {
|
|||
|
||||
describe('createContentGeneratorConfig', () => {
|
||||
const originalEnv = process.env;
|
||||
const mockConfig = {
|
||||
getModel: vi.fn().mockReturnValue('gemini-pro'),
|
||||
setModel: vi.fn(),
|
||||
flashFallbackHandler: vi.fn(),
|
||||
} as unknown as Config;
|
||||
|
||||
beforeEach(() => {
|
||||
// Reset modules to re-evaluate imports and environment variables
|
||||
vi.resetModules();
|
||||
// Restore process.env before each test
|
||||
process.env = { ...originalEnv };
|
||||
vi.clearAllMocks();
|
||||
});
|
||||
|
||||
afterAll(() => {
|
||||
|
@ -80,7 +86,7 @@ describe('createContentGeneratorConfig', () => {
|
|||
it('should configure for Gemini using GEMINI_API_KEY when set', async () => {
|
||||
process.env.GEMINI_API_KEY = 'env-gemini-key';
|
||||
const config = await createContentGeneratorConfig(
|
||||
undefined,
|
||||
mockConfig,
|
||||
AuthType.USE_GEMINI,
|
||||
);
|
||||
expect(config.apiKey).toBe('env-gemini-key');
|
||||
|
@ -90,7 +96,7 @@ describe('createContentGeneratorConfig', () => {
|
|||
it('should not configure for Gemini if GEMINI_API_KEY is empty', async () => {
|
||||
process.env.GEMINI_API_KEY = '';
|
||||
const config = await createContentGeneratorConfig(
|
||||
undefined,
|
||||
mockConfig,
|
||||
AuthType.USE_GEMINI,
|
||||
);
|
||||
expect(config.apiKey).toBeUndefined();
|
||||
|
@ -100,7 +106,7 @@ describe('createContentGeneratorConfig', () => {
|
|||
it('should configure for Vertex AI using GOOGLE_API_KEY when set', async () => {
|
||||
process.env.GOOGLE_API_KEY = 'env-google-key';
|
||||
const config = await createContentGeneratorConfig(
|
||||
undefined,
|
||||
mockConfig,
|
||||
AuthType.USE_VERTEX_AI,
|
||||
);
|
||||
expect(config.apiKey).toBe('env-google-key');
|
||||
|
@ -111,7 +117,7 @@ describe('createContentGeneratorConfig', () => {
|
|||
process.env.GOOGLE_CLOUD_PROJECT = 'env-gcp-project';
|
||||
process.env.GOOGLE_CLOUD_LOCATION = 'env-gcp-location';
|
||||
const config = await createContentGeneratorConfig(
|
||||
undefined,
|
||||
mockConfig,
|
||||
AuthType.USE_VERTEX_AI,
|
||||
);
|
||||
expect(config.vertexai).toBe(true);
|
||||
|
@ -123,7 +129,7 @@ describe('createContentGeneratorConfig', () => {
|
|||
process.env.GOOGLE_CLOUD_PROJECT = '';
|
||||
process.env.GOOGLE_CLOUD_LOCATION = '';
|
||||
const config = await createContentGeneratorConfig(
|
||||
undefined,
|
||||
mockConfig,
|
||||
AuthType.USE_VERTEX_AI,
|
||||
);
|
||||
expect(config.apiKey).toBeUndefined();
|
||||
|
|
|
@ -52,17 +52,17 @@ export type ContentGeneratorConfig = {
|
|||
authType?: AuthType | undefined;
|
||||
};
|
||||
|
||||
export async function createContentGeneratorConfig(
|
||||
model: string | undefined,
|
||||
export function createContentGeneratorConfig(
|
||||
config: Config,
|
||||
authType: AuthType | undefined,
|
||||
): Promise<ContentGeneratorConfig> {
|
||||
): ContentGeneratorConfig {
|
||||
const geminiApiKey = process.env.GEMINI_API_KEY || undefined;
|
||||
const googleApiKey = process.env.GOOGLE_API_KEY || undefined;
|
||||
const googleCloudProject = process.env.GOOGLE_CLOUD_PROJECT || undefined;
|
||||
const googleCloudLocation = process.env.GOOGLE_CLOUD_LOCATION || undefined;
|
||||
|
||||
// Use runtime model from config if available, otherwise fallback to parameter or default
|
||||
const effectiveModel = model || DEFAULT_GEMINI_MODEL;
|
||||
const effectiveModel = config.getModel() || DEFAULT_GEMINI_MODEL;
|
||||
|
||||
const contentGeneratorConfig: ContentGeneratorConfig = {
|
||||
model: effectiveModel,
|
||||
|
@ -80,10 +80,14 @@ export async function createContentGeneratorConfig(
|
|||
if (authType === AuthType.USE_GEMINI && geminiApiKey) {
|
||||
contentGeneratorConfig.apiKey = geminiApiKey;
|
||||
contentGeneratorConfig.vertexai = false;
|
||||
contentGeneratorConfig.model = await getEffectiveModel(
|
||||
getEffectiveModel(
|
||||
contentGeneratorConfig.apiKey,
|
||||
contentGeneratorConfig.model,
|
||||
);
|
||||
).then((newModel) => {
|
||||
if (newModel !== contentGeneratorConfig.model) {
|
||||
config.flashFallbackHandler?.(contentGeneratorConfig.model, newModel);
|
||||
}
|
||||
});
|
||||
|
||||
return contentGeneratorConfig;
|
||||
}
|
||||
|
|
Loading…
Reference in New Issue