Create app.py

app.py

@@ -0,0 +1,269 @@
+import os
+import time
+import json
+import grpc
+import asyncio
+from typing import List, Optional
+from fastapi import FastAPI, HTTPException, Request
+from fastapi.responses import JSONResponse, StreamingResponse
+from pydantic import BaseModel
+from dotenv import load_dotenv
+from grpc_tools import protoc
+import re
+
+# 加载环境变量
+load_dotenv()
+
+# 配置类
+class Config:
+    def __init__(self):
+        self.API_PREFIX = os.getenv('API_PREFIX', '/')
+        self.API_KEY = os.getenv('API_KEY', '')
+        self.MAX_RETRY_COUNT = int(os.getenv('MAX_RETRY_COUNT', 3))
+        self.RETRY_DELAY = int(os.getenv('RETRY_DELAY', 5000))
+        self.COMMON_GRPC = 'runtime-native-io-vertex-inference-grpc-service-lmuw6mcn3q-ul.a.run.app'
+        self.COMMON_PROTO = 'protos/VertexInferenceService.proto'
+        self.GPT_GRPC = 'runtime-native-io-gpt-inference-grpc-service-lmuw6mcn3q-ul.a.run.app'
+        self.GPT_PROTO = 'protos/GPTInferenceService.proto'
+        self.PORT = int(os.getenv('PORT', 8787))
+        self.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"
+        ]
+
+    def is_valid_model(self, model):
+        regex_input = r'^(claude-3-(5-sonnet|haiku|sonnet|opus))-(\d{8})$'
+        match_input = re.match(regex_input, model)
+        normalized_model = f"{match_input.group(1)}@{match_input.group(3)}" if match_input else model
+        return normalized_model in self.SUPPORTED_MODELS
+
+# gRPC处理类
+class GRPCHandler:
+    def __init__(self, proto_file):
+        self.proto_file = proto_file
+        self._compile_proto()
+        self._load_proto()
+
+    def _compile_proto(self):
+        proto_dir = os.path.dirname(self.proto_file)
+        proto_file = os.path.basename(self.proto_file)
+        protoc.main((
+            '',
+            f'-I{proto_dir}',
+            f'--python_out=.',
+            f'--grpc_python_out=.',
+            os.path.join(proto_dir, proto_file)
+        ))
+
+    def _load_proto(self):
+        module_name = os.path.splitext(os.path.basename(self.proto_file))[0] + '_pb2_grpc'
+        proto_module = __import__(module_name)
+        self.stub_class = getattr(proto_module, f"{module_name.split('_')[0]}Stub")
+
+    async def grpc_to_pieces(self, model, content, rules, temperature, top_p):
+        channel = grpc.aio.secure_channel(
+            config.COMMON_GRPC if not model.startswith('gpt') else config.GPT_GRPC,
+            grpc.ssl_channel_credentials()
+        )
+        stub = self.stub_class(channel)
+
+        try:
+            request = self._build_request(model, content, rules, temperature, top_p)
+            response = await stub.Predict(request)
+            return self._process_response(response, model)
+        except grpc.RpcError as e:
+            print(f"RPC failed: {e}")
+            return {"error": str(e)}
+        finally:
+            await channel.close()
+
+    async def grpc_to_pieces_stream(self, model, content, rules, temperature, top_p):
+        channel = grpc.aio.secure_channel(
+            config.COMMON_GRPC if not model.startswith('gpt') else config.GPT_GRPC,
+            grpc.ssl_channel_credentials()
+        )
+        stub = self.stub_class(channel)
+
+        try:
+            request = self._build_request(model, content, rules, temperature, top_p)
+            async for response in stub.PredictWithStream(request):
+                result = self._process_stream_response(response, model)
+                if result:
+                    yield f"data: {json.dumps(result)}\n\n"
+        except grpc.RpcError as e:
+            print(f"Stream RPC failed: {e}")
+            yield f"data: {json.dumps({'error': str(e)})}\n\n"
+        finally:
+            await channel.close()
+
+    def _build_request(self, model, content, rules, temperature, top_p):
+        if model.startswith('gpt'):
+            return self.stub_class.Request(
+                models=model,
+                messages=[
+                    {"role": 0, "message": rules},
+                    {"role": 1, "message": content}
+                ],
+                temperature=temperature or 0.1,
+                top_p=top_p or 1.0
+            )
+        else:
+            return self.stub_class.Request(
+                models=model,
+                args={
+                    "messages": {
+                        "unknown": 1,
+                        "message": content
+                    },
+                    "rules": rules
+                }
+            )
+
+    def _process_response(self, response, model):
+        if response.response_code == 200:
+            if model.startswith('gpt'):
+                message = response.body.message_warpper.message.message
+            else:
+                message = response.args.args.args.message
+            return chat_completion_with_model(message, model)
+        return {"error": f"Invalid response code: {response.response_code}"}
+
+    def _process_stream_response(self, response, model):
+        if response.response_code == 204:
+            return None
+        elif response.response_code == 200:
+            if model.startswith('gpt'):
+                message = response.body.message_warpper.message.message
+            else:
+                message = response.args.args.args.message
+            return chat_completion_stream_with_model(message, model)
+        else:
+            return {"error": f"Invalid response code: {response.response_code}"}
+
+# 工具函数
+def messages_process(messages):
+    rules = ''
+    message = ''
+
+    for msg in messages:
+        role = msg.role
+        content = msg.content
+        
+        if isinstance(content, list):
+            content = ''.join([item.get('text', '') for item in content if item.get('text')])
+        
+        if role == 'system':
+            rules += f"system:{content};\r\n"
+        elif role in ['user', 'assistant']:
+            message += f"{role}:{content};\r\n"
+
+    return rules, message
+
+def chat_completion_with_model(message: str, model: str):
+    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,
+            },
+        ],
+    }
+
+def chat_completion_stream_with_model(text: str, model: str):
+    return {
+        "id": "chatcmpl-Nekohy",
+        "object": "chat.completion.chunk",
+        "created": 0,
+        "model": model,
+        "choices": [
+            {
+                "index": 0,
+                "delta": {
+                    "content": text,
+                },
+                "finish_reason": None,
+            },
+        ],
+    }
+
+# 初始化配置
+config = Config()
+
+# 初始化 FastAPI 应用
+app = FastAPI()
+
+# 定义请求模型
+class ChatMessage(BaseModel):
+    role: str
+    content: str
+
+class ChatCompletionRequest(BaseModel):
+    model: str
+    messages: List[ChatMessage]
+    stream: Optional[bool] = False
+    temperature: Optional[float] = None
+    top_p: Optional[float] = None
+
+# 路由定义
+@app.get("/")
+async def root():
+    return {"message": "API 服务运行中~"}
+
+@app.get("/ping")
+async def ping():
+    return {"message": "pong"}
+
+@app.get(config.API_PREFIX + "/v1/models")
+async def list_models():
+    with open('cloud_model.json', 'r') as f:
+        cloud_models = json.load(f)
+    
+    models = [
+        {"id": model["unique"], "object": "model", "owned_by": "pieces-os"}
+        for model in cloud_models["iterable"]
+    ]
+    
+    return JSONResponse({
+        "object": "list",
+        "data": models
+    })
+
+@app.post(config.API_PREFIX + "/v1/chat/completions")
+async def chat_completions(request: ChatCompletionRequest):
+    if not config.is_valid_model(request.model):
+        raise HTTPException(status_code=404, detail=f"Model '{request.model}' does not exist")
+
+    rules, content = messages_process(request.messages)
+    
+    grpc_handler = GRPCHandler(config.COMMON_PROTO if not request.model.startswith('gpt') else config.GPT_PROTO)
+    
+    if request.stream:
+        return StreamingResponse(
+            grpc_handler.grpc_to_pieces_stream(
+                request.model, content, rules, request.temperature, request.top_p
+            ),
+            media_type="text/event-stream"
+        )
+    else:
+        response = await grpc_handler.grpc_to_pieces(
+            request.model, content, rules, request.temperature, request.top_p
+        )
+        return JSONResponse(content=response)
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=config.PORT)