更新生图功能

.gitignore

@@ -17,3 +17,5 @@ logs/
 Thumbs.db
 
 data/
+test/*.png
+test/*.jpeg

src/api/client.js

@@ -59,7 +59,7 @@ function buildRequesterConfig(headers, body = null) {
 
 // 统一错误处理
 async function handleApiError(error, token) {
-  const status = error.response?.status || error.status;
+  const status = error.response?.status || error.status || 'Unknown';
   let errorBody = error.message;
   
   if (error.response?.data?.readable) {
@@ -160,8 +160,7 @@ export async function generateAssistantResponse(requestBody, callback) {
     try {
       const axiosConfig = { ...buildAxiosConfig(config.api.url, headers, requestBody), responseType: 'stream' };
       const response = await axios(axiosConfig);
-      if (response.status === 403) tokenManager.disableCurrentToken(token);
-
+      
       response.data.on('data', chunk => processChunk(chunk.toString()));
       await new Promise((resolve, reject) => {
         response.data.on('end', resolve);
@@ -171,20 +170,21 @@ export async function generateAssistantResponse(requestBody, callback) {
       await handleApiError(error, token);
     }
   } else {
-    const streamResponse = requester.antigravity_fetchStream(config.api.url, buildRequesterConfig(headers, requestBody));
-    let errorBody = '';
-    let statusCode = null;
+    try {
+      const streamResponse = requester.antigravity_fetchStream(config.api.url, buildRequesterConfig(headers, requestBody));
+      let errorBody = '';
+      let statusCode = null;
 
-    await new Promise((resolve, reject) => {
-      streamResponse
-        .onStart(({ status }) => {
-          statusCode = status;
-          if (status === 403) tokenManager.disableCurrentToken(token);
-        })
-        .onData((chunk) => statusCode !== 200 ? errorBody += chunk : processChunk(chunk))
-        .onEnd(() => statusCode !== 200 ? reject(new Error(`API请求失败 (${statusCode}): ${errorBody}`)) : resolve())
-        .onError(reject);
-    });
+      await new Promise((resolve, reject) => {
+        streamResponse
+          .onStart(({ status }) => { statusCode = status; })
+          .onData((chunk) => statusCode !== 200 ? errorBody += chunk : processChunk(chunk))
+          .onEnd(() => statusCode !== 200 ? reject({ status: statusCode, message: errorBody }) : resolve())
+          .onError(reject);
+      });
+    } catch (error) {
+      await handleApiError(error, token);
+    }
   }
 }
 
@@ -195,9 +195,17 @@ export async function getAvailableModels() {
   const headers = buildHeaders(token);
   
   try {
-    const data = useAxios
-      ? (await axios(buildAxiosConfig(config.api.modelsUrl, headers, {}))).data
-      : await (await requester.antigravity_fetch(config.api.modelsUrl, buildRequesterConfig(headers, {}))).json();
+    let data;
+    if (useAxios) {
+      data = (await axios(buildAxiosConfig(config.api.modelsUrl, headers, {}))).data;
+    } else {
+      const response = await requester.antigravity_fetch(config.api.modelsUrl, buildRequesterConfig(headers, {}));
+      if (response.status !== 200) {
+        const errorBody = await response.text();
+        throw { status: response.status, message: errorBody };
+      }
+      data = await response.json();
+    }
     
     return {
       object: 'list',
@@ -227,8 +235,7 @@ export async function generateAssistantResponseNoStream(requestBody) {
       const response = await requester.antigravity_fetch(config.api.noStreamUrl, buildRequesterConfig(headers, requestBody));
       if (response.status !== 200) {
         const errorBody = await response.text();
-        if (response.status === 403) tokenManager.disableCurrentToken(token);
-        throw new Error(response.status === 403 ? `该账号没有使用权限，已自动禁用。错误详情: ${errorBody}` : `API请求失败 (${response.status}): ${errorBody}`);
+        throw { status: response.status, message: errorBody };
       }
       data = await response.json();
     }
@@ -241,6 +248,7 @@ export async function generateAssistantResponseNoStream(requestBody) {
   let content = '';
   let thinkingContent = '';
   const toolCalls = [];
+  const imageContents = [];
   
   for (const part of parts) {
     if (part.thought === true) {
@@ -249,6 +257,11 @@ export async function generateAssistantResponseNoStream(requestBody) {
       content += part.text;
     } else if (part.functionCall) {
       toolCalls.push(convertToToolCall(part.functionCall));
+    } else if (part.inlineData) {
+      imageContents.push({
+        type: 'image_url',
+        image_url: { url: `data:${part.inlineData.mimeType};base64,${part.inlineData.data}` }
+      });
     }
   }
   
@@ -257,6 +270,13 @@ export async function generateAssistantResponseNoStream(requestBody) {
     content = `<think>\n${thinkingContent}\n</think>\n${content}`;
   }
   
+  // 生图模型：转换为markdown格式
+  if (imageContents.length > 0) {
+    let markdown = content ? content + '\n\n' : '';
+    markdown += imageContents.map(img => `![image](${img.image_url.url})`).join('\n\n');
+    return { content: markdown, toolCalls };
+  }
+  
   return { content, toolCalls };
 }
 

src/config/config.js

@@ -26,7 +26,7 @@ API_KEY=sk-text
 
 # 其他配置
 USE_NATIVE_AXIOS=false
-TIMEOUT=30000
+TIMEOUT=180000
 # PROXY=http://127.0.0.1:7897
 
 # 系统提示词

src/server/index.js

@@ -6,6 +6,40 @@ import config from '../config/config.js';
 
 const app = express();
 
+// 工具函数：生成响应元数据
+const createResponseMeta = () => ({
+  id: `chatcmpl-${Date.now()}`,
+  created: Math.floor(Date.now() / 1000)
+});
+
+// 工具函数：设置流式响应头
+const setStreamHeaders = (res) => {
+  res.setHeader('Content-Type', 'text/event-stream');
+  res.setHeader('Cache-Control', 'no-cache');
+  res.setHeader('Connection', 'keep-alive');
+};
+
+// 工具函数：构建流式数据块
+const createStreamChunk = (id, created, model, delta, finish_reason = null) => ({
+  id,
+  object: 'chat.completion.chunk',
+  created,
+  model,
+  choices: [{ index: 0, delta, finish_reason }]
+});
+
+// 工具函数：写入流式数据
+const writeStreamData = (res, data) => {
+  res.write(`data: ${JSON.stringify(data)}\n\n`);
+};
+
+// 工具函数：结束流式响应
+const endStream = (res, id, created, model, finish_reason) => {
+  writeStreamData(res, createStreamChunk(id, created, model, {}, finish_reason));
+  res.write('data: [DONE]\n\n');
+  res.end();
+};
+
 app.use(express.json({ limit: config.security.maxRequestSize }));
 
 app.use((err, req, res, next) => {
@@ -51,65 +85,55 @@ app.get('/v1/models', async (req, res) => {
 app.post('/v1/chat/completions', async (req, res) => {
   const { messages, model, stream = true, tools, ...params} = req.body;
   try {
-    
     if (!messages) {
       return res.status(400).json({ error: 'messages is required' });
     }
     
+    const isImageModel = model.includes('-image');
     const requestBody = await generateRequestBody(messages, model, params, tools);
-    // console.log(JSON.stringify(requestBody,null,2));
+    if (isImageModel) {
+      requestBody.request.generationConfig={
+        candidateCount: 1,
+        // imageConfig:{
+        //   aspectRatio: "1:1"
+        // }
+      }
+      requestBody.requestType="image_gen";
+      //delete requestBody.request.systemInstruction;
+      delete requestBody.request.tools;
+      delete requestBody.request.toolConfig;
+    }
+    //console.log(JSON.stringify(requestBody,null,2))
     
-    if (stream) {
-      res.setHeader('Content-Type', 'text/event-stream');
-      res.setHeader('Cache-Control', 'no-cache');
-      res.setHeader('Connection', 'keep-alive');
-      
-      const id = `chatcmpl-${Date.now()}`;
-      const created = Math.floor(Date.now() / 1000);
+    const { id, created } = createResponseMeta();
+    
+    if (stream && !isImageModel) {
+      setStreamHeaders(res);
       let hasToolCall = false;
       
       await generateAssistantResponse(requestBody, (data) => {
-        if (data.type === 'tool_calls') {
-          hasToolCall = true;
-          res.write(`data: ${JSON.stringify({
-            id,
-            object: 'chat.completion.chunk',
-            created,
-            model,
-            choices: [{ index: 0, delta: { tool_calls: data.tool_calls }, finish_reason: null }]
-          })}\n\n`);
-        } else {
-          res.write(`data: ${JSON.stringify({
-            id,
-            object: 'chat.completion.chunk',
-            created,
-            model,
-            choices: [{ index: 0, delta: { content: data.content }, finish_reason: null }]
-          })}\n\n`);
-        }
+        const delta = data.type === 'tool_calls' 
+          ? { tool_calls: data.tool_calls } 
+          : { content: data.content };
+        if (data.type === 'tool_calls') hasToolCall = true;
+        writeStreamData(res, createStreamChunk(id, created, model, delta));
       });
       
-      res.write(`data: ${JSON.stringify({
-        id,
-        object: 'chat.completion.chunk',
-        created,
-        model,
-        choices: [{ index: 0, delta: {}, finish_reason: hasToolCall ? 'tool_calls' : 'stop' }]
-      })}\n\n`);
-      res.write('data: [DONE]\n\n');
-      res.end();
+      endStream(res, id, created, model, hasToolCall ? 'tool_calls' : 'stop');
+    } else if (stream && isImageModel) {
+      setStreamHeaders(res);
+      const { content } = await generateAssistantResponseNoStream(requestBody);
+      writeStreamData(res, createStreamChunk(id, created, model, { content }));
+      endStream(res, id, created, model, 'stop');
     } else {
       const { content, toolCalls } = await generateAssistantResponseNoStream(requestBody);
-      
       const message = { role: 'assistant', content };
-      if (toolCalls.length > 0) {
-        message.tool_calls = toolCalls;
-      }
+      if (toolCalls.length > 0) message.tool_calls = toolCalls;
       
       res.json({
-        id: `chatcmpl-${Date.now()}`,
+        id,
         object: 'chat.completion',
-        created: Math.floor(Date.now() / 1000),
+        created,
         model,
         choices: [{
           index: 0,
@@ -121,37 +145,22 @@ app.post('/v1/chat/completions', async (req, res) => {
   } catch (error) {
     logger.error('生成响应失败:', error.message);
     if (!res.headersSent) {
+      const { id, created } = createResponseMeta();
+      const errorContent = `错误: ${error.message}`;
+      
       if (stream) {
-        res.setHeader('Content-Type', 'text/event-stream');
-        res.setHeader('Cache-Control', 'no-cache');
-        res.setHeader('Connection', 'keep-alive');
-        const id = `chatcmpl-${Date.now()}`;
-        const created = Math.floor(Date.now() / 1000);
-        res.write(`data: ${JSON.stringify({
-          id,
-          object: 'chat.completion.chunk',
-          created,
-          model,
-          choices: [{ index: 0, delta: { content: `错误: ${error.message}` }, finish_reason: null }]
-        })}\n\n`);
-        res.write(`data: ${JSON.stringify({
-          id,
-          object: 'chat.completion.chunk',
-          created,
-          model,
-          choices: [{ index: 0, delta: {}, finish_reason: 'stop' }]
-        })}\n\n`);
-        res.write('data: [DONE]\n\n');
-        res.end();
+        setStreamHeaders(res);
+        writeStreamData(res, createStreamChunk(id, created, model, { content: errorContent }));
+        endStream(res, id, created, model, 'stop');
       } else {
         res.json({
-          id: `chatcmpl-${Date.now()}`,
+          id,
           object: 'chat.completion',
-          created: Math.floor(Date.now() / 1000),
+          created,
           model,
           choices: [{
             index: 0,
-            message: { role: 'assistant', content: `错误: ${error.message}` },
+            message: { role: 'assistant', content: errorContent },
             finish_reason: 'stop'
           }]
         });

test/test-image-generation.js

@@ -0,0 +1,81 @@
+import fs from 'fs';
+import path from 'path';
+
+const API_URL = 'http://localhost:8045/v1/chat/completions';
+const API_KEY = 'sk-text';
+
+async function testImageGeneration(stream = true) {
+  console.log(`测试生图模型 (${stream ? '流式' : '非流式'})...\n`);
+  
+  const response = await fetch(API_URL, {
+    method: 'POST',
+    headers: {
+      'Content-Type': 'application/json',
+      'Authorization': `Bearer ${API_KEY}`
+    },
+    body: JSON.stringify({
+      model: 'gemini-2.5-flash-image',
+      messages: [{ role: 'user', content: '画一个二次元美少女' }],
+      stream
+    })
+  });
+
+  let fullContent = '';
+  
+  if (stream) {
+    let buffer = '';
+    const reader = response.body.getReader();
+    const decoder = new TextDecoder();
+
+    while (true) {
+      const { done, value } = await reader.read();
+      if (done) break;
+      
+      buffer += decoder.decode(value, { stream: true });
+      const lines = buffer.split('\n');
+      buffer = lines.pop();
+      
+      for (const line of lines) {
+        if (!line.startsWith('data: ') || line.includes('[DONE]')) continue;
+        try {
+          const data = JSON.parse(line.slice(6));
+          const content = data.choices[0]?.delta?.content;
+          if (content) fullContent = content;
+        } catch (e) {}
+      }
+    }
+  } else {
+    const data = await response.json();
+    fullContent = data.choices[0]?.message?.content || '';
+  }
+
+  console.log('响应内容:\n', fullContent.substring(0, 200), '...\n');
+  
+  // 提取markdown中的图片
+  const imageRegex = /!\[.*?\]\((data:image\/(.*?);base64,([^)]+))\)/g;
+  let match;
+  let imageCount = 0;
+  
+  while ((match = imageRegex.exec(fullContent)) !== null) {
+    imageCount++;
+    const base64Data = match[3];
+    const ext = match[2];
+    const filename = `generated_${Date.now()}_${imageCount}.${ext}`;
+    const filepath = path.join('test', filename);
+    
+    fs.writeFileSync(filepath, Buffer.from(base64Data, 'base64'));
+    console.log(`✓ 图片已保存: ${filepath}`);
+  }
+  
+  if (imageCount === 0) {
+    console.log('✗ 未找到图片');
+  } else {
+    console.log(`\n✓ 共保存 ${imageCount} 张图片`);
+  }
+}
+
+(async () => {
+  // await testImageGeneration(true);
+  // console.log('\n' + '='.repeat(50) + '\n');
+  await testImageGeneration(false);
+})().catch(console.error);