mirror of
https://github.com/HeyPuter/puter.git
synced 2026-01-09 06:30:48 -06:00
Add xAI image generation provider integration
This commit is contained in:
@@ -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,
|
||||
};
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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)!;
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
},
|
||||
},
|
||||
];
|
||||
Reference in New Issue
Block a user