Merge pull request #1 from juzibot/feat/videosupport

谷歌视频上传代理

Author: windmemory

src/google/google-proxy.controller.ts

@@ -1,16 +1,19 @@
 import {
   Body,
   Controller,
+  Get,
   Headers,
   HttpCode,
   HttpException,
   Inject,
+  Param,
   Post,
   Query,
-  Request,
+  Req,
+  Res,
   StreamableFile,
 } from '@nestjs/common';
-import { Request as ExpressRequest } from 'express';
+import { Request, Response } from 'express';
 import { GoogleProxyService } from './google-proxy.service';
 
 @Controller('google')
@@ -24,9 +27,9 @@ export class GoogleProxyController {
     @Body() body: any,
     @Headers() headers: any,
     @Query() query: any,
-    @Request() req: ExpressRequest,
+    @Req() req: Request,
   ) {
-    const params = req.params;
+    const params = req.params as any;
     const reqParams = params.reqParams;
     const [model, method] = reqParams.split(':');
 
@@ -50,4 +53,65 @@ export class GoogleProxyController {
       throw new HttpException('Method not found', 404);
     }
   }
+
+  @Post('upload/v1beta/files')
+  async uploadFileInit(
+    @Query() query: any,
+    @Body() body: any,
+    @Headers() headers: any,
+    @Res({ passthrough: true }) res: Response,
+    @Req() req: any,
+  ) {
+    if (query.upload_id) {
+      const queryString = new URLSearchParams(query).toString();
+      const uploadUrl = `https://generativelanguage.googleapis.com/upload/v1beta/files?${queryString}`;
+      const bufferBody = req.body;
+      return this.service.uploadFileData(uploadUrl, bufferBody, headers);
+    }
+    const result = await this.service.uploadFileInit(body, headers);
+    res.status(result.status);
+    Object.keys(result.headers).forEach(key => {
+      res.setHeader(key, result.headers[key]);
+    });
+    return result.data;
+  }
+
+  @Post('upload/v1beta/files/:path(*)')
+  async uploadFileData(
+    @Param('path') path: string,
+    @Body() body: Buffer,
+    @Headers() headers: any,
+    @Query() query: any,
+  ) {
+    const uploadUrl = `https://generativelanguage.googleapis.com/upload/v1beta/files/${path}`;
+    const chunkIndex = query.chunk_index ? parseInt(query.chunk_index) : 0;
+    const totalChunks = query.total_chunks ? parseInt(query.total_chunks) : 1;
+    if (totalChunks > 1) {
+      return this.service.uploadFileChunk(uploadUrl, body, headers, chunkIndex, totalChunks);
+    } else {
+      return this.service.uploadFileData(uploadUrl, body, headers);
+    }
+  }
+
+  @Post('upload/v1beta/files/chunk/:path(*)')
+  async uploadFileChunk(
+    @Param('path') path: string,
+    @Body() body: Buffer,
+    @Headers() headers: any,
+    @Query() query: any,
+  ) {
+    const uploadUrl = `https://generativelanguage.googleapis.com/upload/v1beta/files/${path}`;
+    const chunkIndex = parseInt(query.chunk_index || '0');
+    const totalChunks = parseInt(query.total_chunks || '1');
+    return this.service.uploadFileChunk(uploadUrl, body, headers, chunkIndex, totalChunks);
+  }
+
+  @Get('v1beta/:path(*)')
+  async getFile(
+    @Param('path') path: string,
+    @Headers() headers: any,
+  ) {
+    const url = `https://generativelanguage.googleapis.com/v1beta/${path}`;
+    return this.service.getFileInfo(url, headers);
+  }
 }

src/google/google-proxy.service.ts

