Update app.py

app.py

@@ -1,114 +1,153 @@
 from __future__ import annotations
 
+import json
 import os
 import sys
-import json
-import subprocess
+import urllib.parse
+import urllib.request
 from pathlib import Path
-from typing import Any, Dict
-
-import requests
-from fastapi import FastAPI
-from fastapi.responses import JSONResponse
-import uvicorn
 
+from huggingface_hub import hf_hub_download
 
-# ======================
-# Config
-# ======================
-MODEL_PATH = os.getenv("MODEL_PATH", "/data/models/model.gguf")
+# 配置项 (将路径重命名为通用命名，避免 gemma/mixtao 混淆)
+MODEL_REPO = os.getenv("MODEL_REPO", "mixtao/MixTAO-7Bx2-MoE-v8.1-GGUF")
+MODEL_FILE = os.getenv("MODEL_FILE", "mixtao-7bx2-moe-v8.1.Q4_K_M.gguf")
+MODEL_DIR = Path(os.getenv("MODEL_DIR", "/data/models/llm"))
+CHAT_TEMPLATE_FILE = Path(os.getenv("CHAT_TEMPLATE_FILE", "/data/models/llm/chat_template.jinja"))
 
 LLAMA_SERVER_BIN = os.getenv("LLAMA_SERVER_BIN", "/opt/llama.cpp/llama-server")
-LLAMA_HOST = os.getenv("LLAMA_HOST", "127.0.0.1")
-LLAMA_PORT = int(os.getenv("LLAMA_PORT", "7860"))
+LLAMA_HOST = os.getenv("LLAMA_HOST", "0.0.0.0")
+LLAMA_PORT = os.getenv("LLAMA_PORT", "7860")
 
-API_HOST = os.getenv("API_HOST", "0.0.0.0")
-API_PORT = int(os.getenv("API_PORT", "8000"))
+# OpenAI API 安全配置
+API_KEY = os.getenv("API_KEY", "")  # 如果设置，客户端请求必须携带 Bearer Token
 
+# 性能及运行参数
 THREADS = os.getenv("THREADS", "4")
-CTX_SIZE = os.getenv("CTX_SIZE", "4096")
+CTX_SIZE = os.getenv("CTX_SIZE", "2048")
+BATCH_SIZE = os.getenv("BATCH_SIZE", "default")
+UBATCH_SIZE = os.getenv("UBATCH_SIZE", "default")
 GPU_LAYERS = os.getenv("GPU_LAYERS", "0")
+FLASH_ATTN = os.getenv("FLASH_ATTN", "false")
+CACHE_TYPE_K = os.getenv("CACHE_TYPE_K", "default")
+CACHE_TYPE_V = os.getenv("CACHE_TYPE_V", "default")
 
+TEMPERATURE = os.getenv("TEMPERATURE", "0.2")
+TOP_P = os.getenv("TOP_P", "0.95")
+TOP_K = os.getenv("TOP_K", "64")
+REPEAT_PENALTY = os.getenv("REPEAT_PENALTY", "1.08")
 
-# ======================
-# Start llama.cpp server
-# ======================
-def start_llama_server():
-    cmd = [
-        LLAMA_SERVER_BIN,
-        "-m", MODEL_PATH,
-        "--host", LLAMA_HOST,
-        "--port", str(LLAMA_PORT),
-        "--threads", str(THREADS),
-        "--ctx-size", str(CTX_SIZE),
-        "--n-gpu-layers", str(GPU_LAYERS),
-        "--parallel", "1",
-        "--cont-batching",
-    ]
-
-    print("[startup] launching llama.cpp server:")
-    print(" ".join(cmd), flush=True)
 
-    return subprocess.Popen(cmd, stdout=sys.stdout, stderr=sys.stderr)
+def log(message: str) -> None:
+    print(f"[startup] {message}", flush=True)
 
 
-# ======================
-# OpenAI-compatible API layer
-# ======================
-app = FastAPI()
+def download_model() -> str:
+    MODEL_DIR.mkdir(parents=True, exist_ok=True)
+    local_file = MODEL_DIR / MODEL_FILE
+    if local_file.exists():
+        log(f"Using cached model: {local_file}")
+        return str(local_file)
 
+    log(f"Downloading {MODEL_REPO}/{MODEL_FILE}")
+    model_path = hf_hub_download(
+        repo_id=MODEL_REPO,
+        filename=MODEL_FILE,
+        local_dir=str(MODEL_DIR),
+        local_dir_use_symlinks=False
+    )
+    log(f"Model ready: {model_path}")
+    return model_path
 
-def llama_chat_completion(payload: Dict[str, Any]) -> Dict[str, Any]:
-    """
-    Forward to llama.cpp OpenAI-compatible endpoint
-    """
-    url = f"http://{LLAMA_HOST}:{LLAMA_PORT}/v1/chat/completions"
 
-    resp = requests.post(url, json=payload, timeout=600)
-    resp.raise_for_status()
-    return resp.json()
+def download_chat_template() -> str | None:
+    if CHAT_TEMPLATE_FILE.exists() and CHAT_TEMPLATE_FILE.stat().st_size > 0:
+        log(f"Using cached chat template: {CHAT_TEMPLATE_FILE}")
+        return str(CHAT_TEMPLATE_FILE)
 
