# main.py
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List
from vllm import LLM
import numpy as np

# Initialize the model
llm = LLM(model='BAAI/bge-m3', task="embed")

# Initialize FastAPI app
app = FastAPI()

# Define request schemas
class DocumentsRequest(BaseModel):
    documents: List[str]

class QueryRequest(BaseModel):
    query: str

@app.get("/", tags=["Home"])
def api_home():
    return {'hello': 'Welcome!'}

# API to embed documents
@app.post("/embed_documents")
def embed_documents(request: DocumentsRequest):
    try:
        docs = request.documents
        docs_embd = llm.encode(docs)
        docs_embd = [doc.outputs.data.numpy().tolist() for doc in docs_embd]
        return {"embeddings": docs_embd}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error embedding documents: {str(e)}")

# API to embed query
@app.post("/embed_query")
def embed_query(request: QueryRequest):
    try:
        query = request.query
        query_embd = llm.encode(query)
        query_embd = query_embd[0].outputs.data.numpy().tolist()
        return {"embedding": query_embd}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error embedding query: {str(e)}")