Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 7 additions & 2 deletions src/openai/lib/_parsing/_embeddings.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
from __future__ import annotations

import sys
import array
import base64
from typing import cast
Expand All @@ -24,12 +25,16 @@ def parse_embedding_response(
data = cast(object, embedding.embedding)
if not isinstance(data, str):
continue
decoded = base64.b64decode(data)
if not has_numpy():
# use array for base64 optimisation
embedding.embedding = array.array("f", base64.b64decode(data)).tolist()
floats = array.array("f", decoded)
if sys.byteorder == "big":
floats.byteswap()
embedding.embedding = floats.tolist()
else:
embedding.embedding = np.frombuffer( # type: ignore[no-untyped-call]
base64.b64decode(data), dtype="float32"
decoded, dtype="<f4"
).tolist()

return obj
35 changes: 33 additions & 2 deletions tests/lib/test_embeddings.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,10 @@
from __future__ import annotations

import json
import array
import base64
import struct
import binascii
from types import SimpleNamespace
from typing import Any, cast
from typing_extensions import Literal

Expand All @@ -18,7 +19,7 @@
from openai.types.create_embedding_response import CreateEmbeddingResponse

VALUES = [0.125, -2.5, 3.75]
ENCODED = base64.b64encode(array.array("f", VALUES).tobytes()).decode("ascii")
ENCODED = base64.b64encode(struct.pack("<3f", *VALUES)).decode("ascii")
EncodingFormat = Literal["float", "base64"] | Omit
ResponseMode = Literal["normal", "raw", "streaming"]

Expand Down Expand Up @@ -63,6 +64,36 @@ def test_decode_preserves_response_and_non_string_vectors(encoding_format: Omit
assert parsed.model == "text-embedding-3-small"


def test_stdlib_decoder_swaps_big_endian_values(monkeypatch: pytest.MonkeyPatch) -> None:
response = make_response(ENCODED)
decoded = base64.b64decode(ENCODED)

class Floats:
swapped = False

def byteswap(self) -> None:
self.swapped = True

def tolist(self) -> list[float]:
assert self.swapped
return VALUES

floats = Floats()

def make_array(typecode: str, initializer: bytes) -> Floats:
assert typecode == "f"
assert initializer == decoded
return floats

monkeypatch.setattr(embeddings_parser, "has_numpy", lambda: False)
monkeypatch.setattr(embeddings_parser, "sys", SimpleNamespace(byteorder="big"))
monkeypatch.setattr(embeddings_parser.array, "array", make_array)

parsed = embeddings_parser.parse_embedding_response(response, encoding_format=omit)

assert parsed.data[0].embedding == VALUES


@pytest.mark.parametrize("encoding_format", ["float", "base64", None])
@pytest.mark.parametrize("vectors", [(ENCODED,), ("abc",), ()], ids=["encoded", "invalid", "empty"])
def test_explicit_format_is_untouched(
Expand Down