from fastapi import FastAPI, Query
from pydantic import BaseModel
from typing import List
from model import get_text_embedding, get_image_embedding

app = FastAPI(title="Jina CLIP V2 API")

class TextInput(BaseModel):
    texts: List[str]

class ImageInput(BaseModel):
    image_urls: List[str]

@app.post("/text-embedding")
def text_embedding(input_data: TextInput, truncate_dim: int = Query(512)):
    embeddings = get_text_embedding(input_data.texts, truncate_dim)
    return {"embeddings": embeddings}

@app.post("/image-embedding")
def image_embedding(input_data: ImageInput, truncate_dim: int = Query(512)):
    embeddings = get_image_embedding(input_data.image_urls, truncate_dim)
    return {"embeddings": embeddings}