This is an automated email from the ASF dual-hosted git repository. wenjin272 pushed a commit to branch main in repository https://gitbox.apache.org/repos/asf/flink-agents.git
commit 1e91a775872d75d13cd31b6d73384cedc05d2553 Author: purshotam shah <[email protected]> AuthorDate: Fri Sep 18 10:50:19 2026 -0700 [integrations][python] Align the OpenAI embedding model with the Java connection additional_kwargs are sent as extra request body properties and may not repeat a typed field; request_timeout 0 disables the timeout as in the OpenAI chat models; request_timeout, max_retries, dimensions, encoding_format and model are validated at construction and per call; blank or null strings mean absent or default; base64 responses are decoded (previously the characters of the base64 string were returned) with padding restored only for completely unpadded values, so partially padded input is rejected as it is by Java's decoder; responses are checked for size, index consistency and malformed vectors; token usage tolerates partial or string counts. The mocked tests run in the unit suite; only the live-API test keeps the integration marker. Compatibility: additional_kwargs entries that are not typed fields were dropped and are now sent; entries that repeat a typed field previously overrode it and now fail validation; model_kwargs nests additional_kwargs; dimensions must be an integer. Documented in embedding_models.md. Generated-by: Claude Code 2.1.273 (Claude Fable 5.1) --- .../embedding_models/openai_embedding_model.py | 341 ++++++++++++++++-- .../tests/test_openai_embedding_model.py | 383 ++++++++++++++++++++- 2 files changed, 692 insertions(+), 32 deletions(-) diff --git a/python/flink_agents/integrations/embedding_models/openai_embedding_model.py b/python/flink_agents/integrations/embedding_models/openai_embedding_model.py index 92d09b76..b8cd0256 100644 --- a/python/flink_agents/integrations/embedding_models/openai_embedding_model.py +++ b/python/flink_agents/integrations/embedding_models/openai_embedding_model.py @@ -15,10 +15,14 @@ # See the License for the specific language governing permissions and # limitations under the License. ################################################################################ -from typing import Any, Dict, Sequence +import array +import base64 +import sys +from collections.abc import Mapping +from typing import Any, Dict, Sequence, cast from openai import NOT_GIVEN, OpenAI -from pydantic import Field +from pydantic import Field, field_validator from typing_extensions import override from flink_agents.api.embedding_models.embedding_model import ( @@ -31,6 +35,168 @@ from flink_agents.api.embedding_models.embedding_model import ( DEFAULT_REQUEST_TIMEOUT = 30.0 DEFAULT_BASE_URL = "https://api.openai.com/v1"; DEFAULT_MAX_RETRIES = 3 +DEFAULT_ENCODING_FORMAT = "float" +# Upper bounds shared with the OpenAI chat models and the Java connection. +MAX_OPENAI_TIMEOUT_SECONDS = 2_147_483.647 +MAX_OPENAI_RETRIES = 2_147_483_647 +MAX_DIMENSIONS = 2_147_483_647 # int range, as the Java connection requires +ENCODING_FORMATS = frozenset({"float", "base64"}) +# Typed request fields; additional_kwargs may not repeat them. +RESERVED_ADDITIONAL_KWARGS = frozenset( + {"model", "input", "encoding_format", "dimensions", "user"} +) + + +def _blank_to_none(value: Any, name: str) -> str | None: + """Require a string or None; blank strings are treated as absent, as in Java.""" + if value is None: + return None + if not isinstance(value, str): + msg = f"{name} must be a string, got: {type(value).__name__}" + raise TypeError(msg) + return value if value.strip() else None + + +def _check_additional_kwargs(additional_kwargs: Any) -> Dict[str, Any]: + """Require a mapping with non-blank string keys and no typed request field.""" + if not isinstance(additional_kwargs, Mapping): + msg = ( + f"additional_kwargs must be a map, got: {type(additional_kwargs).__name__}" + ) + raise TypeError(msg) + # Per-call keys are stringified as Java's Map keys are (the setup field already + # requires string keys); null and blank keys are rejected in both languages. + if any(key is None or not str(key).strip() for key in additional_kwargs): + msg = "additional_kwargs contains an empty key." + raise ValueError(msg) + additional_kwargs = {str(key): value for key, value in additional_kwargs.items()} + collisions = sorted(RESERVED_ADDITIONAL_KWARGS & additional_kwargs.keys()) + if collisions: + msg = ( + f"additional_kwargs must not contain the typed request fields {collisions}; " + "set them through the corresponding setup argument instead." + ) + raise ValueError(msg) + return additional_kwargs + + +def _check_encoding_format(encoding_format: Any) -> str: + """Validate encoding_format; a blank value means the default, as in Java.""" + if _blank_to_none(encoding_format, "encoding_format") is None: + return DEFAULT_ENCODING_FORMAT + if encoding_format not in ENCODING_FORMATS: + msg = ( + f"encoding_format must be one of {sorted(ENCODING_FORMATS)}, " + f"got: {encoding_format!r}" + ) + raise ValueError(msg) + return encoding_format + + +def _check_model(model: Any) -> str: + if not isinstance(model, str) or not model.strip(): + msg = "OpenAI embedding requires a non-empty 'model' (setup or per-call parameter)." + raise ValueError(msg) + return model + + +def _check_dimensions(dimensions: Any) -> int | None: + """Accept a positive integer (an integral float counts, as in Java).""" + if dimensions is None: + return None + if isinstance(dimensions, float) and dimensions.is_integer(): + dimensions = int(dimensions) + if ( + isinstance(dimensions, bool) + or not isinstance(dimensions, int) + or not 0 < dimensions <= MAX_DIMENSIONS + ): + msg = f"dimensions must be a positive integer, got: {dimensions!r}" + raise ValueError(msg) + return dimensions + + +def _usage_count(value: Any) -> int | None: + """A token count, or None when a compatible server sends something unusable.""" + if isinstance(value, str) and value.strip().isdecimal(): + return int(value) + if isinstance(value, bool): + return None + if isinstance(value, int): + return value + if isinstance(value, float) and value.is_integer(): + return int(value) + return None + + +def _decode_embedding(value: Any, *, base64_requested: bool) -> list[float]: + # With encoding_format="base64" the SDK hands the base64 string through untouched. + if isinstance(value, str): + if not base64_requested: + msg = "a string vector was returned although encoding_format was float" + raise TypeError(msg) + # validate=True so corrupted input raises instead of being silently truncated. + # Padding is restored only for a completely unpadded value, which Java's + # decoder accepts too; partial padding such as "AACAPw=" stays invalid, as it + # is for Java. + padded = value if "=" in value else value + "=" * (-len(value) % 4) + floats = array.array("f", base64.b64decode(padded, validate=True)) + if sys.byteorder == "big": + floats.byteswap() # the wire format is little-endian float32 + return floats.tolist() + vector = list(value) + if any(isinstance(x, bool) or not isinstance(x, int | float) for x in vector): + msg = "embedding contains non-numeric values" + raise TypeError(msg) + return vector + + +def _order_embeddings( + data: Sequence[Any], expected: int, *, base64_requested: bool +) -> list[list[float]]: + """Place each vector at its response ``index``; mirrors the Java connection. + + An item without an index keeps its response position. Every slot must be + filled exactly once, so out-of-range, duplicated or contradictory indices + are rejected rather than silently mapped. + """ + if len(data) != expected: + msg = f"OpenAI returned {len(data)} embeddings for {expected} input texts." + raise RuntimeError(msg) + ordered: list[list[float] | None] = [None] * expected + for position, item in enumerate(data): + if item is None: + msg = f"OpenAI returned a malformed embedding at position {position}." + raise RuntimeError(msg) + index = getattr(item, "index", None) + if isinstance(index, float) and index.is_integer(): + index = int(index) # an integral float counts, as Jackson does in Java + if index is None: + target = position + elif isinstance(index, int) and not isinstance(index, bool): + target = index + else: + # A present but non-integer index is malformed, not absent. + msg = ( + f"OpenAI returned a non-integer embedding index {index!r} " + f"at position {position}." + ) + raise RuntimeError(msg) + if not 0 <= target < expected or ordered[target] is not None: + msg = ( + f"OpenAI returned an unexpected embedding index {target} " + f"at position {position}." + ) + raise RuntimeError(msg) + try: + ordered[target] = _decode_embedding( + item.embedding, base64_requested=base64_requested + ) + except (TypeError, ValueError) as e: + msg = f"OpenAI returned a malformed embedding at position {position}: {e}" + raise RuntimeError(msg) from e + # The size and duplicate checks above guarantee every slot is filled. + return cast("list[list[float]]", ordered) class OpenAIEmbeddingModelConnection(BaseEmbeddingModelConnection): @@ -45,7 +211,8 @@ class OpenAIEmbeddingModelConnection(BaseEmbeddingModelConnection): base_url : str Base URL for the OpenAI API (default: https://api.openai.com/v1). request_timeout : float - The timeout for making HTTP requests to OpenAI API. + The timeout for making HTTP requests to OpenAI API. Set to 0 to disable + timeouts. max_retries : int Maximum number of retries for failed requests. organization : Optional[str] @@ -55,17 +222,46 @@ class OpenAIEmbeddingModelConnection(BaseEmbeddingModelConnection): """ api_key: str = Field(description="OpenAI API key for authentication.") + + # A null (YAML ``~``) or blank base_url means the default, as in Java. + @field_validator("base_url", mode="before") + @classmethod + def _default_base_url(cls, value: Any) -> Any: + try: + return ( + DEFAULT_BASE_URL if _blank_to_none(value, "base_url") is None else value + ) + except TypeError as e: # pydantic only converts ValueError + raise ValueError(str(e)) from e + + # A YAML null means the default, as parseRequestTimeout(null) does in Java. + @field_validator("request_timeout", mode="before") + @classmethod + def _default_request_timeout(cls, value: Any) -> Any: + return DEFAULT_REQUEST_TIMEOUT if value is None else value + + @field_validator("max_retries", mode="before") + @classmethod + def _default_max_retries(cls, value: Any) -> Any: + return DEFAULT_MAX_RETRIES if value is None else value + base_url: str = Field( default=DEFAULT_BASE_URL, description="Base URL for the OpenAI API.", ) request_timeout: float = Field( default=DEFAULT_REQUEST_TIMEOUT, - description="The timeout for making HTTP requests to OpenAI API.", + description="The timeout for making HTTP requests to OpenAI API. " + "Set to 0 to disable timeouts.", + ge=0, + le=MAX_OPENAI_TIMEOUT_SECONDS, + allow_inf_nan=False, ) max_retries: int = Field( default=DEFAULT_MAX_RETRIES, description="Maximum number of retries for failed requests.", + ge=0, + le=MAX_OPENAI_RETRIES, ) organization: str | None = Field( default=None, @@ -80,7 +276,7 @@ class OpenAIEmbeddingModelConnection(BaseEmbeddingModelConnection): self, api_key: str, base_url: str = DEFAULT_BASE_URL, - request_timeout: float | None = DEFAULT_REQUEST_TIMEOUT, + request_timeout: float = DEFAULT_REQUEST_TIMEOUT, max_retries: int = DEFAULT_MAX_RETRIES, organization: str | None = None, project: str | None = None, @@ -105,10 +301,12 @@ class OpenAIEmbeddingModelConnection(BaseEmbeddingModelConnection): if self.__client is None: self.__client = OpenAI( api_key=self.api_key, - base_url=self.base_url, - timeout=self.request_timeout, - organization=self.organization, - project=self.project, + # Blank strings mean the default / unset (the SDK may then read its + # OPENAI_ORG_ID / OPENAI_PROJECT_ID environment variables). + base_url=_blank_to_none(self.base_url, "base_url") or DEFAULT_BASE_URL, + timeout=None if self.request_timeout == 0 else self.request_timeout, + organization=_blank_to_none(self.organization, "organization"), + project=_blank_to_none(self.project, "project"), max_retries=self.max_retries, ) return self.__client @@ -122,38 +320,77 @@ class OpenAIEmbeddingModelConnection(BaseEmbeddingModelConnection): def embed_with_usage( self, text: str | Sequence[str], **kwargs: Any ) -> EmbeddingResult[list[float] | list[list[float]]]: - """Generate embeddings and return OpenAI token usage when available.""" + """Generate embeddings and return OpenAI token usage when available. + + ``additional_kwargs`` (from the setup) are sent as extra request body + properties; they may not repeat the typed fields, and ``None`` values are + omitted. Per-call keyword arguments override the setup's entry by entry, so + a per-call ``additional_kwargs`` map replaces the setup's map. Other + keyword arguments are ignored, as in the Java connection. Vectors are + returned in input order and ``base64`` responses are decoded. + """ # Extract OpenAI specific parameters - model = kwargs.pop("model") - encoding_format = kwargs.pop("encoding_format", None) - dimensions = kwargs.pop("dimensions", None) - user = kwargs.pop("user", None) + model = _check_model(kwargs.pop("model", None)) + encoding_format = _check_encoding_format(kwargs.pop("encoding_format", None)) + dimensions = _check_dimensions(kwargs.pop("dimensions", None)) + user = _blank_to_none(kwargs.pop("user", None), "user") + raw_additional_kwargs = kwargs.pop("additional_kwargs", None) + additional_kwargs = _check_additional_kwargs( + {} if raw_additional_kwargs is None else raw_additional_kwargs + ) + extra_body = {k: v for k, v in additional_kwargs.items() if v is not None} + + if text is None: + msg = "text must be a string or a sequence of strings, got: None" + raise TypeError(msg) + if not isinstance(text, str): + if isinstance(text, Mapping | set | frozenset): + msg = f"text must be a string or a sequence of strings, got: {type(text).__name__}" + raise TypeError(msg) + text = list(text) # materialize one-shot iterables before validating + for position, item in enumerate(text): + if not isinstance(item, str): + msg = f"Text at index {position} is not a string; every input must be a string." + raise TypeError(msg) + if len(text) == 0: + return EmbeddingResult(embeddings=[], token_usage=None) # Create the embedding request + # The resolved format is always sent, as the Java connection does; a "base64" + # response is decoded below. response = self.client.embeddings.create( model=model, input=text, - encoding_format=encoding_format - if encoding_format is not None - else NOT_GIVEN, + encoding_format=encoding_format, dimensions=dimensions if dimensions is not None else NOT_GIVEN, user=user if user is not None else NOT_GIVEN, + extra_body=extra_body or None, ) usage = getattr(response, "usage", None) token_usage = None if usage is not None: - prompt_tokens = getattr(usage, "prompt_tokens", None) - total_tokens = getattr(usage, "total_tokens", None) - if prompt_tokens is not None or total_tokens is not None: + prompt_tokens = _usage_count(getattr(usage, "prompt_tokens", None)) + total_tokens = _usage_count(getattr(usage, "total_tokens", None)) + # Embedding requests have no completion tokens, so a missing side of the + # usage equals the other side (as the Tongyi connection does). + if prompt_tokens is None: + prompt_tokens = total_tokens + if prompt_tokens is not None: token_usage = EmbeddingTokenUsage( - prompt_tokens=int(prompt_tokens or 0), + prompt_tokens=int(prompt_tokens), total_tokens=int( total_tokens if total_tokens is not None else prompt_tokens ), ) - embeddings = [list(embedding.embedding) for embedding in response.data] + # A compatible server may answer 200 with no ``data``; report it as a short + # response rather than a bare TypeError. + embeddings = _order_embeddings( + getattr(response, "data", None) or [], + 1 if isinstance(text, str) else len(text), + base64_requested=encoding_format == "base64", + ) return EmbeddingResult( embeddings=embeddings[0] if isinstance(text, str) else embeddings, token_usage=token_usage, @@ -188,12 +425,14 @@ class OpenAIEmbeddingModelSetup(BaseEmbeddingModelSetup): A unique identifier representing your end-user, which can help OpenAI to monitor and detect abuse. additional_kwargs : Dict[str, Any] - Additional parameters for the OpenAI embeddings API. + Extra request body properties sent with every embeddings request. Keys + may not repeat the typed arguments above (model, encoding_format, + dimensions, user). """ encoding_format: str = Field( default="float", - description='The format to return the embeddings in (default: "float").', + description='The format to return the embeddings in: "float" or "base64".', ) dimensions: int | None = Field( default=None, @@ -205,9 +444,51 @@ class OpenAIEmbeddingModelSetup(BaseEmbeddingModelSetup): ) additional_kwargs: Dict[str, Any] = Field( default_factory=dict, - description="Additional parameters for the OpenAI embeddings API.", + description="Extra request body properties for the OpenAI embeddings API.", ) + # "before" mode so non-string keys are stringified as the Java setup does, instead + # of failing pydantic's Dict[str, Any] check first. + @field_validator("additional_kwargs", mode="before") + @classmethod + def _validate_additional_kwargs(cls, value: Any) -> Dict[str, Any]: + try: + return _check_additional_kwargs(value) + except TypeError as e: # pydantic only converts ValueError + raise ValueError(str(e)) from e + + # "before" mode so a null from YAML (``encoding_format: ~``) means the default, + # as in Java, instead of failing the ``str`` field check. A TypeError would + # escape pydantic, so it is reported as a ValueError like the other fields. + @field_validator("encoding_format", mode="before") + @classmethod + def _validate_encoding_format(cls, value: Any) -> str: + try: + return _check_encoding_format(value) + except TypeError as e: + raise ValueError(str(e)) from e + + @field_validator("connection") + @classmethod + def _validate_connection(cls, value: Any) -> Any: + # The field also accepts a connection object; only a name can be blank. + if isinstance(value, str) and not value.strip(): + msg = "OpenAI embedding requires a non-empty 'connection' on the embedding model setup." + raise ValueError(msg) + return value + + # "before" mode so pydantic's lax coercion (True -> 1, "256" -> 256) does not + # run first; the per-call path and the Java setup reject those inputs too. + @field_validator("dimensions", mode="before") + @classmethod + def _validate_dimensions(cls, value: Any) -> int | None: + return _check_dimensions(value) + + @field_validator("model") + @classmethod + def _validate_model(cls, value: str) -> str: + return _check_model(value) + def __init__( self, *, @@ -243,10 +524,10 @@ class OpenAIEmbeddingModelSetup(BaseEmbeddingModelSetup): if self.dimensions is not None: base_kwargs["dimensions"] = self.dimensions - if self.user is not None: + if _blank_to_none(self.user, "user") is not None: base_kwargs["user"] = self.user - return { - **base_kwargs, - **self.additional_kwargs, - } + if self.additional_kwargs: + base_kwargs["additional_kwargs"] = dict(self.additional_kwargs) + + return base_kwargs diff --git a/python/flink_agents/integrations/embedding_models/tests/test_openai_embedding_model.py b/python/flink_agents/integrations/embedding_models/tests/test_openai_embedding_model.py index 9d63c249..1ebfc006 100644 --- a/python/flink_agents/integrations/embedding_models/tests/test_openai_embedding_model.py +++ b/python/flink_agents/integrations/embedding_models/tests/test_openai_embedding_model.py @@ -15,11 +15,16 @@ # See the License for the specific language governing permissions and # limitations under the License. ################################################################################ +import array +import base64 import os from types import SimpleNamespace +from typing import Any from unittest.mock import MagicMock import pytest +from openai import NOT_GIVEN +from pydantic import ValidationError from flink_agents.api.resource import Resource, ResourceType from flink_agents.api.resource_context import ResourceContext @@ -28,12 +33,11 @@ from flink_agents.integrations.embedding_models.openai_embedding_model import ( OpenAIEmbeddingModelSetup, ) -pytestmark = pytest.mark.integration - test_model = os.environ.get("TEST_EMBEDDING_MODEL", "text-embedding-3-small") api_key = os.environ.get("TEST_API_KEY") [email protected] @pytest.mark.skipif(api_key is None, reason="TEST_API_KEY is not set") def test_openai_embedding_model() -> None: connection = OpenAIEmbeddingModelConnection(name="openai", api_key=api_key) @@ -89,3 +93,378 @@ def test_openai_embedding_model_returns_token_usage() -> None: assert result.token_usage is not None assert result.token_usage.prompt_tokens == 5 assert result.token_usage.total_tokens == 5 + + +def test_openai_embedding_connection_zero_timeout_disables_timeout( + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured: dict[str, Any] = {} + + def fake_openai(**kwargs: Any) -> object: + captured.update(kwargs) + return object() + + monkeypatch.setattr( + "flink_agents.integrations.embedding_models.openai_embedding_model.OpenAI", + fake_openai, + ) + zero = OpenAIEmbeddingModelConnection( + name="openai", api_key="k", request_timeout=0, base_url=" ", organization=" " + ) + zero.client + assert captured["timeout"] is None + # Whitespace-only strings mean the default / absent, as in Java. + assert captured["base_url"] == "https://api.openai.com/v1"; + assert captured["organization"] is None + + captured.clear() + ten = OpenAIEmbeddingModelConnection(name="openai", api_key="k", request_timeout=10) + ten.client + assert captured["timeout"] == 10 + + +def test_openai_embedding_connection_rejects_invalid_timeout() -> None: + with pytest.raises(ValidationError): + OpenAIEmbeddingModelConnection(name="openai", api_key="k", request_timeout=-1) + with pytest.raises(ValidationError): + OpenAIEmbeddingModelConnection(name="openai", api_key="k", max_retries=-1) + with pytest.raises(ValidationError): + OpenAIEmbeddingModelConnection(name="openai", api_key="k", max_retries=2**31) + # A blank api_key stays accepted: unauthenticated OpenAI-compatible servers rely on + # it and the SDK then omits the Authorization header. + assert OpenAIEmbeddingModelConnection(name="openai", api_key="").api_key == "" + + +def _setup_with_mock_client( + additional_kwargs: dict[str, Any], +) -> tuple[OpenAIEmbeddingModelSetup, MagicMock]: + connection = OpenAIEmbeddingModelConnection(name="openai", api_key="fake-key") + mock_client = MagicMock() + mock_client.embeddings.create.return_value = SimpleNamespace( + data=[SimpleNamespace(embedding=[0.1, 0.2, 0.3])], + usage=SimpleNamespace(prompt_tokens=5, total_tokens=5), + ) + connection._OpenAIEmbeddingModelConnection__client = mock_client + mock_ctx = MagicMock(spec=ResourceContext) + mock_ctx.get_resource = lambda name, type: connection + setup = OpenAIEmbeddingModelSetup( + name="openai", + model=test_model, + connection="openai", + dimensions=3, + additional_kwargs=additional_kwargs, + resource_context=mock_ctx, + ) + setup.open() + return setup, mock_client + + +def test_openai_embedding_additional_kwargs_sent_as_extra_body() -> None: + setup, mock_client = _setup_with_mock_client({"custom_flag": True, "skip": None}) + + setup.embed("Hello, Flink Agent!") + + kwargs = mock_client.embeddings.create.call_args.kwargs + assert kwargs["model"] == test_model + assert kwargs["dimensions"] == 3 + assert kwargs["extra_body"] == {"custom_flag": True} + + setup_without_extras, mock_client = _setup_with_mock_client({}) + setup_without_extras.embed("Hello, Flink Agent!") + assert mock_client.embeddings.create.call_args.kwargs["extra_body"] is None + + +def test_openai_embedding_additional_kwargs_rejects_typed_fields() -> None: + with pytest.raises(ValidationError, match="typed request fields"): + OpenAIEmbeddingModelSetup( + name="openai", + model=test_model, + connection="openai", + additional_kwargs={"dimensions": 8, "custom": 1}, + ) + connection = OpenAIEmbeddingModelConnection(name="openai", api_key="fake-key") + with pytest.raises(ValueError, match=r"\['model'\]"): + connection.embed("x", model=test_model, additional_kwargs={"model": "other"}) + with pytest.raises(TypeError, match="must be a map"): + connection.embed("x", model=test_model, additional_kwargs="foo") + with pytest.raises(ValueError, match="empty key"): + connection.embed("x", model=test_model, additional_kwargs={" ": 1}) + with pytest.raises(ValueError, match="empty key"): + connection.embed("x", model=test_model, additional_kwargs={None: 1}) + with pytest.raises(TypeError, match="must be a map"): + connection.embed("x", model=test_model, additional_kwargs=[]) + with pytest.raises(ValidationError, match="empty key"): + OpenAIEmbeddingModelSetup( + name="openai", + model=test_model, + connection="openai", + additional_kwargs={"": 1}, + ) + # Non-string keys are stringified as in Java; a non-mapping is a ValidationError. + int_key = OpenAIEmbeddingModelSetup( + name="openai", model=test_model, connection="openai", additional_kwargs={1: "x"} + ) + assert int_key.model_kwargs["additional_kwargs"] == {"1": "x"} + with pytest.raises(ValidationError, match="must be a map"): + OpenAIEmbeddingModelSetup( + name="openai", + model=test_model, + connection="openai", + additional_kwargs="foo", + ) + blank_user = OpenAIEmbeddingModelSetup( + name="openai", model=test_model, connection="openai", user=" " + ) + assert "user" not in blank_user.model_kwargs + + +def test_openai_embedding_setup_validates_encoding_format_and_dimensions() -> None: + with pytest.raises(ValidationError, match="encoding_format"): + OpenAIEmbeddingModelSetup( + name="openai", model=test_model, connection="openai", encoding_format="hex" + ) + # A null (YAML ``~``) or blank encoding_format means the default, as in Java. + for absent in (None, " "): + setup = OpenAIEmbeddingModelSetup( + name="openai", model=test_model, connection="openai", encoding_format=absent + ) + assert setup.encoding_format == "float" + with pytest.raises(ValidationError): + OpenAIEmbeddingModelSetup( + name="openai", model=test_model, connection="openai", dimensions=0 + ) + for bad in (True, "256"): + with pytest.raises(ValidationError, match="dimensions"): + OpenAIEmbeddingModelSetup( + name="openai", model=test_model, connection="openai", dimensions=bad + ) + connection = OpenAIEmbeddingModelConnection(name="openai", api_key="fake-key") + with pytest.raises(ValueError, match="encoding_format"): + connection.embed("x", model=test_model, encoding_format="FLOAT") + with pytest.raises(ValueError, match="dimensions"): + connection.embed("x", model=test_model, dimensions=-1) + with pytest.raises(ValueError, match="dimensions"): + connection.embed("x", model=test_model, dimensions="256") + with pytest.raises(ValueError, match="dimensions"): + connection.embed("x", model=test_model, dimensions=True) + with pytest.raises(ValueError, match="dimensions"): + connection.embed("x", model=test_model, dimensions=2**40) + with pytest.raises(ValueError, match="'model'"): + connection.embed("x", model=" ") + with pytest.raises(TypeError, match="user must be a string"): + connection.embed("x", model=test_model, user=12345) + with pytest.raises(TypeError, match="encoding_format must be a string"): + connection.embed("x", model=test_model, encoding_format=1) + with pytest.raises(ValidationError, match="model"): + OpenAIEmbeddingModelSetup(name="openai", model="", connection="openai") + with pytest.raises(ValidationError, match="connection"): + OpenAIEmbeddingModelSetup(name="openai", model=test_model, connection=" ") + # A connection object (not a name) is accepted, as in the sibling setups. + setup_with_object = OpenAIEmbeddingModelSetup( + name="openai", model=test_model, connection=connection + ) + assert setup_with_object.connection is connection + for absent in (None, " "): + conn = OpenAIEmbeddingModelConnection( + name="openai", api_key="k", base_url=absent + ) + assert conn.base_url == "https://api.openai.com/v1"; + null_defaults = OpenAIEmbeddingModelConnection( + name="openai", api_key="k", request_timeout=None, max_retries=None + ) + assert (null_defaults.request_timeout, null_defaults.max_retries) == (30.0, 3) + with pytest.raises(ValidationError, match="base_url must be a string"): + OpenAIEmbeddingModelConnection(name="openai", api_key="k", base_url=8080) + with pytest.raises(TypeError, match="got: None"): + connection.embed(None, model=test_model) + with pytest.raises(TypeError, match="got: dict"): + connection.embed({"id": "doc-1"}, model=test_model) + with pytest.raises(ValidationError, match="encoding_format must be a string"): + OpenAIEmbeddingModelSetup( + name="openai", model=test_model, connection="openai", encoding_format=1 + ) + with pytest.raises(TypeError, match="Text at index 1 is not a string"): + connection.embed(["ok", None], model=test_model) + # An integral float is accepted as in Java. + setup = OpenAIEmbeddingModelSetup( + name="openai", model=test_model, connection="openai", dimensions=256.0 + ) + assert setup.model_kwargs["dimensions"] == 256 + + +def test_openai_embedding_blank_strings_and_default_format() -> None: + connection, mock_client = _connection_returning([SimpleNamespace(embedding=[0.1])]) + + connection.embed("x", model=test_model, encoding_format=" ", user=" ") + + kwargs = mock_client.embeddings.create.call_args.kwargs + # Blank user is omitted and a blank encoding_format means the default, as in Java. + assert kwargs["user"] is NOT_GIVEN + assert kwargs["encoding_format"] == "float" + + connection.embed("x", model=test_model, encoding_format="base64", user="u1") + kwargs = mock_client.embeddings.create.call_args.kwargs + assert kwargs["user"] == "u1" + assert kwargs["encoding_format"] == "base64" + + +def _connection_returning( + data: list[Any] | None, +) -> tuple[OpenAIEmbeddingModelConnection, MagicMock]: + connection = OpenAIEmbeddingModelConnection(name="openai", api_key="fake-key") + mock_client = MagicMock() + mock_client.embeddings.create.return_value = SimpleNamespace( + data=data, usage=SimpleNamespace(prompt_tokens=None, total_tokens=7) + ) + connection._OpenAIEmbeddingModelConnection__client = mock_client + return connection, mock_client + + +def test_openai_embedding_decodes_base64_and_orders_by_index() -> None: + encoded = base64.b64encode(array.array("f", [0.5, 0.25]).tobytes()).decode() + connection, _ = _connection_returning( + [ + SimpleNamespace(index=1, embedding=[0.9]), + SimpleNamespace(index=0, embedding=encoded), + ] + ) + + result = connection.embed_with_usage( + ["a", "b"], model=test_model, encoding_format="base64" + ) + + assert result.embeddings == [[0.5, 0.25], [0.9]] + + # Unpadded base64 (one float, "AACAPw" instead of "AACAPw==") decodes as in Java. + unpadded, _ = _connection_returning([SimpleNamespace(index=0, embedding="AACAPw")]) + assert unpadded.embed("a", model=test_model, encoding_format="base64") == [1.0] + # Partial padding ("AACAPw=" for "AACAPw==") is rejected, as by Java's decoder. + partial, _ = _connection_returning([SimpleNamespace(index=0, embedding="AACAPw=")]) + with pytest.raises(RuntimeError, match="malformed embedding at position 0"): + partial.embed("a", model=test_model, encoding_format="base64") + # A missing prompt_tokens equals total_tokens for embedding requests. + assert result.token_usage is not None + assert result.token_usage.prompt_tokens == 7 + assert result.token_usage.total_tokens == 7 + + +def test_openai_embedding_rejects_inconsistent_responses() -> None: + short, _ = _connection_returning([SimpleNamespace(index=0, embedding=[0.1])]) + with pytest.raises(RuntimeError, match="1 embeddings for 2 input texts"): + short.embed(["a", "b"], model=test_model) + + contradictory, _ = _connection_returning( + [SimpleNamespace(embedding=[0.1]), SimpleNamespace(index=0, embedding=[0.2])] + ) + with pytest.raises( + RuntimeError, match="unexpected embedding index 0 at position 1" + ): + contradictory.embed(["a", "b"], model=test_model) + + duplicate, _ = _connection_returning( + [ + SimpleNamespace(index=0, embedding=[0.1]), + SimpleNamespace(index=0, embedding=[0.2]), + ] + ) + with pytest.raises(RuntimeError, match="unexpected embedding index"): + duplicate.embed(["a", "b"], model=test_model) + string_index, _ = _connection_returning( + [ + SimpleNamespace(index="1", embedding=[0.1]), + SimpleNamespace(index="0", embedding=[0.2]), + ] + ) + with pytest.raises( + RuntimeError, match="non-integer embedding index '1' at position 0" + ): + string_index.embed(["a", "b"], model=test_model) + float_index, _ = _connection_returning( + [ + SimpleNamespace(index=1.0, embedding=[0.1]), + SimpleNamespace(index=0.0, embedding=[0.2]), + ] + ) + assert float_index.embed(["a", "b"], model=test_model) == [[0.2], [0.1]] + + malformed, _ = _connection_returning( + [SimpleNamespace(index=0, embedding="AACAPwAA")] + ) + with pytest.raises(RuntimeError, match="malformed embedding at position 0"): + malformed.embed("a", model=test_model, encoding_format="base64") + corrupted, _ = _connection_returning([SimpleNamespace(index=0, embedding="@@@@")]) + with pytest.raises(RuntimeError, match="malformed embedding at position 0"): + corrupted.embed("a", model=test_model, encoding_format="base64") + strings, _ = _connection_returning( + [SimpleNamespace(index=0, embedding=["0.1", "0.2"])] + ) + with pytest.raises(RuntimeError, match="malformed embedding at position 0"): + strings.embed("a", model=test_model) + # A string vector is only decoded when base64 was requested. + unexpected_b64, _ = _connection_returning( + [SimpleNamespace(index=0, embedding="AAAAAAAAAAA=")] + ) + with pytest.raises(RuntimeError, match="malformed embedding at position 0"): + unexpected_b64.embed("a", model=test_model) + null_vector, _ = _connection_returning([SimpleNamespace(index=0, embedding=None)]) + with pytest.raises(RuntimeError, match="malformed embedding at position 0"): + null_vector.embed("a", model=test_model) + no_data, _ = _connection_returning(None) + with pytest.raises(RuntimeError, match="0 embeddings for 1 input texts"): + no_data.embed("a", model=test_model) + null_item, _ = _connection_returning([None]) + with pytest.raises(RuntimeError, match="malformed embedding at position 0"): + null_item.embed("a", model=test_model) + + +def test_openai_embedding_empty_batch_skips_request() -> None: + connection, mock_client = _connection_returning([]) + result = connection.embed_with_usage([], model=test_model) + assert result.embeddings == [] + assert result.token_usage is None + mock_client.embeddings.create.assert_not_called() + + # A one-shot iterable is materialized before validation and sent as a list. + connection, mock_client = _connection_returning( + [ + SimpleNamespace(index=0, embedding=[0.1]), + SimpleNamespace(index=1, embedding=[0.2]), + ] + ) + result = connection.embed_with_usage((t for t in ["a", "b"]), model=test_model) + assert result.embeddings == [[0.1], [0.2]] + assert mock_client.embeddings.create.call_args.kwargs["input"] == ["a", "b"] + + +def test_openai_embedding_unusable_usage_yields_none() -> None: + connection = OpenAIEmbeddingModelConnection(name="openai", api_key="fake-key") + mock_client = MagicMock() + mock_client.embeddings.create.return_value = SimpleNamespace( + data=[SimpleNamespace(index=0, embedding=[0.1])], + usage=SimpleNamespace(prompt_tokens="n/a", total_tokens=None), + ) + connection._OpenAIEmbeddingModelConnection__client = mock_client + + result = connection.embed_with_usage("a", model=test_model) + + assert result.embeddings == [0.1] + assert result.token_usage is None + + # A huge integer count is passed through rather than overflowing a float. + mock_client.embeddings.create.return_value = SimpleNamespace( + data=[SimpleNamespace(index=0, embedding=[0.1])], + usage=SimpleNamespace(prompt_tokens=10**400, total_tokens=None), + ) + assert ( + connection.embed_with_usage("a", model=test_model).token_usage.prompt_tokens + == 10**400 + ) + + # Numeric strings, which some compatible servers send, still count. + mock_client.embeddings.create.return_value = SimpleNamespace( + data=[SimpleNamespace(index=0, embedding=[0.1])], + usage=SimpleNamespace(prompt_tokens="5", total_tokens="5"), + ) + result = connection.embed_with_usage("a", model=test_model) + assert result.token_usage is not None + assert (result.token_usage.prompt_tokens, result.token_usage.total_tokens) == (5, 5)
