Update api/index.js

api/index.js

@@ -1,4 +1,4 @@
-import grpc from '@grpc/grpc-js';
+import grpc from '@huayue/grpc-js';
 import protoLoader from '@grpc/proto-loader';
 import {AutoRouter, cors, error, json} from 'itty-router';
 import dotenv from 'dotenv';
@@ -12,6 +12,7 @@ dotenv.config();
 // 获取当前文件的目录路径（ESM 方式）
 const __dirname = dirname(fileURLToPath(import.meta.url));
 // 初始化配置
+// 初始化配置
 class Config {
         constructor() {
                 this.API_PREFIX = process.env.API_PREFIX || '/';
@@ -23,6 +24,32 @@ class Config {
                 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 {
@@ -59,167 +86,80 @@ const withAuth = (request) => {
 		}
 	}
 };
-
 // 返回运行信息
 const logger = (res, req) => {
 	console.log(req.method, res.status, req.url, Date.now() - req.start, 'ms');
 };
-
-// 定义模型映射信息
-const MODEL_INFO = {
-    "claude-3-sonnet-20240229": {
-        "provider": "anthropic",
-        "mapping": "claude-3-sonnet@20240229"
-    },
-    "claude-3-opus-20240229": {
-        "provider": "anthropic",
-        "mapping": "claude-3-opus@20240229"
-    },
-    "claude-3-haiku-20240307": {
-        "provider": "anthropic",
-        "mapping": "claude-3-haiku@20240307"
-    },
-    "claude-3-5-sonnet-20240620": {
-        "provider": "anthropic",
-        "mapping": "claude-3-5-sonnet@20240620"
-    },
-    "gpt-4o-mini": {
-        "provider": "openai",
-        "mapping": "gpt-4o-mini"
-    },
-    "gpt-4o": {
-        "provider": "openai",
-        "mapping": "gpt-4o"
-    },
-    "gpt-4-turbo": {
-        "provider": "openai",
-        "mapping": "gpt-4-turbo"
-    },
-    "gpt-4": {
-        "provider": "openai",
-        "mapping": "gpt-4"
-    },
-    "gpt-3.5-turbo": {
-        "provider": "openai",
-        "mapping": "gpt-3.5-turbo"
-    },
-    "gemini-1.5-pro": {
-        "provider": "google",
-        "mapping": "gemini-1.5-pro"
-    },
-    "gemini-1.5-flash": {
-        "provider": "google",
-        "mapping": "gemini-1.5-flash"
-    },
-    "chat-bison": {
-        "provider": "pieces-os",
-        "mapping": "chat-bison"
-    },
-    "codechat-bison": {
-        "provider": "pieces-os",
-        "mapping": "codechat-bison"
-    }
-};
-
-// 定义路由
 const router = AutoRouter({
-  before: [preflight], // 只保留 CORS preflight 检查
-  missing: () => error(404, '404 not found.'),
-  finally: [corsify, logger],
+	before: [preflight, withAuth],
+	missing: () => error(404, '404 not found.'),
+	finally: [corsify, logger],
 });
-
-// 根路由
-router.get('/', () => json({
-    service: "AI Chat Completion Proxy",
-    usage: {
-        endpoint: "/v1/chat/completions", 
-        method: "POST",
-        headers: {
-            "Content-Type": "application/json",
-            "Authorization": "Bearer YOUR_API_KEY"
-        },
-        body: {
-            model: "One of: " + Object.keys(MODEL_INFO).join(", "),
-            messages: [
-                { role: "system", content: "You are a helpful assistant." },
-                { role: "user", content: "Hello, who are you?" }
-            ],
-            stream: false,
-            temperature: 0.7,
-            top_p: 1
-        }
-    },
-    note: "Replace YOUR_API_KEY with your actual API key."
-}));
-
-// models 路由
-router.get(config.API_PREFIX + '/v1/models', withAuth, () =>
+// Router路径
+router.get('/', () => json({ message: 'API 服务运行中~' }));
+router.get('/ping', () => json({ message: 'pong' }));
+router.get(config.API_PREFIX + '/v1/models', () =>
     json({
-        object: "list",
-        data: Object.entries(MODEL_INFO).map(([modelId, info]) => ({
-            id: modelId,
-            object: "model",
-            created: Date.now(),
-            owned_by: "pieces-os",
-            permission: [],
-            root: modelId,
-            parent: null,
-            mapping: info.mapping,
-            provider: info.provider
-        }))
+            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));
 
-// chat 路由
-router.post(config.API_PREFIX + '/v1/chat/completions', withAuth, (req) => handleCompletion(req));
-
-async function GrpcToPieces(models, message, rules, stream, temperature, top_p) {
-    const credentials = grpc.credentials.createSsl();
-
-    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')) {
-        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,
-        };
-        const GRPCobjects = grpc.loadPackageDefinition(packageDefinition).runtime.aot.machine_learning.parents.gpt;
-        client = new GRPCobjects.GPTInferenceService(config.GPT_GRPC, credentials, options);
-    } else {
-        const packageDefinition = new GRPCHandler(config.COMMON_PROTO).packageDefinition;
-        request = {
-            models: models,
-            args: {
-                messages: {
-                    unknown: 1,
-                    message: message
-                },
-                rules: rules
-            }
-        };
-        const GRPCobjects = grpc.loadPackageDefinition(packageDefinition).runtime.aot.machine_learning.parents.vertex;
-        client = new GRPCobjects.VertexInferenceService(config.COMMON_GRPC, credentials, options);
-    }
-
-    console.log('Request:', JSON.stringify(request, null, 2));
-    return await ConvertOpenai(client, request, models, stream, metadata);
+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) {
@@ -246,89 +186,120 @@ async function messagesProcess(messages) {
         return { rules, 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) => {
-                            console.log('Stream response:', JSON.stringify(response, null, 2));
-                            let response_code = Number(response.response_code);
-                            if (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;
-                                }
-                                if (response_message) {
-                                    controller.enqueue(encoder.encode(`data: ${JSON.stringify(ChatCompletionStreamWithModel(response_message, model))}\n\n`));
+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 {
+                                                                        console.error(`Invalid response code: ${response_code}`);
+                                                                        controller.error(error);
+                                                                }
+                                                        } 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();
+                                                });
+
+                                                // 处理取消
+                                                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 {
-                                controller.error(new Error(`Error: stream chunk response code ${response_code}`));
-                                controller.close();
-                            }
-                        });
-                        call.on('error', (error) => {
-                            console.error('Stream 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);
-                    });
-                });
-                console.log('Non-stream response:', JSON.stringify(call, null, 2));
-                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;
-                    }
-                    if (response_message) {
-                        return new Response(JSON.stringify(ChatCompletionWithModel(response_message, model)), {
-                            headers: {
-                                'Content-Type': 'application/json',
-                            },
-                        });
-                    } else {
-                        throw new Error('Response message is empty or undefined');
-                    }
-                } else {
-                    throw new Error(`Error: response code ${response_code}`);
+                        }
+                } catch (err) {
+                        console.error(`Attempt ${i + 1} failed:`, err);
+                        if (i === config.MAX_RETRY_COUNT - 1) {
+                                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'
+                                        }
+                                });
+                        }
+                        await new Promise((resolve) => setTimeout(resolve, config.RETRY_DELAY));
                 }
