|
|
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 |
