Update main.ts

main.ts

@@ -1,19 +1,17 @@
 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;
+const MAX_DOCUMENT_SIZE_MB = 20; // 设置最大文档大小限制（单位：MB）
 const MAX_DOCUMENT_SIZE_BYTES = MAX_DOCUMENT_SIZE_MB * 1024 * 1024;
-const MODELS_CACHE_DURATION = 60000;
+const MODELS_CACHE_DURATION = 60000; // 1分钟模型缓存
 
-// --- 接口定义 ---
 interface OpenAIMessage {
   role: "system" | "user" | "assistant";
   content: string | Array<{
     type: string;
     text?: string;
     image_url?: { url: string };
-    document?: { url: string; type: string };
+    document?: { url: string; type: string }; // 支持多种文档类型
   }>;
 }
 
@@ -25,12 +23,6 @@ interface OpenAIRequest {
   stream?: boolean;
 }
 
-interface OpenAITTSRequest {
-  model: string;
-  input: string;
-  voice: 'alloy' | 'echo' | 'fable' | 'onyx' | 'shimmer' | 'nova' | string;
-}
-
 class GoogleAIService {
   public apiKeys: string[];
   public currentKeyIndex = 0;
@@ -58,63 +50,7 @@ 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 };
-      }
-    });
-  }
-
-  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 }
-    };
-    
-    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 方法保持不变, 这里为了简洁省略，请保留您文件中的这些方法)
   async fetchOfficialModels(): Promise<any[]> {
     const now = Date.now();
     if (this.cachedModels.length > 0 && (now - this.modelsLastFetch) < MODELS_CACHE_DURATION) {
@@ -148,68 +84,20 @@ 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 },
       { name: "models/gemini-1.5-flash", displayName: "Gemini 1.5 Flash", description: "Fast and versatile multimodal model for diverse tasks, supports images and documents (PDF, TXT, MD)", supportedGenerationMethods: ["generateContent"], maxTokens: 1000000, supportsDocuments: true },
-      { name: "models/gemini-2.0-flash-preview-image-generation", displayName: "Gemini 2.0 Flash Image Generation", description: "Advanced model for generating and editing high-quality images with text and image outputs", supportedGenerationMethods: ["generateContent"], maxTokens: 100000, capabilities: ["text", "image_generation", "image_editing"] },
-      { 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"] },
+      { name: "models/gemini-2.0-flash-preview-image-generation", displayName: "Gemini 2.0 Flash Image Generation", description: "Advanced model for generating and editing high-quality images with text and image outputs", supportedGenerationMethods: ["generateContent"], maxTokens: 100000, capabilities: ["text", "image_generation", "image_editing"] }
     ];
   }
+
   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}`;
-
-    console.log(`Generating speech with model: ${fullModelName}, voice: ${voiceName}`);
-
-    const requestBody = {
-      contents: [{
-        parts: [{ "text": text }]
-      }],
-      generationConfig: {
-        responseModalities: ["AUDIO"],
-        speechConfig: {
-          voiceConfig: {
-            prebuiltVoiceConfig: {
-              voiceName: voiceName
-            }
-          }
-        }
-      },
-      model: fullModelName,
-    };
-
-    const response = await fetch(
-      `https://generativelanguage.googleapis.com/v1beta/${fullModelName}:generateContent?key=${apiKey}`,
-      {
-        method: "POST",
-        headers: { "Content-Type": "application/json" },
-        body: JSON.stringify(requestBody),
-      }
-    );
 
-    if (!response.ok) {
-      const errorBody = await response.json().catch(() => response.text());
-      const errorMessage = errorBody?.error?.message || JSON.stringify(errorBody);
-      console.error(`Google TTS API Error: ${response.status} - ${errorMessage}`);
-      throw new Error(`Google TTS API request failed with status ${response.status}: ${errorMessage}`);
-    }
-
-    const data = await response.json();
-    const audioData = data.candidates?.[0]?.content?.parts?.[0]?.inlineData?.data;
-    
-    if (!audioData) {
-      console.error("Invalid TTS response from Google AI:", JSON.stringify(data));
-      throw new Error("No audio data received from Google AI TTS service.");
-    }
-
-    return audioData;
-  }
   private getDocumentType(url: string): string {
     const lowerUrl = url.toLowerCase();
     if (lowerUrl.startsWith('data:application/pdf') || lowerUrl.includes('.pdf')) return 'pdf';
@@ -219,6 +107,10 @@ 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);
 
