Andreas/LibreChat
1
0
Fork 0
mirror of https://github.com/danny-avila/LibreChat.git synced 2026-02-24 19:34:08 +01:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

💎 fix: Gemini Image Gen Tool Vertex AI Auth and File Storage (#11923)

Browse source 
* chore: saveToCloudStorage function and enhance error handling

- Removed unnecessary parameters and streamlined the logic for saving images to cloud storage.
- Introduced buffer handling for base64 image data and improved the integration with file strategy functions.
- Enhanced error handling during local image saving to ensure robustness.
- Updated the createGeminiImageTool function to reflect changes in the saveToCloudStorage implementation.

* refactor: streamline image persistence logic in GeminiImageGen

- Consolidated image saving functionality by renaming and refactoring the saveToCloudStorage function to persistGeneratedImage.
- Improved error handling and logging for image persistence operations.
- Enhanced the replaceUnwantedChars function to better sanitize input strings.
- Updated createGeminiImageTool to reflect changes in image handling and ensure consistent behavior across storage strategies.

* fix: clean up GeminiImageGen by removing unused functions and improving logging

- Removed the getSafeFormat and persistGeneratedImage functions to streamline image handling.
- Updated logging in createGeminiImageTool for clarity and consistency.
- Consolidated imports by eliminating unused dependencies, enhancing code maintainability.

* chore: update environment configuration and manifest for unused GEMINI_VERTEX_ENABLED

- Removed the Vertex AI configuration option from .env.example to simplify setup.
- Updated the manifest.json to reflect the removal of the Vertex AI dependency in the authentication field.
- Cleaned up the createGeminiImageTool function by eliminating unused fields related to Vertex AI, streamlining the code.

* fix: update loadAuthValues call in loadTools function for GeminiImageGen tool

- Modified the loadAuthValues function call to include throwError: false, preventing exceptions on authentication failures.
- Removed the unused processFileURL parameter from the tool context object, streamlining the code.

* refactor: streamline GoogleGenAI initialization in GeminiImageGen

- Removed unused file system access check for Google application credentials, simplifying the environment setup.
- Added googleAuthOptions to the GoogleGenAI instantiation, enhancing the configuration for authentication.

* fix: update Gemini API Key label and description in manifest.json

- Changed the label to indicate that the Gemini API Key is optional.
- Revised the description to clarify usage with Vertex AI and service accounts, enhancing user guidance.

* fix: enhance abort signal handling in createGeminiImageTool

- Introduced derivedSignal to manage abort events during image generation, improving responsiveness to cancellation requests.
- Added an abortHandler to log when image generation is aborted, enhancing debugging capabilities.
- Ensured proper cleanup of event listeners in the finally block to prevent memory leaks.

* fix: update authentication handling for plugins to support optional fields

- Added support for optional authentication fields in the manifest and PluginAuthForm.
- Updated the checkPluginAuth function to correctly validate plugins with optional fields.
- Enhanced tests to cover scenarios with optional authentication fields, ensuring accurate validation logic.
This commit is contained in:
Danny Avila 2026-02-24 08:21:02 -05:00 • committed by GitHub
parent 1d0a4c501f
commit f3eb197675
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 136 additions and 181 deletions
Download patch file Download diff file Expand all files Collapse all files

4
.env.example
View file

@ -243,10 +243,6 @@ GOOGLE_KEY=user_provided
# Option A: Use dedicated Gemini API key for image generation
# GEMINI_API_KEY=your-gemini-api-key
# Option B: Use Vertex AI (no API key needed, uses service account)
# Set this to enable Vertex AI and allow tool without requiring API keys
# GEMINI_VERTEX_ENABLED=true
# Vertex AI model for image generation (defaults to gemini-2.5-flash-image)
# GEMINI_IMAGE_MODEL=gemini-2.5-flash-image

7
api/app/clients/tools/manifest.json
View file

@ -161,9 +161,10 @@
"icon": "assets/gemini_image_gen.svg",
"authConfig": [
{
"authField": "GEMINI_API_KEY||GOOGLE_KEY||GEMINI_VERTEX_ENABLED",
"label": "Gemini API Key (Optional if Vertex AI is configured)",
"description": "Your Google Gemini API Key from <a href='https://aistudio.google.com/app/apikey' target='_blank'>Google AI Studio</a>. Leave blank if using Vertex AI with service account."
"authField": "GEMINI_API_KEY||GOOGLE_KEY||GOOGLE_SERVICE_KEY_FILE",
"label": "Gemini API Key (optional)",
"description": "Your Google Gemini API Key from <a href='https://aistudio.google.com/app/apikey' target='_blank'>Google AI Studio</a>. Leave blank to use Vertex AI with a service account (GOOGLE_SERVICE_KEY_FILE or api/data/auth.json).",
"optional": true
}
]
}

