mirror of
https://github.com/stackblitz-labs/bolt.diy
synced 2025-06-26 18:26:38 +00:00
- Remove hardcoded static models from all provider classes to simplify maintenance and improve flexibility. Introduce a new `formatTokens` utility function to standardize token count display across providers. This change ensures consistency and reduces redundancy in model definitions. - add anthropic-beta header for increased output limit
83 lines
2.4 KiB
TypeScript
83 lines
2.4 KiB
TypeScript
import { BaseProvider } from '~/lib/modules/llm/base-provider';
|
|
import type { ModelInfo } from '~/lib/modules/llm/types';
|
|
import type { IProviderSetting } from '~/types/model';
|
|
import type { LanguageModelV1 } from 'ai';
|
|
import { createGoogleGenerativeAI } from '@ai-sdk/google';
|
|
import { formatTokens } from './tokenFormat';
|
|
|
|
export default class GoogleProvider extends BaseProvider {
|
|
name = 'Google';
|
|
getApiKeyLink = 'https://aistudio.google.com/app/apikey';
|
|
|
|
config = {
|
|
apiTokenKey: 'GOOGLE_GENERATIVE_AI_API_KEY',
|
|
};
|
|
|
|
staticModels: ModelInfo[] = [];
|
|
|
|
async getDynamicModels(
|
|
apiKeys?: Record<string, string>,
|
|
settings?: IProviderSetting,
|
|
serverEnv?: Record<string, string>,
|
|
): Promise<ModelInfo[]> {
|
|
const { apiKey } = this.getProviderBaseUrlAndKey({
|
|
apiKeys,
|
|
providerSettings: settings,
|
|
serverEnv: serverEnv as any,
|
|
defaultBaseUrlKey: '',
|
|
defaultApiTokenKey: 'GOOGLE_GENERATIVE_AI_API_KEY',
|
|
});
|
|
|
|
if (!apiKey) {
|
|
throw `Missing Api Key configuration for ${this.name} provider`;
|
|
}
|
|
|
|
const response = await fetch(`https://generativelanguage.googleapis.com/v1beta/models?key=${apiKey}`, {
|
|
headers: {
|
|
['Content-Type']: 'application/json',
|
|
},
|
|
});
|
|
|
|
const res = (await response.json()) as any;
|
|
|
|
const data = res.models.filter((model: any) => model.outputTokenLimit > 8000);
|
|
|
|
return data.map((m: any) => {
|
|
const totalTokens = m.inputTokenLimit + m.outputTokenLimit;
|
|
return {
|
|
name: m.name.replace('models/', ''),
|
|
label: `${m.displayName} - context: ${formatTokens(totalTokens)} tokens`,
|
|
provider: this.name,
|
|
maxTokenAllowed: totalTokens || 8000,
|
|
};
|
|
});
|
|
}
|
|
|
|
getModelInstance(options: {
|
|
model: string;
|
|
serverEnv: any;
|
|
apiKeys?: Record<string, string>;
|
|
providerSettings?: Record<string, IProviderSetting>;
|
|
}): LanguageModelV1 {
|
|
const { model, serverEnv, apiKeys, providerSettings } = options;
|
|
|
|
const { apiKey } = this.getProviderBaseUrlAndKey({
|
|
apiKeys,
|
|
providerSettings: providerSettings?.[this.name],
|
|
serverEnv: serverEnv as any,
|
|
defaultBaseUrlKey: '',
|
|
defaultApiTokenKey: 'GOOGLE_GENERATIVE_AI_API_KEY',
|
|
});
|
|
|
|
if (!apiKey) {
|
|
throw new Error(`Missing API key for ${this.name} provider`);
|
|
}
|
|
|
|
const google = createGoogleGenerativeAI({
|
|
apiKey,
|
|
});
|
|
|
|
return google(model);
|
|
}
|
|
}
|