"""
E5-small-v2 embedding server — OpenAI-compatible /v1/embeddings endpoint.
Minimal, CPU-only, ~100MB RAM.
"""
from fastapi import FastAPI
from pydantic import BaseModel
from sentence_transformers import SentenceTransformer
import numpy as np

app = FastAPI()
model = SentenceTransformer("intfloat/e5-small-v2", device="cpu")


class EmbedRequest(BaseModel):
    model: str = "intfloat/e5-small-v2"
    input: list[str]


class EmbedData(BaseModel):
    object: str = "embedding"
    index: int
    embedding: list[float]


class EmbedResponse(BaseModel):
    object: str = "list"
    data: list[EmbedData]
    model: str


@app.post("/v1/embeddings")
def embed(req: EmbedRequest):
    # E5 requires "query: " or "passage: " prefix
    prefixed = [f"passage: {t}" if "\n" not in t else t for t in req.input]
    embs = model.encode(prefixed, normalize_embeddings=True, show_progress_bar=False)
    data = [
        EmbedData(index=i, embedding=emb.tolist())
        for i, emb in enumerate(embs)
    ]
    return EmbedResponse(data=data, model=req.model)


@app.get("/health")
def health():
    return {"status": "ok"}


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8081)