Upload 6 files

Dockerfile

@@ -1,8 +1,8 @@
-# 使用 node:20.16.0-alpine 作为基础镜像
-FROM node:20.16.0-alpine
+# 使用 Node.js 18 作为基础镜像
+FROM node:18-slim
 
 # 设置工作目录
-WORKDIR /app
+WORKDIR /usr/src/app
 
 # 复制 package.json 和 package-lock.json 到工作目录
 COPY package*.json ./
@@ -15,7 +15,7 @@ ENV API_PREFIX=/api \
   PORT=8787
 
 # 安装项目依赖
-RUN npm install
+RUN npm ci --only=production
 
 # 复制项目的源代码到工作目录
 COPY . .

api/index.js

@@ -1,303 +1,391 @@
-import grpc from '@grpc/grpc-js';
-import protoLoader from '@grpc/proto-loader';
-import {AutoRouter, cors, error, json} from 'itty-router';
-import dotenv from 'dotenv';
-import path,{ dirname } from 'path';
-import { fileURLToPath } from 'url';
-import {createServerAdapter} from '@whatwg-node/server';
-import {createServer} from 'http';
+import grpc from '@huayue/grpc-js'
+import protoLoader from '@grpc/proto-loader'
+import { AutoRouter, cors, error, json } from 'itty-router'
+import dotenv from 'dotenv'
+import path, { dirname } from 'path'
+import { fileURLToPath } from 'url'
+import { createServerAdapter } from '@whatwg-node/server'
+import { createServer } from 'http'
 
 // 加载环境变量
-dotenv.config();
+dotenv.config()
 // 获取当前文件的目录路径（ESM 方式）
-const __dirname = dirname(fileURLToPath(import.meta.url));
+const __dirname = dirname(fileURLToPath(import.meta.url))
+// 初始化配置
 // 初始化配置
 class Config {
-        constructor() {
-                this.API_PREFIX = process.env.API_PREFIX || '/';
-                this.API_KEY = process.env.API_KEY || '';
-                this.MAX_RETRY_COUNT = process.env.MAX_RETRY_COUNT || 3;
-                this.RETRY_DELAY = process.env.RETRY_DELAY || 5000;
-                this.COMMON_GRPC = 'runtime-native-io-vertex-inference-grpc-service-lmuw6mcn3q-ul.a.run.app';
-                this.COMMON_PROTO = path.join(__dirname,'..', 'protos', 'VertexInferenceService.proto')
-                this.GPT_GRPC = 'runtime-native-io-gpt-inference-grpc-service-lmuw6mcn3q-ul.a.run.app';
-                this.GPT_PROTO = path.join(__dirname,'..', 'protos', 'GPTInferenceService.proto')
-                this.PORT = process.env.PORT || 8787;
-        }
+  constructor() {
+    this.API_PREFIX = process.env.API_PREFIX || '/'
+    this.API_KEY = process.env.API_KEY || ''
+    this.MAX_RETRY_COUNT = process.env.MAX_RETRY_COUNT || 3
+    this.RETRY_DELAY = process.env.RETRY_DELAY || 5000
+    this.COMMON_GRPC = 'runtime-native-io-vertex-inference-grpc-service-lmuw6mcn3q-ul.a.run.app'
+    this.COMMON_PROTO = path.join(__dirname, '..', 'protos', 'VertexInferenceService.proto')
+    this.GPT_GRPC = 'runtime-native-io-gpt-inference-grpc-service-lmuw6mcn3q-ul.a.run.app'
+    this.GPT_PROTO = path.join(__dirname, '..', 'protos', 'GPTInferenceService.proto')
+    this.PORT = process.env.PORT || 8787
+    // 添加支持的模型列表
+    this.SUPPORTED_MODELS = process.env.SUPPORTED_MODELS || [
+      'gpt-4o-mini',
+      'gpt-4o',
+      'gpt-4-turbo',
+      'gpt-4',
+      'gpt-3.5-turbo',
+      'claude-3-sonnet@20240229',
+      'claude-3-opus@20240229',
+      'claude-3-haiku@20240307',
+      'claude-3-5-sonnet@20240620',
+      'gemini-1.5-flash',
+      'gemini-1.5-pro',
+      'chat-bison',
+      'codechat-bison',
+    ]
+  }
+
+  // 添加模型验证方法
+  isValidModel(model) {
+    // 处理 Claude 模型的特殊格式
+    const RegexInput = /^(claude-3-(5-sonnet|haiku|sonnet|opus))-(\d{8})$/
+    const matchInput = model.match(RegexInput)
+    const normalizedModel = matchInput ? `${matchInput[1]}@${matchInput[3]}` : model
+
+    return this.SUPPORTED_MODELS.includes(normalizedModel)
+  }
 }
 class GRPCHandler {
-        constructor(protoFilePath) {
-                // 动态加载传入的 .proto 文件路径
-                this.packageDefinition = protoLoader.loadSync(protoFilePath, {
-                        keepCase: true,
-                        longs: String,
-                        enums: String,
-                        defaults: true,
-                        oneofs: true
-                });
-        }
+  constructor(protoFilePath) {
+    // 动态加载传入的 .proto 文件路径
+    this.packageDefinition = protoLoader.loadSync(protoFilePath, {
+      keepCase: true,
+      longs: String,
+      enums: String,
+      defaults: true,
+      oneofs: true,
+    })
+  }
 }
