import os
from typing import Literal

import pytest

from openai import OpenAI
from openai.types import CreateEmbeddingResponse


@pytest.fixture
def client():
    return OpenAI(
        base_url=os.environ.get("HFENDPOINTS_BASE_URL", "http://localhost:8000/api/v1/"),
        api_key=os.environ.get("HFENDPOINTS_API_KEY", os.environ.get("HF_TOKEN", "hf_notdefined")),
        organization="hfendpoints-images",
    )


@pytest.mark.parametrize("encoding_format", ["float", "base64"])
def test_openai_embedding_single(client, encoding_format: Literal["float", "base64"]):
    response: CreateEmbeddingResponse = client.embeddings.create(
        input="Hello, how are you?",
        model="all-MiniLM-L6-v2",
        encoding_format=encoding_format,
    )

    assert len(response.model)

    assert response.object == "list"
    assert len(response.data) == 1

    assert response.data[0].object == "embedding"
    assert len(response.data[0].embedding) == 384


@pytest.mark.parametrize("encoding_format", ["float", "base64"])
def test_openai_embedding_batch(client, encoding_format: Literal["float", "base64"]):
    response: CreateEmbeddingResponse = client.embeddings.create(
        input=["Hello, how are you?", "How are you?"],
        model="all-MiniLM-L6-v2",
        encoding_format=encoding_format,
    )

    assert len(response.model)

    assert response.object == "list"
    assert len(response.data) == 2

    assert response.data[0].object == "embedding"
    assert len(response.data[0].embedding) == 384

    assert response.data[1].object == "embedding"
    assert len(response.data[1].embedding) == 384