from fastapi import FastAPI
from sentence_transformers import SentenceTransformer


app = FastAPI(
    title="Generate Embeddings API",
    description="""
## Generate Embeddings API
A simple API to generate sentence embeddings using **Sentence Transformers**.

### 🔗 Useful Links
- [Hugging Face Code](https://huggingface.co/spaces/TheMihirNaik/generate-embeddings-api/tree/main)
- [SEO Workflow Automation Newsletter](https://mihirnaik.substack.com/)
- [Mihir Naik](https://www.mihirnaik.com)

### How It Works
1. Send a **POST request** to `/generate-embeddings/` with a list of sentences.
2. Get **embeddings** as JSON response.

Built by **Mihir Naik** 🚀
""",
    version="1.0.0",
    contact={
        "name": "Mihir Naik",
        "url": "https://www.mihirnaik.com",
    }
)


try:
    all_MiniLM_L6_V2_model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
except Exception as e:
    raise RuntimeError("Failed to load the all-MiniLM-L6-v2 model.") from e

try:
    intfloat_e5_large_v2_model = SentenceTransformer('intfloat/e5-large-v2')
except Exception as e:
    raise RuntimeError("Failed to load the intfloat/e5-large-v2 model.") from e

@app.get("/", include_in_schema=False)
def redirect_to_docs():
    """
    Redirects to the FastAPI documentation.

    Returns:
        RedirectResponse: A response that redirects to the documentation URL.
    """
    from fastapi.responses import RedirectResponse
    return RedirectResponse(url="/docs")

from typing import List

@app.post("/generate-embeddings/sentence-transformers-all-MiniLM-L6-v2/")
def generate_embeddings_all_MiniLM_L6_V2_model(sentences: List[str]):
    """
    Generates embeddings for a list of sentences using the MiniLM model.

    Args:
        sentences (List[str]): A list of sentences to generate embeddings for.

    Returns:
        dict: A dictionary containing the embeddings as a JSON-compatible list.
    """
    if not sentences:
        raise ValueError("The input list of sentences must not be empty.")
    
    embeddings = all_MiniLM_L6_V2_model.encode(sentences)
    return {"embeddings": embeddings.tolist()}  # Return embeddings as a JSON-compatible list


@app.post("/generate-embeddings/intfloat-e5-large-v2/")
def generate_embeddings_intfloat_e5_large_v2_model(sentences: List[str]):
    """
    Generates embeddings for a list of sentences using the intfloat/e5-large-v2 model.

    Args:
        sentences (List[str]): A list of sentences to generate embeddings for.

    Returns:
        dict: A dictionary containing the embeddings as a JSON-compatible list.
    """

    if not sentences:
        raise ValueError("The input list of sentences must not be empty.")
    
    embeddings = intfloat_e5_large_v2_model.encode(sentences)
    return {"embeddings": embeddings.tolist()}