-const config = new Config();
+const config = new Config()
 // 中间件
 // 添加运行回源
 const { preflight, corsify } = cors({
-	origin: '*',
-	allowMethods: '*',
-	exposeHeaders: '*',
-});
+  origin: '*',
+  allowMethods: '*',
+  exposeHeaders: '*',
+})
 
 // 添加认证
 const withAuth = (request) => {
-	if (config.API_KEY) {
-		const authHeader = request.headers.get('Authorization');
-		if (!authHeader || !authHeader.startsWith('Bearer ')) {
-			return error(401, 'Unauthorized: Missing or invalid Authorization header');
-		}
-		const token = authHeader.substring(7);
-		if (token !== config.API_KEY) {
-			return error(403, 'Forbidden: Invalid API key');
-		}
-	}
-};
+  if (config.API_KEY) {
+    const authHeader = request.headers.get('Authorization')
+    if (!authHeader || !authHeader.startsWith('Bearer ')) {
+      return error(401, 'Unauthorized: Missing or invalid Authorization header')
+    }
+    const token = authHeader.substring(7)
+    if (token !== config.API_KEY) {
+      return error(403, 'Forbidden: Invalid API key')
+    }
+  }
+}
 // 返回运行信息
 const logger = (res, req) => {
-	console.log(req.method, res.status, req.url, Date.now() - req.start, 'ms');
-};
+  console.log(req.method, res.status, req.url, Date.now() - req.start, 'ms')
+}
 const router = AutoRouter({
-	before: [preflight, withAuth],
-	missing: () => error(404, '404 not found.'),
-	finally: [corsify, logger],
-});
+  before: [preflight, withAuth],
+  missing: () => error(404, '404 not found.'),
+  finally: [corsify, logger],
+})
 // Router路径