@@ -13,7 +13,6 @@ export class GoogleProxyService {
 
   async generateContent(body: any, headers: any, query: any, model: string) {
     const url = `https://generativelanguage.googleapis.com/v1beta/models/${model}:generateContent`;
-
     return this.makeRequest(url, headers, body, query, false);
   }
 
@@ -24,55 +23,188 @@ export class GoogleProxyService {
     model: string,
   ) {
     const url = `https://generativelanguage.googleapis.com/v1beta/models/${model}:streamGenerateContent`;
-
     return this.makeRequest(url, headers, body, query, true);
   }
 
+  async uploadFileInit(body: any, headers: any): Promise<{ status: number; headers: any; data: any }> {
+    const url = 'https://generativelanguage.googleapis.com/upload/v1beta/files';
+    const result = await this.makeRequest(url, headers, body, {}, false, {
+      customHeaders: {
+        'X-Goog-Upload-Protocol': headers['x-goog-upload-protocol'],
+        'X-Goog-Upload-Command': headers['x-goog-upload-command'],
+        'X-Goog-Upload-Header-Content-Length': headers['x-goog-upload-header-content-length'],
+        'X-Goog-Upload-Header-Content-Type': headers['x-goog-upload-header-content-type'],
+      },
+      validateStatus: (status) => status === 200 || status === 308,
+      timeout: 2 * MINUTE,
+      returnFullResponse: true,
+    });
+    return result;
+  }
+
+  async uploadFileData(uploadUrl: string, body: Buffer, headers: any) {
+    const contentLength = headers['content-length'] || body.length;
+    const uploadOffset = headers['x-goog-upload-offset'] || '0';
+    const uploadCommand = headers['x-goog-upload-command'] || 'upload, finalize';
+    const result = await this.makeRequest(uploadUrl, headers, body, {}, false, {
+      customHeaders: {
+        'Content-Length': contentLength.toString(),
+        'X-Goog-Upload-Offset': uploadOffset,
+        'X-Goog-Upload-Command': uploadCommand,
+      },
+      timeout: 5 * MINUTE,
+      maxContentLength: Infinity,
+      maxBodyLength: Infinity,
+      validateStatus: (status) => status === 200 || status === 308,
+      isBinaryData: true,
+    });
+    return result.data;
+  }
+
+  async uploadFileChunk(
+    uploadUrl: string,
+    chunkData: Buffer,
+    headers: any,
+    chunkIndex: number,
+    totalChunks: number
+  ) {
+    const contentLength = chunkData.length;
+    const uploadOffset = headers['x-goog-upload-offset'] || '0';
+    const isLastChunk = chunkIndex === totalChunks - 1;
+    const uploadCommand = isLastChunk ? 'upload, finalize' : 'upload';
+
+    const result = await this.makeRequest(uploadUrl, headers, chunkData, {}, false, {
+      customHeaders: {
+        'Content-Length': contentLength.toString(),
+        'X-Goog-Upload-Offset': uploadOffset,
+        'X-Goog-Upload-Command': uploadCommand,
+      },
+      timeout: 5 * MINUTE,
+      maxContentLength: Infinity,
+      maxBodyLength: Infinity,
+      validateStatus: (status) => status === 200 || status === 308,
+      isBinaryData: true,
+    });
+    return result.data;
+  }
+
+  async getFileInfo(url: string, headers: any) {
+    return this.makeRequest(url, headers, null, {}, false, {
+      method: 'GET',
+      timeout: 30000,
+    });
+  }
+
   private async makeRequest(
     url: string,
     headers: any,
     body: any,
     query: any,
     stream?: boolean,
+    options?: {
+      method?: string;
+      customHeaders?: Record<string, string>;
+      validateStatus?: (status: number) => boolean;
+      timeout?: number;
+      returnFullResponse?: boolean;
+      maxContentLength?: number;
+      maxBodyLength?: number;
+      isBinaryData?: boolean;
+    },
   ) {
     const { httpAgent, httpsAgent } = this.getAgents();
+    const {
+      method = 'POST',
+      customHeaders = {},
+      validateStatus = (status) => status === 200,
+      timeout = 10 * MINUTE,
+      returnFullResponse = false,
+      maxContentLength,
+      maxBodyLength,
+      isBinaryData = false,
+    } = options || {};
+
+    const requestHeaders: Record<string, string> = {
+      ...customHeaders,
+    };
+    if (!isBinaryData) {
+      requestHeaders['Content-Type'] = 'application/json';
+    }
+    if (headers['anthropic-version']) {
+      requestHeaders['anthropic-version'] = headers['anthropic-version'];
+    }
+    if (headers['x-api-key']) {
+      requestHeaders['x-api-key'] = headers['x-api-key'];
+    }
+    if (headers['x-goog-api-key']) {
+      requestHeaders['x-goog-api-key'] = headers['x-goog-api-key'];
+    }
+    const axiosConfig: any = {
+      httpAgent,
+      httpsAgent,
+      method,
+      headers: requestHeaders,
+      params: query,
+      responseType: stream ? 'stream' : 'json',
+      data: body,
+      timeout,
+      validateStatus,
+    };
+    if (maxContentLength !== undefined) {
+      axiosConfig.maxContentLength = maxContentLength;
+    }
+    if (maxBodyLength !== undefined) {
+      axiosConfig.maxBodyLength = maxBodyLength;
+    }
     let response: any;
     try {
-      response = await axios(url, {
-        httpAgent,
-        httpsAgent,
-        method: 'POST',
-        headers: {
-          'Content-Type': 'application/json',
-          'anthropic-version': headers['anthropic-version'],
-          'x-api-key': headers['x-api-key'],
-        },
-        params: query,
-        responseType: stream ? 'stream' : 'json',
-        data: body,
-        timeout: 10 * MINUTE,
-      });
+      response = await axios(url, axiosConfig);
     } catch (e) {
       if (e.response) {
         if (stream) {
           return e.response.data;
         }
-        throw new HttpException(e.response.data, e.response.status);
+        throw new HttpException({
+          message: e.response.data?.message || e.message,
+          data: e.response.data,
+          status: e.response.status,
+          statusText: e.response.statusText,
+          headers: e.response.headers,
+        }, e.response.status);
       } else if (e.request) {
-        console.log(e.message);
-        throw new Error(
-          `Failed to send message. error message: ${e.message}, request: ${e.request}`,
-        );
+        throw new HttpException({
+          message: `Network request failed: ${e.message}`,
+          code: e.code,
+          errno: e.errno,
+          syscall: e.syscall,
+          hostname: e.hostname,
+          port: e.port,
+          path: e.path,
+        }, 500);
       } else {
-        throw e;
+        throw new HttpException({
+          message: e.message,
+          stack: e.stack,
+          name: e.name,
+        }, 500);
       }
     }
-
-    if (response.status !== 200) {
-      const error = new HttpException(response.data, response.status);
-      throw error;
+    if (!validateStatus(response.status)) {
+      throw new HttpException({
+        message: response.data?.message || `HTTP ${response.status} Error`,
+        data: response.data,
+        status: response.status,
+        statusText: response.statusText,
+        headers: response.headers,
+      }, response.status);
     }
-    return response.data;
+    return returnFullResponse
+      ? {
+          status: response.status,
+          headers: response.headers,
+          data: response.data,
+        }
+      : response.data;
   }
 
   private getAgents() {

src/main.ts

@@ -2,18 +2,25 @@ import { NestFactory } from '@nestjs/core';
 import { ConfigService } from '@nestjs/config';
 import { AppModule } from './app.module';
 import { Logger, ValidationPipe } from '@nestjs/common';
-import { urlencoded, json } from 'express';
+import { urlencoded, json, raw } from 'express';
 
 async function bootstrap() {
   const app = await NestFactory.create(AppModule);
   app.enableShutdownHooks();
+
+  app.use((req, res, next) => {
+    if (req.path === '/google/upload/v1beta/files' && req.query.upload_id) {
+      raw({ type: '*/*', limit: '100mb' })(req, res, next);
+    } else {
+      next();
+    }
+  });
+
   app.use(json({ limit: '50mb' }));
   app.useGlobalPipes(new ValidationPipe({ transform: true }));
-  app.use(urlencoded({ extended: true, limit: '50mb' }));
-
+  app.use(urlencoded({ extended: true, limit: '50mb', parameterLimit: 50000 }));
   const configService = app.get(ConfigService);
   const port = configService.get<number>('port');
-
   await app.listen(port, () => {
     Logger.log(`Server started listening on port: ${port}`);
   });