worsarise's picture
Update main.py
0179583 verified
Raw
History Blame Contribute Delete
1.51 kB
import os
from fastapi import FastAPI, Request, HTTPException
from sentence_transformers import SentenceTransformer
os.environ["OMP_NUM_THREADS"] = "2"
os.environ["MKL_NUM_THREADS"] = "2"
app = FastAPI(title="OpenAI Compatible API")
# Keep the model you specifically chose to host
MODEL_ID = "Snowflake/snowflake-arctic-embed-l"
model = SentenceTransformer(MODEL_ID, device="cpu")
@app.get("/health")
def health_check():
return {"status": "healthy"}
@app.post("/v1/embeddings")
async def create_embeddings(request: Request):
try:
data = await request.json()
except Exception:
raise HTTPException(status_code=400, detail="Invalid JSON format")
inputs = data.get("input")
# 🛡️ THE SAFEGUARD: If n8n sends an empty string, force a valid text
# so the model generates real numbers instead of an array of zeros.
if not inputs or inputs == [""] or inputs == "":
inputs = ["dummy text to prevent zero vector database crash"]
if isinstance(inputs, str):
inputs = [inputs]
# Generate vectors
embeddings = model.encode(inputs, normalize_embeddings=True).tolist()
response_data = [
{
"object": "embedding",
"embedding": emb,
"index": i
}
for i, emb in enumerate(embeddings)
]
return {
"object": "list",
"data": response_data,
"model": MODEL_ID,
"usage": {"prompt_tokens": 0, "total_tokens": 0}
}