-router.get('/', () => json({ message: 'API 服务运行中~' }));
-router.get('/ping', () => json({ message: 'pong' }));
+router.get('/', () => json({ message: 'API 服务运行中~' }))
+router.get('/ping', () => json({ message: 'pong' }))
 router.get(config.API_PREFIX + '/v1/models', () =>
-    json({
-            object: 'list',
-            data: [
-                    { id: "gpt-4o-mini", object: "model", owned_by: "pieces-os" },
-                    { id: "gpt-4o", object: "model", owned_by: "pieces-os" },
-                    { id: "gpt-4-turbo", object: "model", owned_by: "pieces-os" },
-                    { id: "gpt-4", object: "model", owned_by: "pieces-os" },
-                    { id: "gpt-3.5-turbo", object: "model", owned_by: "pieces-os" },
-                    { id: "claude-3-sonnet@20240229", object: "model", owned_by: "pieces-os" },
-                    { id: "claude-3-opus@20240229", object: "model", owned_by: "pieces-os" },
-                    { id: "claude-3-haiku@20240307", object: "model", owned_by: "pieces-os" },
-                    { id: "claude-3-5-sonnet@20240620", object: "model", owned_by: "pieces-os" },
-                    { id: "gemini-1.5-flash", object: "model", owned_by: "pieces-os" },
-                    { id: "gemini-1.5-pro", object: "model", owned_by: "pieces-os" },
-                    { id: "chat-bison", object: "model", owned_by: "pieces-os" },
-                    { id: "codechat-bison", object: "model", owned_by: "pieces-os" },
-            ],
-    })
-);
-router.post(config.API_PREFIX + '/v1/chat/completions', (req) => handleCompletion(req));
+  json({
+    object: 'list',
+    data: [
+      { id: 'gpt-4o-mini', object: 'model', owned_by: 'pieces-os' },
+      { id: 'gpt-4o', object: 'model', owned_by: 'pieces-os' },
+      { id: 'gpt-4-turbo', object: 'model', owned_by: 'pieces-os' },
+      { id: 'gpt-4', object: 'model', owned_by: 'pieces-os' },
+      { id: 'gpt-3.5-turbo', object: 'model', owned_by: 'pieces-os' },
+      { id: 'claude-3-sonnet@20240229', object: 'model', owned_by: 'pieces-os' },
+      { id: 'claude-3-opus@20240229', object: 'model', owned_by: 'pieces-os' },
+      { id: 'claude-3-haiku@20240307', object: 'model', owned_by: 'pieces-os' },
+      { id: 'claude-3-5-sonnet@20240620', object: 'model', owned_by: 'pieces-os' },
+      { id: 'gemini-1.5-flash', object: 'model', owned_by: 'pieces-os' },
+      { id: 'gemini-1.5-pro', object: 'model', owned_by: 'pieces-os' },
+      { id: 'chat-bison', object: 'model', owned_by: 'pieces-os' },
+      { id: 'codechat-bison', object: 'model', owned_by: 'pieces-os' },
+    ],
+  }),
+)
+router.post(config.API_PREFIX + '/v1/chat/completions', (req) => handleCompletion(req))
 
