Update main.ts

main.ts

@@ -254,7 +254,7 @@ class GoogleAIService {
     }
   }
 
-  async generateContentWithDocument(messages: OpenAIMessage[], modelName: string): Promise<string> {
+  async generateContentWithDocument(messages: OpenAIMessage[], modelName: string, maxTokens?: number): Promise<string> {
     const apiKey = this.getNextApiKey();
     const fullModelName = modelName.startsWith('models/') ? modelName : `models/${modelName}`;
     const documentModel = this.isDocumentModel(fullModelName) ? fullModelName : 'models/gemini-1.5-pro-latest';
@@ -299,7 +299,7 @@ class GoogleAIService {
 
     const requestBody = {
       contents,
-      generationConfig: { temperature: 0.7, maxOutputTokens: 8192 }
+      generationConfig: { temperature: 0.7, maxOutputTokens: maxTokens || 8192 }
     };
 
     const response = await fetch(
@@ -342,10 +342,10 @@ class GoogleAIService {
   }
   
   // The rest of the original methods from the user's code
-  async generateContent(messages: OpenAIMessage[], modelName: string, enableSearch: boolean = false): Promise<string> {
+  async generateContent(messages: OpenAIMessage[], modelName: string, maxTokens?: number, 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);
+      return await this.generateContentWithDocument(messages, modelName, maxTokens);
     }
 
     const apiKey = this.getNextApiKey();
@@ -375,7 +375,7 @@ class GoogleAIService {
 
     const requestBody: any = {
       contents,
-      generationConfig: { temperature: 0.7, maxOutputTokens: 4096 }
+      generationConfig: { temperature: 0.7, maxOutputTokens: maxTokens || 8192 }
     };
     if (enableSearch) {
       requestBody.tools = [{ googleSearchRetrieval: {} }];
@@ -456,7 +456,7 @@ class GoogleAIService {
     return result;
   }
     
-  async generateContentWithGrounding(messages: OpenAIMessage[], modelName: string): Promise<string> {
+  async generateContentWithGrounding(messages: OpenAIMessage[], modelName: string, maxTokens?: number): 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 : '' }] }));
@@ -464,7 +464,7 @@ class GoogleAIService {
     const requestBody = {
       contents,
       tools: [{ googleSearch: {} }],
-      generationConfig: { temperature: 0.7, maxOutputTokens: 4096 }
+      generationConfig: { temperature: 0.7, maxOutputTokens: maxTokens || 8192 }
     };
 
     const response = await fetch(
@@ -474,12 +474,12 @@ class GoogleAIService {
 
     if (!response.ok) {
       console.warn(`Google Search API failed: ${response.status}, trying alternative.`);
-      return await this.generateContentWithSearchPrompt(messages, modelName);
+      return await this.generateContentWithSearchPrompt(messages, modelName, maxTokens);
     }
 
     const data = await response.json();
     if (!data.candidates || data.candidates.length === 0) {
-      return await this.generateContentWithSearchPrompt(messages, modelName);
+      return await this.generateContentWithSearchPrompt(messages, modelName, maxTokens);
     }
 
     const candidate = data.candidates[0];
@@ -489,13 +489,13 @@ class GoogleAIService {
     return candidate.content?.parts[0]?.text || "No response generated";
   }
 
-  async generateContentWithSearchPrompt(messages: OpenAIMessage[], modelName: string): Promise<string> {
+  async generateContentWithSearchPrompt(messages: OpenAIMessage[], modelName: string, maxTokens?: number): 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}.`;
     }
-    return await this.generateContent(enhancedMessages, modelName, false);
+    return await this.generateContent(enhancedMessages, modelName, maxTokens, false);
   }
 
   async generateOrEditImage(prompt: string, modelName: string, inputImages?: any[]): Promise<string> {
@@ -602,7 +602,8 @@ class OpenAICompatibleServer {
       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 maxTokens = body.max_tokens || 1024000;
+      console.log(`Request for model: ${requestedModel}, stream: ${stream}, max_tokens: ${maxTokens}`);
 
       const lastMessage = body.messages[body.messages.length - 1];
       const content = typeof lastMessage.content === "string"
@@ -631,16 +632,16 @@ class OpenAICompatibleServer {
 
       // Routing logic based on keywords and content types
       if (hasDocument) {
-        responseText = await this.googleAI.generateContentWithDocument(body.messages, requestedModel);
+        responseText = await this.googleAI.generateContentWithDocument(body.messages, requestedModel, maxTokens);
       } 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 (requestedModel.endsWith("-search")) {
         const searchMessages = [{ ...lastMessage, content: content }];
-        responseText = await this.googleAI.generateContentWithGrounding(searchMessages, requestedModel.slice(0, -"-search".length));
+        responseText = await this.googleAI.generateContentWithGrounding(searchMessages, requestedModel.slice(0, -"-search".length), maxTokens);
       } else {
-        responseText = await this.googleAI.generateContent(body.messages, requestedModel, false);
+        responseText = await this.googleAI.generateContent(body.messages, requestedModel, maxTokens, false);
       }
 
       if (stream) {