@@ -226,6 +118,8 @@ class GoogleAIService {
       if (documentUrl.startsWith("http")) {
         throw new Error("Document URL downloads are not supported. Please provide base64 encoded data URLs.");
       }
+      // 如果不是data url或http url，则假定为纯base64数据，但这是一种不推荐的格式
+      // 为了健壮性，我们强制要求使用标准的 data URL
       throw new Error("Document must be provided as a standard base64 data URL (e.g., 'data:application/pdf;base64,...').");
     }
 
@@ -235,6 +129,8 @@ class GoogleAIService {
     }
     const [mimeInfo, base64Data] = parts;
 
+    // **改进1: 检查文件大小**
+    // Base64 字符串的长度约是原始数据的 4/3。
     const approxSizeInBytes = base64Data.length * 0.75;
     if (approxSizeInBytes > MAX_DOCUMENT_SIZE_BYTES) {
         throw new Error(`Document size (${(approxSizeInBytes / 1024 / 1024).toFixed(2)}MB) exceeds the ${MAX_DOCUMENT_SIZE_MB}MB limit.`);
@@ -252,9 +148,11 @@ class GoogleAIService {
       }
     }
     
+    // 自动识别PDF的MIME类型
     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(",");
@@ -266,6 +164,7 @@ 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}`;
@@ -273,57 +172,260 @@ class GoogleAIService {
 
     console.log(`Processing document with model: ${documentModel}`);
 
-    const contents = this._buildContents(messages.slice(0, -1));
-    const lastMessage = messages[messages.length - 1];
-    
-    // Complex last message handling
-    if (typeof lastMessage.content === "string") {
-      contents.push({ role: "user", parts: [{text: lastMessage.content}] });
-    } else {
-        const docParts = [];
-        const textParts = [];
-        for (const part of lastMessage.content) {
-            if (part.type === "document" && part.document) {
-                const docData = this.extractDocumentData(part.document.url);
-                docParts.push({ inlineData: { mimeType: docData.mimeType, data: docData.data } });
-            } else if (part.type === "text") {
-                textParts.push({text: part.text});
-            }
+    let contents;
+    try {
+      contents = messages.map(msg => {
+        if (typeof msg.content === "string") {
+          return { role: msg.role === "assistant" ? "model" : "user", parts: [{ text: msg.content }] };
         }
-        contents.push({ role: "user", parts: [...docParts, ...textParts] });
+
+        const messageParts = msg.content.map(part => {
+          if (part.type === "text") return { text: part.text };
+
+          if (part.type === "image_url" && part.image_url) {
+            const { mimeType, data } = this.extractImageData(part.image_url.url);
+            return { inlineData: { mimeType, data } };
+          }
+
+          if (part.type === "document" && part.document) {
+            const docData = this.extractDocumentData(part.document.url);
+            console.log(`Processing document: ${docData.docType}, mime: ${docData.mimeType}, size: ${(docData.data.length * 0.75 / 1024).toFixed(2)} KB`);
+
+            if (docData.docType === 'txt' || docData.docType === 'md') {
+              const prefix = docData.docType === 'md' ? 'Markdown document content:\n' : 'Text document content:\n';
+              return { text: `${prefix}${docData.text}` };
+            }
+            if (docData.docType === 'pdf') {
+              return { inlineData: { mimeType: docData.mimeType, data: docData.data } };
+            }
+            return { text: `[Document type '${docData.docType}' is not supported for direct processing. Please convert to PDF, TXT, or MD.]` };
+          }
+          return { text: "" };
+        });
+        return { role: msg.role === "assistant" ? "model" : "user", parts: messageParts.filter(p => p.text || p.inlineData) };
+      });
+    } catch (error) {
+      throw error;
     }
 
-    const requestBody = { contents, generationConfig: { temperature: 0.7, maxOutputTokens: 8192 } };
-    const response = await fetch( `https://generativelanguage.googleapis.com/v1beta/${documentModel}:generateContent?key=${apiKey}`, { method: "POST", headers: { "Content-Type": "application/json" }, body: JSON.stringify(requestBody) } );
+    const requestBody = {
+      contents,
+      generationConfig: { temperature: 0.7, maxOutputTokens: 8192 }
+    };
+
+    const response = await fetch(
+      `https://generativelanguage.googleapis.com/v1beta/${documentModel}:generateContent?key=${apiKey}`,
+      {
+        method: "POST",
+        headers: { "Content-Type": "application/json" },
+        body: JSON.stringify(requestBody),
+      }
+    );
 
     if (!response.ok) {
-        const errorBody = await response.json().catch(() => response.text());
-        const errorMessage = errorBody?.error?.message || JSON.stringify(errorBody);
-        throw new Error(`Google API request failed: ${errorMessage}`);
+      const errorBody = await response.json().catch(() => response.text());
+      const errorMessage = errorBody?.error?.message || JSON.stringify(errorBody);
+      console.error(`Google API Error: ${response.status} - ${errorMessage}`);
+      throw new Error(`Google API request failed with status ${response.status}: ${errorMessage}`);
     }
