Add xAI image generation provider integration

This commit is contained in:
jelveh
2025-12-14 15:46:28 -08:00
parent 54dd60d44b
commit a030fcd9b8
4 changed files with 137 additions and 1 deletions

View File

@@ -50,4 +50,7 @@ export const XAI_COST_MAP = {
// Grok 2
'xai:grok-2:prompt_tokens': 200,
'xai:grok-2:completion-tokens': 1000,
};
// Grok Image
'xai:grok-2-image:output': 7_000_000,
};

View File

@@ -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<string, IImageProvider>;
await this.eventService.emit('ai.image.registerProviders', extensionProviders);

View File

@@ -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 <https://www.gnu.org/licenses/>.
*/
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<string> {
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)!;
}
}

View File

@@ -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
},
},
];