diff --git a/src/backend/src/services/MeteringService/costMaps/xaiCostMap.ts b/src/backend/src/services/MeteringService/costMaps/xaiCostMap.ts index 2653ef54..ee431c26 100644 --- a/src/backend/src/services/MeteringService/costMaps/xaiCostMap.ts +++ b/src/backend/src/services/MeteringService/costMaps/xaiCostMap.ts @@ -50,4 +50,7 @@ export const XAI_COST_MAP = { // Grok 2 'xai:grok-2:prompt_tokens': 200, 'xai:grok-2:completion-tokens': 1000, -}; \ No newline at end of file + + // Grok Image + 'xai:grok-2-image:output': 7_000_000, +}; diff --git a/src/backend/src/services/ai/image/AIImageGenerationService.ts b/src/backend/src/services/ai/image/AIImageGenerationService.ts index 2954fab9..0c5e0d9a 100644 --- a/src/backend/src/services/ai/image/AIImageGenerationService.ts +++ b/src/backend/src/services/ai/image/AIImageGenerationService.ts @@ -30,6 +30,7 @@ import { MeteringService } from '../../MeteringService/MeteringService.js'; import { GeminiImageGenerationProvider } from './providers/GeminiImageGenerationProvider/GeminiImageGenerationProvider.js'; import { OpenAiImageGenerationProvider } from './providers/OpenAiImageGenerationProvider/OpenAiImageGenerationProvider.js'; import { TogetherImageGenerationProvider } from './providers/TogetherImageGenerationProvider/TogetherImageGenerationProvider.js'; +import { XAIImageGenerationProvider } from './providers/XAIImageGenerationProvider/XAIImageGenerationProvider.js'; import { IGenerateParams, IImageModel, IImageProvider } from './providers/types.js'; export class AIImageGenerationService extends BaseService { @@ -108,6 +109,11 @@ export class AIImageGenerationService extends BaseService { this.#providers['together-image-generation'] = new TogetherImageGenerationProvider({ apiKey: togetherConfig.apiKey || togetherConfig.secret_key }, this.meteringService, this.errorService, this.eventService); } + const xaiConfig = this.config.providers?.['xai-image-generation'] || this.config.providers?.['xai'] || this.global_config?.services?.['xai']; + if ( xaiConfig && (xaiConfig.apiKey || xaiConfig.secret_key) ) { + this.#providers['xai-image-generation'] = new XAIImageGenerationProvider({ apiKey: xaiConfig.apiKey || xaiConfig.secret_key }, this.meteringService, this.errorService); + } + // emit event for extensions to add providers const extensionProviders = {} as Record; await this.eventService.emit('ai.image.registerProviders', extensionProviders); diff --git a/src/backend/src/services/ai/image/providers/XAIImageGenerationProvider/XAIImageGenerationProvider.ts b/src/backend/src/services/ai/image/providers/XAIImageGenerationProvider/XAIImageGenerationProvider.ts new file mode 100644 index 00000000..2e507a0b --- /dev/null +++ b/src/backend/src/services/ai/image/providers/XAIImageGenerationProvider/XAIImageGenerationProvider.ts @@ -0,0 +1,112 @@ +/* + * Copyright (C) 2024-present Puter Technologies Inc. + * + * This file is part of Puter. + * + * Puter is free software: you can redistribute it and/or modify + * it under the terms of the GNU Affero General Public License as published + * by the Free Software Foundation, either version 3 of the License, or + * (at your option) any later version. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * GNU Affero General Public License for more details. + * + * You should have received a copy of the GNU Affero General Public License + * along with this program. If not, see . + */ + +import { OpenAI } from 'openai'; +import APIError from '../../../../../api/APIError.js'; +import { ErrorService } from '../../../../../modules/core/ErrorService.js'; +import { Context } from '../../../../../util/context.js'; +import { MeteringService } from '../../../../MeteringService/MeteringService.js'; +import { IGenerateParams, IImageModel, IImageProvider } from '../types.js'; +import { XAI_IMAGE_GENERATION_MODELS } from './models.js'; + +const DEFAULT_MODEL = 'grok-2-image'; +const PRICE_KEY = 'output'; + +export class XAIImageGenerationProvider implements IImageProvider { + #client: OpenAI; + #meteringService: MeteringService; + #errors: ErrorService; + + constructor (config: { apiKey: string }, meteringService: MeteringService, errorService: ErrorService) { + if ( ! config.apiKey ) { + throw new Error('xAI image generation requires an API key'); + } + + this.#meteringService = meteringService; + this.#errors = errorService; + this.#client = new OpenAI({ + apiKey: config.apiKey, + baseURL: 'https://api.x.ai/v1', + }); + } + + models (): IImageModel[] { + return XAI_IMAGE_GENERATION_MODELS; + } + + getDefaultModel (): string { + return DEFAULT_MODEL; + } + + async generate (params: IGenerateParams): Promise { + const { prompt, test_mode } = params; + let { model } = params; + + const selectedModel = this.#getModel(model); + + if ( test_mode ) { + return 'https://puter-sample-data.puter.site/image_example.png'; + } + + if ( typeof prompt !== 'string' || prompt.trim().length === 0 ) { + throw new Error('`prompt` must be a non-empty string'); + } + + const actor = Context.get('actor'); + const user_private_uid = actor?.private_uid ?? 'UNKNOWN'; + if ( user_private_uid === 'UNKNOWN' ) { + this.#errors.report('xai-image-generation:unknown-user', { + message: 'failed to get a user ID for an xAI request', + alarm: true, + trace: true, + }); + } + + const priceInCents = selectedModel.costs[PRICE_KEY]; + const costInMicroCents = priceInCents * 1_000_000; + const usageAllowed = await this.#meteringService.hasEnoughCredits(actor, costInMicroCents); + + if ( ! usageAllowed ) { + throw APIError.create('insufficient_funds'); + } + + const response = await this.#client.images.generate({ + model: selectedModel.id, + prompt, + user: user_private_uid, + }); + + const first = response.data?.[0] as { url?: string; b64_json?: string } | undefined; + const url = first?.url || (first?.b64_json ? `data:image/png;base64,${ first.b64_json}` : undefined); + + if ( ! url ) { + throw new Error('Failed to extract image URL from xAI response'); + } + + this.#meteringService.incrementUsage(actor, `xai:${selectedModel.id}:${PRICE_KEY}`, 1, costInMicroCents); + + return url; + } + + #getModel (model?: string) { + const models = this.models(); + const found = models.find(m => m.id === model || m.aliases?.includes(model ?? '')); + return found || models.find(m => m.id === DEFAULT_MODEL)!; + } +} diff --git a/src/backend/src/services/ai/image/providers/XAIImageGenerationProvider/models.ts b/src/backend/src/services/ai/image/providers/XAIImageGenerationProvider/models.ts new file mode 100644 index 00000000..34055d2b --- /dev/null +++ b/src/backend/src/services/ai/image/providers/XAIImageGenerationProvider/models.ts @@ -0,0 +1,15 @@ +import { IImageModel } from '../types'; + +export const XAI_IMAGE_GENERATION_MODELS: IImageModel[] = [ + { + id: 'grok-2-image', + aliases: ['grok-image'], + name: 'Grok 2 Image', + version: '1.0', + costs_currency: 'usd-cents', + index_cost_key: 'output', + costs: { + output: 7, // $0.07 per image + }, + }, +];