+
     const data = await response.json();
-    return data.candidates?.[0]?.content?.parts?.[0]?.text || "No response text.";
+    const promptFeedback = data.promptFeedback;
+    if (promptFeedback && promptFeedback.blockReason) {
+      const reason = promptFeedback.blockReason;
+      const safetyRatings = promptFeedback.safetyRatings?.map((r: any) => `${r.category}: ${r.probability}`).join(', ') || 'N/A';
+      throw new Error(`Request blocked by Google API. Reason: ${reason}. Safety Ratings: [${safetyRatings}]`);
+    }
+
+    if (!data.candidates || data.candidates.length === 0) {
+      throw new Error("No response generated for document content. The content might be empty or unreadable.");
+    }
+
+    const candidate = data.candidates[0];
+    if (candidate.finishReason === "SAFETY") {
+        throw new Error("Response blocked due to safety filters. Check content for sensitive topics.");
+    }
+    if (candidate.finishReason === "RECITATION") {
+        throw new Error("Response blocked due to recitation policy. The model's output was too similar to a copyrighted source.");
+    }
+
+    return candidate.content?.parts[0]?.text || "Document processed, but no text response was generated.";
   }
-  async generateContent(messages: OpenAIMessage[], modelName: string): Promise<string> {
+  
+  // The rest of the original methods from the user's code
+  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) {
+      return await this.generateContentWithDocument(messages, modelName);
+    }
+
     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 } };
 