-            }
-        } catch (err) {
-            console.error(`Attempt ${i + 1} failed:`, err);
-            if (i === config.MAX_RETRY_COUNT - 1) {
-                return error(500, `All retry attempts failed. Last error: ${err.message}`);
-            }
-            await new Promise((resolve) => setTimeout(resolve, config.RETRY_DELAY));
         }
-    }
 }
 
+
 function ChatCompletionWithModel(message, model) {
         return {
                 id: 'Chat-Nekohy',
@@ -371,25 +342,41 @@ function ChatCompletionStreamWithModel(text, model) {
 }
 
 async function handleCompletion(request) {
-    try {
-        const { model: inputModel, messages, stream, temperature, top_p } = await request.json();
-        
-        // 获取模型映射
-        const modelInfo = MODEL_INFO[inputModel];
-        if (!modelInfo) {
-            return error(400, `Unsupported model: ${inputModel}`);
+        try {
+                // todo stream逆向接口
+                // 解析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);
         }
-        
-        const mappedModel = modelInfo.mapping;
-        
-        // 解析 system 和 user/assistant 消息
-        const { rules, message: content } = await messagesProcess(messages);
-        
-        // 使用映射后的模型名称
-        return await GrpcToPieces(mappedModel, content, rules, stream, temperature, top_p);
-    } catch (err) {
-        return error(500, err.message);
-    }
 }
 
 (async () => {