from fastapi import FastAPI, Request, HTTPException
from pydantic import BaseModel, Field
from sentence_transformers import SentenceTransformer
import numpy as np
import logging

# 设置日志记录
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI()

try:
    # Load the Sentence Transformer model
    model = SentenceTransformer("BAAI/bge-large-zh-v1.5")
except Exception as e:
    logger.error(f"Failed to load model: {e}")
    raise HTTPException(status_code=500, detail="Model loading failed")

class EmbeddingRequest(BaseModel):
    input: str = Field(..., min_length=1, max_length=1000)

@app.post("/embeddings")
async def embeddings(request: EmbeddingRequest):
    input_text = request.input

    try:
        if not input_text:
            return {
                "object": "list",
                "data": [],
                "model": "BAAI/bge-large-zh-v1.5",
                "usage": {
                    "prompt_tokens": 0,
                    "total_tokens": 0
                }
            }

        # Calculate embeddings
        embeddings = model.encode(input_text)

        # Format the embeddings in OpenAI compatible format
        data = {
            "object": "list",
            "data": [
                {
                    "object": "embedding",
                    "embedding": embeddings.tolist(),
                    "index": 0
                }
            ],
            "model": "BAAI/bge-large-zh-v1.5",
            "usage": {
                "prompt_tokens": len(input_text),
                "total_tokens": len(input_text)
            }
        }

        return data
    except Exception as e:
        logger.error(f"Error processing embeddings: {e}")
        raise HTTPException(status_code=500, detail="Internal Server Error")