-    const response = await fetch( `https://generativelanguage.googleapis.com/v1beta/${fullModelName}:generateContent?key=${apiKey}`, { method: "POST", headers: { "Content-Type": "application/json" }, body: JSON.stringify(requestBody) } );
+    const contents = 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 };
+      }
+    });
+
+    const requestBody: any = {
+      contents,
+      generationConfig: { temperature: 0.7, maxOutputTokens: 4096 }
+    };
+    if (enableSearch) {
+      requestBody.tools = [{ googleSearchRetrieval: {} }];
+    }
+
+    const response = await fetch(
+      `https://generativelanguage.googleapis.com/v1beta/${fullModelName}:generateContent?key=${apiKey}`,
+      { method: "POST", headers: { "Content-Type": "application/json" }, body: JSON.stringify(requestBody) }
+    );
+
+    if (!response.ok) {
+      const errorText = await response.text();
+      throw new Error(`Google AI API error: ${response.status} - ${errorText}`);
+    }
+    const data = await response.json();
+    if (!data.candidates || data.candidates.length === 0) {
+      throw new Error("No response generated from Google AI");
+    }
+    const candidate = data.candidates[0];
+    if (candidate.finishReason === "SAFETY") {
+      throw new Error("Response blocked due to safety filters");
+    }
+    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}`;
+    const requestParts: any[] = [{ text: prompt }];
+
+    if (inputImage) {
+      requestParts.push({ inline_data: { mime_type: inputImage.mimeType, data: inputImage.data } });
+      console.log(`Editing image with model: ${fullModelName}`);
+    } else {
+      console.log(`Generating image with model: ${fullModelName}`);
+    }
+
+    const requestBody = {
+      contents: [{ parts: requestParts }],
+      generationConfig: { responseModalities: ["TEXT", "IMAGE"], temperature: 0.7 }
+    };
+
+    const response = await fetch(
+      `https://generativelanguage.googleapis.com/v1beta/${fullModelName}:generateContent?key=${apiKey}`,
+      { method: "POST", headers: { "Content-Type": "application/json" }, body: JSON.stringify(requestBody) }
+    );
+
+    if (!response.ok) {
+      const errorText = await response.text();
+      throw new Error(`Image ${inputImage ? 'editing' : 'generation'} failed: ${response.status} - ${errorText}`);
+    }
+    const data = await response.json();
+    if (!data.candidates || data.candidates.length === 0) {
+      throw new Error(`No ${inputImage ? 'edited' : 'generated'} image returned`);
+    }
+
+    const candidate = data.candidates[0];
+    if (candidate.finishReason === "SAFETY") {
+      throw new Error(`Image ${inputImage ? 'editing' : 'generation'} blocked due to safety filters`);
+    }
+
+    const responseParts = candidate.content?.parts || [];
+    let textResponse = "";
+    let imageBase64 = "";
+
+    for (const part of responseParts) {
+      if (part.text) textResponse += part.text;
+      if (part.inlineData?.data) imageBase64 = part.inlineData.data;
+      if (part.inline_data?.data) imageBase64 = part.inline_data.data;
+    }
+
+    const result: { text?: string; imageBase64?: string; imageUrl?: string } = {};
+    if (textResponse) result.text = textResponse;
+    if (imageBase64) {
+      result.imageBase64 = imageBase64;
+      result.imageUrl = `data:image/png;base64,${imageBase64}`;
+    }
+    return result;
+  }
     
-    if (!response.ok) { throw new Error(`Google AI API error: ${await response.text()}`); }
+  async generateContentWithGrounding(messages: OpenAIMessage[], modelName: string): Promise<string> {
+    const apiKey = this.getNextApiKey();
+    const fullModelName = modelName.startsWith('models/') ? modelName : `models/${modelName}`;
+    const contents = messages.map(msg => ({ role: msg.role === 'assistant' ? 'model' : 'user', parts: [{ text: typeof msg.content === 'string' ? msg.content : '' }] }));
+
+    const requestBody = {
+      contents,
+      tools: [{ googleSearch: {} }],
+      generationConfig: { temperature: 0.7, maxOutputTokens: 4096 }
+    };
+
+    const response = await fetch(
+      `https://generativelanguage.googleapis.com/v1beta/${fullModelName}:generateContent?key=${apiKey}`,
+      { method: "POST", headers: { "Content-Type": "application/json" }, body: JSON.stringify(requestBody) }
+    );
+
+    if (!response.ok) {
+      console.warn(`Google Search API failed: ${response.status}, trying alternative.`);
+      return await this.generateContentWithSearchPrompt(messages, modelName);
+    }
 
     const data = await response.json();
-    if (data.promptFeedback?.blockReason) {
-        throw new Error(`Request blocked by Google: ${data.promptFeedback.blockReason}`);
+    if (!data.candidates || data.candidates.length === 0) {
+      return await this.generateContentWithSearchPrompt(messages, modelName);
+    }
+
+    const candidate = data.candidates[0];
+    if (candidate.finishReason === "SAFETY") {
+      throw new Error("Response blocked due to safety filters");
+    }
+    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];
+    if (typeof lastMessage.content === "string") {
+      lastMessage.content = `Please provide the most current and accurate information available about: ${lastMessage.content}.`;
     }
-    const candidate = data.candidates?.[0];
-    if (!candidate) { throw new Error("No response generated from Google AI."); }
-    if (candidate.finishReason !== "STOP") {
-        console.warn(`Stream finished with reason: ${candidate.finishReason}`);
+    return await this.generateContent(enhancedMessages, modelName, false);
+  }
+
+  async generateOrEditImage(prompt: string, modelName: string, inputImages?: any[]): Promise<string> {
+    if (this.isImageGenerationModel(modelName)) {
+      try {
+        let inputImage: { mimeType: string; data: string } | undefined;
+        if (inputImages && inputImages.length > 0) {
+          inputImage = this.extractImageData(inputImages[0].url);
+        }
+        const result = await this.generateOrEditImageWithGemini(prompt, modelName, inputImage);
+        let response = "";
+        if (result.text) response += result.text + "\\\\n\\\\n";
+        if (result.imageUrl) response += `${inputImage ? 'Edited' : 'Generated'} image:\\\\n${result.imageUrl}`;
+        return response || `Image processing complete.`;
+      } catch (error) {
+        return `Image processing failed: ${error.message}`;
+      }
     }
-    return candidate.content?.parts?.[0]?.text || "";
+    return `Model ${modelName} does not support image generation. Use a model like gemini-2.0-flash-preview-image-generation.`;
   }
 }
 
