Update api/index.js

api/index.js

@@ -173,46 +173,66 @@ router.get(config.API_PREFIX + '/v1/models', withAuth, () =>
 // chat 路由
 router.post(config.API_PREFIX + '/v1/chat/completions', withAuth, (req) => handleCompletion(req));
 
+function getMetadata() {
+  const metadata = new grpc.Metadata();
+  metadata.set('user-agent', 'dart-grpc/2.0.0');
+  return metadata;
+}
+
 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);
+    // 使用系统的根证书
+    const credentials = grpc.credentials.createSsl();
+
+    // 创建自定义metadata
+    function getMetadata() {
+        const metadata = new grpc.Metadata();
+        metadata.set('user-agent', 'dart-grpc/2.0.0');
+        return metadata;
+    }
+
+    const metadata = getMetadata();
+    const options = {
+        'grpc.primary_user_agent': 'dart-grpc/2.0.0'
+    };
+
+    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,
         }
-        return await ConvertOpenai(client,request,models,stream);
+        // 获取gRPC对象
+        const GRPCobjects = grpc.loadPackageDefinition(packageDefinition).runtime.aot.machine_learning.parents.gpt;
+        client = new GRPCobjects.GPTInferenceService(config.GPT_GRPC, credentials, options);
+    } 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, options);
+    }
+
+    return await ConvertOpenai(client, request, models, stream, metadata);
 }
 
 async function messagesProcess(messages) {
@@ -239,69 +259,80 @@ async function messagesProcess(messages) {
         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',
-                                        },
-                                })
-                } 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
+async function ConvertOpenai(client, request, model, stream, metadata) {
+    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) => {
+                            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 = call.body.message_warpper.message.message;
+                                    response_message = response.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',
-                                                },
-                                        });
+                                    response_message = response.args.args.args.message;
                                 }
-                        }
-                } catch (err) {
-                        console.error(err);
-                        await new Promise((resolve) => setTimeout(resolve, config.RETRY_DELAY));
+                                // 否则，将数据块加入流中
+                                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();
+                            }
+                        });
+                        call.on('error', (error) => {
+                            controller.error(error);
+                            controller.close();
+                        });
+                        call.on('end', () => {
+                            controller.close();
+                        });
+                    }
+                });
+                return new Response(ReturnStream, {
+                    headers: {
+                        'Content-Type': 'text/event-stream',
+                    },
+                });
+            } 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 (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',
+                        },
+                    });
+                } else {
+                    throw new Error(`Error: response code ${response_code}`);
                 }
+            }
+        } catch (err) {
+            console.error(err);
+            if (i === config.MAX_RETRY_COUNT - 1) {
+                return error(500, err.message);
+            }
+            await new Promise((resolve) => setTimeout(resolve, config.RETRY_DELAY));
         }
-        return error(500, err.message);
+    }
 }
 
 function ChatCompletionWithModel(message, model) {