Update Dockerfile

Dockerfile

@@ -2,65 +2,124 @@ FROM python:3.11-slim
 
 WORKDIR /app
 
-# 安装依赖
+# 安装必要库
 RUN pip install fastapi uvicorn google-genai
 
-# 生成 main.py（支持自动获取模型列表）
+# 生成全新的 main.py，支持 Chat Completions、Embeddings、Models
 RUN cat > main.py << 'PYEOF'
 from fastapi import FastAPI, Request
-from fastapi.responses import JSONResponse
+from fastapi.responses import JSONResponse, StreamingResponse
 from google import genai
-import os, uvicorn
+import os, uvicorn, json, time
 
 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 端点
+# ---------- 模型名称映射（短名 -> 完整 resources 名）----------
 MODEL_MAP = {
-    "text-embedding-004": "models/text-embedding-004",
+    # 聊天模型
+    "gemini-2.5-flash": "models/gemini-2.5-flash",
+    "gemini-2.5-pro": "models/gemini-2.5-pro",
+    "gemini-2.5-pro-preview-03-25": "models/gemini-2.5-pro-preview-03-25",
+    "gemini-1.5-flash": "models/gemini-1.5-flash",
+    "gemini-1.5-pro": "models/gemini-1.5-pro",
+    # 嵌入模型
     "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
+# 预定义的所有可用模型列表（用于 /v1/models）
+AVAILABLE_MODELS = [
+    {"id": "gemini-2.5-flash", "object": "model", "owned_by": "gemini"},
+    {"id": "gemini-2.5-pro", "object": "model", "owned_by": "gemini"},
+    {"id": "gemini-1.5-flash", "object": "model", "owned_by": "gemini"},
+    {"id": "gemini-1.5-pro", "object": "model", "owned_by": "gemini"},
+    {"id": "gemini-embedding-001", "object": "model", "owned_by": "gemini"},
+    {"id": "gemini-embedding-2-preview", "object": "model", "owned_by": "gemini"},
+]
+
+# ---------- 辅助函数 ----------
+def get_full_model_name(short_name: str) -> str:
+    """将短模型名转换为完整的 resources 路径，若未找到则自动添加 models/ 前缀"""
+    return MODEL_MAP.get(short_name, f"models/{short_name}")
 
+def convert_openai_messages(messages: list) -> list:
+    """将 OpenAI 格式的 messages 转换为 Gemini 要求的 contents 格式"""
+    contents = []
+    for msg in messages:
+        role = msg.get("role", "user")
+        text = msg.get("content", "")
+        # Gemini 使用 "user" 和 "model" 角色，将 "assistant" 映射为 "model"
+        if role == "assistant":
+            role = "model"
+        elif role not in ("user", "model"):
+            role = "user"  # 其他角色一律视为 user
+        contents.append({"role": role, "parts": [{"text": text}]})
+    return contents
+
+# ---------- API 端点 ----------
 @app.get("/v1/models")
 async def list_models():
-    models = fetch_embedding_models()
-    return JSONResponse({"object": "list", "data": models})
+    """返回可用模型列表（可手动扩展）"""
+    return JSONResponse({"object": "list", "data": AVAILABLE_MODELS})
+
+@app.post("/v1/chat/completions")
+async def chat_completions(request: Request):
+    """OpenAI 兼容的聊天补全端点（非流式）"""
+    try:
+        body = await request.json()
+        model_short = body.get("model", "gemini-2.5-flash")
+        messages = body.get("messages", [])
+        temperature = body.get("temperature", 0.7)
+        max_tokens = body.get("max_tokens", 1000)
+        
+        model_full = get_full_model_name(model_short)
+        contents = convert_openai_messages(messages)
+        
+        # 调用 Gemini API
+        response = client.models.generate_content(
+            model=model_full,
+            contents=contents,
+            config={
+                "temperature": temperature,
+                "max_output_tokens": max_tokens,
+            }
+        )
+        
+        # 构造 OpenAI 兼容响应
+        choice = {
+            "index": 0,
+            "message": {
+                "role": "assistant",
+                "content": response.text
+            },
+            "finish_reason": "stop"
+        }
+        return JSONResponse({
+            "id": f"chatcmpl-{int(time.time())}",
+            "object": "chat.completion",
+            "created": int(time.time()),
+            "model": model_short,
+            "choices": [choice],
+            "usage": {
+                "prompt_tokens": response.usage_metadata.prompt_token_count,
+                "completion_tokens": response.usage_metadata.candidates_token_count,
+                "total_tokens": response.usage_metadata.total_token_count
+            }
+        })
+    except Exception as e:
+        return JSONResponse({"error": str(e)}, status_code=500)
 
 @app.post("/v1/embeddings")
 async def create_embeddings(request: Request):
+    """OpenAI 兼容的嵌入端点"""
     try:
         body = await request.json()
         model_short = body.get("model", "gemini-embedding-001")
-        # 自动补全模型名称为完整路径
-        model_full = MODEL_MAP.get(model_short, f"models/{model_short}")
         input_text = body.get("input", "")
+        model_full = get_full_model_name(model_short)
+        
         result = client.models.embed_content(
             model=model_full,
             contents=input_text