Upload 3 files

Procfile

@@ -1 +1 @@
-web: gunicorn app:app
+web: gunicorn app:app -w 2 -k uvicorn.workers.UvicornWorker

app.py

@@ -1,108 +1,105 @@
-from flask import Flask, request, jsonify
-from sentence_transformers import SentenceTransformer
-import numpy as np
+from fastapi import FastAPI, HTTPException, Depends
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
 import os
-import time
+import numpy as np
+from sentence_transformers import SentenceTransformer
+from typing import List, Optional
 
-app = Flask(__name__)
+app = FastAPI()
 
-# 加载模型
-model_name = "BAAI/bge-small-en-v1.5"
-model = SentenceTransformer(model_name)
+# 允许跨域请求
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
 
-# 支持的模型列表
-SUPPORTED_MODELS = {
-    "text-embedding-3-small": model,
-    "bge-small-en-v1.5": model
+# 模型映射：OpenAI模型名 → 开源模型名
+MODEL_MAPPING = {
+    "text-embedding-3-small": "BAAI/bge-small-en-v1.5",
+    "text-embedding-3-large": "BAAI/bge-large-en-v1.5"  # 新增大模型映射
 }
 
-# 简单的API密钥验证（可选）
-API_KEY = os.getenv("API_KEY", "your-default-api-key")
+# 加载模型（懒加载，首次请求时加载）
+models = {}
+
+def get_model(model_name: str):
+    if model_name not in models:
+        # 检查是否支持该模型
+        if model_name not in MODEL_MAPPING:
+            raise HTTPException(status_code=400, detail=f"不支持的模型: {model_name}")
+        # 加载模型
+        models[model_name] = SentenceTransformer(MODEL_MAPPING[model_name])
+    return models[model_name]
+
+# 验证API密钥
+def verify_api_key(authorization: Optional[str] = None):
+    if not authorization or not authorization.startswith("Bearer "):
+        raise HTTPException(status_code=401, detail="未提供有效的API密钥")
+    api_key = authorization[len("Bearer "):]
+    if api_key != os.getenv("API_KEY"):
+        raise HTTPException(status_code=401, detail="无效的API密钥")
+    return True
+
+# 请求体模型（对齐OpenAI格式）
+class EmbeddingRequest(BaseModel):
+    input: str or List[str]
+    model: str
+    encoding_format: Optional[str] = "float"  # 仅支持float，忽略base64
+
+# 响应体模型（对齐OpenAI格式）
+class EmbeddingData(BaseModel):
+    object: str = "embedding"
+    embedding: List[float]
+    index: int
 
-def verify_api_key(headers):
-    """验证API密钥"""
-    auth_header = headers.get("Authorization")
-    if not auth_header or not auth_header.startswith("Bearer "):
-        return False
-    return auth_header.split("Bearer ")[1] == API_KEY
+class EmbeddingResponse(BaseModel):
+    object: str = "list"
+    data: List[EmbeddingData]
+    model: str
+    usage: dict = {"prompt_tokens": 0, "total_tokens": 0}
 
-@app.route('/v1/embeddings', methods=['POST'])
-def create_embedding():
-    """生成嵌入向量，兼容OpenAI API格式"""
-    # 验证API密钥
-    if not verify_api_key(request.headers):
-        return jsonify({
-            "error": {
-                "message": "Invalid API key",
-                "type": "invalid_request_error",
-                "param": None,
-                "code": "invalid_api_key"
-            }
-        }), 401
-    
-    # 解析请求
-    data = request.json
-    if not data or "input" not in data:
-        return jsonify({
-            "error": {
-                "message": "Missing input",
-                "type": "invalid_request_error",
-                "param": None,
-                "code": "missing_input"
-            }
-        }), 400
-    
-    # 获取模型（默认为text-embedding-3-small）
-    model_name = data.get("model", "text-embedding-3-small")
-    if model_name not in SUPPORTED_MODELS:
-        return jsonify({
-            "error": {
-                "message": f"Model {model_name} not found",
-                "type": "invalid_request_error",
-                "param": None,
-                "code": "model_not_found"
-            }
-        }), 404
-    
-    # 处理输入（支持单文本或文本列表）
-    inputs = data["input"]
-    if isinstance(inputs, str):
-        inputs = [inputs]
-    
-    # 计算嵌入向量
-    start_time = time.time()
-    embeddings = model.encode(inputs, normalize_embeddings=True)
-    processing_time = time.time() - start_time
-    
-    # 准备响应数据
-    response_data = {
-        "object": "list",
-        "data": [
-            {
-                "object": "embedding",
-                "embedding": embedding.tolist(),
-                "index": i
-            } for i, embedding in enumerate(embeddings)
-        ],
-        "model": model_name,
-        "usage": {
-            "prompt_tokens": sum(len(text.split()) for text in inputs),  # 简单估算
-            "total_tokens": sum(len(text.split()) for text in inputs)
-        }
-    }
-    
-    return jsonify(response_data)
+@app.post("/v1/embeddings", response_model=EmbeddingResponse)
+async def create_embedding(
+    request: EmbeddingRequest,
+    _: bool = Depends(verify_api_key)
+):
+    try:
+        # 获取模型
+        model = get_model(request.model)
+        
+        # 处理输入（支持单文本或文本列表）
+        inputs = [request.input] if isinstance(request.input, str) else request.input
+        
+        # 计算嵌入
+        embeddings = model.encode(inputs, normalize_embeddings=True)
+        
+        # 构建响应
+        data = [
+            EmbeddingData(embedding=embedding.tolist(), index=i)
+            for i, embedding in enumerate(embeddings)
+        ]
+        
+        # 估算token数（简单近似：每个单词约1 token）
+        prompt_tokens = sum(len(text.split()) for text in inputs)
+        
+        return EmbeddingResponse(
+            data=data,
+            model=request.model,
+            usage={"prompt_tokens": prompt_tokens, "total_tokens": prompt_tokens}
+        )
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
 
-@app.route('/health', methods=['GET'])
-def health_check():
-    """健康检查接口"""
-    return jsonify({
-        "status": "healthy",
-        "model": model_name,
-        "supported_models": list(SUPPORTED_MODELS.keys())
-    })
+# 健康检查接口
+@app.get("/health")
+async def health_check():
+    return {"status": "healthy", "models": list(MODEL_MAPPING.keys())}
 
-if __name__ == '__main__':
-    # 生产环境应使用Gunicorn等WSGI服务器
-    app.run(host='0.0.0.0', port=int(os.getenv('PORT', 7860)))
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)
     

requirements.txt

@@ -1,5 +1,7 @@
-flask==2.3.3
+fastapi==0.110.0
+uvicorn==0.29.0
+gunicorn==21.2.0
 sentence-transformers==2.7.0
 torch==2.2.2
 numpy==1.26.4
-gunicorn==21.2.0  # 用于生产环境部署
+pydantic==2.6.4