+    encoded_repo = urllib.parse.quote(MODEL_REPO, safe="/")
+    api_url = f"https://huggingface.co{encoded_repo}"
+    log("Fetching chat template from model metadata")
 
-@app.post("/v1/chat/completions")
-def chat_completions(body: Dict[str, Any]):
-    """
-    OpenAI-compatible endpoint
-    """
     try:
-        result = llama_chat_completion(body)
-        return JSONResponse(content=result)
-    except Exception as e:
-        return JSONResponse(
-            status_code=500,
-            content={
-                "error": {
-                    "message": str(e),
-                    "type": "server_error",
-                }
-            },
-        )
+        req = urllib.request.Request(api_url, headers={"User-Agent": "Mozilla/5.0"})
+        with urllib.request.urlopen(req, timeout=30) as response:
+            metadata = json.loads(response.read().decode("utf-8"))
+    except Exception as exc:
+        log(f"Could not fetch chat template metadata: {exc}")
+        return None
+
+    template = (metadata.get("gguf") or {}).get("chat_template")
+    if not template:
+        log("No chat template found in model metadata; llama-server will use GGUF metadata")
+        return None
+
+    CHAT_TEMPLATE_FILE.parent.mkdir(parents=True, exist_ok=True)
+    CHAT_TEMPLATE_FILE.write_text(template, encoding="utf-8")
+    log(f"Chat template ready: {CHAT_TEMPLATE_FILE}")
+    return str(CHAT_TEMPLATE_FILE)
 
 
-@app.get("/health")
-def health():
-    return {"status": "ok"}
+def build_command(model_path: str, template_path: str | None) -> list[str]:
+    def is_valid(value: str) -> bool:
+        return value.strip().lower() not in {"", "default", "auto", "none", "off", "false"}
 
+    cmd = [
+        LLAMA_SERVER_BIN,
+        "-m", model_path,
+        "--host", LLAMA_HOST,
+        "--port", LLAMA_PORT,
+        "--threads", THREADS,
+        "--ctx-size", CTX_SIZE,
+        "--n-gpu-layers", GPU_LAYERS,
+        "--parallel", "4",          # 开启多并发槽位以支持多用户同时调用 API
+        "--cont-batching",          # 启用连续批处理以优化 API 吞吐量
+        "--temp", TEMPERATURE,
+        "--top-p", TOP_P,
+        "--top-k", TOP_K,
+        "--repeat-penalty", REPEAT_PENALTY,
+    ]
+
+    # 核心改造：注入 OpenAI API 专属控制参数
+    cmd.append("--embedding")       # 允许对外提供 OpenAI 兼容的 /v1/embeddings 接口
+    
+    if API_KEY.strip():
+        cmd.extend(["--api-key", API_KEY.strip()]) # 启用 API 鉴权机制
+        log("OpenAI API authentication enabled with custom API_KEY.")
+    else:
+        log("Warning: Running without API_KEY. Server is publicly accessible.")
 
-# ======================
-# Main
-# ======================
-def main():
-    # 1. start llama.cpp backend
-    proc = start_llama_server()
+    # 动态参数注入
+    if is_valid(BATCH_SIZE): cmd.extend(["--batch-size", BATCH_SIZE])
+    if is_valid(UBATCH_SIZE): cmd.extend(["--ubatch-size", UBATCH_SIZE])
+    if is_valid(CACHE_TYPE_K): cmd.extend(["--cache-type-k", CACHE_TYPE_K])
+    if is_valid(CACHE_TYPE_V): cmd.extend(["--cache-type-v", CACHE_TYPE_V])
+    if is_valid(FLASH_ATTN): cmd.append("--flash-attn")
+    if template_path: cmd.extend(["--chat-template-file", template_path])
 
-    # 2. wait a bit for server init
-    import time
-    time.sleep(3)
+    return cmd
 
-    # 3. start OpenAI API gateway
-    print(f"[startup] OpenAI API server running on http://{API_HOST}:{API_PORT}")
-    uvicorn.run(app, host=API_HOST, port=API_PORT)
+
+def main() -> None:
+    binary_dir = str(Path(LLAMA_SERVER_BIN).parent)
+    existing_lib = os.environ.get("LD_LIBRARY_PATH")
+    os.environ["LD_LIBRARY_PATH"] = binary_dir if not existing_lib else f"{binary_dir}:{existing_lib}"
+
+    for env_var in ["OMP_NUM_THREADS", "OPENBLAS_NUM_THREADS", "MKL_NUM_THREADS"]:
+        os.environ.setdefault(env_var, THREADS)
+
+    model_path = download_model()
+    template_path = download_chat_template()
+    cmd = build_command(model_path, template_path)
+
+    log("Starting OpenAI-compatible llama.cpp API server")
+    log(" ".join(cmd))
+    os.execvpe(cmd[0], cmd, os.environ)
 
 
 if __name__ == "__main__":
-    main()
+    try:
+        main()
+    except Exception as exc:
+        print(f"[fatal] {exc}", file=sys.stderr, flush=True)
+        raise