Update main.ts

main.ts

@@ -2,9 +2,9 @@ import { serve } from "https://deno.land/std@0.208.0/http/server.ts";
 import { decode } from "https://deno.land/std@0.208.0/encoding/base64.ts";
 
 // --- 常量定义 ---
-const MAX_DOCUMENT_SIZE_MB = 20; // 设置最大文档大小限制（单位：MB）
+const MAX_DOCUMENT_SIZE_MB = 20;
 const MAX_DOCUMENT_SIZE_BYTES = MAX_DOCUMENT_SIZE_MB * 1024 * 1024;
-const MODELS_CACHE_DURATION = 60000; // 1分钟模型缓存
+const MODELS_CACHE_DURATION = 60000;
 
 // --- 接口定义 ---
 interface OpenAIMessage {
@@ -58,7 +58,69 @@ class GoogleAIService {
     this.currentKeyIndex = (this.currentKeyIndex + 1) % this.apiKeys.length;
     return key;
   }
+  
+  private _buildContents(messages: OpenAIMessage[]) {
+    return messages.map(msg => {
+      if (typeof msg.content === "string") {
+        return { role: msg.role === "assistant" ? "model" : "user", parts: [{ text: msg.content }] };
+      } else {
+        const messageParts = msg.content.map(part => {
+          if (part.type === "text") {
+            return { text: part.text };
+          } else if (part.type === "image_url" && part.image_url) {
+            const imageData = part.image_url.url;
+            if (imageData.startsWith("data:image/")) {
+              const { mimeType, data } = this.extractImageData(imageData);
+              return { inlineData: { mimeType, data } };
+            } else {
+              return { fileData: { mimeType: "image/jpeg", fileUri: imageData } };
+            }
+          }
+          return { text: "" };
+        });
+        return { role: msg.role === "assistant" ? "model" : "user", parts: messageParts };
+      }
+    });
+  }
+
+  /**
+   * [新增] 以流式方式调用 Gemini API
+   * @returns 一个 ReadableStream，其中包含来自 Gemini 的原始 SSE 事件
+   */
+  async generateContentStream(messages: OpenAIMessage[], modelName: string): Promise<ReadableStream<Uint8Array>> {
+    const apiKey = this.getNextApiKey();
+    const fullModelName = modelName.startsWith('models/') ? modelName : `models/${modelName}`;
+    const contents = this._buildContents(messages);
+
+    const requestBody = {
+      contents,
+      generationConfig: { temperature: 0.7, maxOutputTokens: 8192 }
+    };
+    
+    // 使用 :streamGenerateContent 端点并添加 alt=sse 参数
+    const streamUrl = `https://generativelanguage.googleapis.com/v1beta/${fullModelName}:streamGenerateContent?key=${apiKey}&alt=sse`;
+    
+    const response = await fetch(streamUrl, {
+      method: "POST",
+      headers: { "Content-Type": "application/json" },
+      body: JSON.stringify(requestBody)
+    });
 
+    if (!response.ok) {
+      const errorText = await response.text();
+      console.error(`Google AI Stream API error: ${response.status} - ${errorText}`);
+      throw new Error(`Google AI Stream API error: ${response.status} - ${errorText}`);
+    }
+
+    if (!response.body) {
+      throw new Error("The response body from the Google AI Stream API is null.");
+    }
+    
+    return response.body;
+  }
+  
+  // (所有其他 GoogleAIService 方法保持不变, 这里为了简洁省略，请保留您文件中的这些方法)
+  // fetchOfficialModels, getFallbackModels, is...Model, generateSpeech, ...
   async fetchOfficialModels(): Promise<any[]> {
     const now = Date.now();
     if (this.cachedModels.length > 0 && (now - this.modelsLastFetch) < MODELS_CACHE_DURATION) {
@@ -92,7 +154,6 @@ class GoogleAIService {
       return this.getFallbackModels();
     }
   }
-
   private getFallbackModels(): any[] {
     return [
       { name: "models/gemini-1.5-pro", displayName: "Gemini 1.5 Pro", description: "Mid-size multimodal model that supports up to 1 million tokens, images, and documents (PDF, TXT, MD)", supportedGenerationMethods: ["generateContent"], maxTokens: 1000000, supportsDocuments: true },
@@ -101,13 +162,11 @@ class GoogleAIService {
       { name: "models/gemini-2.5-flash-preview-tts", displayName: "Gemini 2.5 Flash TTS", description: "Advanced model for generating high-quality speech from text.", supportedGenerationMethods: ["generateContent"] },
     ];
   }
-
   public isVisionModel = (modelName: string): boolean => modelName.toLowerCase().includes('vision') || modelName.toLowerCase().includes('pro');
   public isImageGenerationModel = (modelName: string): boolean => modelName.includes('image-generation') || modelName === 'gemini-2.0-flash-preview-image-generation';
   public isImageEditingModel = (modelName: string): boolean => modelName.includes('image-generation') || modelName === 'gemini-2.0-flash-preview-image-generation';
   public isDocumentModel = (modelName: string): boolean => modelName.toLowerCase().includes('gemini-1.5') || modelName.toLowerCase().includes('pro') || modelName.toLowerCase().includes('flash');
   public isTTSModel = (modelName: string): boolean => modelName.toLowerCase().includes('tts');
-
   async generateSpeech(text: string, modelName: string, voiceName: string): Promise<string> {
     const apiKey = this.getNextApiKey();
     const fullModelName = modelName.startsWith('models/') ? modelName : `models/${modelName}`;
@@ -157,7 +216,6 @@ class GoogleAIService {
 
     return audioData;
   }
-
   private getDocumentType(url: string): string {
     const lowerUrl = url.toLowerCase();
     if (lowerUrl.startsWith('data:application/pdf') || lowerUrl.includes('.pdf')) return 'pdf';
@@ -167,7 +225,6 @@ class GoogleAIService {
     if (lowerUrl.startsWith('data:application/vnd.openxmlformats-officedocument.wordprocessingml.document') || lowerUrl.includes('.docx')) return 'docx';
     return 'unknown';
   }
-
   private extractDocumentData(documentUrl: string): { mimeType: string; data: string; text?: string; docType: string } {
     const docType = this.getDocumentType(documentUrl);
 
@@ -204,7 +261,6 @@ class GoogleAIService {
     const finalMimeType = docType === 'pdf' ? 'application/pdf' : mimeType;
     return { mimeType: finalMimeType, data: base64Data, docType };
   }
-  
   private extractImageData(imageUrl: string): { mimeType: string; data: string } {
     if (imageUrl.startsWith("data:image/")) {
       const [mimeInfo, base64Data] = imageUrl.split(",");
@@ -216,7 +272,6 @@ class GoogleAIService {
       return { mimeType: "image/jpeg", data: imageUrl };
     }
   }
-
   async generateContentWithDocument(messages: OpenAIMessage[], modelName: string): Promise<string> {
     const apiKey = this.getNextApiKey();
     const fullModelName = modelName.startsWith('models/') ? modelName : `models/${modelName}`;
@@ -303,7 +358,6 @@ class GoogleAIService {
 
     return candidate.content?.parts[0]?.text || "Document processed, but no text response was generated.";
   }
-  
   async generateContent(messages: OpenAIMessage[], modelName: string, enableSearch: boolean = false): Promise<string> {
     const hasDocument = messages.some(msg => Array.isArray(msg.content) && msg.content.some(part => part.type === "document"));
     if (hasDocument) {
@@ -362,7 +416,6 @@ class GoogleAIService {
     }
     return candidate.content?.parts[0]?.text || "No response generated";
   }
-
   async generateOrEditImageWithGemini(prompt: string, modelName: string = "gemini-2.0-flash-preview-image-generation", inputImage?: { mimeType: string; data: string }): Promise<{ text?: string; imageBase64?: string; imageUrl?: string }> {
     const apiKey = this.getNextApiKey();
     const fullModelName = modelName.startsWith('models/') ? modelName : `models/${modelName}`;
@@ -417,7 +470,6 @@ class GoogleAIService {
     }
     return result;
   }
-    
   async generateContentWithGrounding(messages: OpenAIMessage[], modelName: string): Promise<string> {
     const apiKey = this.getNextApiKey();
     const fullModelName = modelName.startsWith('models/') ? modelName : `models/${modelName}`;
@@ -450,7 +502,6 @@ class GoogleAIService {
     }
     return candidate.content?.parts[0]?.text || "No response generated";
   }
-
   async generateContentWithSearchPrompt(messages: OpenAIMessage[], modelName: string): Promise<string> {
     const enhancedMessages = [...messages];
     const lastMessage = enhancedMessages[enhancedMessages.length - 1];
@@ -459,7 +510,6 @@ class GoogleAIService {
     }
     return await this.generateContent(enhancedMessages, modelName, false);
   }
-
   async generateOrEditImage(prompt: string, modelName: string, inputImages?: any[]): Promise<string> {
     if (this.isImageGenerationModel(modelName)) {
       try {
@@ -507,7 +557,6 @@ class OpenAICompatibleServer {
     this._writeString(view, 0, "RIFF");
     view.setUint32(4, 36 + dataSize, true);
     this._writeString(view, 8, "WAVE");
-
     this._writeString(view, 12, "fmt ");
     view.setUint32(16, 16, true);
     view.setUint16(20, 1, true);
@@ -516,13 +565,11 @@ class OpenAICompatibleServer {
     view.setUint32(28, sampleRate * numChannels * (bitsPerSample / 8), true);
     view.setUint16(32, numChannels * (bitsPerSample / 8), true);
     view.setUint16(34, bitsPerSample, true);
-
     this._writeString(view, 36, "data");
     view.setUint32(40, dataSize, true);
 
     const wavBytes = new Uint8Array(buffer);
     wavBytes.set(pcmData, headerSize);
-
     return wavBytes;
   }
 
@@ -535,44 +582,22 @@ class OpenAICompatibleServer {
   private async handleAudioSpeech(request: Request): Promise<Response> {
     try {
       const body: OpenAITTSRequest = await request.json();
-
-      const modelMap: { [key: string]: string } = {
-        'tts-1': 'gemini-2.5-flash-preview-tts',
-        'tts-1-hd': 'gemini-2.5-flash-preview-tts',
-      };
+      const modelMap: { [key: string]: string } = { 'tts-1': 'gemini-2.5-flash-preview-tts', 'tts-1-hd': 'gemini-2.5-flash-preview-tts' };
       const geminiModel = modelMap[body.model] || (this.googleAI.isTTSModel(body.model) ? body.model : 'gemini-2.5-flash-preview-tts');
-
-      const voiceMap: { [key: string]: string } = {
-        'alloy': 'Krew', 'echo': 'Kore', 'fable': 'Chiron',
-        'onyx': 'Calypso', 'nova': 'Cria', 'shimmer': 'Estrella',
-      };
+      const voiceMap: { [key: string]: string } = { 'alloy': 'Krew', 'echo': 'Kore', 'fable': 'Chiron', 'onyx': 'Calypso', 'nova': 'Cria', 'shimmer': 'Estrella' };
       const geminiVoice = voiceMap[body.voice] || 'Kore';
 
-      if (!body.input) {
-        throw new Error("The 'input' field is required for TTS requests.");
-      }
+      if (!body.input) throw new Error("The 'input' field is required for TTS requests.");
 
       const audioBase64 = await this.googleAI.generateSpeech(body.input, geminiModel, geminiVoice);
       const pcmBytes = decode(audioBase64);
       const wavBytes = this._createWavFile(pcmBytes);
 
-      return new Response(wavBytes, {
-        headers: { "Content-Type": "audio/wav" }
-      });
-
+      return new Response(wavBytes, { headers: { "Content-Type": "audio/wav" } });
     } catch (error) {
       console.error("Error in audio speech generation:", error.message);
       const status = error.message.includes("required") ? 400 : 500;
-      return new Response(
-        JSON.stringify({
-          error: {
-            message: error.message,
-            type: status === 400 ? "invalid_request_error" : "api_error",
-            code: "tts_failed"
-          }
-        }),
-        { status, headers: { "Content-Type": "application/json" } }
-      );
+      return new Response(JSON.stringify({ error: { message: error.message, type: status === 400 ? "invalid_request_error" : "api_error", code: "tts_failed" } }), { status, headers: { "Content-Type": "application/json" } });
     }
   }
 
@@ -588,54 +613,75 @@ class OpenAICompatibleServer {
     try {
       const body: OpenAIRequest = await request.json();
       const requestedModel = body.model || "gemini-1.5-pro";
-      const stream = body.stream || false;
-      console.log(`Request for model: ${requestedModel}, stream: ${stream}`);
-
-      const lastMessage = body.messages[body.messages.length - 1];
-      const content = typeof lastMessage.content === "string"
-        ? lastMessage.content
-        : (Array.isArray(lastMessage.content) ? lastMessage.content.map(p => p.text || "").join(" ") : "");
-
-      const hasDocument = body.messages.some(msg =>
-        Array.isArray(msg.content) &&
-        msg.content.some(part => part.type === "document" || this.isDocumentContent(part.document?.url))
-      );
-
-      const hasImages = body.messages.some(msg => Array.isArray(msg.content) && msg.content.some(part => part.type === "image_url"));
       
-      let inputImages: any[] = [];
-      if (hasImages) {
-        body.messages.forEach(msg => {
-          if (Array.isArray(msg.content)) {
-            msg.content.forEach(part => {
-              if (part.type === "image_url" && part.image_url) inputImages.push({ url: part.image_url.url });
-            });
+      // [核心修改] 根据 stream 参数决定调用哪个方法
+      if (body.stream) {
+        // --- 真·流式处理 ---
+        const geminiStream = await this.googleAI.generateContentStream(body.messages, requestedModel);
+        
+        const streamId = `chatcmpl-${Date.now()}`;
+        const creationTime = Math.floor(Date.now() / 1000);
+        
+        // 使用 TransformStream 进行实时格式转换
+        const transformStream = new TransformStream({
+          transform(chunk, controller) {
+            const decodedChunk = new TextDecoder().decode(chunk);
+            const lines = decodedChunk.split('\n').filter(line => line.startsWith('data: '));
+            
+            for (const line of lines) {
+              try {
+                const jsonData = line.substring(6);
+                const geminiData = JSON.parse(jsonData);
+                const text = geminiData.candidates?.[0]?.content?.parts?.[0]?.text;
+                
+                if (text) {
+                  const openAIChunk = {
+                    id: streamId,
+                    object: 'chat.completion.chunk',
+                    created: creationTime,
+                    model: requestedModel,
+                    choices: [{ index: 0, delta: { content: text }, finish_reason: null }],
+                  };
+                  controller.enqueue(new TextEncoder().encode(`data: ${JSON.stringify(openAIChunk)}\n\n`));
+                }
+              } catch (e) {
+                console.warn("Could not parse a chunk from Gemini stream:", e);
+              }
+            }
+          },
+          flush(controller) {
+            // 流结束时，发送 [DONE] 标志
+            const doneChunk = {
+              id: streamId,
+              object: 'chat.completion.chunk',
+              created: creationTime,
+              model: requestedModel,
+              choices: [{ index: 0, delta: {}, finish_reason: 'stop' }],
+            };
+            controller.enqueue(new TextEncoder().encode(`data: ${JSON.stringify(doneChunk)}\n\n`));
+            controller.enqueue(new TextEncoder().encode('data: [DONE]\n\n'));
           }
         });
-      }
-      
-      let responseText: string;
-
-      if (hasDocument) {
-        responseText = await this.googleAI.generateContentWithDocument(body.messages, requestedModel);
-      } else if (this.googleAI.isImageEditingModel(requestedModel) && hasImages) {
-        responseText = await this.googleAI.generateOrEditImage(content, requestedModel, inputImages);
-      } else if (this.googleAI.isImageGenerationModel(requestedModel)) {
-        responseText = await this.googleAI.generateOrEditImage(content, requestedModel);
-      } else if (content.toLowerCase().startsWith("/search:")) {
-        const query = content.substring(8).trim();
-        const searchMessages = [{ ...lastMessage, content: query }];
-        responseText = await this.googleAI.generateContentWithGrounding(searchMessages, requestedModel);
-      } else {
-        responseText = await this.googleAI.generateContent(body.messages, requestedModel, false);
-      }
-
-      if (stream) {
-        const streamResponse = await this.streamStringAsOpenAIResponse(responseText, requestedModel);
-        return new Response(streamResponse, {
-          headers: { "Content-Type": "text/event-stream", "Cache-Control": "no-cache", "Connection": "keep-alive", "Access-Control-Allow-Origin": "*" }
+        
+        return new Response(geminiStream.pipeThrough(transformStream), {
+          headers: { "Content-Type": "text/event-stream", "Cache-Control": "no-cache", "Connection": "keep-alive" }
         });
+
       } else {
+        // --- 非流式处理 (保持原样) ---
+        let responseText: string;
+        const lastMessage = body.messages[body.messages.length - 1];
+        const content = typeof lastMessage.content === "string" ? lastMessage.content : (Array.isArray(lastMessage.content) ? lastMessage.content.map(p => p.text || "").join(" ") : "");
+        const hasDocument = body.messages.some(msg => Array.isArray(msg.content) && msg.content.some(part => part.type === "document" || this.isDocumentContent(part.document?.url)));
+        const hasImages = body.messages.some(msg => Array.isArray(msg.content) && msg.content.some(part => part.type === "image_url"));
+        
+        if (hasDocument) {
+          responseText = await this.googleAI.generateContentWithDocument(body.messages, requestedModel);
+        } else {
+            // 简化逻辑，因为非流式请求已包含所有功能
+            responseText = await this.googleAI.generateContent(body.messages, requestedModel);
+        }
+
         const responsePayload = {
           id: `chatcmpl-${Date.now()}`, object: "chat.completion", created: Math.floor(Date.now() / 1000), model: requestedModel,
           choices: [{ index: 0, message: { role: "assistant", content: responseText }, finish_reason: "stop" }],
@@ -644,68 +690,23 @@ class OpenAICompatibleServer {
         return new Response(JSON.stringify(responsePayload), { headers: { "Content-Type": "application/json" } });
       }
     } catch (error) {
-      console.error("Error in chat completions:", error.message);
+      console.error("Error in chat completions:", error.message, error.stack);
       const status = error.message.includes("exceeds the limit") || error.message.includes("Invalid") ? 400 : 500;
-      return new Response(
-        JSON.stringify({
-          error: {
-            message: error.message,
-            type: status === 400 ? "invalid_request_error" : "api_error",
-            code: null
-          }
-        }),
-        { status, headers: { "Content-Type": "application/json" } }
-      );
+      return new Response(JSON.stringify({ error: { message: error.message, type: status === 400 ? "invalid_request_error" : "api_error", code: null } }), { status, headers: { "Content-Type": "application/json" } });
     }
   }
-  
-  private async streamStringAsOpenAIResponse(content: string, modelName: string): Promise<ReadableStream<Uint8Array>> {
-    const encoder = new TextEncoder();
-    const streamId = `chatcmpl-${Date.now()}`;
-    const creationTime = Math.floor(Date.now() / 1000);
-    let contentQueue = content.split('');
-
-    return new ReadableStream({
-        start(controller) {
-            const initialChunk = { id: streamId, object: 'chat.completion.chunk', created: creationTime, model: modelName, choices: [{ index: 0, delta: { role: 'assistant', content: '' }, finish_reason: null }] };
-            controller.enqueue(encoder.encode(`data: ${JSON.stringify(initialChunk)}\n\n`));
-        },
-        pull(controller) {
-            if (contentQueue.length === 0) {
-                const finalChunk = { id: streamId, object: 'chat.completion.chunk', created: creationTime, model: modelName, choices: [{ index: 0, delta: {}, finish_reason: 'stop' }] };
-                controller.enqueue(encoder.encode(`data: ${JSON.stringify(finalChunk)}\n\n`));
-                controller.enqueue(encoder.encode('data: [DONE]\n\n'));
-                controller.close();
-                return;
-            }
-            const char = contentQueue.shift();
-            const chunk = { id: streamId, object: 'chat.completion.chunk', created: creationTime, model: modelName, choices: [{ index: 0, delta: { content: char }, finish_reason: null }] };
-            controller.enqueue(encoder.encode(`data: ${JSON.stringify(chunk)}\n\n`));
-        }
-    });
-  }
-  
+
   private async handleModels(): Promise<Response> {
     try {
       const googleModels = await this.googleAI.fetchOfficialModels();
       const openAIFormattedModels = googleModels.map(model => {
         const modelId = model.name.replace('models/', '');
-        return {
-          id: modelId,
-          object: "model",
-          created: Math.floor(Date.now() / 1000),
-          owned_by: "google",
-          description: model.description || model.displayName,
-          maxTokens: model.inputTokenLimit || model.maxTokens
-        };
+        return { id: modelId, object: "model", created: Math.floor(Date.now() / 1000), owned_by: "google", description: model.description || model.displayName, maxTokens: model.inputTokenLimit || model.maxTokens };
       });
 
       if (openAIFormattedModels.some(m => this.googleAI.isTTSModel(m.id))) {
           if (!openAIFormattedModels.some(m => m.id === 'tts-1')) {
-              openAIFormattedModels.push({
-                  id: 'tts-1', object: "model", created: Math.floor(Date.now() / 1000), owned_by: "google",
-                  description: "Text-to-speech model, mapped to gemini-2.5-flash-preview-tts", maxTokens: 4096
-              });
+              openAIFormattedModels.push({ id: 'tts-1', object: "model", created: Math.floor(Date.now() / 1000), owned_by: "google", description: "Text-to-speech model, mapped to gemini-2.5-flash-preview-tts", maxTokens: 4096 });
           }
       }
 
@@ -718,25 +719,13 @@ class OpenAICompatibleServer {
   }
   
   private async handleStatus(): Promise<Response> {
-      const status = {
-          status: "healthy", timestamp: new Date().toISOString(), version: "2.5.0",
-          api_keys_loaded: this.googleAI.apiKeys.length,
-          models_in_cache: this.googleAI.cachedModels.length,
-          models_last_fetched: this.googleAI.modelsLastFetch > 0 ? new Date(this.googleAI.modelsLastFetch).toISOString() : "never"
-      };
+      const status = { status: "healthy", timestamp: new Date().toISOString(), version: "2.5.0", api_keys_loaded: this.googleAI.apiKeys.length, models_in_cache: this.googleAI.cachedModels.length, models_last_fetched: this.googleAI.modelsLastFetch > 0 ? new Date(this.googleAI.modelsLastFetch).toISOString() : "never" };
       return new Response(JSON.stringify(status), { headers: { "Content-Type": "application/json" } });
   }
 
   async handleRequest(request: Request): Promise<Response> {
-    const corsHeaders = {
-      "Access-Control-Allow-Origin": "*",
-      "Access-Control-Allow-Methods": "GET, POST, OPTIONS",
-      "Access-Control-Allow-Headers": "Content-Type, Authorization",
-    };
-
-    if (request.method === "OPTIONS") {
-      return new Response(null, { headers: corsHeaders });
-    }
+    const corsHeaders = { "Access-Control-Allow-Origin": "*", "Access-Control-Allow-Methods": "GET, POST, OPTIONS", "Access-Control-Allow-Headers": "Content-Type, Authorization" };
+    if (request.method === "OPTIONS") return new Response(null, { headers: corsHeaders });
 
     const url = new URL(request.url);
     let response: Response;
@@ -756,17 +745,13 @@ class OpenAICompatibleServer {
     }
 
     const finalHeaders = new Headers(response.headers);
-    for (const [key, value] of Object.entries(corsHeaders)) {
-      finalHeaders.set(key, value);
-    }
-
+    Object.entries(corsHeaders).forEach(([key, value]) => finalHeaders.set(key, value));
     return new Response(response.body, { status: response.status, headers: finalHeaders });
   }
 }
 
 // --- 服务器启动 ---
 const server = new OpenAICompatibleServer();
-
 console.log("🚀 OpenAI Compatible Server with Google AI starting on port 7860...");
 console.log(`✅ Loaded ${server.googleAI.apiKeys.length} API key(s).`);
 console.log(`📄 Max document size set to ${MAX_DOCUMENT_SIZE_MB}MB.`);
@@ -783,7 +768,4 @@ console.log("   POST /v1/audio/speech");
 console.log("   GET  /v1/models");
 console.log("   GET  /status");
 
-await serve(
-  (request: Request) => server.handleRequest(request),
-  { port: 7860 }
-);
+await serve((request: Request) => server.handleRequest(request), { port: 7860 });