@@ -336,62 +438,11 @@ class OpenAICompatibleServer {
     this.authKey = Deno.env.get("AUTH_KEY") || "";
   }
 
-  private _writeString(view: DataView, offset: number, str: string) {
-    for (let i = 0; i < str.length; i++) {
-      view.setUint8(offset + i, str.charCodeAt(i));
-    }
-  }
-  
-  private _createWavFile(pcmData: Uint8Array): Uint8Array {
-    const numChannels = 1, sampleRate = 24000, bitsPerSample = 16, dataSize = pcmData.length;
-    const headerSize = 44, buffer = new ArrayBuffer(headerSize + dataSize), view = new DataView(buffer);
-    
-    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);
-    view.setUint16(22, numChannels, true);
-    view.setUint32(24, sampleRate, true);
-    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;
-  }
-
   private authenticate(request: Request): boolean {
     if (!this.authKey) return true;
     const authHeader = request.headers.get("Authorization");
     return authHeader ? authHeader.replace("Bearer ", "") === this.authKey : false;
   }
-  
-  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 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 geminiVoice = voiceMap[body.voice] || 'Kore';
-
-      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" } });
-    } 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" } });
-    }
-  }
 
   private isDocumentContent(url?: string): boolean {
     if (!url) return false;
@@ -405,87 +456,55 @@ 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"));
       
-      if (body.stream) {
-        // --- [增强版] 真·流式处理 ---
-        const geminiStream = await this.googleAI.generateContentStream(body.messages, requestedModel);
-        
-        const streamId = `chatcmpl-${Date.now()}`;
-        const creationTime = Math.floor(Date.now() / 1000);
-        
-        let buffer = "";
-        const textDecoder = new TextDecoder();
-        
-        const transformStream = new TransformStream({
-          transform(chunk, controller) {
-            buffer += textDecoder.decode(chunk);
-            const lines = buffer.split('\n');
-            buffer = lines.pop() || ""; // 保留不完整的行到下一次处理
-
-            for (const line of lines) {
-              if (!line.startsWith('data: ')) continue;
-              const jsonData = line.substring(6);
-              if (!jsonData) continue;
-              
-              try {
-                const geminiData = JSON.parse(jsonData);
-
-                // --- 核心逻辑：检查内容和终止原因 ---
-                const candidate = geminiData.candidates?.[0];
-                if (!candidate) continue;
-
-                // 1. 提取文本内容
-                const text = candidate.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`));
-                }
-
-                // 2. 检查终止原因
-                const finishReason = candidate.finishReason;
-                if (finishReason && finishReason !== "STOP") {
-                    const reasonText = `\n\n[STREAM INTERRUPTED BY GOOGLE: ${finishReason}]`;
-                    const openAIChunk = {
-                        id: streamId, object: 'chat.completion.chunk', created: creationTime, model: requestedModel,
-                        choices: [{ index: 0, delta: { content: reasonText }, finish_reason: finishReason.toLowerCase() }],
-                    };
-                    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 finalChunk = {
-              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(finalChunk)}\n\n`));
-            controller.enqueue(new TextEncoder().encode('data: [DONE]\n\n'));
+      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 });
+            });
           }
         });
