#!/usr/bin/env python3
"""轻量 bge-m3 embedding server — OpenAI 兼容 /v1/embeddings 接口"""
import os, sys, json, time
from typing import List

# 环境变量
HOST = os.environ.get("EMBED_HOST", "0.0.0.0")
PORT = int(os.environ.get("EMBED_PORT", "8000"))
MODEL = os.environ.get("EMBED_MODEL", "BAAI/bge-m3")
DEVICE = os.environ.get("EMBED_DEVICE", "cpu")

from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
from pydantic import BaseModel, Field
import uvicorn

app = FastAPI(title="bge-m3 Embedding Server", version="1.0")

class EmbedRequest(BaseModel):
    input: str | List[str]
    model: str = MODEL

class EmbeddingObject(BaseModel):
    object: str = "embedding"
    index: int
    embedding: List[float]

class UsageInfo(BaseModel):
    prompt_tokens: int
    total_tokens: int

class EmbedResponse(BaseModel):
    object: str = "list"
    data: List[EmbeddingObject]
    model: str = MODEL
    usage: UsageInfo

# 延迟加载
model_pipe = None

def load_model():
    global model_pipe
    if model_pipe is not None:
        return
    print(f"[embed] Loading {MODEL} on {DEVICE}...", flush=True)
    t0 = time.time()
    from sentence_transformers import SentenceTransformer
    model_pipe = SentenceTransformer(MODEL, device=DEVICE)
    elapsed = time.time() - t0
    print(f"[embed] Loaded in {elapsed:.1f}s. Dimension: {model_pipe.get_sentence_embedding_dimension()}", flush=True)

@app.on_event("startup")
async def startup():
    load_model()

@app.get("/health")
async def health():
    return {"status": "ok", "model": MODEL, "device": DEVICE,
            "loaded": model_pipe is not None}

@app.post("/v1/embeddings")
async def embed(req: EmbedRequest):
    t0 = time.time()
    texts = req.input if isinstance(req.input, list) else [req.input]
    
    # bge-m3 需要加 prefix
    prefixed = [f"为这个句子生成表示以用于检索相关文章：{t}" for t in texts]
    
    vecs = model_pipe.encode(prefixed, normalize_embeddings=True,
                             show_progress_bar=False)
    vecs = vecs.tolist()
    
    total_tokens = sum(len(t) for t in texts)
    elapsed = time.time() - t0
    
    data = [
        EmbeddingObject(index=i, embedding=vecs[i])
        for i in range(len(vecs))
    ]
    return EmbedResponse(data=data, usage=UsageInfo(prompt_tokens=total_tokens, total_tokens=total_tokens))

if __name__ == "__main__":
    load_model()
    uvicorn.run(app, host=HOST, port=PORT, log_level="info")