feat(API): 集成 Gemini 图片生成能力
- 添加 Gemini 模型类型检测和初始化 - 实现图片生成请求处理逻辑 - 支持流式返回生成的图片数据 - 将生成的图片保存到数据库
This commit is contained in:
gaoziman
parent
4c43fb4471
commit
c72b4ce3e2
@ -1,6 +1,6 @@
|
||||
import { NextResponse } from 'next/server';
|
||||
import { db } from '@/drizzle/db';
|
||||
import { conversations, messages, userSettings } from '@/drizzle/schema';
|
||||
import { conversations, messages, userSettings, type GeneratedImageData } from '@/drizzle/schema';
|
||||
import { eq } from 'drizzle-orm';
|
||||
import { nanoid } from 'nanoid';
|
||||
import { executeTool } from '@/services/tools';
|
||||
@ -94,6 +94,19 @@ function isCodexModel(modelId: string): boolean {
|
||||
return modelId.startsWith('gpt-') && modelId.includes('codex');
|
||||
}
|
||||
|
||||
// 判断是否为 Gemini 图片生成模型
|
||||
function isGeminiImageModel(modelId: string): boolean {
|
||||
// 支持的图片生成模型列表
|
||||
const imageModels = [
|
||||
'gemini-2.0-flash-preview-image-generation',
|
||||
'gemini-3-pro-image-preview',
|
||||
'imagen-3.0-generate-002',
|
||||
];
|
||||
return imageModels.some(model => modelId.includes(model)) ||
|
||||
modelId.includes('image-generation') ||
|
||||
modelId.includes('imagen');
|
||||
}
|
||||
|
||||
// 规范化 URL（移除末尾斜杠，避免拼接时出现双斜杠）
|
||||
function normalizeBaseUrl(url: string): string {
|
||||
return url.replace(/\/+$/, '');
|
||||
@ -359,6 +372,7 @@ export async function POST(request: Request) {
|
||||
// 判断使用的模型类型
|
||||
const useModel = model || conversation.model;
|
||||
const isCodex = isCodexModel(useModel);
|
||||
const isGeminiImage = isGeminiImageModel(useModel);
|
||||
|
||||
// 创建 SSE 响应
|
||||
const encoder = new TextEncoder();
|
||||
@ -420,6 +434,7 @@ export async function POST(request: Request) {
|
||||
let totalInputTokens = 0;
|
||||
let totalOutputTokens = 0;
|
||||
let usedTools: string[] = []; // 收集使用过的工具名称
|
||||
let generatedImages: GeneratedImageData[] = []; // Gemini 生成的图片
|
||||
|
||||
// 如果有文档解析失败，将警告添加到内容开头
|
||||
if (documentParseErrors.length > 0) {
|
||||
@ -433,13 +448,35 @@ export async function POST(request: Request) {
|
||||
}
|
||||
|
||||
// 【重要】处理器选择优先级说明：
|
||||
// 1. 首先检查 apiFormat === 'openai'：如果用户选择了 "OpenAI 兼容" 格式，
|
||||
// 0. 首先检查 isGeminiImage：Gemini 图片生成模型有专门的处理逻辑
|
||||
// 1. 然后检查 apiFormat === 'openai'：如果用户选择了 "OpenAI 兼容" 格式，
|
||||
// 则所有模型（包括 Codex 模型）都统一使用 /v1/chat/completions 端点
|
||||
// 这是因为第三方中转站通常只支持 OpenAI 兼容的 /v1/chat/completions 端点
|
||||
// 2. 然后检查 isCodex：如果是 Claude 原生格式 + Codex 模型，才使用 /v1/responses 端点
|
||||
// 3. 最后是普通的 Claude 原生格式，使用 /v1/messages 端点
|
||||
|
||||
if (apiFormat === 'openai') {
|
||||
if (isGeminiImage) {
|
||||
// ==================== Gemini 图片生成模型处理 ====================
|
||||
// Gemini 图片生成模型使用专门的 Gemini API 端点
|
||||
console.log('[API/chat] 使用 Gemini 图片生成模型:', useModel);
|
||||
const result = await handleGeminiImageChat({
|
||||
cchUrl,
|
||||
apiKey: decryptedApiKey,
|
||||
model: useModel,
|
||||
systemPrompt,
|
||||
temperature,
|
||||
historyMessages,
|
||||
message,
|
||||
controller,
|
||||
encoder,
|
||||
images,
|
||||
});
|
||||
|
||||
fullContent = result.fullContent;
|
||||
generatedImages = result.generatedImages;
|
||||
totalInputTokens = result.inputTokens;
|
||||
totalOutputTokens = result.outputTokens;
|
||||
} else if (apiFormat === 'openai') {
|
||||
// ==================== OpenAI 兼容格式处理 ====================
|
||||
// 当用户选择 "OpenAI 兼容" 时，无论什么模型都走这个分支
|
||||
// 第三方中转站统一使用 /v1/chat/completions 端点
|
||||
@ -528,6 +565,7 @@ export async function POST(request: Request) {
|
||||
content: fullContent,
|
||||
thinkingContent: thinkingContent || null,
|
||||
usedTools: usedTools.length > 0 ? usedTools : null,
|
||||
generatedImages: generatedImages.length > 0 ? generatedImages : null,
|
||||
inputTokens: totalInputTokens,
|
||||
outputTokens: totalOutputTokens,
|
||||
status: 'completed',
|
||||
@ -1974,6 +2012,203 @@ async function handleOpenAICompatibleChat(params: OpenAICompatibleChatParams): P
|
||||
};
|
||||
}
|
||||
|
||||
// ==================== Gemini 图片生成模型处理函数 ====================
|
||||
interface GeminiImageChatParams {
|
||||
cchUrl: string;
|
||||
apiKey: string;
|
||||
model: string;
|
||||
systemPrompt: string;
|
||||
temperature: number;
|
||||
historyMessages: { role: string; content: string }[];
|
||||
message: string;
|
||||
controller: ReadableStreamDefaultController;
|
||||
encoder: TextEncoder;
|
||||
// 用户上传的图片
|
||||
images?: {
|
||||
type: 'image';
|
||||
media_type: string;
|
||||
data: string;
|
||||
}[];
|
||||
}
|
||||
|
||||
async function handleGeminiImageChat(params: GeminiImageChatParams): Promise<{
|
||||
fullContent: string;
|
||||
generatedImages: GeneratedImageData[];
|
||||
inputTokens: number;
|
||||
outputTokens: number;
|
||||
}> {
|
||||
const {
|
||||
cchUrl,
|
||||
apiKey,
|
||||
model,
|
||||
systemPrompt,
|
||||
temperature,
|
||||
historyMessages,
|
||||
message,
|
||||
controller,
|
||||
encoder,
|
||||
images,
|
||||
} = params;
|
||||
|
||||
// 创建安全的 stream 写入器
|
||||
const safeWriter = createSafeStreamWriter(controller, encoder);
|
||||
|
||||
// 发送生成开始事件
|
||||
safeWriter.write({
|
||||
type: 'image_generation_start',
|
||||
model,
|
||||
});
|
||||
|
||||
// 构建 Gemini API 请求内容
|
||||
// Gemini 使用 contents 数组格式
|
||||
const contents: Array<{
|
||||
role: 'user' | 'model';
|
||||
parts: Array<{ text?: string; inlineData?: { mimeType: string; data: string } }>;
|
||||
}> = [];
|
||||
|
||||
// 添加历史消息
|
||||
for (const msg of historyMessages) {
|
||||
if (msg.content && msg.content.trim() !== '') {
|
||||
contents.push({
|
||||
role: msg.role === 'assistant' ? 'model' : 'user',
|
||||
parts: [{ text: msg.content }],
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// 构建当前用户消息（支持多模态）
|
||||
const currentParts: Array<{ text?: string; inlineData?: { mimeType: string; data: string } }> = [];
|
||||
|
||||
// 如果有图片，先添加图片
|
||||
if (images && images.length > 0) {
|
||||
for (const img of images) {
|
||||
currentParts.push({
|
||||
inlineData: {
|
||||
mimeType: img.media_type,
|
||||
data: img.data,
|
||||
},
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// 添加文本消息
|
||||
if (message) {
|
||||
currentParts.push({ text: message });
|
||||
}
|
||||
|
||||
contents.push({
|
||||
role: 'user',
|
||||
parts: currentParts,
|
||||
});
|
||||
|
||||
// 构建请求体
|
||||
const requestBody = {
|
||||
contents,
|
||||
systemInstruction: systemPrompt ? { parts: [{ text: systemPrompt }] } : undefined,
|
||||
generationConfig: {
|
||||
temperature,
|
||||
responseModalities: ['TEXT', 'IMAGE'], // 请求同时返回文本和图片
|
||||
},
|
||||
};
|
||||
|
||||
console.log('[handleGeminiImageChat] Sending request to Gemini API:', {
|
||||
model,
|
||||
contentsCount: contents.length,
|
||||
hasImages: images && images.length > 0,
|
||||
});
|
||||
|
||||
// 发送请求到 Gemini API
|
||||
// Gemini API 端点格式: /v1beta/models/{model}:generateContent
|
||||
const geminiEndpoint = `${normalizeBaseUrl(cchUrl)}/v1beta/models/${model}:generateContent`;
|
||||
|
||||
const response = await fetch(geminiEndpoint, {
|
||||
method: 'POST',
|
||||
headers: {
|
||||
'Content-Type': 'application/json',
|
||||
'x-goog-api-key': apiKey,
|
||||
},
|
||||
body: JSON.stringify(requestBody),
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
const errorText = await response.text();
|
||||
throw new Error(`Gemini API error: ${response.status} - ${errorText}`);
|
||||
}
|
||||
|
||||
// 解析 Gemini 响应
|
||||
const geminiResponse = await response.json();
|
||||
|
||||
let fullContent = '';
|
||||
const generatedImages: GeneratedImageData[] = [];
|
||||
let totalInputTokens = 0;
|
||||
let totalOutputTokens = 0;
|
||||
|
||||
// 提取 usage 信息
|
||||
if (geminiResponse.usageMetadata) {
|
||||
totalInputTokens = geminiResponse.usageMetadata.promptTokenCount || 0;
|
||||
totalOutputTokens = geminiResponse.usageMetadata.candidatesTokenCount || 0;
|
||||
}
|
||||
|
||||
// 解析候选响应
|
||||
if (geminiResponse.candidates && geminiResponse.candidates.length > 0) {
|
||||
const candidate = geminiResponse.candidates[0];
|
||||
|
||||
if (candidate.content && candidate.content.parts) {
|
||||
for (const part of candidate.content.parts) {
|
||||
// 处理文本内容
|
||||
if (part.text) {
|
||||
fullContent += part.text;
|
||||
safeWriter.write({
|
||||
type: 'text',
|
||||
content: part.text,
|
||||
});
|
||||
}
|
||||
|
||||
// 处理生成的图片
|
||||
if (part.inlineData) {
|
||||
const imageData: GeneratedImageData = {
|
||||
mimeType: part.inlineData.mimeType,
|
||||
data: part.inlineData.data,
|
||||
};
|
||||
generatedImages.push(imageData);
|
||||
|
||||
// 发送生成图片事件
|
||||
safeWriter.write({
|
||||
type: 'generated_image',
|
||||
image: imageData,
|
||||
index: generatedImages.length - 1,
|
||||
});
|
||||
|
||||
console.log('[handleGeminiImageChat] Generated image:', {
|
||||
mimeType: imageData.mimeType,
|
||||
dataLength: imageData.data.length,
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// 发送生成完成事件
|
||||
safeWriter.write({
|
||||
type: 'image_generation_complete',
|
||||
imageCount: generatedImages.length,
|
||||
});
|
||||
|
||||
console.log('[handleGeminiImageChat] Response processed:', {
|
||||
textLength: fullContent.length,
|
||||
imageCount: generatedImages.length,
|
||||
inputTokens: totalInputTokens,
|
||||
outputTokens: totalOutputTokens,
|
||||
});
|
||||
|
||||
return {
|
||||
fullContent,
|
||||
generatedImages,
|
||||
inputTokens: totalInputTokens,
|
||||
outputTokens: totalOutputTokens,
|
||||
};
|
||||
}
|
||||
|
||||
// 构建 Claude 工具定义
|
||||
function buildClaudeToolDefinitions(toolIds: string[]) {
|
||||
const toolMap: Record<string, object> = {
|
||||
|
||||
Loading…
Reference in New Issue
Block a user