-        
-        return new Response(geminiStream.pipeThrough(transformStream), {
-          headers: { "Content-Type": "text/event-stream", "Cache-Control": "no-cache", "Connection": "keep-alive" }
-        });
-
+      }
+      
+      let responseText: string;
+
+      // Routing logic based on keywords and content types
+      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 {
-        // --- 非流式处理 ---
-        const hasDocument = body.messages.some(msg => Array.isArray(msg.content) && msg.content.some(part => part.type === "document" || this.isDocumentContent(part.document?.url)));
-        let responseText;
-
-        if (hasDocument) {
-          responseText = await this.googleAI.generateContentWithDocument(body.messages, requestedModel);
-        } else {
-          responseText = await this.googleAI.generateContent(body.messages, requestedModel);
-        }
+        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": "*" }
+        });
+      } else {
         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" }],
@@ -494,27 +513,60 @@ class OpenAICompatibleServer {
         return new Response(JSON.stringify(responsePayload), { headers: { "Content-Type": "application/json" } });
       }
     } catch (error) {
-      console.error("Error in chat completions:", error.message, error.stack);
+      console.error("Error in chat completions:", error.message);
       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 };
-      });
-
-      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 });
-          }
-      }
-
-      const models = { object: "list", data: openAIFormattedModels };
+      const models = {
+        object: "list",
+        data: 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 new Response(JSON.stringify(models), { headers: { "Content-Type": "application/json" } });
     } catch (error) {
       console.error("Error fetching models:", error);
@@ -523,23 +575,34 @@ 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;
 
+    // Handle routes
     if (url.pathname === "/health" || url.pathname === "/status") {
       response = await this.handleStatus();
     } else if (!this.authenticate(request)) {
       response = new Response(JSON.stringify({ error: { message: "Unauthorized" } }), { status: 401 });
-    } else if (url.pathname === "/v1/audio/speech" && request.method === "POST") {
-      response = await this.handleAudioSpeech(request);
     } else if (url.pathname === "/v1/chat/completions" && request.method === "POST") {
       response = await this.handleChatCompletions(request);
     } else if (url.pathname === "/v1/models" && request.method === "GET") {
@@ -548,18 +611,24 @@ class OpenAICompatibleServer {
       response = new Response("Not Found", { status: 404 });
     }
 
+    // Add CORS headers to all responses
     const finalHeaders = new Headers(response.headers);
-    Object.entries(corsHeaders).forEach(([key, value]) => finalHeaders.set(key, value));
+    for (const [key, value] of Object.entries(corsHeaders)) {
+      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("🚀 OpenAI Compatible Server with Google AI starting on port 8000...");
 console.log(`✅ Loaded ${server.googleAI.apiKeys.length} API key(s).`);
 console.log(`📄 Max document size set to ${MAX_DOCUMENT_SIZE_MB}MB.`);
 
+// Pre-fetch models at startup
 server.googleAI.fetchOfficialModels().then(models => {
   console.log(`✅ Successfully fetched ${models.length} models from Google AI.`);
 }).catch(error => {
@@ -568,8 +637,10 @@ server.googleAI.fetchOfficialModels().then(models => {
 
 console.log("\n🔗 Endpoints:");
 console.log("   POST /v1/chat/completions");
-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 }
+);