litellm/tests/integration/providers/test_openai_embedding_wire.py

import base64
import json
import struct
from typing import Final, Literal

import httpx
import openai
import pytest
from integration._support.client import Gateway
from integration._support.wire import Reply, Request, wire_server
from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter

_MODEL: Final = "openai/text-embedding-3-small"
_SDK_KEY: Final = "synthetic-openai-key"
_INPUTS: Final = ("embedding alpha marker", "embedding bravo marker")
_RAW_INPUT: Final = "single raw embedding marker"
_VECTORS: Final = ((0.25, 0.5, 0.75), (-1.0, 0.125, 2.0))
_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])


class _Embedding(BaseModel):
    model_config = ConfigDict(extra="forbid")

    object: Literal["embedding"]
    index: int
    embedding: list[float]


class _Usage(BaseModel):
    prompt_tokens: int
    total_tokens: int


class _EmbeddingResponse(BaseModel):
    data: tuple[_Embedding, ...]
    usage: _Usage


def _encoded(vector: tuple[float, ...]) -> str:
    return base64.b64encode(struct.pack("<3f", *vector)).decode("ascii")


def _embedding_reply(embeddings: list[str | list[float]]) -> Reply:
    return Reply(
        body=json.dumps(
            {
                "object": "list",
                "data": [
                    {"object": "embedding", "index": index, "embedding": embedding}
                    for index, embedding in enumerate(embeddings)
                ],
                "model": "text-embedding-3-small",
                "usage": {"prompt_tokens": 4, "total_tokens": 4},
            }
        ).encode()
    )


# openai-python 2.33.0 defaults omitted encoding_format to base64: https://github.com/openai/openai-python/blob/v2.33.0/src/openai/resources/embeddings.py
def test_openai_sdk_embedding_preserves_base64_and_decodes_vectors(gateway: Gateway) -> None:
    def respond(request: Request) -> Reply:
        assert request.method == "POST"
        assert request.target == "/v1/embeddings"
        assert request.headers["authorization"] == f"Bearer {_SDK_KEY}"
        assert request.headers["authorization"] != f"Bearer {virtual_key}"
        assert _JSON_OBJECT.validate_json(request.body) == {
            "model": "text-embedding-3-small",
            "input": list(_INPUTS),
            "encoding_format": "base64",
            "dimensions": 3,
            "user": "u1",
        }
        return _embedding_reply([_encoded(vector) for vector in _VECTORS])

    with wire_server(respond) as wire, gateway.scenario() as scenario:
        model: Final = scenario.model(
            model=_MODEL,
            api_base=f"{wire.url}/v1",
            api_key=_SDK_KEY,
        )
        virtual_key: Final = scenario.key()
        with (
            httpx.Client(trust_env=False) as http_client,
            openai.OpenAI(
                base_url=f"{str(gateway.client.base_url).rstrip('/')}/v1",
                api_key=virtual_key,
                max_retries=0,
                http_client=http_client,
            ) as client,
        ):
            response: Final = client.embeddings.create(
                model=model,
                input=list(_INPUTS),
                dimensions=3,
                user="u1",
            )
        assert [item.embedding for item in response.data] == [list(vector) for vector in _VECTORS], (
            response.model_dump_json()
        )
        assert [item.index for item in response.data] == [0, 1], response.model_dump_json()
        assert (response.usage.prompt_tokens, response.usage.total_tokens) == (4, 4), response.model_dump_json()
        assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/v1/embeddings")]


@pytest.mark.parametrize("encoding_format", (None, "float"), ids=("omitted", "float"))
def test_openai_embedding_forwards_omitted_or_float_encoding_format(
    gateway: Gateway,
    encoding_format: str | None,
) -> None:
    expected_body: Final = (
        {"model": "text-embedding-3-small", "input": _RAW_INPUT}
        if encoding_format is None
        else {"model": "text-embedding-3-small", "input": _RAW_INPUT, "encoding_format": "float"}
    )

    def respond(request: Request) -> Reply:
        assert request.method == "POST"
        assert request.target == "/v1/embeddings"
        assert request.headers["authorization"] == f"Bearer {_SDK_KEY}"
        body: Final = _JSON_OBJECT.validate_json(request.body)
        if encoding_format is None and "encoding_format" in body:
            return Reply(status=400, body=b'{"error":"encoding_format must be omitted"}')
        assert body == expected_body, body
        return _embedding_reply([[0.5, -0.25, 1.0]])

    with wire_server(respond) as wire, gateway.scenario() as scenario:
        model: Final = scenario.model(
            model=_MODEL,
            api_base=f"{wire.url}/v1",
            api_key=_SDK_KEY,
        )
        virtual_key: Final = scenario.key()
        request_body: Final = (
            {"model": model, "input": _RAW_INPUT}
            if encoding_format is None
            else {"model": model, "input": _RAW_INPUT, "encoding_format": encoding_format}
        )
        response: Final = gateway.request("POST", "/v1/embeddings", request_body, key=virtual_key)
        assert response.status_code == 200, response.text
        payload: Final = _EmbeddingResponse.model_validate_json(response.content)
        assert [item.model_dump() for item in payload.data] == [
            {"object": "embedding", "index": 0, "embedding": [0.5, -0.25, 1.0]}
        ], response.text
        assert (payload.usage.prompt_tokens, payload.usage.total_tokens) == (4, 4), response.text
        assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/v1/embeddings")]


def test_openai_embedding_relays_requested_base64_string_unchanged(gateway: Gateway) -> None:
    encoded: Final = _encoded(_VECTORS[0])

    def respond(request: Request) -> Reply:
        assert request.method == "POST"
        assert request.target == "/v1/embeddings"
        assert request.headers["authorization"] == f"Bearer {_SDK_KEY}"
        body: Final = _JSON_OBJECT.validate_json(request.body)
        assert body == {"model": "text-embedding-3-small", "input": _RAW_INPUT, "encoding_format": "base64"}, body
        return _embedding_reply([encoded])

    with wire_server(respond) as wire, gateway.scenario() as scenario:
        model: Final = scenario.model(
            model=_MODEL,
            api_base=f"{wire.url}/v1",
            api_key=_SDK_KEY,
        )
        virtual_key: Final = scenario.key()
        response: Final = gateway.request(
            "POST",
            "/v1/embeddings",
            {"model": model, "input": _RAW_INPUT, "encoding_format": "base64"},
            key=virtual_key,
        )
        assert response.status_code == 200, response.text
        payload: Final = _JSON_OBJECT.validate_json(response.content)
        assert payload["data"] == [{"object": "embedding", "index": 0, "embedding": encoded}], response.text
        assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/v1/embeddings")]