-async function GrpcToPieces(models, message, rules, stream, temperature, top_p) {
-        // 在非GPT类型的模型中，temperature和top_p是无效的
-        // 使用系统的根证书
-        const credentials = grpc.credentials.createSsl();
-        let client,request;
-        if (models.includes('gpt')){
-                // 加载proto文件
-                const packageDefinition = new GRPCHandler(config.GPT_PROTO).packageDefinition;
-                // 构建请求消息
-                request = {
-                        models: models,
-                        messages: [
-                                {role: 0, message: rules}, // system
-                                {role: 1, message: message} // user
-                        ],
-                        temperature:temperature || 0.1,
-                        top_p:top_p ?? 1,
-                }
-                // 获取gRPC对象
-                const GRPCobjects = grpc.loadPackageDefinition(packageDefinition).runtime.aot.machine_learning.parents.gpt;
-                client = new GRPCobjects.GPTInferenceService(config.GPT_GRPC, credentials);
-        } else {
-                // 加载proto文件
-                const packageDefinition = new GRPCHandler(config.COMMON_PROTO).packageDefinition;
-                // 构建请求消息
-                request = {
-                        models: models,
-                        args: {
-                                messages: {
-                                        unknown: 1,
-                                        message: message
-                                },
-                                rules: rules
-                        }
-                };
-                // ���取gRPC对象
-                const GRPCobjects = grpc.loadPackageDefinition(packageDefinition).runtime.aot.machine_learning.parents.vertex;
-                client = new GRPCobjects.VertexInferenceService(config.COMMON_GRPC, credentials);
-        }
-        return await ConvertOpenai(client,request,models,stream);
+async function GrpcToPieces(inputModel, OriginModel, message, rules, stream, temperature, top_p) {
+  // 在非GPT类型的模型中，temperature和top_p是无效的
+  // 使用系统的根证书
+  const credentials = grpc.credentials.createSsl()
+  let client, request
+  if (inputModel.includes('gpt')) {
+    // 加载proto文件
+    const packageDefinition = new GRPCHandler(config.GPT_PROTO).packageDefinition
+    // 构建请求消息
+    request = {
+      models: inputModel,
+      messages: [
+        { role: 0, message: rules }, // system
+        { role: 1, message: message }, // user
+      ],
+      temperature: temperature || 0.1,
+      top_p: top_p ?? 1,
+    }
+    // 获取gRPC对象
+    const GRPCobjects = grpc.loadPackageDefinition(packageDefinition).runtime.aot.machine_learning.parents.gpt
+    client = new GRPCobjects.GPTInferenceService(config.GPT_GRPC, credentials)
+  } else {
+    // 加载proto文件
+    const packageDefinition = new GRPCHandler(config.COMMON_PROTO).packageDefinition
+    // 构建请求消息
+    request = {
+      models: inputModel,
+      args: {
+        messages: {
+          unknown: 1,
+          message: message,
+        },
+        rules: rules,
+      },
+    }
+    // 获取gRPC对象
+    const GRPCobjects = grpc.loadPackageDefinition(packageDefinition).runtime.aot.machine_learning.parents.vertex
+    client = new GRPCobjects.VertexInferenceService(config.COMMON_GRPC, credentials)
+  }
+  return await ConvertOpenai(client, request, inputModel, OriginModel, stream)
 }
 
 async function messagesProcess(messages) {
-        let rules = '';
-        let message = '';
+  let rules = ''
+  let message = ''
 
-        for (const msg of messages) {
-                let role = msg.role;
-                // 格式化为字符串
-                const contentStr = Array.isArray(msg.content)
-                    ? msg.content
-                    .filter((item) => item.text)
-                    .map((item) => item.text)
-                    .join('') || ''
-                    : msg.content;
-                // 判断身份
-                if (role === 'system') {
-                        rules += `system:${contentStr};\r\n`;
-                } else if (['user', 'assistant'].includes(role)) {
-                        message += `${role}:${contentStr};\r\n`;
-                }
-        }
+  for (const msg of messages) {
+    let role = msg.role
+    // 格式化为字符串
+    const contentStr = Array.isArray(msg.content)
+      ? msg.content
+          .filter((item) => item.text)
+          .map((item) => item.text)
+          .join('') || ''
+      : msg.content
+    // 判断身份
+    if (role === 'system') {
+      rules += `system:${contentStr};\r\n`
+    } else if (['user', 'assistant'].includes(role)) {
+      message += `${role}:${contentStr};\r\n`
+    }
+  }
 
-        return { rules, message };
+  return { rules, message }
 }
 
-async function ConvertOpenai(client,request,model,stream) {
-        for (let i = 0; i < config.MAX_RETRY_COUNT; i++) {
-                try {
-                        if (stream) {
-                                const call = client.PredictWithStream(request);
-                                const encoder = new TextEncoder();
-                                const ReturnStream = new ReadableStream({
-                                    start(controller) {
-                                            call.on('data', (response) => {
-                                                    let response_code = Number(response.response_code);
-                                                    if (response_code === 204) {
-                                                            // 如果 response_code 是 204，关闭流
-                                                            controller.close()
-                                                            call.destroy()
-                                                    } else if (response_code === 200) {
-                                                            let response_message
-                                                            if (model.includes('gpt')) {
-                                                                    response_message = response.body.message_warpper.message.message;
-                                                            } else {
-                                                                    response_message = response.args.args.args.message;
-                                                            }
-                                                            // 否则，将数据块加入流中
-                                                            controller.enqueue(encoder.encode(`data: ${JSON.stringify(ChatCompletionStreamWithModel(response_message, model))}\n\n`));
-                                                    } else {
-                                                            controller.error(new Error(`Error: stream chunk is not success`));
-                                                            controller.close()
-                                                    }
-                                            })
-                                    }
-                                    });
-                                return new Response(ReturnStream, {
-                                        headers: {
-                                                'Content-Type': 'text/event-stream',
-                                        },
-                                })
+async function ConvertOpenai(client, request, inputModel, OriginModel, stream) {
+  const metadata = new grpc.Metadata()
+  metadata.set('User-Agent', 'dart-grpc/2.0.0')
+  for (let i = 0; i < config.MAX_RETRY_COUNT; i++) {
+    try {
+      if (stream) {
+        const call = client.PredictWithStream(request, metadata)
+        const encoder = new TextEncoder()
+        const ReturnStream = new ReadableStream({
+          start(controller) {
+            // 处理数据
+            call.on('data', (response) => {
+              try {
+                let response_code = Number(response.response_code)
+                if (response_code === 204) {
+                  controller.close()
+                  call.destroy()
+                } else if (response_code === 200) {
+                  let response_message
+                  if (inputModel.includes('gpt')) {
+                    response_message = response.body.message_warpper.message.message
+                  } else {
+                    response_message = response.args.args.args.message
+                  }
+                  controller.enqueue(
+                    encoder.encode(`data: ${JSON.stringify(ChatCompletionStreamWithModel(response_message, OriginModel))}\n\n`),
+                  )
                 } else {
-                        const call = await new Promise((resolve, reject) => {
-                                client.Predict(request, (err, response) => {
-                                        if (err) reject(err);
-                                        else resolve(response);
-                                });
-                        });
-                        let response_code = Number(call.response_code);
-                        if (response_code === 200) {
-                                let response_message
-                                if (model.includes('gpt')) {
-                                        response_message = call.body.message_warpper.message.message;
-                                } else {
-                                        response_message = call.args.args.args.message;
-                                }
-                                return new Response(JSON.stringify(ChatCompletionWithModel(response_message, model)), {
-                                                headers: {
-                                                        'Content-Type': 'application/json',
-                                                },
-                                        });
-                                }
-                        }
-                } catch (err) {
-                        console.error(err);
-                        await new Promise((resolve) => setTimeout(resolve, config.RETRY_DELAY));
+                  console.error(`Invalid response code: ${response_code}`)
+                  controller.error(error)
                 }
-        }
-        return error(500, err.message);
-}
+              } catch (error) {
+                console.error('Error processing stream data:', error)
+                controller.error(error)
+              }
+            })
+
+            // 处理错误
+            call.on('error', (error) => {
+              console.error('Stream error:', error)
+              // 如果是 INTERNAL 错误且包含 RST_STREAM，可能是正常的流结束
+              if (error.code === 13 && error.details.includes('RST_STREAM')) {
+                controller.close()
+              } else {
+                controller.error(error)
+              }
+              call.destroy()
+            })
+
+            // 处理结束
+            call.on('end', () => {
+              controller.close()
+            })
 
-function renameIfNeeded(input) {
-        // 替换的正则表达式
-        const regex = /^(claude-3-(5-sonnet|haiku|sonnet|opus))-(\d{8})$/;
-        const match = input.match(regex);
-        if (match) {
-                return `${match[1]}@${match[3]}`;
+            // 处理取消
+            return () => {
+              call.destroy()
+            }
+          },
+        })
+        return new Response(ReturnStream, {
+          headers: {
+            'Content-Type': 'text/event-stream',
+            Connection: 'keep-alive',
+            'Cache-Control': 'no-cache',
+            'Transfer-Encoding': 'chunked',
+          },
+        })
+      } else {
+        // 非流式调用保持不变
+        const call = await new Promise((resolve, reject) => {
+          client.Predict(request, metadata, (err, response) => {
+            if (err) reject(err)
+            else resolve(response)
+          })
+        })
+        let response_code = Number(call.response_code)
+        if (response_code === 200) {
+          let response_message
+          if (inputModel.includes('gpt')) {
+            response_message = call.body.message_warpper.message.message
+          } else {
+            response_message = call.args.args.args.message
+          }
+          return new Response(JSON.stringify(ChatCompletionWithModel(response_message, OriginModel)), {
+            headers: {
+              'Content-Type': 'application/json',
+            },
+          })
+        } else {
+          throw new Error(`Invalid response code: ${response_code}`)
         }
-        return input;
+      }
+    } catch (err) {
+      console.error(`Attempt ${i + 1} failed:`, err)
+      await new Promise((resolve) => setTimeout(resolve, config.RETRY_DELAY))
+    }
+  }
+  return new Response(
+    JSON.stringify({
+      error: {
+        message: 'An error occurred while processing your request',
+        type: 'server_error',
+        code: 'internal_error',
+        param: null,
+      },
+    }),
+    {
+      status: 500,
+      headers: {
+        'Content-Type': 'application/json',
+      },
+    },
+  )
 }
 
 function ChatCompletionWithModel(message, model) {
-        return {
-                id: 'Chat-Nekohy',
-                object: 'chat.completion',
-                created: Date.now(),
-                model,
-                usage: {
-                        prompt_tokens: 0,
-                        completion_tokens: 0,
-                        total_tokens: 0,
-                },
-                choices: [
-                        {
-                                message: {
-                                        content: message,
-                                        role: 'assistant',
-                                },
-                                index: 0,
-                        },
-                ],
-        };
+  return {
+    id: 'Chat-Nekohy',
+    object: 'chat.completion',
+    created: Date.now(),
+    model,
+    usage: {
+      prompt_tokens: 0,
+      completion_tokens: 0,
+      total_tokens: 0,
+    },
+    choices: [
+      {
+        message: {
+          content: message,
+          role: 'assistant',
+        },
+        index: 0,
+      },
+    ],
+  }
 }
 
 function ChatCompletionStreamWithModel(text, model) {
-        return {
-                id: 'chatcmpl-Nekohy',
-                object: 'chat.completion.chunk',
-                created: 0,
-                model,
-                choices: [
-                        {
-                                index: 0,
-                                delta: {
-                                        content: text,
-                                },
-                                finish_reason: null,
-                        },
-                ],
-        };
+  return {
+    id: 'chatcmpl-Nekohy',
+    object: 'chat.completion.chunk',
+    created: 0,
+    model,
+    choices: [
+      {
+        index: 0,
+        delta: {
+          content: text,
+        },
+        finish_reason: null,
+      },
+    ],
+  }
 }
 
 async function handleCompletion(request) {
-        try {
-                // todo stream逆向接口
-                // 解析openai格式API请求
-                const { model: inputModel, messages, stream,temperature,top_p} = await request.json();
-                console.log(inputModel,messages,stream)
-                const model = renameIfNeeded(inputModel);
-                // 解析system和user/assistant消息
-                const { rules, message:content } = await messagesProcess(messages);
-                console.log(rules,content)
-                // 响应码，回复的消息
-                return await GrpcToPieces(model, content, rules, stream, temperature, top_p);
-        } catch (err) {
-                return error(500, err.message);
-        }
+  try {
+    // 解析openai格式API请求
+    const { model: OriginModel, messages, stream, temperature, top_p } = await request.json()
+    const RegexInput = /^(claude-3-(5-sonnet|haiku|sonnet|opus))-(\d{8})$/
+    const matchInput = OriginModel.match(RegexInput)
+    const inputModel = matchInput ? `${matchInput[1]}@${matchInput[3]}` : OriginModel
+    // 添加模型验证
+    if (!config.isValidModel(inputModel)) {
+      return new Response(
+        JSON.stringify({
+          error: {
+            message: `Model '${OriginModel}' does not exist`,
+            type: 'invalid_request_error',
+            param: 'model',
+            code: 'model_not_found',
+          },
+        }),
+        {
+          status: 404,
+          headers: {
+            'Content-Type': 'application/json',
+          },
+        },
+      )
+    }
+    console.log(inputModel, messages, stream)
+    // 解析system和user/assistant消息
+    const { rules, message: content } = await messagesProcess(messages)
+    console.log(rules, content)
+    // 响应码，回复的消息
+    return await GrpcToPieces(inputModel, OriginModel, content, rules, stream, temperature, top_p)
+  } catch (err) {
+    return error(500, err.message)
+  }
 }
 
-(async () => {
-	//For Cloudflare Workers
-	if (typeof addEventListener === 'function') return;
-	// For Nodejs
-	const ittyServer = createServerAdapter(router.fetch);
-	console.log(`Listening on http://localhost:${config.PORT}`);
-	const httpServer = createServer(ittyServer);
-	httpServer.listen(config.PORT);
-})();
+;(async () => {
+  //For Cloudflare Workers
+  if (typeof addEventListener === 'function') return
+  // For Nodejs
+  const ittyServer = createServerAdapter(router.fetch)
+  console.log(`Listening on http://localhost:${config.PORT}`)
+  const httpServer = createServer(ittyServer)
+  httpServer.listen(config.PORT)
+})()

package.json

@@ -10,7 +10,7 @@
   "author": "Nekohy",
   "private": true,
   "dependencies": {
-    "@grpc/grpc-js": "^1.12.2",
+    "@huayue/grpc-js": "^1.12.1",
     "@grpc/proto-loader": "^0.7.13",
     "@whatwg-node/server": "^0.9.50",
     "dotenv": "^16.4.5",