202
api/app/clients/tools/structured/GeminiImageGen.js
View file

@ -1,4 +1,3 @@
const fs = require('fs');
const path = require('path');
const sharp = require('sharp');
const { v4 } = require('uuid');
@ -6,12 +5,7 @@ const { ProxyAgent } = require('undici');
const { GoogleGenAI } = require('@google/genai');
const { tool } = require('@langchain/core/tools');
const { logger } = require('@librechat/data-schemas');
const {
FileContext,
ContentTypes,
FileSources,
EImageOutputType,
} = require('librechat-data-provider');
const { ContentTypes, EImageOutputType } = require('librechat-data-provider');
const {
geminiToolkit,
loadServiceKey,
@ -59,17 +53,12 @@ const displayMessage =
* @returns {string} - The processed string
*/
function replaceUnwantedChars(inputString) {
return inputString?.replace(/[^\w\s\-_.,!?()]/g, '') || '';
}
/**
* Validate and sanitize image format
* @param {string} format - The format to validate
* @returns {string} - Safe format
*/
function getSafeFormat(format) {
const allowedFormats = ['png', 'jpg', 'jpeg', 'webp', 'gif'];
return allowedFormats.includes(format?.toLowerCase()) ? format.toLowerCase() : 'png';
return (
inputString
?.replace(/\r\n|\r|\n/g, ' ')
.replace(/"/g, '')
.trim() || ''
);
}
/**
@ -117,11 +106,8 @@ async function initializeGeminiClient(options = {}) {
return new GoogleGenAI({ apiKey: googleKey });
}
// Fall back to Vertex AI with service account
logger.debug('[GeminiImageGen] Using Vertex AI with service account');
const credentialsPath = getDefaultServiceKeyPath();
// Use loadServiceKey for consistent loading (supports file paths, JSON strings, base64)
const serviceKey = await loadServiceKey(credentialsPath);
if (!serviceKey || !serviceKey.project_id) {
@ -131,75 +117,14 @@ async function initializeGeminiClient(options = {}) {
);
}
// Set GOOGLE_APPLICATION_CREDENTIALS for any Google Cloud SDK dependencies
try {
await fs.promises.access(credentialsPath);
process.env.GOOGLE_APPLICATION_CREDENTIALS = credentialsPath;
} catch {
// File doesn't exist, skip setting env var
}
return new GoogleGenAI({
vertexai: true,
project: serviceKey.project_id,
location: process.env.GOOGLE_LOC || process.env.GOOGLE_CLOUD_LOCATION || 'global',
googleAuthOptions: { credentials: serviceKey },
});
}
/**
* Save image to local filesystem
* @param {string} base64Data - Base64 encoded image data
* @param {string} format - Image format
* @param {string} userId - User ID
* @returns {Promise<string>} - The relative URL
*/
async function saveImageLocally(base64Data, format, userId) {
const safeFormat = getSafeFormat(format);
const safeUserId = userId ? path.basename(userId) : 'default';
const imageName = `gemini-img-${v4()}.${safeFormat}`;
const userDir = path.join(process.cwd(), 'client/public/images', safeUserId);
await fs.promises.mkdir(userDir, { recursive: true });
const filePath = path.join(userDir, imageName);
await fs.promises.writeFile(filePath, Buffer.from(base64Data, 'base64'));
logger.debug('[GeminiImageGen] Image saved locally to:', filePath);
return `/images/${safeUserId}/${imageName}`;
}
/**
* Save image to cloud storage
* @param {Object} params - Parameters
* @returns {Promise<string|null>} - The storage URL or null
*/
async function saveToCloudStorage({ base64Data, format, processFileURL, fileStrategy, userId }) {
if (!processFileURL || !fileStrategy || !userId) {
return null;
}
try {
const safeFormat = getSafeFormat(format);
const safeUserId = path.basename(userId);
const dataURL = `data:image/${safeFormat};base64,${base64Data}`;
const imageName = `gemini-img-${v4()}.${safeFormat}`;
const result = await processFileURL({
URL: dataURL,
basePath: 'images',
userId: safeUserId,
fileName: imageName,
fileStrategy,
context: FileContext.image_generation,
});
return result.filepath;
} catch (error) {
logger.error('[GeminiImageGen] Error saving to cloud storage:', error);
return null;
}
}
/**
* Convert image files to Gemini inline data format
* @param {Object} params - Parameters
@ -390,34 +315,18 @@ function createGeminiImageTool(fields = {}) {
throw new Error('This tool is only available for agents.');
}
// Skip validation during tool creation - validation happens at runtime in initializeGeminiClient
// This allows the tool to be added to agents when using Vertex AI without requiring API keys
// The actual credentials check happens when the tool is invoked
const {
req,
imageFiles = [],
processFileURL,
userId,
fileStrategy,
GEMINI_API_KEY,
GOOGLE_KEY,
// GEMINI_VERTEX_ENABLED is used for auth validation only (not used in code)
// When set as env var, it signals Vertex AI is configured and bypasses API key requirement
} = fields;
const { req, imageFiles = [], userId, fileStrategy, GEMINI_API_KEY, GOOGLE_KEY } = fields;
const imageOutputType = fields.imageOutputType || EImageOutputType.PNG;
const geminiImageGenTool = tool(
async ({ prompt, image_ids, aspectRatio, imageSize }, _runnableConfig) => {
async ({ prompt, image_ids, aspectRatio, imageSize }, runnableConfig) => {
if (!prompt) {
throw new Error('Missing required field: prompt');
}
logger.debug('[GeminiImageGen] Generating image with prompt:', prompt?.substring(0, 100));
logger.debug('[GeminiImageGen] Options:', { aspectRatio, imageSize });
logger.debug('[GeminiImageGen] Generating image', { aspectRatio, imageSize });
// Initialize Gemini client with user-provided credentials
let ai;
try {
ai = await initializeGeminiClient({
@ -432,10 +341,8 @@ function createGeminiImageTool(fields = {}) {
];
}
// Build request contents
const contents = [{ text: replaceUnwantedChars(prompt) }];
// Add context images if provided
if (image_ids?.length > 0) {
const contextImages = await convertImagesToInlineData({
imageFiles,
@ -447,28 +354,34 @@ function createGeminiImageTool(fields = {}) {
logger.debug('[GeminiImageGen] Added', contextImages.length, 'context images');
}
// Generate image
let apiResponse;
const geminiModel = process.env.GEMINI_IMAGE_MODEL || 'gemini-2.5-flash-image';
try {
// Build config with optional imageConfig
const config = {
responseModalities: ['TEXT', 'IMAGE'],
};
const config = {
responseModalities: ['TEXT', 'IMAGE'],
};
// Add imageConfig if aspectRatio or imageSize is specified
// Note: gemini-2.5-flash-image doesn't support imageSize
const supportsImageSize = !geminiModel.includes('gemini-2.5-flash-image');
if (aspectRatio || (imageSize && supportsImageSize)) {
config.imageConfig = {};
if (aspectRatio) {
config.imageConfig.aspectRatio = aspectRatio;
}
if (imageSize && supportsImageSize) {
config.imageConfig.imageSize = imageSize;
}
const supportsImageSize = !geminiModel.includes('gemini-2.5-flash-image');
if (aspectRatio || (imageSize && supportsImageSize)) {
config.imageConfig = {};
if (aspectRatio) {
config.imageConfig.aspectRatio = aspectRatio;
}
if (imageSize && supportsImageSize) {
config.imageConfig.imageSize = imageSize;
}
}
let derivedSignal = null;
let abortHandler = null;
if (runnableConfig?.signal) {
derivedSignal = AbortSignal.any([runnableConfig.signal]);
abortHandler = () => logger.debug('[GeminiImageGen] Image generation aborted');
derivedSignal.addEventListener('abort', abortHandler, { once: true });
config.abortSignal = derivedSignal;
}
try {
apiResponse = await ai.models.generateContent({
model: geminiModel,
contents,
@ -480,9 +393,12 @@ function createGeminiImageTool(fields = {}) {
[{ type: ContentTypes.TEXT, text: `Image generation failed: ${error.message}` }],
{ content: [], file_ids: [] },
];
} finally {
if (abortHandler && derivedSignal) {
derivedSignal.removeEventListener('abort', abortHandler);
}
}
// Check for safety blocks
const safetyBlock = checkForSafetyBlock(apiResponse);
if (safetyBlock) {
logger.warn('[GeminiImageGen] Safety block:', safetyBlock);
@ -509,46 +425,7 @@ function createGeminiImageTool(fields = {}) {
const imageData = convertedBuffer.toString('base64');
const mimeType = outputFormat === 'jpeg' ? 'image/jpeg' : `image/${outputFormat}`;
logger.debug('[GeminiImageGen] Image format:', { outputFormat, mimeType });
let imageUrl;
const useLocalStorage = !fileStrategy || fileStrategy === FileSources.local;
if (useLocalStorage) {
try {
imageUrl = await saveImageLocally(imageData, outputFormat, userId);
} catch (error) {
logger.error('[GeminiImageGen] Local save failed:', error);
imageUrl = `data:${mimeType};base64,${imageData}`;
}
} else {
const cloudUrl = await saveToCloudStorage({
base64Data: imageData,
format: outputFormat,
processFileURL,
fileStrategy,
userId,
});
if (cloudUrl) {
imageUrl = cloudUrl;
} else {
// Fallback to local
try {
imageUrl = await saveImageLocally(imageData, outputFormat, userId);
} catch (_error) {
imageUrl = `data:${mimeType};base64,${imageData}`;
}
}
}
logger.debug('[GeminiImageGen] Image URL:', imageUrl);
// For the artifact, we need a data URL (same as OpenAI)
// The local file save is for persistence, but the response needs a data URL
const dataUrl = `data:${mimeType};base64,${imageData}`;
// Return in content_and_artifact format (same as OpenAI)
const file_ids = [v4()];
const content = [
{
@ -567,8 +444,7 @@ function createGeminiImageTool(fields = {}) {
},
];
// Record token usage for balance tracking (don't await to avoid blocking response)
const conversationId = _runnableConfig?.configurable?.thread_id;
const conversationId = runnableConfig?.configurable?.thread_id;
recordTokenUsage({
usageMetadata: apiResponse.usageMetadata,
req,

3
api/app/clients/tools/util/handleTools.js
View file

@ -207,7 +207,7 @@ const loadTools = async ({
},
gemini_image_gen: async (toolContextMap) => {
const authFields = getAuthFields('gemini_image_gen');
const authValues = await loadAuthValues({ userId: user, authFields });
const authValues = await loadAuthValues({ userId: user, authFields, throwError: false });
const imageFiles = options.tool_resources?.[EToolResources.image_edit]?.files ?? [];
const toolContext = buildImageToolContext({
imageFiles,
@ -222,7 +222,6 @@ const loadTools = async ({
isAgent: !!agent,
req: options.req,
imageFiles,
processFileURL: options.processFileURL,
userId: user,
fileStrategy,
});

25
client/src/components/Plugins/Store/PluginAuthForm.tsx
View file

@ -20,6 +20,7 @@ function PluginAuthForm({ plugin, onSubmit, isEntityTool }: TPluginAuthFormProps
const localize = useLocalize();
const authConfig = plugin?.authConfig ?? [];
const allFieldsOptional = authConfig.length > 0 && authConfig.every((c) => c.optional === true);
return (
<div className="flex w-full flex-col items-center gap-2">
@ -38,6 +39,7 @@ function PluginAuthForm({ plugin, onSubmit, isEntityTool }: TPluginAuthFormProps
>
{authConfig.map((config: TPluginAuthConfig, i: number) => {
const authField = config.authField.split('||')[0];
const isOptional = config.optional === true;
return (
<div key={`${authField}-${i}`} className="flex w-full flex-col gap-1">
<label
@ -55,19 +57,24 @@ function PluginAuthForm({ plugin, onSubmit, isEntityTool }: TPluginAuthFormProps
aria-invalid={!!errors[authField]}
aria-describedby={`${authField}-error`}
aria-label={config.label}
aria-required="true"
aria-required={!isOptional}
/* autoFocus is generally disabled due to the fact that it can disorient users,
* but in this case, the required field must be navigated to anyways, and the component's functionality
* emulates that of a new modal opening, where users would expect focus to be shifted to the new content */
// eslint-disable-next-line jsx-a11y/no-autofocus
autoFocus={i === 0}
{...register(authField, {
required: `${config.label} is required.`,
minLength: {
value: 1,
message: `${config.label} must be at least 1 character long`,
},
})}
{...register(
authField,
isOptional
? {}
: {
required: `${config.label} is required.`,
minLength: {
value: 1,
message: `${config.label} must be at least 1 character long`,
},
},
)}
className="flex h-10 max-h-10 w-full resize-none rounded-md border border-gray-200 bg-transparent px-3 py-2 text-sm text-gray-700 shadow-[0_0_10px_rgba(0,0,0,0.05)] outline-none placeholder:text-gray-400 focus:border-gray-400 focus:bg-gray-50 focus:outline-none focus:ring-0 focus:ring-gray-400 focus:ring-opacity-0 focus:ring-offset-2 disabled:cursor-not-allowed disabled:opacity-50 dark:border-gray-500 dark:bg-gray-700 dark:text-gray-50 dark:shadow-[0_0_15px_rgba(0,0,0,0.10)] dark:focus:border-gray-400 focus:dark:bg-gray-600 dark:focus:outline-none dark:focus:ring-0 dark:focus:ring-gray-400 dark:focus:ring-offset-0"
/>
</HoverCardTrigger>
@ -82,7 +89,7 @@ function PluginAuthForm({ plugin, onSubmit, isEntityTool }: TPluginAuthFormProps
);
})}
<button
disabled={!isDirty || !isValid || isSubmitting}
disabled={allFieldsOptional ? isSubmitting : !isDirty || !isValid || isSubmitting}
type="button"
className="btn btn-primary relative"
onClick={() => {

71
packages/api/src/tools/format.spec.ts
View file

@ -189,6 +189,77 @@ describe('format.ts helper functions', () => {
const result = checkPluginAuth(plugin);
expect(result).toBe(true);
});
it('should return true when auth field is marked as optional with no env vars', () => {
const plugin: TPlugin = {
name: 'Test',
pluginKey: 'test',
description: 'Test plugin',
authConfig: [
{ authField: 'MISSING_KEY', label: 'API Key', description: 'API Key', optional: true },
],
};
const result = checkPluginAuth(plugin);
expect(result).toBe(true);
});
it('should return true when all auth fields are optional', () => {
const plugin: TPlugin = {
name: 'Test',
pluginKey: 'test',
description: 'Test plugin',
authConfig: [
{ authField: 'MISSING_KEY_A', label: 'Key A', description: 'Key A', optional: true },
{ authField: 'MISSING_KEY_B', label: 'Key B', description: 'Key B', optional: true },
],
};
const result = checkPluginAuth(plugin);
expect(result).toBe(true);
});
it('should return false when required field is missing even if optional field exists', () => {
const plugin: TPlugin = {
name: 'Test',
pluginKey: 'test',
description: 'Test plugin',
authConfig: [
{ authField: 'MISSING_KEY', label: 'Required Key', description: 'Required' },
{
authField: 'OPTIONAL_KEY',
label: 'Optional Key',
description: 'Optional',
optional: true,
},
],
};
const result = checkPluginAuth(plugin);
expect(result).toBe(false);
});
it('should return true when optional field has no env var but required field is satisfied', () => {
process.env.REQUIRED_KEY = 'valid-key';
const plugin: TPlugin = {
name: 'Test',
pluginKey: 'test',
description: 'Test plugin',
authConfig: [
{ authField: 'REQUIRED_KEY', label: 'Required Key', description: 'Required' },
{
authField: 'OPTIONAL_KEY',
label: 'Optional Key',
description: 'Optional',
optional: true,
},
],
};
const result = checkPluginAuth(plugin);
expect(result).toBe(true);
});
});
describe('getToolkitKey', () => {

4
packages/api/src/tools/format.ts
View file

@ -31,6 +31,10 @@ export const checkPluginAuth = (plugin?: TPlugin): boolean => {
}
return plugin.authConfig.every((authFieldObj) => {
if (authFieldObj.optional === true) {
return true;
}
const authFieldOptions = authFieldObj.authField.split('||');
let isFieldAuthenticated = false;

1
packages/data-provider/src/schemas.ts
View file

@ -544,6 +544,7 @@ export const tPluginAuthConfigSchema = z.object({
authField: z.string(),
label: z.string(),
description: z.string(),
optional: z.boolean().optional(),
});
export type TPluginAuthConfig = z.infer<typeof tPluginAuthConfigSchema>;