from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, root_validator
from transformers import AutoModel
from typing import List
import os

if platform.system() == "Windows":
    print("Windows detected. Assigning cache directory to Transformers in AppData\Local.")
    transformers_cache_directory = os.path.join(os.getenv('LOCALAPPDATA'), 'transformers_cache')
    if not os.path.exists(transformers_cache_directory):
        try:
            os.mkdir(transformers_cache_directory)
            print(f"First launch. Directory '{transformers_cache_directory}' created successfully.")
        except OSError as e:
            print(f"Error creating directory '{transformers_cache_directory}': {e}")
    else:
        print(f"Directory '{transformers_cache_directory}' already exists.")
    os.environ['TRANSFORMERS_CACHE'] = transformers_cache_directory
    print("Environment variable assigned.")
    del transformers_cache_directory

else:
    print("Windows not detected. Assignment of Transformers cache directory not necessary.")


model = AutoModel.from_pretrained('jinaai/jina-embeddings-v2-base-en',token = "hf_GkUomApayMBJteRvrjvslfyLRvfp QRckba".replace(" ", ""), trust_remote_code=True)

app = FastAPI()

class Validation(BaseModel):
    prompt: List[str]


#Endpoint
@app.post("/jina_embedding")
async def chaatie_agent(item: Validation):
    # Assuming model.encode returns a list of numpy arrays (one for each prompt)
    embeddings = model.encode(item.prompt)

    # Convert each numpy array in the list to a list
    embeddings_list = [embedding.tolist() for embedding in embeddings]

    return {"embeddings": embeddings_list}