Update app.py

app.py

@@ -4,7 +4,7 @@ from pydantic import BaseModel
 import os
 import numpy as np
 from sentence_transformers import SentenceTransformer
-from typing import List, Optional
+from typing import List, Optional, Union  # 导入Union
 import logging
 
 # 配置日志
@@ -12,8 +12,8 @@ logging.basicConfig(
     level=logging.INFO,
     format="%(asctime)s-%(name)s-%(levelname)s-%(message)s",
     handlers=[
-        logging.FileHandler("embedding_service.log"),  # 日志写入文件
-        logging.StreamHandler()  # 同时输出到控制台
+        logging.FileHandler("embedding_service.log"),
+        logging.StreamHandler()
     ]
 )
 logger = logging.getLogger("embedding_service")
@@ -35,19 +35,15 @@ MODEL_MAPPING = {
     "text-embedding-3-large": "BAAI/bge-large-en-v1.5"
 }
 
-# 加载模型（懒加载，首次请求时加载）
+# 加载模型（懒加载）
 models = {}
 
 def get_model(model_name: str):
     logger.info(f"尝试获取模型: {model_name}")
+    model_to_load = MODEL_MAPPING.get(model_name, model_name)  # 兼容直接用开源模型名
     if model_name not in models:
-        if model_name not in MODEL_MAPPING:
-            error_msg = f"不支持的模型: {model_name}"
-            logger.error(error_msg)
-            raise HTTPException(status_code=400, detail=error_msg)
-        logger.info(f"开始加载模型: {MODEL_MAPPING[model_name]}")
         try:
-            models[model_name] = SentenceTransformer(MODEL_MAPPING[model_name])
+            models[model_name] = SentenceTransformer(model_to_load)
             logger.info(f"模型 {model_name} 加载成功")
         except Exception as e:
             error_msg = f"加载模型 {model_name} 失败: {str(e)}"
@@ -57,21 +53,18 @@ def get_model(model_name: str):
 
 # 验证API密钥
 def verify_api_key(authorization: Optional[str] = Header(None)):
-    logger.info("执行API密钥验证")
     logger.info(f"Authorization头部内容: {authorization}")
     if not authorization or not authorization.startswith("Bearer "):
-        logger.warning("未提供有效的API密钥格式")
         raise HTTPException(status_code=401, detail="未提供有效的API密钥")
     api_key = authorization[len("Bearer "):]
     if api_key != os.getenv("API_KEY"):
-        logger.warning("无效的API密钥")
         raise HTTPException(status_code=401, detail="无效的API密钥")
     logger.info("API密钥验证通过")
     return True
 
 # 请求体模型
 class EmbeddingRequest(BaseModel):
-    input: str or List[str]
+    input: Union[str, List[str]]  # 支持str或List[str]
     model: str
     encoding_format: Optional[str] = "float"
 
@@ -90,9 +83,10 @@ class EmbeddingResponse(BaseModel):
 @app.post("/v1/embeddings", response_model=EmbeddingResponse)
 async def create_embedding(
     request: Request,
-    req: EmbeddingRequest
+    req: EmbeddingRequest,
+    _: bool = Depends(verify_api_key)
 ):
-    # 先打印完整请求信息（在验证之前）
+    # 打印请求信息
     logger.info("\n===== 接收到的完整请求信息 =====")
     logger.info(f"请求方法: {request.method}")
     logger.info(f"请求URL: {request.url}")
@@ -102,11 +96,7 @@ async def create_embedding(
     logger.info(f"请求体: {await request.body()}")
     logger.info("===============================\n")
     
-    # 手动执行验证（在打印日志之后）
-    authorization = request.headers.get("Authorization")
-    verify_api_key(authorization)
-    
-    # 原有嵌入处理逻辑
+    # 嵌入生成逻辑
     logger.info(f"收到嵌入请求，模型: {req.model}, 输入类型: {type(req.input)}")
     try:
         model = get_model(req.model)
@@ -131,9 +121,7 @@ async def create_embedding(
             usage={"prompt_tokens": prompt_tokens, "total_tokens": prompt_tokens}
         )
     except Exception as e:
-        error_msg = f"处理嵌入请求时发生错误: {str(e)}"
-        logger.error(error_msg)
-        raise HTTPException(status_code=500, detail=error_msg)
+        raise HTTPException(status_code=500, detail=f"处理嵌入请求时发生错误: {str(e)}")
 
 @app.get("/health")
 async def health_check(request: Request):
@@ -144,9 +132,9 @@ async def health_check(request: Request):
     for name, value in request.headers.items():
         logger.info(f"  {name}: {value}")
     logger.info("===============================\n")
-    return {"status": "healthy", "models": list(MODEL_MAPPING.keys())}
+    return {"status": "healthy", "models": list(MODEL_MAPPING.keys()) + list(models.keys())}
 
 if __name__ == "__main__":
     import uvicorn
     logger.info("启动服务")
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=7860)