Delete app.py

app.py

@@ -1,478 +0,0 @@
-import os
-import time
-import asyncio
-import grpc
-from dotenv import load_dotenv
-from fastapi import FastAPI, HTTPException, Depends
-from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import StreamingResponse
-from pydantic import BaseModel
-from typing import List, Optional
-from google.protobuf import descriptor_pb2
-from google.protobuf import descriptor
-from google.protobuf import descriptor_pool
-from google.protobuf import symbol_database
-from google.protobuf.compiler import plugin_pb2
-from google.protobuf.json_format import MessageToDict
-
-# 加载环境变量
-load_dotenv()
-
-# 内嵌 Proto 定义
-GPT_PROTO = """
-syntax = "proto3";
-package runtime.aot.machine_learning.parents.gpt;
-
-service GPTInferenceService {
-    rpc Predict (Request) returns (Response);
-    rpc PredictWithStream (Request) returns (stream Response);
-}
-
-message Message {
-    uint64 role = 1;
-    string message = 2;
-}
-
-message Request {
-    string models = 1;
-    repeated Message messages = 2;
-    double temperature = 3;
-    double top_p = 4;
-}
-
-message Response {
-    uint64 response_code = 2;
-    optional Body body = 4;
-}
-
-message Body {
-    string id = 1;
-    string object = 2;
-    uint64 time = 3;
-    MessageWarpper message_warpper = 4;
-    Unknown unknown = 5;
-}
-
-message MessageWarpper {
-    int64 arg1 = 1;
-    Message message = 2;
-}
-
-message Unknown {
-    int64 arg1 = 1;
-    int64 arg2 = 2;
-    int64 arg3 = 3;
-}
-"""
-
-VERTEX_PROTO = """
-syntax = "proto3";
-package runtime.aot.machine_learning.parents.vertex;
-
-service VertexInferenceService {
-    rpc Predict (Requests) returns (Response);
-    rpc PredictWithStream (Requests) returns (stream Response);
-}
-
-message Messages {
-    int64 unknown = 1;
-    string message = 2;
-}
-
-message Requests {
-    string models = 1;
-    Args args = 2;
-}
-
-message Args {
-    Messages messages = 2;
-    string rules = 3;
-}
-
-message Response {
-    int64 response_code = 2;
-    Args1 args = 4;
-}
-
-message Args1 {
-    Args2 args = 1;
-}
-
-message Args2 {
-    Messages args = 2;
-}
-"""
-
-# 配置类
-class Config:
-    API_PREFIX = os.getenv('API_PREFIX', '/')
-    API_KEY = os.getenv('API_KEY', '')
-    MAX_RETRY_COUNT = int(os.getenv('MAX_RETRY_COUNT', 3))
-    RETRY_DELAY = int(os.getenv('RETRY_DELAY', 5000))
-    COMMON_GRPC = 'runtime-native-io-vertex-inference-grpc-service-lmuw6mcn3q-ul.a.run.app'
-    GPT_GRPC = 'runtime-native-io-gpt-inference-grpc-service-lmuw6mcn3q-ul.a.run.app'
-    PORT = int(os.getenv('PORT', 8787))
-
-config = Config()
-
-# 动态生成 Proto
-def generate_proto_classes():
-    pool = descriptor_pool.DescriptorPool()
-    
-    # 为 GPT 服务创建文件描述符
-    gpt_file = descriptor_pb2.FileDescriptorProto()
-    gpt_file.name = "gpt_service.proto"
-    gpt_file.package = "runtime.aot.machine_learning.parents.gpt"
-    gpt_file.syntax = "proto3"
-
-    # GPT 服务消息定义
-    message = gpt_file.message_type.add()
-    message.name = "Message"
-    field = message.field.add()
-    field.name = "role"
-    field.number = 1
-    field.type = descriptor.FieldDescriptor.TYPE_UINT64
-    field = message.field.add()
-    field.name = "message"
-    field.number = 2
-    field.type = descriptor.FieldDescriptor.TYPE_STRING
-
-    # Request 消息
-    request = gpt_file.message_type.add()
-    request.name = "Request"
-    field = request.field.add()
-    field.name = "models"
-    field.number = 1
-    field.type = descriptor.FieldDescriptor.TYPE_STRING
-    field = request.field.add()
-    field.name = "messages"
-    field.number = 2
-    field.type = descriptor.FieldDescriptor.TYPE_MESSAGE
-    field.type_name = ".runtime.aot.machine_learning.parents.gpt.Message"
-    field.label = descriptor.FieldDescriptor.LABEL_REPEATED
-    field = request.field.add()
-    field.name = "temperature"
-    field.number = 3
-    field.type = descriptor.FieldDescriptor.TYPE_DOUBLE
-    field = request.field.add()
-    field.name = "top_p"
-    field.number = 4
-    field.type = descriptor.FieldDescriptor.TYPE_DOUBLE
-
-    # Response 消息
-    response = gpt_file.message_type.add()
-    response.name = "Response"
-    field = response.field.add()
-    field.name = "response_code"
-    field.number = 2
-    field.type = descriptor.FieldDescriptor.TYPE_UINT64
-    field = response.field.add()
-    field.name = "body"
-    field.number = 4
-    field.type = descriptor.FieldDescriptor.TYPE_MESSAGE
-    field.type_name = ".runtime.aot.machine_learning.parents.gpt.Body"
-    field.label = descriptor.FieldDescriptor.LABEL_OPTIONAL
-
-    # Body 消息
-    body = gpt_file.message_type.add()
-    body.name = "Body"
-    field = body.field.add()
-    field.name = "id"
-    field.number = 1
-    field.type = descriptor.FieldDescriptor.TYPE_STRING
-    field = body.field.add()
-    field.name = "object"
-    field.number = 2
-    field.type = descriptor.FieldDescriptor.TYPE_STRING
-    field = body.field.add()
-    field.name = "time"
-    field.number = 3
-    field.type = descriptor.FieldDescriptor.TYPE_UINT64
-    field = body.field.add()
-    field.name = "message_warpper"
-    field.number = 4
-    field.type = descriptor.FieldDescriptor.TYPE_MESSAGE
-    field.type_name = ".runtime.aot.machine_learning.parents.gpt.MessageWarpper"
-
-    # MessageWarpper 消息
-    message_wrapper = gpt_file.message_type.add()
-    message_wrapper.name = "MessageWarpper"
-    field = message_wrapper.field.add()
-    field.name = "arg1"
-    field.number = 1
-    field.type = descriptor.FieldDescriptor.TYPE_INT64
-    field = message_wrapper.field.add()
-    field.name = "message"
-    field.number = 2
-    field.type = descriptor.FieldDescriptor.TYPE_MESSAGE
-    field.type_name = ".runtime.aot.machine_learning.parents.gpt.Message"
-
-    # GPT 服务定义
-    service = gpt_file.service.add()
-    service.name = "GPTInferenceService"
-    method = service.method.add()
-    method.name = "Predict"
-    method.input_type = ".runtime.aot.machine_learning.parents.gpt.Request"
-    method.output_type = ".runtime.aot.machine_learning.parents.gpt.Response"
-    method = service.method.add()
-    method.name = "PredictWithStream"
-    method.input_type = ".runtime.aot.machine_learning.parents.gpt.Request"
-    method.output_type = ".runtime.aot.machine_learning.parents.gpt.Response"
-    method.server_streaming = True
-
-    # 将文件描述符添加到池中
-    pool.Add(gpt_file)
-
-    # Vertex 服务的定义类似...
-    # 这里省略 Vertex 服务的定义，原理相同
-
-    return pool
-
-# 生成 Proto 类
-proto_pool = generate_proto_classes()
-
-# FastAPI 应用
-app = FastAPI()
-
-# CORS 中间件
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-# 设置 API 前缀
-API_PREFIX = "/ai"
-
-# 创建一个路由器
-router = APIRouter(prefix=API_PREFIX)
-
-# 认证依赖
-def verify_api_key(authorization: str = None):
-    if config.API_KEY:
-        if not authorization or not authorization.startswith('Bearer '):
-            raise HTTPException(status_code=401, detail='Unauthorized: Missing or invalid Authorization header')
-        token = authorization.split(' ')[1]
-        if token != config.API_KEY:
-            raise HTTPException(status_code=403, detail='Forbidden: Invalid API key')
-
-# 模型列表
-MODELS = [
-    {"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"},
-]
-
-# API 路由
-@app.get("/")
-async def root():
-    return {"message": "API 服务运行中~"}
-
-@app.get("/ping")
-async def ping():
-    return {"message": "pong"}
-
-@app.get(f"{config.API_PREFIX}/v1/models")
-async def get_models():
-    return {"object": "list", "data": MODELS}
-
-# 请求模型
-class Message(BaseModel):
-    role: str
-    content: str
-
-class ChatCompletionRequest(BaseModel):
-    model: str
-    messages: List[Message]
-    stream: Optional[bool] = False
-    temperature: Optional[float] = 0.1
-    top_p: Optional[float] = 1.0
-
-@router.post("/v1/chat/completions")
-async def chat_completions(request: ChatCompletionRequest):
-    try:
-        rules, content = process_messages(request.messages)
-        return await grpc_to_pieces(request.model, content, rules, request.stream, request.temperature, request.top_p)
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
-
-app.include_router(router)
-
-def process_messages(messages):
-    rules = ''
-    content = ''
-    for msg in messages:
-        if msg.role == 'system':
-            rules += f"system:{msg.content};\r\n"
-        elif msg.role in ['user', 'assistant']:
-            content += f"{msg.role}:{msg.content};\r\n"
-    return rules, content
-
-async def grpc_to_pieces(model, message, rules, stream, temperature, top_p):
-    credentials = grpc.ssl_channel_credentials()
-    
-    try:
-        if 'gpt' in model:
-            channel = grpc.secure_channel(config.GPT_GRPC, credentials)
-            stub = GPTInferenceServiceStub(channel)
-            
-            # 创建 GPT 请求
-            request = {
-                'models': model,
-                'messages': [
-                    {'role': 0, 'message': rules},
-                    {'role': 1, 'message': message}
-                ],
-                'temperature': temperature,
-                'top_p': top_p
-            }
-        else:
-            channel = grpc.secure_channel(config.COMMON_GRPC, credentials)
-            stub = VertexInferenceServiceStub(channel)
-            
-            # 创建 Vertex 请求
-            request = {
-                'models': model,
-                'args': {
-                    'messages': {'unknown': 1, 'message': message},
-                    'rules': rules
-                }
-            }
-
-        for _ in range(config.MAX_RETRY_COUNT):
-            try:
-                if stream:
-                    return await stream_response(stub, request, model)
-                else:
-                    return await single_response(stub, request, model)
-            except Exception as e:
-                print(f"Error: {e}")
-                await asyncio.sleep(config.RETRY_DELAY / 1000)
-                continue
-        
-        raise HTTPException(status_code=500, detail="Max retry count reached")
-    
-    finally:
-        channel.close()
-
-async def stream_response(stub, request, model):
-    async def generate():
-        try:
-            responses = stub.PredictWithStream(request)
-            for response in responses:
-                response_code = response.response_code
-                if response_code == 204:
-                    break
-                elif response_code == 200:
-                    if 'gpt' in model:
-                        message = response.body.message_warpper.message.message
-                    else:
-                        message = response.args.args.args.message
-                    
-                    chunk = {
-                        "id": "chatcmpl-Nekohy",
-                        "object": "chat.completion.chunk",
-                        "created": 0,
-                        "model": model,
-                        "choices": [{
-                            "index": 0,
-                            "delta": {
-                                "content": message,
-                            },
-                            "finish_reason": None,
-                        }],
-                    }
-                    yield f"data: {json.dumps(chunk)}\n\n"
-        except Exception as e:
-            raise HTTPException(status_code=500, detail=str(e))
-        
-        yield "data: [DONE]\n\n"
-
-    return StreamingResponse(
-        generate(),
-        media_type="text/event-stream",
-        headers={
-            "Cache-Control": "no-cache",
-            "Connection": "keep-alive",
-        }
-    )
-
-async def single_response(stub, request, model):
-    response = await asyncio.get_event_loop().run_in_executor(
-        None, stub.Predict, request
-    )
-    
-    if response.response_code == 200:
-        if 'gpt' in model:
-            message = response.body.message_warpper.message.message
-        else:
-            message = response.args.args.args.message
-            
-        return {
-            "id": "Chat-Nekohy",
-            "object": "chat.completion",
-            "created": int(time.time()),
-            "model": model,
-            "usage": {
-                "prompt_tokens": 0,
-                "completion_tokens": 0,
-                "total_tokens": 0,
-            },
-            "choices": [{
-                "message": {
-                    "content": message,
-                    "role": "assistant",
-                },
-                "index": 0,
-            }],
-        }
-    else:
-        raise HTTPException(
-            status_code=500,
-            detail=f"Error response code: {response.response_code}"
-        )
-
-# gRPC 服务存根类
-class GPTInferenceServiceStub:
-    def __init__(self, channel):
-        self.channel = channel
-        self.stub = grpc.ProtoRPC(channel)
-
-    def Predict(self, request):
-        return self.stub.Predict(request)
-
-    def PredictWithStream(self, request):
-        return self.stub.PredictWithStream(request)
-
-class VertexInferenceServiceStub:
-    def __init__(self, channel):
-        self.channel = channel
-        self.stub = grpc.ProtoRPC(channel)
-
-    def Predict(self, request):
-        return self.stub.Predict(request)
-
-    def PredictWithStream(self, request):
-        return self.stub.PredictWithStream(request)
-
-if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(
-        app,
-        host="0.0.0.0",
-        port=config.PORT,
-        log_level="info"
-    )
-