Update Dockerfile

Dockerfile

@@ -1,7 +1,11 @@
 FROM python:3.11-slim
+
 WORKDIR /app
-# 核心修正：安装新的 google-genai 库
+
+# 安装依赖
 RUN pip install fastapi uvicorn google-genai
+
+# 生成 main.py（支持自动获取模型列表）
 RUN cat > main.py << 'PYEOF'
 from fastapi import FastAPI, Request
 from fastapi.responses import JSONResponse
@@ -12,28 +16,50 @@ app = FastAPI()
 GOOGLE_API_KEY = os.environ["GOOGLE_API_KEY"]
 client = genai.Client(api_key=GOOGLE_API_KEY)
 
-# 已更新的模型列表，包含当前可用的模型
+# 静态模型列表：作为 API 调用失败时的后备方案
+STATIC_MODELS = [
+    {"id": "gemini-embedding-001", "object": "model", "owned_by": "gemini"},
+    {"id": "gemini-embedding-2-preview", "object": "model", "owned_by": "gemini"},
+]
+
+# 模型名称映射（短名称 -> 完整路径），用于 /v1/embeddings 端点
+MODEL_MAP = {
+    "text-embedding-004": "models/text-embedding-004",
+    "gemini-embedding-001": "models/gemini-embedding-001",
+    "gemini-embedding-2-preview": "models/gemini-embedding-2-preview",
+}
+
+def fetch_embedding_models():
+    """动态获取支持 embedContent 的 Gemini Embedding 模型列表"""
+    try:
+        all_models = client.models.list()
+        embedding_models = []
+        for m in all_models:
+            if hasattr(m, 'supported_actions') and 'embedContent' in m.supported_actions:
+                # 提取短名称 (e.g., "models/gemini-embedding-001" -> "gemini-embedding-001")
+                model_id = m.name.split('/')[-1]
+                embedding_models.append({
+                    "id": model_id,
+                    "object": "model",
+                    "owned_by": "gemini"
+                })
+        return embedding_models
+    except Exception as e:
+        print(f"动态获取模型列表失败，回退到静态列表: {e}")
+        return STATIC_MODELS
+
 @app.get("/v1/models")
 async def list_models():
-    return JSONResponse({
-        "object": "list",
-        "data": [
-            {"id": "gemini-embedding-001", "object": "model", "owned_by": "gemini"},
-            {"id": "gemini-embedding-2-preview", "object": "model", "owned_by": "gemini"}
-        ]
-    })
+    models = fetch_embedding_models()
+    return JSONResponse({"object": "list", "data": models})
 
 @app.post("/v1/embeddings")
 async def create_embeddings(request: Request):
     try:
         body = await request.json()
-        # 默认使用可用的 gemini-embedding-001 模型
-        model = body.get("model", "gemini-embedding-001")
-        # 自动补全模型名的完整路径
-        if not model.startswith("models/"):
-            model_full = f"models/{model}"
-        else:
-            model_full = model
+        model_short = body.get("model", "gemini-embedding-001")
+        # 自动补全模型名称为完整路径
+        model_full = MODEL_MAP.get(model_short, f"models/{model_short}")
         input_text = body.get("input", "")
         result = client.models.embed_content(
             model=model_full,
@@ -46,7 +72,7 @@ async def create_embeddings(request: Request):
                 "embedding": result.embeddings[0].values,
                 "index": 0
             }],
-            "model": model
+            "model": model_short
         })
     except Exception as e:
         return JSONResponse({"error": str(e)}, status_code=500)