Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 2 additions & 3 deletions src/renderer/src/aiCore/plugins/PluginBuilder.ts
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,6 @@ import { SystemProviderIds } from '@renderer/types'
import { isOllamaProvider, isSupportEnableThinkingProvider } from '@renderer/utils/provider'

import type { AiSdkMiddlewareConfig } from '../types/middlewareConfig'
import { isOpenRouterGeminiGenerateImageModel } from '../utils/image'
import { getReasoningTagName } from '../utils/reasoning'
import { createAnthropicCachePlugin } from './anthropicCachePlugin'
import { createNoThinkPlugin } from './noThinkPlugin'
Expand Down Expand Up @@ -96,8 +95,8 @@ export function buildPlugins({ provider, model, config }: BuildPluginsContext):
plugins.push(createQwenThinkingPlugin(enableThinking))
}

// 0.6 OpenRouter Gemini image generation
if (isOpenRouterGeminiGenerateImageModel(model, provider)) {
// 0.6 OpenRouter image generation
if (provider.id === SystemProviderIds.openrouter && config.enableGenerateImage) {
plugins.push(createOpenrouterGenerateImagePlugin())
}

Expand Down
169 changes: 169 additions & 0 deletions src/renderer/src/aiCore/plugins/__tests__/PluginBuilder.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,169 @@
import type { Assistant, Model, Provider } from '@renderer/types'
import { SystemProviderIds } from '@renderer/types'
import { describe, expect, it, vi } from 'vitest'

import { buildPlugins } from '../PluginBuilder'

vi.mock('@logger', () => ({
loggerService: {
withContext: () => ({
debug: vi.fn(),
error: vi.fn(),
warn: vi.fn(),
info: vi.fn()
})
}
}))

vi.mock('@renderer/hooks/useSettings', () => ({
getEnableDeveloperMode: vi.fn(() => false)
}))

vi.mock('@renderer/config/models', () => ({
isGemini3Model: vi.fn(() => false),
isQwen35Model: vi.fn(() => false),
isSupportedThinkingTokenQwenModel: vi.fn(() => false)
}))

vi.mock('@renderer/utils/provider', () => ({
isOllamaProvider: vi.fn(() => false),
isSupportEnableThinkingProvider: vi.fn(() => false)
}))

vi.mock('@cherrystudio/ai-core/built-in/plugins', () => ({
createPromptToolUsePlugin: vi.fn(() => ({ name: 'promptToolUse' })),
webSearchPlugin: vi.fn(() => ({ name: 'webSearch' }))
}))

vi.mock('../anthropicCachePlugin', () => ({
createAnthropicCachePlugin: vi.fn(() => ({ name: 'anthropicCache' }))
}))

vi.mock('../noThinkPlugin', () => ({
createNoThinkPlugin: vi.fn(() => ({ name: 'noThink' }))
}))

vi.mock('../openrouterGenerateImagePlugin', () => ({
createOpenrouterGenerateImagePlugin: vi.fn(() => ({ name: 'openrouterGenerateImage' }))
}))

vi.mock('../openrouterReasoningPlugin', () => ({
createOpenrouterReasoningPlugin: vi.fn(() => ({ name: 'openrouterReasoning' }))
}))

vi.mock('../qwenThinkingPlugin', () => ({
createQwenThinkingPlugin: vi.fn(() => ({ name: 'qwenThinking' }))
}))

vi.mock('../reasoningExtractionPlugin', () => ({
createReasoningExtractionPlugin: vi.fn(() => ({ name: 'reasoningExtraction' }))
}))

vi.mock('../searchOrchestrationPlugin', () => ({
searchOrchestrationPlugin: vi.fn(() => ({ name: 'searchOrchestration' }))
}))

vi.mock('../simulateStreamingPlugin', () => ({
createSimulateStreamingPlugin: vi.fn(() => ({ name: 'simulateStreaming' }))
}))

vi.mock('../skipGeminiThoughtSignaturePlugin', () => ({
createSkipGeminiThoughtSignaturePlugin: vi.fn(() => ({ name: 'skipGeminiThoughtSignature' }))
}))

vi.mock('../telemetryPlugin', () => ({
createTelemetryPlugin: vi.fn(() => ({ name: 'telemetry' }))
}))

vi.mock('../../utils/reasoning', () => ({
getReasoningTagName: vi.fn(() => 'think')
}))

function createAssistant(): Assistant {
return {
id: 'assistant-1',
name: 'Test Assistant',
prompt: '',
topics: [],
type: 'assistant',
settings: {}
}
}

function createModel(provider = SystemProviderIds.openrouter): Model {
return {
id: 'google/gemini-2.5-flash-image-preview',
name: 'Gemini 2.5 Flash Image',
provider
} as Model
}

function createProvider(id = SystemProviderIds.openrouter): Provider {
return {
id,
name: 'Test Provider',
type: id === SystemProviderIds.openrouter ? 'openai' : 'google'
} as Provider
}

describe('PluginBuilder', () => {
it('mounts openrouterGenerateImage plugin when provider is openrouter and generate image is enabled', () => {
const plugins = buildPlugins({
provider: createProvider(),
model: createModel(),
config: {
assistant: createAssistant(),
streamOutput: true,
enableReasoning: false,
isPromptToolUse: false,
isSupportedToolUse: false,
isImageGenerationEndpoint: false,
enableWebSearch: false,
enableGenerateImage: true,
enableUrlContext: false
}
})

expect(plugins.map((plugin) => plugin.name)).toContain('openrouterGenerateImage')
})

it('does not mount openrouterGenerateImage plugin when generate image is disabled', () => {
const plugins = buildPlugins({
provider: createProvider(),
model: createModel(),
config: {
assistant: createAssistant(),
streamOutput: true,
enableReasoning: false,
isPromptToolUse: false,
isSupportedToolUse: false,
isImageGenerationEndpoint: false,
enableWebSearch: false,
enableGenerateImage: false,
enableUrlContext: false
}
})

expect(plugins.map((plugin) => plugin.name)).not.toContain('openrouterGenerateImage')
})

it('does not mount openrouterGenerateImage plugin for non-openrouter providers', () => {
const plugins = buildPlugins({
provider: { ...createProvider(SystemProviderIds.openrouter), id: 'gemini' } as any,
model: createModel(SystemProviderIds.openrouter),
config: {
assistant: createAssistant(),
streamOutput: true,
enableReasoning: false,
isPromptToolUse: false,
isSupportedToolUse: false,
isImageGenerationEndpoint: false,
enableWebSearch: false,
enableGenerateImage: true,
enableUrlContext: false
}
})

expect(plugins.map((plugin) => plugin.name)).not.toContain('openrouterGenerateImage')
})
})
95 changes: 1 addition & 94 deletions src/renderer/src/aiCore/utils/__tests__/image.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -3,11 +3,9 @@
* Tests for Gemini image generation utilities
*/

import type { Model, Provider } from '@renderer/types'
import { SystemProviderIds } from '@renderer/types'
import { describe, expect, it } from 'vitest'

import { buildGeminiGenerateImageParams, isOpenRouterGeminiGenerateImageModel } from '../image'
import { buildGeminiGenerateImageParams } from '../image'

describe('image utils', () => {
describe('buildGeminiGenerateImageParams', () => {
Expand All @@ -27,95 +25,4 @@ describe('image utils', () => {
expect(result.responseModalities).toHaveLength(2)
})
})

describe('isOpenRouterGeminiGenerateImageModel', () => {
const mockOpenRouterProvider: Provider = {
id: SystemProviderIds.openrouter,
name: 'OpenRouter',
apiKey: 'test-key',
apiHost: 'https://openrouter.ai/api/v1',
isSystem: true
} as Provider

const mockOtherProvider: Provider = {
id: SystemProviderIds.openai,
name: 'OpenAI',
apiKey: 'test-key',
apiHost: 'https://api.openai.com/v1',
isSystem: true
} as Provider

it('should return true for OpenRouter Gemini 2.5 Flash Image model', () => {
const model: Model = {
id: 'google/gemini-2.5-flash-image-preview',
name: 'Gemini 2.5 Flash Image',
provider: SystemProviderIds.openrouter
} as Model

const result = isOpenRouterGeminiGenerateImageModel(model, mockOpenRouterProvider)
expect(result).toBe(true)
})

it('should return false for non-Gemini model on OpenRouter', () => {
const model: Model = {
id: 'openai/gpt-4',
name: 'GPT-4',
provider: SystemProviderIds.openrouter
} as Model

const result = isOpenRouterGeminiGenerateImageModel(model, mockOpenRouterProvider)
expect(result).toBe(false)
})

it('should return false for Gemini model on non-OpenRouter provider', () => {
const model: Model = {
id: 'gemini-2.5-flash-image-preview',
name: 'Gemini 2.5 Flash Image',
provider: SystemProviderIds.gemini
} as Model

const result = isOpenRouterGeminiGenerateImageModel(model, mockOtherProvider)
expect(result).toBe(false)
})

it('should return false for Gemini model without image suffix', () => {
const model: Model = {
id: 'google/gemini-2.5-flash',
name: 'Gemini 2.5 Flash',
provider: SystemProviderIds.openrouter
} as Model

const result = isOpenRouterGeminiGenerateImageModel(model, mockOpenRouterProvider)
expect(result).toBe(false)
})

it('should handle model ID with partial match', () => {
const model: Model = {
id: 'google/gemini-2.5-flash-image-generation',
name: 'Gemini Image Gen',
provider: SystemProviderIds.openrouter
} as Model

const result = isOpenRouterGeminiGenerateImageModel(model, mockOpenRouterProvider)
expect(result).toBe(true)
})

it('should return false for custom provider', () => {
const customProvider: Provider = {
id: 'custom-provider-123',
name: 'Custom Provider',
apiKey: 'test-key',
apiHost: 'https://custom.com'
} as Provider

const model: Model = {
id: 'gemini-2.5-flash-image-preview',
name: 'Gemini 2.5 Flash Image',
provider: 'custom-provider-123'
} as Model

const result = isOpenRouterGeminiGenerateImageModel(model, customProvider)
expect(result).toBe(false)
})
})
})
11 changes: 0 additions & 11 deletions src/renderer/src/aiCore/utils/image.ts
Original file line number Diff line number Diff line change
@@ -1,16 +1,5 @@
import type { Model, Provider } from '@renderer/types'
import { isSystemProvider, SystemProviderIds } from '@renderer/types'

export function buildGeminiGenerateImageParams(): Record<string, any> {
return {
responseModalities: ['TEXT', 'IMAGE']
}
}

export function isOpenRouterGeminiGenerateImageModel(model: Model, provider: Provider): boolean {
return (
model.id.includes('gemini-2.5-flash-image') &&
isSystemProvider(provider) &&
provider.id === SystemProviderIds.openrouter
)
}
12 changes: 12 additions & 0 deletions src/renderer/src/config/models/__tests__/vision.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -99,6 +99,7 @@ describe('vision helpers', () => {

providerMock.mockReturnValue({ type: 'custom' } as any)
expect(isGenerateImageModel(createModel({ id: 'gemini-2.5-flash-image' }))).toBe(true)
expect(isGenerateImageModel(createModel({ id: 'gemini-3.1-flash-image-preview' }))).toBe(true)
})

it('returns false when openai-response model is not on allow list', () => {
Expand All @@ -111,6 +112,7 @@ describe('vision helpers', () => {
expect(isPureGenerateImageModel(createModel({ id: 'gpt-image-1' }))).toBe(true)
expect(isPureGenerateImageModel(createModel({ id: 'gpt-4o' }))).toBe(false)
expect(isPureGenerateImageModel(createModel({ id: 'gemini-2.5-flash-image-preview' }))).toBe(true)
expect(isPureGenerateImageModel(createModel({ id: 'gemini-3.1-flash-image-preview' }))).toBe(false)
})
})

Expand All @@ -122,12 +124,14 @@ describe('vision helpers', () => {

it('detects models with restricted image size support and enhancement', () => {
expect(isImageEnhancementModel(createModel({ id: 'qwen-image-edit' }))).toBe(true)
expect(isImageEnhancementModel(createModel({ id: 'gemini-3.1-flash-image-preview' }))).toBe(true)
expect(isImageEnhancementModel(createModel({ id: 'gpt-4o' }))).toBe(false)
})

it('identifies dedicated and auto-enabled image generation models', () => {
expect(isDedicatedImageGenerationModel(createModel({ id: 'grok-2-image-1212' }))).toBe(true)
expect(isAutoEnableImageGenerationModel(createModel({ id: 'gemini-2.5-flash-image-ultra' }))).toBe(true)
expect(isAutoEnableImageGenerationModel(createModel({ id: 'gemini-3.1-flash-image-preview' }))).toBe(true)
})

it('returns false when models are not in dedicated or auto-enable sets', () => {
Expand Down Expand Up @@ -314,6 +318,14 @@ describe('isVisionModel', () => {
group: ''
})
).toBe(true)
expect(
isVisionModel({
id: 'gemini-3.1-flash-image-preview',
name: '',
provider: '',
group: ''
})
).toBe(true)
})

it('should return true for gemini exp models', () => {
Expand Down
12 changes: 12 additions & 0 deletions src/renderer/src/config/models/default.ts
Original file line number Diff line number Diff line change
Expand Up @@ -360,6 +360,12 @@ export const SYSTEM_MODELS: Record<SystemProviderId | 'defaultModel', Model[]> =
name: 'Gemini 2.5 Flash Image',
group: 'Gemini 2.5'
},
{
id: 'gemini-3.1-flash-image-preview',
provider: 'gemini',
name: 'Gemini 3.1 Flash Image Preview',
group: 'Gemini 3'
},
{
id: 'gemini-3-pro-image-preview',
provider: 'gemini',
Expand Down Expand Up @@ -1346,6 +1352,12 @@ export const SYSTEM_MODELS: Record<SystemProviderId | 'defaultModel', Model[]> =
name: 'Google: Gemini 2.5 Flash Image',
group: 'google'
},
{
id: 'google/gemini-3.1-flash-image-preview',
provider: 'openrouter',
name: 'Google: Gemini 3.1 Flash Image Preview',
group: 'google'
},
{
id: 'google/gemini-2.5-flash-preview',
provider: 'openrouter',
Expand Down
Loading
Loading