Files
Gemini_Draw/server/index.ts

459 lines
13 KiB
TypeScript

import express from 'express';
import fs from 'node:fs';
import path from 'node:path';
import dotenv from 'dotenv';
import multer from 'multer';
import { GoogleGenAI } from '@google/genai';
dotenv.config({ path: '.env.local' });
dotenv.config();
type InlineInput = {
name?: string;
mimeType?: string;
data?: string;
base64?: string;
dataUrl?: string;
url?: string;
};
type GenerateRequest = {
apiKey?: string;
prompt?: string;
instruction?: string;
model?: string;
imageSize?: '1K' | '2K' | '4K';
aspectRatio?: '1:1' | '4:3' | '3:4' | '16:9' | '9:16';
images?: InlineInput[];
documents?: InlineInput[];
files?: InlineInput[];
};
type GeminiPart = {
text?: string;
inlineData?: {
data: string;
mimeType: string;
};
};
const app = express();
const upload = multer({
storage: multer.memoryStorage(),
limits: {
fileSize: Number(process.env.API_MAX_FILE_MB || 25) * 1024 * 1024,
files: Number(process.env.API_MAX_FILES || 20),
},
});
let runtimeApiKey = process.env.GEMINI_API_KEY || process.env.API_KEY || '';
const apiPort = Number(process.env.API_PORT || 3002);
const defaultImageModel = process.env.GEMINI_IMAGE_MODEL || 'gemini-3.1-flash-image-preview';
const defaultTextModel = process.env.GEMINI_TEXT_MODEL || 'gemini-2.5-flash';
const apiAuthToken = process.env.API_AUTH_TOKEN || '';
const apiAuthDisabled = process.env.API_AUTH_DISABLED === 'true';
if (!runtimeApiKey) {
console.warn('GEMINI_API_KEY/API_KEY is not set. API calls will fail until it is configured.');
}
app.use((req, res, next) => {
res.setHeader('Access-Control-Allow-Origin', process.env.API_CORS_ORIGIN || '*');
res.setHeader('Access-Control-Allow-Methods', 'GET,POST,OPTIONS');
res.setHeader('Access-Control-Allow-Headers', 'Content-Type,Authorization,x-api-key,x-gemini-api-key');
if (req.method === 'OPTIONS') {
res.status(204).end();
return;
}
next();
});
app.use(express.json({ limit: process.env.API_JSON_LIMIT || '50mb' }));
function requireAuth(req: express.Request, res: express.Response, next: express.NextFunction) {
if (apiAuthDisabled) {
next();
return;
}
if (!apiAuthToken) {
res.status(503).json({
ok: false,
error: 'API_AUTH_TOKEN is required. Set it in .env.local, or set API_AUTH_DISABLED=true for local-only development.',
});
return;
}
const authHeader = req.header('authorization') || '';
const bearerToken = authHeader.startsWith('Bearer ') ? authHeader.slice(7) : '';
const headerToken = req.header('x-api-key') || '';
if (bearerToken === apiAuthToken || headerToken === apiAuthToken) {
next();
return;
}
res.status(401).json({
ok: false,
error: 'Unauthorized. Send Authorization: Bearer YOUR_API_AUTH_TOKEN or x-api-key.',
});
}
function getRequestApiKey(req: express.Request, body?: GenerateRequest) {
const headerKey = req.header('x-gemini-api-key') || '';
return body?.apiKey || headerKey || runtimeApiKey;
}
function persistApiKeyToEnvLocal(apiKey: string) {
const envPath = path.resolve(process.cwd(), '.env.local');
let content = '';
if (fs.existsSync(envPath)) {
content = fs.readFileSync(envPath, 'utf8');
}
const escaped = apiKey.replace(/\\/g, '\\\\').replace(/"/g, '\\"');
const line = `GEMINI_API_KEY="${escaped}"`;
if (/^GEMINI_API_KEY=.*$/m.test(content)) {
content = content.replace(/^GEMINI_API_KEY=.*$/m, line);
} else {
content = content ? `${content.trimEnd()}\n${line}\n` : `${line}\n`;
}
fs.writeFileSync(envPath, content, 'utf8');
}
function stripDataPrefix(value: string) {
const match = value.match(/^data:([^;]+);base64,(.+)$/);
if (!match) {
return null;
}
return {
mimeType: match[1],
data: match[2],
};
}
async function fetchAsBase64(url: string) {
const response = await fetch(url);
if (!response.ok) {
throw new Error(`Failed to fetch ${url}: ${response.status} ${response.statusText}`);
}
const contentLength = Number(response.headers.get('content-length') || 0);
const maxBytes = Number(process.env.API_MAX_URL_MB || 25) * 1024 * 1024;
if (contentLength > maxBytes) {
throw new Error(`Remote file is too large. Max size is ${process.env.API_MAX_URL_MB || 25}MB.`);
}
const buffer = Buffer.from(await response.arrayBuffer());
if (buffer.length > maxBytes) {
throw new Error(`Remote file is too large. Max size is ${process.env.API_MAX_URL_MB || 25}MB.`);
}
return {
data: buffer.toString('base64'),
mimeType: response.headers.get('content-type')?.split(';')[0] || 'application/octet-stream',
};
}
async function inlineInputToPart(input: InlineInput): Promise<GeminiPart> {
if (input.dataUrl) {
const parsed = stripDataPrefix(input.dataUrl);
if (!parsed) {
throw new Error(`Invalid dataUrl${input.name ? ` for ${input.name}` : ''}.`);
}
return {
inlineData: {
data: parsed.data,
mimeType: input.mimeType || parsed.mimeType,
},
};
}
if (input.url) {
const fetched = await fetchAsBase64(input.url);
return {
inlineData: {
data: fetched.data,
mimeType: input.mimeType || fetched.mimeType,
},
};
}
const rawData = input.base64 || input.data;
if (!rawData) {
throw new Error(`File input${input.name ? ` ${input.name}` : ''} is missing data, base64, dataUrl, or url.`);
}
const parsed = stripDataPrefix(rawData);
return {
inlineData: {
data: parsed?.data || rawData,
mimeType: input.mimeType || parsed?.mimeType || 'application/octet-stream',
},
};
}
function uploadedFileToPart(file: Express.Multer.File): GeminiPart {
return {
inlineData: {
data: file.buffer.toString('base64'),
mimeType: inferMimeType(file),
},
};
}
function normalizeTextInput(value: unknown) {
if (typeof value === 'string') {
return value;
}
if (Array.isArray(value)) {
return value.map((item) => normalizeTextInput(item)).join('\n');
}
if (value && typeof value === 'object' && 'value' in value) {
return normalizeTextInput((value as { value: unknown }).value);
}
return value == null ? '' : String(value);
}
function inferMimeType(file: Express.Multer.File) {
if (file.mimetype && file.mimetype !== 'application/octet-stream') {
return file.mimetype;
}
const ext = path.extname(file.originalname || '').toLowerCase();
const mimeByExt: Record<string, string> = {
'.png': 'image/png',
'.jpg': 'image/jpeg',
'.jpeg': 'image/jpeg',
'.webp': 'image/webp',
'.gif': 'image/gif',
'.pdf': 'application/pdf',
'.txt': 'text/plain',
'.md': 'text/plain',
'.csv': 'text/csv',
'.json': 'application/json',
'.html': 'text/html',
'.htm': 'text/html',
};
return mimeByExt[ext] || 'application/octet-stream';
}
async function buildParts(body: GenerateRequest, files: Express.Multer.File[] = []) {
const prompt = normalizeTextInput(body.prompt || body.instruction).trim();
if (!prompt) {
throw new Error('prompt or instruction is required.');
}
const inlineInputs = [
...(body.images || []),
...(body.documents || []),
...(body.files || []),
];
const parts: GeminiPart[] = [];
for (const input of inlineInputs) {
parts.push(await inlineInputToPart(input));
}
for (const file of files) {
parts.push(uploadedFileToPart(file));
}
parts.push({ text: prompt });
return parts;
}
function normalizeMultipartBody(req: express.Request): GenerateRequest {
return {
apiKey: req.body.apiKey,
prompt: String(req.body.prompt || req.body.instruction || ''),
model: req.body.model,
imageSize: req.body.imageSize,
aspectRatio: req.body.aspectRatio,
};
}
function collectResponseParts(response: any) {
const images: Array<{ mimeType: string; data: string; dataUrl: string }> = [];
const texts: string[] = [];
for (const candidate of response.candidates || []) {
for (const part of candidate.content?.parts || []) {
if (part.inlineData?.data) {
const mimeType = part.inlineData.mimeType || 'image/png';
images.push({
mimeType,
data: part.inlineData.data,
dataUrl: `data:${mimeType};base64,${part.inlineData.data}`,
});
}
if (part.text) {
texts.push(part.text);
}
}
}
return { images, texts };
}
async function runGemini(req: express.Request, body: GenerateRequest, files: Express.Multer.File[] = [], textOnly = false) {
const parts = await buildParts(body, files);
const requestApiKey = getRequestApiKey(req, body);
if (!requestApiKey) {
throw new Error('Gemini API key is required. Set GEMINI_API_KEY, call POST /api/config/api-key, or send x-gemini-api-key.');
}
const ai = new GoogleGenAI({ apiKey: requestApiKey });
const model = body.model || (textOnly ? defaultTextModel : defaultImageModel);
const config: Record<string, unknown> = {};
if (!textOnly) {
config.imageConfig = {
imageSize: body.imageSize || '1K',
aspectRatio: body.aspectRatio || '1:1',
};
}
const response = await ai.models.generateContent({
model,
contents: { parts } as any,
config,
} as any);
const { images, texts } = collectResponseParts(response);
return {
ok: true,
model,
usedRequestApiKey: Boolean(body.apiKey || req.header('x-gemini-api-key')),
input: {
fileCount: files.length + (body.images?.length || 0) + (body.documents?.length || 0) + (body.files?.length || 0),
partCount: parts.length,
imageSize: textOnly ? undefined : body.imageSize || '1K',
aspectRatio: textOnly ? undefined : body.aspectRatio || '1:1',
},
images,
texts,
usageMetadata: response.usageMetadata,
};
}
function handleError(res: express.Response, error: unknown) {
const message = error instanceof Error ? error.message : String(error);
res.status(400).json({
ok: false,
error: message,
});
}
app.get('/api', (_req, res) => {
res.json({
ok: true,
name: 'Gemini Draw API',
endpoints: [
'GET /api/health',
'GET /api/config',
'POST /api/config/api-key',
'POST /api/generate',
'POST /api/generate/upload',
'POST /api/edit-image',
'POST /api/analyze-document',
],
});
});
app.get('/api/health', (_req, res) => {
res.json({
ok: true,
apiPort,
hasGeminiApiKey: Boolean(runtimeApiKey),
authEnabled: Boolean(apiAuthToken) && !apiAuthDisabled,
authRequired: !apiAuthDisabled,
acceptsPerRequestApiKey: true,
defaultImageModel,
defaultTextModel,
});
});
app.get('/api/config', requireAuth, (_req, res) => {
res.json({
ok: true,
apiPort,
hasGeminiApiKey: Boolean(runtimeApiKey),
apiKeyPreview: runtimeApiKey ? `${runtimeApiKey.slice(0, 6)}...${runtimeApiKey.slice(-4)}` : '',
authEnabled: Boolean(apiAuthToken) && !apiAuthDisabled,
authRequired: !apiAuthDisabled,
defaultImageModel,
defaultTextModel,
});
});
app.post('/api/config/api-key', requireAuth, async (req, res) => {
try {
const nextApiKey = String(req.body.apiKey || '').trim();
const persist = Boolean(req.body.persist);
if (!nextApiKey) {
throw new Error('apiKey is required.');
}
runtimeApiKey = nextApiKey;
process.env.GEMINI_API_KEY = nextApiKey;
if (persist) {
persistApiKeyToEnvLocal(nextApiKey);
}
res.json({
ok: true,
persisted: persist,
apiKeyPreview: `${runtimeApiKey.slice(0, 6)}...${runtimeApiKey.slice(-4)}`,
});
} catch (error) {
handleError(res, error);
}
});
app.post('/api/generate', requireAuth, async (req, res) => {
try {
res.json(await runGemini(req, req.body));
} catch (error) {
handleError(res, error);
}
});
app.post('/api/generate/upload', requireAuth, upload.any(), async (req, res) => {
try {
res.json(await runGemini(req, normalizeMultipartBody(req), (req.files as Express.Multer.File[]) || []));
} catch (error) {
handleError(res, error);
}
});
app.post('/api/edit-image', requireAuth, upload.any(), async (req, res) => {
try {
const files = (req.files as Express.Multer.File[]) || [];
const body = req.is('multipart/form-data') ? normalizeMultipartBody(req) : req.body;
res.json(await runGemini(req, body, files));
} catch (error) {
handleError(res, error);
}
});
app.post('/api/analyze-document', requireAuth, upload.any(), async (req, res) => {
try {
const files = (req.files as Express.Multer.File[]) || [];
const body = req.is('multipart/form-data') ? normalizeMultipartBody(req) : req.body;
res.json(await runGemini(req, body, files, true));
} catch (error) {
handleError(res, error);
}
});
app.listen(apiPort, '0.0.0.0', () => {
console.log(`Gemini Draw API listening on http://localhost:${apiPort}`);
});