Compare commits

...

3 Commits

Author SHA1 Message Date
suyao
7f8d0b06ee Merge branch 'main' into fix/check-api-key 2025-12-01 16:37:43 +08:00
suyao
4be5fedeec fix 2025-12-01 00:07:43 +08:00
suyao
163e016759 fix: enhance provider handling and API key rotation logic in AiProvider 2025-12-01 00:01:01 +08:00
4 changed files with 81 additions and 57 deletions

View File

@@ -120,9 +120,12 @@ export default class ModernAiProvider {
throw new Error('Model is required for completions. Please use constructor with model parameter.') throw new Error('Model is required for completions. Please use constructor with model parameter.')
} }
// 每次请求时重新生成配置以确保API key轮换生效 // Config is now set in constructor, ApiService handles key rotation before passing provider
this.config = providerToAiSdkConfig(this.actualProvider, this.model) if (!this.config) {
logger.debug('Generated provider config for completions', this.config) // If config wasn't set in constructor (when provider only), generate it now
this.config = providerToAiSdkConfig(this.actualProvider, this.model!)
}
logger.debug('Using provider config for completions', this.config)
// 检查 config 是否存在 // 检查 config 是否存在
if (!this.config) { if (!this.config) {

View File

@@ -29,32 +29,6 @@ import { azureAnthropicProviderCreator } from './config/azure-anthropic'
import { COPILOT_DEFAULT_HEADERS } from './constants' import { COPILOT_DEFAULT_HEADERS } from './constants'
import { getAiSdkProviderId } from './factory' import { getAiSdkProviderId } from './factory'
/**
* 获取轮询的API key
* 复用legacy架构的多key轮询逻辑
*/
function getRotatedApiKey(provider: Provider): string {
const keys = provider.apiKey.split(',').map((key) => key.trim())
const keyName = `provider:${provider.id}:last_used_key`
if (keys.length === 1) {
return keys[0]
}
const lastUsedKey = window.keyv.get(keyName)
if (!lastUsedKey) {
window.keyv.set(keyName, keys[0])
return keys[0]
}
const currentIndex = keys.indexOf(lastUsedKey)
const nextIndex = (currentIndex + 1) % keys.length
const nextKey = keys[nextIndex]
window.keyv.set(keyName, nextKey)
return nextKey
}
/** /**
* 处理特殊provider的转换逻辑 * 处理特殊provider的转换逻辑
*/ */
@@ -161,7 +135,7 @@ export function providerToAiSdkConfig(actualProvider: Provider, model: Model): A
const { baseURL, endpoint } = routeToEndpoint(actualProvider.apiHost) const { baseURL, endpoint } = routeToEndpoint(actualProvider.apiHost)
const baseConfig = { const baseConfig = {
baseURL: baseURL, baseURL: baseURL,
apiKey: getRotatedApiKey(actualProvider) apiKey: actualProvider.apiKey
} }
const isCopilotProvider = actualProvider.id === SystemProviderIds.copilot const isCopilotProvider = actualProvider.id === SystemProviderIds.copilot

View File

@@ -8,8 +8,8 @@ import { isDedicatedImageGenerationModel, isEmbeddingModel, isFunctionCallingMod
import { getStoreSetting } from '@renderer/hooks/useSettings' import { getStoreSetting } from '@renderer/hooks/useSettings'
import i18n from '@renderer/i18n' import i18n from '@renderer/i18n'
import store from '@renderer/store' import store from '@renderer/store'
import type { FetchChatCompletionParams } from '@renderer/types'
import type { Assistant, MCPServer, MCPTool, Model, Provider } from '@renderer/types' import type { Assistant, MCPServer, MCPTool, Model, Provider } from '@renderer/types'
import { type FetchChatCompletionParams, isSystemProvider } from '@renderer/types'
import type { StreamTextParams } from '@renderer/types/aiCoreTypes' import type { StreamTextParams } from '@renderer/types/aiCoreTypes'
import { type Chunk, ChunkType } from '@renderer/types/chunk' import { type Chunk, ChunkType } from '@renderer/types/chunk'
import type { Message, ResponseError } from '@renderer/types/newMessage' import type { Message, ResponseError } from '@renderer/types/newMessage'
@@ -22,7 +22,8 @@ import { purifyMarkdownImages } from '@renderer/utils/markdown'
import { isPromptToolUse, isSupportedToolUse } from '@renderer/utils/mcp-tools' import { isPromptToolUse, isSupportedToolUse } from '@renderer/utils/mcp-tools'
import { findFileBlocks, getMainTextContent } from '@renderer/utils/messageUtils/find' import { findFileBlocks, getMainTextContent } from '@renderer/utils/messageUtils/find'
import { containsSupportedVariables, replacePromptVariables } from '@renderer/utils/prompt' import { containsSupportedVariables, replacePromptVariables } from '@renderer/utils/prompt'
import { isEmpty, takeRight } from 'lodash' import { NOT_SUPPORT_API_KEY_PROVIDERS } from '@renderer/utils/provider'
import { cloneDeep, isEmpty, takeRight } from 'lodash'
import type { ModernAiProviderConfig } from '../aiCore/index_new' import type { ModernAiProviderConfig } from '../aiCore/index_new'
import AiProviderNew from '../aiCore/index_new' import AiProviderNew from '../aiCore/index_new'
@@ -43,6 +44,8 @@ import {
// } from './MessagesService' // } from './MessagesService'
// import WebSearchService from './WebSearchService' // import WebSearchService from './WebSearchService'
// FIXME: 这里太多重复逻辑,需要重构
const logger = loggerService.withContext('ApiService') const logger = loggerService.withContext('ApiService')
export async function fetchMcpTools(assistant: Assistant) { export async function fetchMcpTools(assistant: Assistant) {
@@ -95,7 +98,15 @@ export async function fetchChatCompletion({
modelId: assistant.model?.id, modelId: assistant.model?.id,
modelName: assistant.model?.name modelName: assistant.model?.name
}) })
const AI = new AiProviderNew(assistant.model || getDefaultModel())
// Get base provider and apply API key rotation
const baseProvider = getProviderByModel(assistant.model || getDefaultModel())
const providerWithRotatedKey = {
...cloneDeep(baseProvider),
apiKey: getRotatedApiKey(baseProvider)
}
const AI = new AiProviderNew(assistant.model || getDefaultModel(), providerWithRotatedKey)
const provider = AI.getActualProvider() const provider = AI.getActualProvider()
const mcpTools: MCPTool[] = [] const mcpTools: MCPTool[] = []
@@ -172,7 +183,13 @@ export async function fetchMessagesSummary({ messages, assistant }: { messages:
return null return null
} }
const AI = new AiProviderNew(model) // Apply API key rotation
const providerWithRotatedKey = {
...cloneDeep(provider),
apiKey: getRotatedApiKey(provider)
}
const AI = new AiProviderNew(model, providerWithRotatedKey)
const topicId = messages?.find((message) => message.topicId)?.topicId || '' const topicId = messages?.find((message) => message.topicId)?.topicId || ''
@@ -271,7 +288,13 @@ export async function fetchNoteSummary({ content, assistant }: { content: string
return null return null
} }
const AI = new AiProviderNew(model) // Apply API key rotation
const providerWithRotatedKey = {
...cloneDeep(provider),
apiKey: getRotatedApiKey(provider)
}
const AI = new AiProviderNew(model, providerWithRotatedKey)
// only 2000 char and no images // only 2000 char and no images
const truncatedContent = content.substring(0, 2000) const truncatedContent = content.substring(0, 2000)
@@ -359,7 +382,13 @@ export async function fetchGenerate({
return '' return ''
} }
const AI = new AiProviderNew(model) // Apply API key rotation
const providerWithRotatedKey = {
...cloneDeep(provider),
apiKey: getRotatedApiKey(provider)
}
const AI = new AiProviderNew(model, providerWithRotatedKey)
const assistant = getDefaultAssistant() const assistant = getDefaultAssistant()
assistant.model = model assistant.model = model
@@ -404,28 +433,44 @@ export async function fetchGenerate({
export function hasApiKey(provider: Provider) { export function hasApiKey(provider: Provider) {
if (!provider) return false if (!provider) return false
if (['ollama', 'lmstudio', 'vertexai', 'cherryai'].includes(provider.id)) return true if (isSystemProvider(provider) && NOT_SUPPORT_API_KEY_PROVIDERS.includes(provider.id)) return true
return !isEmpty(provider.apiKey) return !isEmpty(provider.apiKey)
} }
/** /**
* Get the first available embedding model from enabled providers * 获取轮询的API key
* 复用legacy架构的多key轮询逻辑
*/ */
// function getFirstEmbeddingModel() { function getRotatedApiKey(provider: Provider): string {
// const providers = store.getState().llm.providers.filter((p) => p.enabled) const keys = provider.apiKey.split(',').map((key) => key.trim())
const keyName = `provider:${provider.id}:last_used_key`
// for (const provider of providers) { if (keys.length === 1) {
// const embeddingModel = provider.models.find((model) => isEmbeddingModel(model)) return keys[0]
// if (embeddingModel) { }
// return embeddingModel
// }
// }
// return undefined const lastUsedKey = window.keyv.get(keyName)
// } if (!lastUsedKey) {
window.keyv.set(keyName, keys[0])
return keys[0]
}
const currentIndex = keys.indexOf(lastUsedKey)
const nextIndex = (currentIndex + 1) % keys.length
const nextKey = keys[nextIndex]
window.keyv.set(keyName, nextKey)
return nextKey
}
export async function fetchModels(provider: Provider): Promise<SdkModel[]> { export async function fetchModels(provider: Provider): Promise<SdkModel[]> {
const AI = new AiProviderNew(provider) // Apply API key rotation
const providerWithRotatedKey = {
...cloneDeep(provider),
apiKey: getRotatedApiKey(provider)
}
const AI = new AiProviderNew(providerWithRotatedKey)
try { try {
return await AI.models() return await AI.models()
@@ -435,12 +480,7 @@ export async function fetchModels(provider: Provider): Promise<SdkModel[]> {
} }
export function checkApiProvider(provider: Provider): void { export function checkApiProvider(provider: Provider): void {
if ( if (isSystemProvider(provider) && !NOT_SUPPORT_API_KEY_PROVIDERS.includes(provider.id)) {
provider.id !== 'ollama' &&
provider.id !== 'lmstudio' &&
provider.type !== 'vertexai' &&
provider.id !== 'copilot'
) {
if (!provider.apiKey) { if (!provider.apiKey) {
window.toast.error(i18n.t('message.error.enter.api.label')) window.toast.error(i18n.t('message.error.enter.api.label'))
throw new Error(i18n.t('message.error.enter.api.label')) throw new Error(i18n.t('message.error.enter.api.label'))
@@ -461,8 +501,7 @@ export function checkApiProvider(provider: Provider): void {
export async function checkApi(provider: Provider, model: Model, timeout = 15000): Promise<void> { export async function checkApi(provider: Provider, model: Model, timeout = 15000): Promise<void> {
checkApiProvider(provider) checkApiProvider(provider)
// Don't pass in provider parameter. We need auto-format URL const ai = new AiProviderNew(model, provider)
const ai = new AiProviderNew(model)
const assistant = getDefaultAssistant() const assistant = getDefaultAssistant()
assistant.model = model assistant.model = model

View File

@@ -183,3 +183,11 @@ export const isSupportAPIVersionProvider = (provider: Provider) => {
} }
return provider.apiOptions?.isNotSupportAPIVersion !== false return provider.apiOptions?.isNotSupportAPIVersion !== false
} }
export const NOT_SUPPORT_API_KEY_PROVIDERS: readonly SystemProviderId[] = [
'ollama',
'lmstudio',
'vertexai',
'aws-bedrock',
'copilot'
]