This is an automated email from the ASF dual-hosted git repository.
dianfu pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git
The following commit(s) were added to refs/heads/master by this push:
new e0785f05542 [FLINK-40439][python] Add DataFrame model provider
management
e0785f05542 is described below
commit e0785f055420c6290c8a5c8ff7ad8b058ba83c8b
Author: auroflow <[email protected]>
AuthorDate: Fri Oct 9 16:02:59 2026 +0800
[FLINK-40439][python] Add DataFrame model provider management
This closes #29454.
---
.../model/triton/TritonModelProviderFactory.java | 6 +
.../triton/TritonModelProviderFactoryTest.java | 41 +-
.../docs/reference/pyflink.dataframe/index.rst | 1 +
.../docs/reference/pyflink.dataframe/llm.rst | 62 +++
flink-python/pyflink/dataframe/__init__.py | 16 +
flink-python/pyflink/dataframe/llm.py | 612 +++++++++++++++++++++
.../pyflink/dataframe/tests/test_model_provider.py | 481 ++++++++++++++++
flink-python/pyflink/dataframe/validation.py | 33 +-
8 files changed, 1243 insertions(+), 9 deletions(-)
diff --git
a/flink-models/flink-model-triton/src/main/java/org/apache/flink/model/triton/TritonModelProviderFactory.java
b/flink-models/flink-model-triton/src/main/java/org/apache/flink/model/triton/TritonModelProviderFactory.java
index 4fd6139a245..951c4c76153 100644
---
a/flink-models/flink-model-triton/src/main/java/org/apache/flink/model/triton/TritonModelProviderFactory.java
+++
b/flink-models/flink-model-triton/src/main/java/org/apache/flink/model/triton/TritonModelProviderFactory.java
@@ -75,6 +75,12 @@ public class TritonModelProviderFactory implements
ModelProviderFactory {
set.add(TritonOptions.RETRY_INITIAL_BACKOFF);
set.add(TritonOptions.RETRY_MAX_BACKOFF);
set.add(TritonOptions.DEFAULT_VALUE);
+ set.add(TritonOptions.HEALTH_CHECK_ENABLED);
+ set.add(TritonOptions.HEALTH_CHECK_INTERVAL);
+ set.add(TritonOptions.CIRCUIT_BREAKER_ENABLED);
+ set.add(TritonOptions.CIRCUIT_BREAKER_FAILURE_THRESHOLD);
+ set.add(TritonOptions.CIRCUIT_BREAKER_TIMEOUT);
+ set.add(TritonOptions.CIRCUIT_BREAKER_HALF_OPEN_REQUESTS);
return set;
}
diff --git
a/flink-models/flink-model-triton/src/test/java/org/apache/flink/model/triton/TritonModelProviderFactoryTest.java
b/flink-models/flink-model-triton/src/test/java/org/apache/flink/model/triton/TritonModelProviderFactoryTest.java
index 7fab9ce7446..b3bbcffc771 100644
---
a/flink-models/flink-model-triton/src/test/java/org/apache/flink/model/triton/TritonModelProviderFactoryTest.java
+++
b/flink-models/flink-model-triton/src/test/java/org/apache/flink/model/triton/TritonModelProviderFactoryTest.java
@@ -17,8 +17,16 @@
package org.apache.flink.model.triton;
+import org.apache.flink.table.api.DataTypes;
+import org.apache.flink.table.catalog.Column;
+import org.apache.flink.table.catalog.ResolvedSchema;
+import org.apache.flink.table.factories.utils.FactoryMocks;
+import org.apache.flink.table.ml.AsyncPredictRuntimeProvider;
+
import org.junit.jupiter.api.Test;
+import java.util.Map;
+
import static org.assertj.core.api.Assertions.assertThat;
/** Test for {@link TritonModelProviderFactory}. */
@@ -42,7 +50,7 @@ class TritonModelProviderFactoryTest {
void testOptionalOptions() {
TritonModelProviderFactory factory = new TritonModelProviderFactory();
assertThat(factory.optionalOptions())
- .hasSize(14)
+ .hasSize(20)
.containsExactlyInAnyOrder(
TritonOptions.MODEL_VERSION,
TritonOptions.TIMEOUT,
@@ -57,6 +65,35 @@ class TritonModelProviderFactoryTest {
TritonOptions.MAX_RETRIES,
TritonOptions.RETRY_INITIAL_BACKOFF,
TritonOptions.RETRY_MAX_BACKOFF,
- TritonOptions.DEFAULT_VALUE);
+ TritonOptions.DEFAULT_VALUE,
+ TritonOptions.HEALTH_CHECK_ENABLED,
+ TritonOptions.HEALTH_CHECK_INTERVAL,
+ TritonOptions.CIRCUIT_BREAKER_ENABLED,
+ TritonOptions.CIRCUIT_BREAKER_FAILURE_THRESHOLD,
+ TritonOptions.CIRCUIT_BREAKER_TIMEOUT,
+ TritonOptions.CIRCUIT_BREAKER_HALF_OPEN_REQUESTS);
+ }
+
+ @Test
+ void testCreateProviderWithHealthCheckAndCircuitBreaker() {
+ final TritonModelProviderFactory factory = new
TritonModelProviderFactory();
+ final ResolvedSchema schema =
+ ResolvedSchema.of(Column.physical("text", DataTypes.STRING()));
+ final Map<String, String> options =
+ Map.of(
+ "provider", "triton",
+ "endpoint", "http://localhost:8000";,
+ "model-name", "test-model",
+ "health-check-enabled", "true",
+ "health-check-interval", "10 s",
+ "circuit-breaker-enabled", "true",
+ "circuit-breaker-failure-threshold", "0.25",
+ "circuit-breaker-timeout", "5 s",
+ "circuit-breaker-half-open-requests", "2");
+
+ assertThat(
+ factory.createModelProvider(
+ FactoryMocks.createModelContext(schema,
schema, options)))
+ .isInstanceOf(AsyncPredictRuntimeProvider.class);
}
}
diff --git a/flink-python/docs/reference/pyflink.dataframe/index.rst
b/flink-python/docs/reference/pyflink.dataframe/index.rst
index d92c806a027..24059215551 100644
--- a/flink-python/docs/reference/pyflink.dataframe/index.rst
+++ b/flink-python/docs/reference/pyflink.dataframe/index.rst
@@ -34,3 +34,4 @@ This page gives an overview of all public PyFlink DataFrame
APIs.
datatype
environment
config
+ llm
diff --git a/flink-python/docs/reference/pyflink.dataframe/llm.rst
b/flink-python/docs/reference/pyflink.dataframe/llm.rst
new file mode 100644
index 00000000000..b51f9c45eff
--- /dev/null
+++ b/flink-python/docs/reference/pyflink.dataframe/llm.rst
@@ -0,0 +1,62 @@
+..
################################################################################
+ Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
+
################################################################################
+
+========
+AI / LLM
+========
+
+Model Providers
+---------------
+
+Model providers configure connections and options for AI functions.
+``OpenAIProvider`` and ``TritonProvider`` provide typed configurations for
Flink's
+built-in providers. See their API references below for parameters and defaults.
+Use ``options`` to supply Java provider option names and string values
directly.
+
+Use ``set_model_provider`` to register a configuration. Pass a provider alone
to
+use its identifier as the name, or supply a name to distinguish configurations.
+Use ``list_model_providers`` to inspect registered names:
+
+.. code-block:: python
+
+ import pyflink.dataframe as pf
+
+ provider = pf.TritonProvider(endpoint="http://localhost:8000";)
+ pf.set_model_provider("inference", provider)
+ pf.list_model_providers() # ["inference"]
+
+A single registered provider is selected automatically. When using several
+providers, choose a default with ``set_default_model_provider("inference")``.
+Registrations are shared within the Python process.
+
+For other providers available in the application's classpath, use
``GenericProvider``
+with provider-specific option names and string values, or subclass
``ModelProvider``
+to create a reusable configuration.
+
+.. currentmodule:: pyflink.dataframe
+
+.. autosummary::
+ :toctree: api/
+
+ ModelProvider
+ OpenAIProvider
+ TritonProvider
+ GenericProvider
+ set_model_provider
+ set_default_model_provider
+ list_model_providers
diff --git a/flink-python/pyflink/dataframe/__init__.py
b/flink-python/pyflink/dataframe/__init__.py
index 174a3e4df79..5c0b3af3941 100644
--- a/flink-python/pyflink/dataframe/__init__.py
+++ b/flink-python/pyflink/dataframe/__init__.py
@@ -73,6 +73,15 @@ from pyflink.dataframe.io import (
read_parquet,
)
from pyflink.dataframe.iteration import CloseableIterator
+from pyflink.dataframe.llm import (
+ GenericProvider,
+ ModelProvider,
+ OpenAIProvider,
+ TritonProvider,
+ list_model_providers,
+ set_default_model_provider,
+ set_model_provider,
+)
from pyflink.dataframe.over_window import following, preceding
from pyflink.dataframe.sql import sql
from pyflink.dataframe.udf import udf
@@ -83,6 +92,13 @@ __all__ = [
"DataFrame",
"GroupedDataFrame",
"DataType",
+ "ModelProvider",
+ "OpenAIProvider",
+ "TritonProvider",
+ "GenericProvider",
+ "set_model_provider",
+ "set_default_model_provider",
+ "list_model_providers",
"col",
"lit",
"udf",
diff --git a/flink-python/pyflink/dataframe/llm.py
b/flink-python/pyflink/dataframe/llm.py
new file mode 100644
index 00000000000..122707dafb4
--- /dev/null
+++ b/flink-python/pyflink/dataframe/llm.py
@@ -0,0 +1,612 @@
+################################################################################
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+################################################################################
+
+from abc import ABC, abstractmethod
+from collections.abc import Mapping as MappingABC
+import json
+from typing import Dict, List, Literal, Mapping, Optional, Tuple, overload
+
+from pyflink.dataframe.validation import _require_int, _require_number
+from pyflink.util.api_stability_decorators import PublicEvolving
+
+__all__ = [
+ "ModelProvider", "OpenAIProvider", "TritonProvider", "GenericProvider",
+ "set_model_provider", "set_default_model_provider", "list_model_providers",
+]
+
+
+def _validate_name(name: object, argument: str) -> str:
+ if not isinstance(name, str):
+ raise TypeError(f"{argument} must be a string")
+ if not name.strip():
+ raise ValueError(f"{argument} must not be empty or whitespace")
+ return name
+
+
+def _stringify_options(options: Dict[str, object]) -> Dict[str, str]:
+ return {
+ name: str(value).lower() if isinstance(value, bool) else str(value)
+ for name, value in options.items() if value is not None
+ }
+
+
+def _merge_options(
+ options: Optional[Mapping[str, str]], parameters: Mapping[str, str],
+) -> Dict[str, str]:
+ if options is not None and not isinstance(options, MappingABC):
+ raise TypeError("options must be a mapping of strings to strings")
+ merged: Dict[str, str] = {}
+ for source in (options, parameters):
+ if source is None:
+ continue
+ for key, value in source.items():
+ if not isinstance(key, str):
+ raise TypeError("Option keys must be strings")
+ if not isinstance(value, str):
+ raise TypeError(f"Option {key!r} must have a string value")
+ if key in merged and merged[key] != value:
+ raise ValueError(f"conflicting values for option {key!r}")
+ merged[key] = value
+ return merged
+
+
+def _validate_type(name: str, value: object, expected: type) -> None:
+ if value is not None and not isinstance(value, expected):
+ raise TypeError(f"{name} must be a {expected.__name__}")
+
+
+def _enum_option(name: str, value: Optional[str], choices: Tuple[str, ...]) ->
Optional[str]:
+ _validate_type(name, value, str)
+ if value is None:
+ return None
+ for choice in choices:
+ if value.lower() == choice.lower():
+ return choice
+ raise ValueError(f"{name} must be one of {choices}")
+
+
+def _header_option(headers: Optional[Mapping[str, str]]) -> Optional[str]:
+ if headers is None:
+ return None
+ if not isinstance(headers, MappingABC):
+ raise TypeError("custom_headers must be a mapping of strings to
strings")
+ if any(not isinstance(key, str) or not isinstance(value, str)
+ for key, value in headers.items()):
+ raise TypeError("custom_headers must contain string keys and values")
+ return json.dumps(dict(headers))
+
+
+@PublicEvolving()
+class ModelProvider(ABC):
+ """
+ Base class for model provider configurations.
+
+ Subclass this to configure a provider available in your application's
classpath.
+ Implement :meth:`provider_identifier` to identify the provider and
:meth:`to_options`
+ to supply its options as string keys and values.
+
+ Example::
+
+ >>> class MyProvider(ModelProvider):
+ ... def provider_identifier(self):
+ ... return "my-provider"
+ ... def to_options(self):
+ ... return {"endpoint": "https://example.test/inference"}
+ >>> MyProvider().provider_identifier()
+ 'my-provider'
+
+ .. versionadded:: 2.4.0
+ """
+
+ @abstractmethod
+ def provider_identifier(self) -> str:
+ """Return the provider identifier, such as ``openai`` or ``triton``.
+
+ .. versionadded:: 2.4.0
+ """
+
+ @abstractmethod
+ def to_options(self) -> Dict[str, str]:
+ """Return provider options as string keys and values.
+
+ .. versionadded:: 2.4.0
+ """
+
+ def model_option_key(self) -> str:
+ """Return the option key for a model name, normally ``model``.
+
+ Override this for providers whose model option has a different name.
+
+ .. versionadded:: 2.4.0
+ """
+ return "model"
+
+
+@PublicEvolving()
+class OpenAIProvider(ModelProvider):
+ """
+ Configure an OpenAI-compatible service for chat completions or embeddings.
+
+ The full endpoint URL selects the task: use a URL ending in
``/chat/completions``
+ for chat or ``/embeddings`` for embedding vectors. Chat options such as
+ ``system_prompt`` and ``temperature`` apply to chat requests;
``dimension`` applies
+ to embeddings. Supply a model name here or when calling an AI function.
+
+ All optional parameters default to ``None``. Omitted options use the Flink
defaults
+ described below, or the service's defaults where Flink defines none.
+
+ :param endpoint: Required full API URL, for example
+ ``https://api.openai.com/v1/chat/completions``.
+ :param api_key: Required API key for authenticating requests.
+ :param model: Model name to use. Default: ``None``; if omitted, supply it
when
+ calling an AI function.
+ :param system_prompt: System message that guides the chat response. If
omitted,
+ Flink uses ``"You are a helpful assistant."``. Set
``""``
+ for an empty system message.
+ :param temperature: Controls randomness in chat responses; typical values
are
+ between 0.0 and 1.0. Default: ``None`` (service
default).
+ :param top_p: Probability cutoff for token selection, between 0 and 1.
Usually
+ set either this or ``temperature``. Default: ``None``
(service default).
+ :param max_tokens: Maximum number of tokens generated in a chat completion.
+ Default: ``None`` (service default).
+ :param stop: Comma-separated strings that stop generation when encountered.
+ Default: ``None`` (no stop sequences supplied).
+ :param presence_penalty: Value between -2 and 2. Positive values
discourage tokens
+ already used in the response, encouraging new
topics.
+ Default: ``None`` (service default).
+ :param n: Number of chat completion choices generated per input. Setting
this to
+ 1 limits token usage. Default: ``None`` (service default).
+ :param seed: Seed for best-effort repeatable sampling. Default: ``None``
(no seed
+ supplied).
+ :param response_format: Chat response format: ``text`` or ``json_object``.
+ Default: ``None`` (service default).
+ :param dimension: Number of elements in each embedding vector. Default:
``None``
+ (the embedding model's default dimension).
+ :param max_context_size: Maximum number of input context tokens before
applying
+ ``context_overflow_action``. Default: ``None``
(no Flink
+ context limit).
+ :param context_overflow_action: ``truncated-tail`` removes excess tokens
from the
+ end, ``truncated-head`` removes them from
the start,
+ and ``skipped`` skips the input. Each has a
``-log``
+ variant that logs the action. If omitted,
Flink uses
+ ``truncated-tail``.
+ :param error_handling_strategy: ``RETRY`` retries failed requests,
``FAILOVER``
+ fails the job, and ``IGNORE`` skips the
failed input.
+ If omitted, Flink uses ``RETRY``.
+ :param retry_num: Number of retries when ``error_handling_strategy`` is
``RETRY``.
+ If omitted, Flink uses ``100``.
+ :param retry_fallback_strategy: Action after retries are exhausted:
``FAILOVER``
+ fails the job and ``IGNORE`` skips the
failed input.
+ If omitted, Flink uses ``FAILOVER``.
+ :param options: Additional Java provider options as string keys and values.
+
+ Example::
+
+ >>> import pyflink.dataframe as pf
+ >>> chat = pf.OpenAIProvider(
+ ... endpoint="https://api.openai.com/v1/chat/completions";,
api_key="key")
+ >>> pf.set_model_provider("chat", chat)
+ >>> embeddings = pf.OpenAIProvider(
+ ... endpoint="https://api.openai.com/v1/embeddings";, api_key="key")
+ >>> pf.set_model_provider("embed", embeddings)
+
+ .. versionadded:: 2.4.0
+ """
+
+ def __init__(
+ self, endpoint: str, api_key: str, *, model: Optional[str] = None,
+ system_prompt: Optional[str] = None,
+ temperature: Optional[float] = None,
+ top_p: Optional[float] = None,
+ max_tokens: Optional[int] = None,
+ stop: Optional[str] = None,
+ presence_penalty: Optional[float] = None,
+ n: Optional[int] = None,
+ seed: Optional[int] = None,
+ response_format: Optional[Literal["text", "json_object"]] = None,
+ dimension: Optional[int] = None,
+ max_context_size: Optional[int] = None,
+ context_overflow_action: Optional[Literal[
+ "truncated-tail", "truncated-tail-log", "truncated-head",
"truncated-head-log",
+ "skipped", "skipped-log",
+ ]] = None,
+ error_handling_strategy: Optional[Literal["RETRY", "FAILOVER",
"IGNORE"]] = None,
+ retry_num: Optional[int] = None,
+ retry_fallback_strategy: Optional[Literal["FAILOVER", "IGNORE"]] =
None,
+ options: Optional[Mapping[str, str]] = None,
+ ):
+ _validate_name(endpoint, "endpoint")
+ _validate_name(api_key, "api_key")
+ typed_options: Dict[str, object] = {
+ "endpoint": endpoint, "api-key": api_key, "model": model,
+ "system-prompt": system_prompt, "temperature": temperature,
"top-p": top_p,
+ "max-tokens": max_tokens, "stop": stop, "presence-penalty":
presence_penalty,
+ "n": n, "seed": seed, "dimension": dimension, "max-context-size":
max_context_size,
+ "retry-num": retry_num,
+ "response-format": _enum_option("response_format", response_format,
+ ("text", "json_object")),
+ "context-overflow-action": _enum_option(
+ "context_overflow_action", context_overflow_action,
+ ("truncated-tail", "truncated-tail-log", "truncated-head",
"truncated-head-log",
+ "skipped", "skipped-log")),
+ "error-handling-strategy": _enum_option(
+ "error_handling_strategy", error_handling_strategy,
+ ("RETRY", "FAILOVER", "IGNORE")),
+ "retry-fallback-strategy": _enum_option(
+ "retry_fallback_strategy", retry_fallback_strategy,
("FAILOVER", "IGNORE")),
+ }
+ for name in ("model", "system-prompt", "stop"):
+ _validate_type(name, typed_options[name], str)
+ for name, integer, minimum in (
+ ("max_tokens", max_tokens, 1), ("n", n, 1), ("dimension",
dimension, 1),
+ ("max_context_size", max_context_size, 1), ("retry_num",
retry_num, 0),
+ ("seed", seed, None),
+ ):
+ if integer is not None:
+ _require_int(integer, name, minimum)
+ for name, number, lower_bound, upper_bound in (
+ ("temperature", temperature, None, None), ("top_p", top_p, 0, 1),
+ ("presence_penalty", presence_penalty, -2, 2),
+ ):
+ if number is not None:
+ _require_number(number, name, lower_bound, maximum=upper_bound)
+ self._options = _merge_options(options,
_stringify_options(typed_options))
+
+ def provider_identifier(self) -> str:
+ """Return ``openai``.
+
+ .. versionadded:: 2.4.0
+ """
+ return "openai"
+
+ def to_options(self) -> Dict[str, str]:
+ """Return a fresh dictionary of provider options.
+
+ .. versionadded:: 2.4.0
+ """
+ return dict(self._options)
+
+
+@PublicEvolving()
+class TritonProvider(ModelProvider):
+ """
+ Configure inference requests to a model served by NVIDIA Triton Inference
Server.
+
+ Supply a model name here or when calling an AI function. For array inputs,
use
+ ``flatten_batch_dim`` to match the model's expected shape. Retry and
+ fallback options control how failed requests are handled; health checks
and the
+ circuit breaker can reduce requests to an unavailable server.
+
+ All optional parameters default to ``None``. Omitted options use the Flink
defaults
+ described below. Specify durations as strings such as ``"30 s"`` or ``"100
ms"``.
+
+ :param endpoint: Required Triton server URL, for example
``http://localhost:8000``.
+ :param model_name: Name of the model to invoke. Default: ``None``; if
omitted,
+ supply it when calling an AI function.
+ :param model_version: Model version to invoke. If omitted, Flink uses
``"latest"``.
+ :param timeout: HTTP timeout for each request, separate from Flink's
asynchronous
+ prediction timeout. If omitted, Flink uses ``"30 s"``.
+ :param flatten_batch_dim: Convert the array input shape from ``[1, N]`` to
``[N]``
+ when the model expects a vector without a batch
dimension.
+ If omitted, Flink uses ``False``.
+ :param priority: Request priority between 0 and 255. Default: ``None`` (no
priority
+ supplied).
+ :param sequence_id: Identifier shared by requests in a stateful model
sequence.
+ Default: ``None`` (no sequence identifier supplied).
+ :param sequence_start: Mark requests as starting a stateful sequence. If
omitted,
+ Flink uses ``False``.
+ :param sequence_end: Mark requests as ending a stateful sequence. If
omitted,
+ Flink uses ``False``.
+ :param compression: Compress request bodies using ``gzip``. Default:
``None``
+ (no compression).
+ :param auth_token: Authentication token sent as a Bearer token. Default:
``None``
+ (no token supplied).
+ :param custom_headers: Additional HTTP headers as a mapping, for example
+ ``{"X-Trace-Id": "abc"}``. Default: ``None`` (no
extra headers).
+ :param max_retries: Additional attempts for transient failures, such as
network
+ errors and server errors. If omitted, Flink uses ``0``
(no retries).
+ :param retry_initial_backoff: Initial delay between retries; delays
increase
+ exponentially up to ``retry_max_backoff``.
If omitted,
+ Flink uses ``"100 ms"``.
+ :param retry_max_backoff: Maximum delay between retries. If omitted, Flink
uses
+ ``"30 s"``.
+ :param default_value: Fallback value when inference fails, expressed as a
string
+ matching the output type: plain text for strings, a
numeric
+ string for numbers, a JSON array for arrays, or
``"null"`` for
+ SQL NULL. Default: ``None`` (propagate failures).
+ :param health_check_enabled: Enable periodic server health checks. If
omitted,
+ Flink uses ``False``.
+ :param health_check_interval: Time between health checks when enabled. If
omitted,
+ Flink uses ``"30 s"``.
+ :param circuit_breaker_enabled: Temporarily stop sending requests when the
server
+ has a high failure rate. If omitted, Flink
uses ``False``.
+ :param circuit_breaker_failure_threshold: Failure rate in ``(0, 1]`` that
opens the
+ circuit breaker; ``0.5`` means
50% failures.
+ If omitted, Flink uses ``0.5``.
+ :param circuit_breaker_timeout: Time to wait before probing recovery after
the
+ circuit breaker opens. If omitted, Flink
uses ``"60 s"``.
+ :param circuit_breaker_half_open_requests: Successful recovery probes
needed to
+ close the circuit breaker. If
omitted,
+ Flink uses ``3``.
+ :param options: Additional Java provider options as string keys and values.
+
+ Example::
+
+ >>> import pyflink.dataframe as pf
+ >>> provider = pf.TritonProvider(
+ ... endpoint="http://localhost:8000";, model_name="classifier",
+ ... flatten_batch_dim=True, max_retries=2, default_value="-1")
+ >>> pf.set_model_provider("classifier", provider)
+
+ .. versionadded:: 2.4.0
+ """
+
+ def __init__(
+ self, endpoint: str, *, model_name: Optional[str] = None,
+ model_version: Optional[str] = None,
+ timeout: Optional[str] = None,
+ flatten_batch_dim: Optional[bool] = None,
+ priority: Optional[int] = None,
+ sequence_id: Optional[str] = None,
+ sequence_start: Optional[bool] = None,
+ sequence_end: Optional[bool] = None,
+ compression: Optional[Literal["gzip"]] = None,
+ auth_token: Optional[str] = None,
+ custom_headers: Optional[Mapping[str, str]] = None,
+ max_retries: Optional[int] = None,
+ retry_initial_backoff: Optional[str] = None,
+ retry_max_backoff: Optional[str] = None,
+ default_value: Optional[str] = None,
+ health_check_enabled: Optional[bool] = None,
+ health_check_interval: Optional[str] = None,
+ circuit_breaker_enabled: Optional[bool] = None,
+ circuit_breaker_failure_threshold: Optional[float] = None,
+ circuit_breaker_timeout: Optional[str] = None,
+ circuit_breaker_half_open_requests: Optional[int] = None,
+ options: Optional[Mapping[str, str]] = None,
+ ):
+ _validate_name(endpoint, "endpoint")
+ typed_options: Dict[str, object] = {
+ "endpoint": endpoint, "model-name": model_name, "model-version":
model_version,
+ "timeout": timeout, "flatten-batch-dim": flatten_batch_dim,
"priority": priority,
+ "sequence-id": sequence_id, "sequence-start": sequence_start,
+ "sequence-end": sequence_end,
+ "compression": _enum_option("compression", compression, ("gzip",)),
+ "auth-token": auth_token, "custom-headers":
_header_option(custom_headers),
+ "max-retries": max_retries, "retry-initial-backoff":
retry_initial_backoff,
+ "retry-max-backoff": retry_max_backoff, "default-value":
default_value,
+ "health-check-enabled": health_check_enabled,
+ "health-check-interval": health_check_interval,
+ "circuit-breaker-enabled": circuit_breaker_enabled,
+ "circuit-breaker-failure-threshold":
circuit_breaker_failure_threshold,
+ "circuit-breaker-timeout": circuit_breaker_timeout,
+ "circuit-breaker-half-open-requests":
circuit_breaker_half_open_requests,
+ }
+ for name in ("model-name", "model-version", "timeout", "sequence-id",
"auth-token",
+ "retry-initial-backoff", "retry-max-backoff",
"default-value",
+ "health-check-interval", "circuit-breaker-timeout"):
+ _validate_type(name, typed_options[name], str)
+ for name in ("flatten-batch-dim", "sequence-start", "sequence-end",
"health-check-enabled",
+ "circuit-breaker-enabled"):
+ _validate_type(name, typed_options[name], bool)
+ for name, integer, minimum, maximum in (
+ ("priority", priority, 0, 255), ("max_retries", max_retries, 0,
None),
+ ("circuit_breaker_half_open_requests",
circuit_breaker_half_open_requests, 1, None),
+ ):
+ if integer is not None:
+ _require_int(integer, name, minimum, maximum=maximum)
+ if circuit_breaker_failure_threshold is not None:
+ _require_number(circuit_breaker_failure_threshold,
"circuit_breaker_failure_threshold",
+ 0, maximum=1, include_minimum=False)
+ self._options = _merge_options(options,
_stringify_options(typed_options))
+
+ def provider_identifier(self) -> str:
+ """Return ``triton``.
+
+ .. versionadded:: 2.4.0
+ """
+ return "triton"
+
+ def model_option_key(self) -> str:
+ """Return ``model-name``, the option key for a model name.
+
+ .. versionadded:: 2.4.0
+ """
+ return "model-name"
+
+ def to_options(self) -> Dict[str, str]:
+ """Return a fresh dictionary of provider options.
+
+ .. versionadded:: 2.4.0
+ """
+ return dict(self._options)
+
+
+@PublicEvolving()
+class GenericProvider(ModelProvider):
+ """
+ Configure an installed model provider with its option names and string
values.
+
+ Use this for a provider without a dedicated Python configuration class.
Option
+ names and values follow that provider's documentation. Supply options as a
+ mapping, keyword arguments, or both. Conflicting values raise
``ValueError``.
+
+ :param identifier: Java model provider factory identifier.
+ :param options: Provider option names and string values, supplied as the
second
+ positional argument. Default: ``None``.
+ :param kwargs: Provider options supplied as keyword arguments.
+
+ Example::
+
+ >>> provider = GenericProvider(
+ ... "my-provider",
+ ... {"endpoint": "https://example.test/inference";, "api-key":
"key"})
+ >>> provider.to_options()
+ {'endpoint': 'https://example.test/inference', 'api-key': 'key'}
+
+ .. versionadded:: 2.4.0
+ """
+
+ def __init__(
+ self, identifier: str, options: Optional[Mapping[str, str]] = None, /,
**kwargs: str,
+ ):
+ _validate_name(identifier, "identifier")
+ self._identifier = identifier
+ self._options = _merge_options(options, kwargs)
+
+ def provider_identifier(self) -> str:
+ """Return the configured provider identifier.
+
+ .. versionadded:: 2.4.0
+ """
+ return self._identifier
+
+ def to_options(self) -> Dict[str, str]:
+ """Return a fresh dictionary of provider options.
+
+ .. versionadded:: 2.4.0
+ """
+ return dict(self._options)
+
+
+_provider_registry: Dict[str, ModelProvider] = {}
+_default_provider: Optional[str] = None
+_UNSET = object()
+
+
+@overload
+def set_model_provider(provider: ModelProvider, /) -> None:
+ ...
+
+
+@overload
+def set_model_provider(*, provider: ModelProvider) -> None:
+ ...
+
+
+@overload
+def set_model_provider(name: str, provider: ModelProvider) -> None:
+ ...
+
+
+@PublicEvolving()
+def set_model_provider(
+ *args: object, name: object = _UNSET, provider: object = _UNSET,
+) -> None:
+ """
+ Register a model provider configuration.
+
+ Pass a provider alone to register it under its identifier, such as
``openai``.
+ Supply a name to keep several configurations of the same provider.
+
+ Registering an existing name updates its configuration. Registered
providers are
+ shared across DataFrame environments in the current Python process.
+
+ :param name: Name used to look up the provider. Omit it to use the
provider identifier.
+ :param provider: Provider configuration to register.
+
+ Example::
+
+ >>> import pyflink.dataframe as pf
+ >>> provider = pf.OpenAIProvider(
+ ... "https://api.openai.com/v1/chat/completions";, api_key="key")
+ >>> pf.set_model_provider(provider)
+ >>> pf.set_model_provider("chat", provider)
+
+ .. versionadded:: 2.4.0
+ """
+ if len(args) > 2:
+ raise TypeError("set_model_provider() takes at most 2 positional
arguments")
+ if args:
+ if name is not _UNSET:
+ raise TypeError("set_model_provider() got multiple values for
argument 'name'")
+ if len(args) == 2:
+ if provider is not _UNSET:
+ raise TypeError("set_model_provider() got multiple values for
argument 'provider'")
+ name, provider = args
+ elif provider is _UNSET:
+ provider = args[0]
+ else:
+ name = args[0]
+ if not isinstance(provider, ModelProvider):
+ raise TypeError("provider must be a ModelProvider")
+ if name is _UNSET:
+ name = provider.provider_identifier()
+ name = _validate_name(name, "name")
+ _provider_registry[name] = provider
+
+
+@PublicEvolving()
+def list_model_providers() -> List[str]:
+ """
+ Return registered provider names in registration order.
+
+ Example::
+
+ >>> import pyflink.dataframe as pf
+ >>> pf.set_model_provider("custom", pf.GenericProvider("my-provider"))
+ >>> pf.list_model_providers()
+ ['custom']
+
+ .. versionadded:: 2.4.0
+ """
+ return list(_provider_registry)
+
+
+@PublicEvolving()
+def set_default_model_provider(name: str) -> None:
+ """
+ Select the registered provider to use by default.
+
+ A single registered provider is selected automatically. If you register
several
+ providers, select a default or choose a provider for each operation.
+
+ :param name: An already registered provider lookup name.
+ :raises ValueError: If the name is not registered.
+
+ Example::
+
+ >>> import pyflink.dataframe as pf
+ >>> pf.set_model_provider("custom", pf.GenericProvider("my-provider"))
+ >>> pf.set_default_model_provider("custom")
+
+ .. versionadded:: 2.4.0
+ """
+ global _default_provider
+ _validate_name(name, "name")
+ if name not in _provider_registry:
+ raise ValueError(f"Model provider {name!r} is not registered")
+ _default_provider = name
+
+
+def _resolve_provider(name: Optional[str] = None) -> ModelProvider:
+ if name is not None:
+ _validate_name(name, "name")
+ if name not in _provider_registry:
+ raise ValueError(f"Model provider {name!r} is not registered")
+ return _provider_registry[name]
+ if _default_provider is not None:
+ return _provider_registry[_default_provider]
+ if len(_provider_registry) == 1:
+ return next(iter(_provider_registry.values()))
+ if not _provider_registry:
+ raise ValueError("No model provider is registered")
+ raise ValueError(
+ "Multiple model providers are registered; call
set_default_model_provider() "
+ "or select a provider explicitly"
+ )
diff --git a/flink-python/pyflink/dataframe/tests/test_model_provider.py
b/flink-python/pyflink/dataframe/tests/test_model_provider.py
new file mode 100644
index 00000000000..71dd1d62a98
--- /dev/null
+++ b/flink-python/pyflink/dataframe/tests/test_model_provider.py
@@ -0,0 +1,481 @@
+################################################################################
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+################################################################################
+
+import json
+import unittest
+from types import MappingProxyType
+from typing import get_type_hints
+
+import pyflink.dataframe as pf
+from pyflink.dataframe import llm
+from pyflink.table import TableEnvironment
+
+
+class ModelProviderTests(unittest.TestCase):
+ def test_triton_options_follow_community_option_names(self):
+
self.assertEqual(pf.TritonProvider("http://localhost:8000";).to_options(),
+ {"endpoint": "http://localhost:8000"})
+ headers = {"X-Request": "comma, colon: quote'", "X-Empty": ""}
+ provider = pf.TritonProvider(
+ "http://localhost:8000";, model_name="model", model_version="2",
timeout="10 s",
+ flatten_batch_dim=False, priority=0,
sequence_id="18446744073709551615",
+ sequence_start=False, sequence_end=True, compression="gzip",
auth_token="token",
+ custom_headers=headers, max_retries=0, retry_initial_backoff="100
ms",
+ retry_max_backoff="1 s", default_value="null",
health_check_enabled=True,
+ health_check_interval="5 s", circuit_breaker_enabled=False,
+ circuit_breaker_failure_threshold=0.5, circuit_breaker_timeout="30
s",
+ circuit_breaker_half_open_requests=2,
+ )
+ self.assertEqual(provider.provider_identifier(), "triton")
+ self.assertEqual(provider.model_option_key(), "model-name")
+ options = provider.to_options()
+ self.assertEqual(json.loads(options.pop("custom-headers")), headers)
+ expected = {
+ "endpoint": "http://localhost:8000";, "model-name": "model",
"model-version": "2",
+ "timeout": "10 s", "flatten-batch-dim": "false", "priority": "0",
+ "sequence-id": "18446744073709551615", "sequence-start": "false",
+ "sequence-end": "true", "compression": "gzip", "auth-token":
"token",
+ "max-retries": "0", "retry-initial-backoff": "100 ms",
"retry-max-backoff": "1 s",
+ "default-value": "null", "health-check-enabled": "true",
"health-check-interval": "5 s",
+ "circuit-breaker-enabled": "false",
"circuit-breaker-failure-threshold": "0.5",
+ "circuit-breaker-timeout": "30 s",
"circuit-breaker-half-open-requests": "2",
+ }
+ self.assertEqual(options, expected)
+ headers["X-Request"] = "changed"
+ self.assertEqual(json.loads(provider.to_options()["custom-headers"]),
+ {"X-Request": "comma, colon: quote'", "X-Empty": ""})
+ provider.to_options().clear()
+ self.assertEqual(set(provider.to_options()), set(expected) |
{"custom-headers"})
+
+ def test_triton_rejects_invalid_typed_options(self):
+ for options, error in [
+ ({"model_name": 1}, TypeError),
+ ({"timeout": 10}, TypeError),
+ ({"flatten_batch_dim": "false"}, TypeError),
+ ({"sequence_start": 1}, TypeError),
+ ({"priority": True}, TypeError),
+ ({"priority": 0.5}, TypeError),
+ ({"priority": -1}, ValueError),
+ ({"priority": 256}, ValueError),
+ ({"max_retries": -1}, ValueError),
+ ({"compression": "deflate"}, ValueError),
+ ({"default_value": [0.0]}, TypeError),
+ ({"custom_headers": [("X-Header", "value")]}, TypeError),
+ ({"custom_headers": {1: "value"}}, TypeError),
+ ({"custom_headers": {"X-Header": 1}}, TypeError),
+ ({"health_check_enabled": 1}, TypeError),
+ ({"circuit_breaker_enabled": "false"}, TypeError),
+ ({"circuit_breaker_failure_threshold": 0.0}, ValueError),
+ ({"circuit_breaker_failure_threshold": 1.1}, ValueError),
+ ({"circuit_breaker_failure_threshold": True}, TypeError),
+ ({"circuit_breaker_failure_threshold": float("inf")}, ValueError),
+ ({"circuit_breaker_half_open_requests": 0}, ValueError),
+ ({"retry_initial_backoff": 100}, TypeError),
+ ({"extra_body": "{}"}, TypeError),
+ ]:
+ with self.subTest(options=options):
+ with self.assertRaises(error):
+ pf.TritonProvider("http://localhost:8000";, **options)
+
+ def test_openai_connection_options_allow_deferred_model_selection(self):
+ endpoint = "https://example.test/v1/embeddings";
+ provider = pf.OpenAIProvider(endpoint=endpoint, api_key="key")
+ self.assertEqual(provider.provider_identifier(), "openai")
+ self.assertEqual(provider.model_option_key(), "model")
+ self.assertEqual(provider.to_options(), {"endpoint": endpoint,
"api-key": "key"})
+ self.assertEqual(
+ pf.OpenAIProvider(endpoint, "key",
model="embed-model").to_options(),
+ {"endpoint": endpoint, "api-key": "key", "model": "embed-model"},
+ )
+
+ def test_openai_options_follow_community_option_names(self):
+ provider = pf.OpenAIProvider(
+ "https://example.test/v1/chat/completions";, "key",
model="chat-model",
+ system_prompt="", temperature=0.0, top_p=1.0, max_tokens=100,
+ stop="END,STOP", presence_penalty=-1.0, n=2, seed=0,
+ response_format="json_object", dimension=3, max_context_size=1000,
+ context_overflow_action="truncated-head-log",
error_handling_strategy="IGNORE",
+ retry_num=0, retry_fallback_strategy="IGNORE",
+ )
+ expected = {
+ "endpoint": "https://example.test/v1/chat/completions";, "api-key":
"key",
+ "model": "chat-model", "system-prompt": "", "temperature": "0.0",
+ "top-p": "1.0", "max-tokens": "100", "stop": "END,STOP",
+ "presence-penalty": "-1.0", "n": "2", "seed": "0",
+ "response-format": "json_object", "dimension": "3",
"max-context-size": "1000",
+ "context-overflow-action": "truncated-head-log",
"error-handling-strategy": "IGNORE",
+ "retry-num": "0", "retry-fallback-strategy": "IGNORE",
+ }
+ self.assertEqual(provider.to_options(), expected)
+ provider.to_options().clear()
+ self.assertEqual(provider.to_options(), expected)
+
+ def test_openai_rejects_invalid_typed_options(self):
+ for overrides, error in [
+ ({"endpoint": None}, TypeError),
+ ({"api_key": ""}, ValueError),
+ ({"model": 1}, TypeError),
+ ({"system_prompt": 1}, TypeError),
+ ({"max_tokens": True}, TypeError),
+ ({"dimension": "3"}, TypeError),
+ ({"temperature": True}, TypeError),
+ ({"temperature": float("nan")}, ValueError),
+ ({"temperature": float("-inf")}, ValueError),
+ ({"seed": True}, TypeError),
+ ({"top_p": -0.1}, ValueError),
+ ({"top_p": 1.1}, ValueError),
+ ({"presence_penalty": -3.0}, ValueError),
+ ({"presence_penalty": 3.0}, ValueError),
+ ({"retry_num": -1}, ValueError),
+ ({"n": 0}, ValueError),
+ ({"max_context_size": 0}, ValueError),
+ ({"stop": ["STOP"]}, TypeError),
+ ({"response_format": "json_schema"}, ValueError),
+ ({"context_overflow_action": "unknown"}, ValueError),
+ ({"error_handling_strategy": "unknown"}, ValueError),
+ ({"retry_fallback_strategy": "RETRY"}, ValueError),
+ ({"task": "chat/completions"}, TypeError),
+ ]:
+ with self.subTest(overrides=overrides):
+ arguments = {"endpoint":
"https://example.test/chat/completions";, "api_key": "key"}
+ arguments.update(overrides)
+ with self.assertRaises(error):
+ pf.OpenAIProvider(**arguments)
+
+ def test_numeric_options_accept_boundaries_and_unbounded_values(self):
+ for top_p, presence_penalty in [(0, -2), (1, 2)]:
+ with self.subTest(top_p=top_p, presence_penalty=presence_penalty):
+ provider = pf.OpenAIProvider(
+ "https://example.test/chat/completions";, "key",
+ temperature=-0.5, seed=-100, top_p=top_p,
presence_penalty=presence_penalty,
+ max_tokens=1, n=1, dimension=1, max_context_size=1,
retry_num=0,
+ )
+ self.assertEqual(provider.to_options(), {
+ "endpoint": "https://example.test/chat/completions";,
"api-key": "key",
+ "temperature": "-0.5", "seed": "-100", "top-p": str(top_p),
+ "presence-penalty": str(presence_penalty), "max-tokens":
"1", "n": "1",
+ "dimension": "1", "max-context-size": "1", "retry-num":
"0",
+ })
+ provider = pf.TritonProvider(
+ "http://localhost:8000";, priority=255, max_retries=0,
+ circuit_breaker_failure_threshold=1,
circuit_breaker_half_open_requests=1,
+ )
+ self.assertEqual(provider.to_options(), {
+ "endpoint": "http://localhost:8000";, "priority": "255",
"max-retries": "0",
+ "circuit-breaker-failure-threshold": "1",
"circuit-breaker-half-open-requests": "1",
+ })
+
+ def test_generic_provider_preserves_raw_options(self):
+ expected = {"endpoint": "HTTPS://example.test/v1", "api-key": "key",
+ "custom_option": "", "enabled": "false", "identifier":
"raw",
+ "self": "raw", "options": "raw"}
+ for form in ("keywords", "mapping", "mixed", "none"):
+ with self.subTest(form=form):
+ options = dict(expected)
+ if form == "keywords":
+ provider = pf.GenericProvider("custom", **options)
+ elif form == "mapping":
+ provider = pf.GenericProvider("custom",
MappingProxyType(options))
+ elif form == "mixed":
+ provider = pf.GenericProvider(
+ "custom", {"endpoint": options.pop("endpoint")},
**options)
+ else:
+ provider = pf.GenericProvider("custom", None, **options)
+ options.clear()
+ self.assertIsInstance(provider, pf.ModelProvider)
+ self.assertEqual(provider.provider_identifier(), "custom")
+ self.assertEqual(provider.model_option_key(), "model")
+ self.assertEqual(provider.to_options(), expected)
+ provider.to_options().clear()
+ self.assertEqual(provider.to_options(), expected)
+
+ def test_typed_providers_accept_raw_options(self):
+ for provider_type, connection, raw, typed, expected in [
+ (pf.OpenAIProvider,
+ {"endpoint": "https://example.test/chat/completions";, "api_key":
"key"},
+ {"model": "chat-model", "max-tokens": "100", "temperature": "0.0",
+ "system-prompt": "", "future-option": "Mixed CASE"},
+ {"max_tokens": 100, "temperature": 0.0, "system_prompt": ""},
+ {"endpoint": "https://example.test/chat/completions";, "api-key":
"key",
+ "model": "chat-model", "max-tokens": "100", "temperature": "0.0",
+ "system-prompt": "", "future-option": "Mixed CASE"}),
+ (pf.TritonProvider,
+ {"endpoint": "http://localhost:8000"},
+ {"model-name": "classifier", "max-retries": "0",
"flatten-batch-dim": "false",
+ "future_option": ""},
+ {"max_retries": 0, "flatten_batch_dim": False},
+ {"endpoint": "http://localhost:8000";, "model-name": "classifier",
+ "max-retries": "0", "flatten-batch-dim": "false",
"future_option": ""}),
+ ]:
+ for parameters in ({}, typed):
+ with self.subTest(provider=provider_type.__name__,
parameters=parameters):
+ options = dict(raw)
+ provider = provider_type(
+ **connection, **parameters,
options=MappingProxyType(options))
+ options.clear()
+ self.assertEqual(provider.to_options(), expected)
+ provider.to_options().clear()
+ self.assertEqual(provider.to_options(), expected)
+
+ def test_public_contracts_have_resolvable_type_hints(self):
+ for api in (pf.set_model_provider, pf.set_default_model_provider,
pf.list_model_providers,
+ pf.ModelProvider.provider_identifier,
pf.ModelProvider.to_options,
+ pf.ModelProvider.model_option_key,
pf.OpenAIProvider.__init__,
+ pf.TritonProvider.__init__, pf.GenericProvider.__init__):
+ with self.subTest(api=api.__qualname__):
+ self.assertTrue(get_type_hints(api))
+
+ def test_providers_reject_conflicting_options(self):
+ for construct, key, value in [
+ (lambda options: pf.OpenAIProvider(
+ "https://example.test/chat/completions";, "key", max_tokens=10,
options=options),
+ "max-tokens", "10"),
+ (lambda options: pf.OpenAIProvider(
+ "https://example.test/chat/completions";, "key",
options=options),
+ "api-key", "key"),
+ (lambda options: pf.OpenAIProvider(
+ "https://example.test/chat/completions";, "key",
system_prompt="", options=options),
+ "system-prompt", ""),
+ (lambda options: pf.TritonProvider(
+ "http://localhost:8000";, flatten_batch_dim=False,
options=options),
+ "flatten-batch-dim", "false"),
+ (lambda options: pf.TritonProvider(
+ "http://localhost:8000";, max_retries=0, options=options),
+ "max-retries", "0"),
+ (lambda options: pf.GenericProvider("custom", options, timeout="10
s"),
+ "timeout", "10 s"),
+ ]:
+ with self.subTest(key=key):
+ options = {key: value}
+ self.assertEqual(construct(options).to_options()[key], value)
+ options[key] = "different"
+ with self.assertRaisesRegex(ValueError, "conflicting values
for option"):
+ construct(options)
+ self.assertEqual(options, {key: "different"})
+
+ def test_generic_provider_rejects_invalid_configuration(self):
+ for identifier, options, error in [
+ (1, {}, TypeError),
+ ("", {}, ValueError),
+ (" ", {}, ValueError),
+ ("custom", {"enabled": True}, TypeError),
+ ("custom", {"retries": 1}, TypeError),
+ ("custom", {"endpoint": None}, TypeError),
+ ]:
+ with self.subTest(identifier=identifier, options=options):
+ with self.assertRaises(error):
+ pf.GenericProvider(identifier, **options)
+
+ def test_providers_reject_invalid_raw_options(self):
+ for name, construct, expected in (
+ ("openai", lambda options: pf.OpenAIProvider(
+ "https://example.test/chat/completions";, "key",
options=options),
+ {"endpoint": "https://example.test/chat/completions";, "api-key":
"key"}),
+ ("triton", lambda options:
pf.TritonProvider("http://localhost:8000";, options=options),
+ {"endpoint": "http://localhost:8000"}),
+ ("generic", lambda options: pf.GenericProvider("custom", options),
{}),
+ ):
+ for options in ([], "options", False, {1: "value"}, {"option":
None},
+ {"option": False}, {"option": 1}):
+ with self.subTest(provider=name, options=options):
+ with self.assertRaises(TypeError):
+ construct(options)
+ self.assertEqual(construct({}).to_options(), expected)
+
+ with self.assertRaises(TypeError):
+ pf.OpenAIProvider("https://example.test/chat/completions";, "key",
+ max_tokens=True, options={"max-tokens": "1"})
+ with self.assertRaises(ValueError):
+ pf.TritonProvider("http://localhost:8000";, priority=256,
options={"priority": "1"})
+
+
+class ModelProviderRegistryTests(unittest.TestCase):
+ def setUp(self):
+ previous = dict(llm._provider_registry)
+ self.addCleanup(llm._provider_registry.update, previous)
+ self.addCleanup(llm._provider_registry.clear)
+ llm._provider_registry.clear()
+ previous_default = llm._default_provider
+ self.addCleanup(setattr, llm, "_default_provider", previous_default)
+ llm._default_provider = None
+
+ def test_registration_is_lazy_and_names_are_ordered(self):
+ class UnserializedProvider(pf.ModelProvider):
+ def provider_identifier(self):
+ raise AssertionError("Registration must not inspect the
provider")
+
+ def to_options(self):
+ raise AssertionError("Registration must not serialize the
provider")
+
+ pf.set_model_provider(name="chat", provider=UnserializedProvider())
+ pf.set_model_provider("embed", provider=pf.GenericProvider("custom"))
+
+ names = pf.list_model_providers()
+ self.assertEqual(names, ["chat", "embed"])
+ names.clear()
+ self.assertEqual(pf.list_model_providers(), ["chat", "embed"])
+
+ def
test_single_provider_registration_derives_names_without_serializing(self):
+ class UnserializedProvider(pf.ModelProvider):
+ def provider_identifier(self):
+ return "custom"
+
+ def to_options(self):
+ raise AssertionError("Registration must not serialize the
provider")
+
+ def model_option_key(self):
+ raise AssertionError("Registration must not inspect model
options")
+
+ provider = UnserializedProvider()
+ pf.set_model_provider(provider)
+ self.assertEqual(pf.list_model_providers(), ["custom"])
+ self.assertIs(llm._resolve_provider(), provider)
+
+
pf.set_model_provider(provider=pf.OpenAIProvider("https://example.test/embeddings";,
"key"))
+ with self.assertRaisesRegex(ValueError, "Multiple model providers"):
+ llm._resolve_provider()
+
+ pf.set_default_model_provider("custom")
+ pf.set_model_provider(pf.TritonProvider("http://localhost:8000";))
+ self.assertEqual(pf.list_model_providers(), ["custom", "openai",
"triton"])
+ self.assertIs(llm._resolve_provider(), provider)
+
+ def test_invalid_registration_preserves_registered_names(self):
+ class InvalidIdentifierProvider(pf.ModelProvider):
+ def __init__(self, identifier):
+ self.identifier = identifier
+
+ def provider_identifier(self):
+ if isinstance(self.identifier, RuntimeError):
+ raise self.identifier
+ return self.identifier
+
+ def to_options(self):
+ raise AssertionError("Registration must not serialize the
provider")
+
+ provider = pf.GenericProvider("custom")
+ other = pf.GenericProvider("other")
+ pf.set_model_provider("chat", provider)
+ pf.set_model_provider("embed", other)
+ pf.set_default_model_provider("chat")
+ invalid_calls = [
+ (("", provider), {}, ValueError),
+ ((" ", provider), {}, ValueError),
+ ((1, provider), {}, TypeError),
+ ((None, provider), {}, TypeError),
+ (("chat", None), {}, TypeError),
+ (("embed", "custom"), {}, TypeError),
+ ((), {}, TypeError),
+ (("chat",), {}, TypeError),
+ ((), {"provider": None}, TypeError),
+ ((), {"name": provider}, TypeError),
+ ((provider, provider), {}, TypeError),
+ ((provider, None), {}, TypeError),
+ (("chat", provider, provider), {}, TypeError),
+ (("chat",), {"name": "embed", "provider": provider}, TypeError),
+ (("chat", provider), {"provider": other}, TypeError),
+ (("chat", provider), {"endpoint": "https://example.test"},
TypeError),
+ ]
+ for identifier, error in [
+ ("", ValueError), (" ", ValueError), (1, TypeError), (None,
TypeError),
+ (RuntimeError("Identifier unavailable"), RuntimeError),
+ ]:
+ invalid_calls.append(((), {"provider":
InvalidIdentifierProvider(identifier)}, error))
+ for args, kwargs, error in invalid_calls:
+ with self.subTest(args=args, kwargs=kwargs):
+ with self.assertRaises(error):
+ pf.set_model_provider(*args, **kwargs)
+ self.assertEqual(pf.list_model_providers(), ["chat", "embed"])
+ self.assertIs(llm._resolve_provider(), provider)
+ self.assertIs(llm._resolve_provider("chat"), provider)
+ self.assertIs(llm._resolve_provider("embed"), other)
+
+ def test_replacement_preserves_order_and_defaults_bound_to_names(self):
+ original = pf.GenericProvider("custom", token="old")
+ pf.set_model_provider(original)
+ replacement = pf.GenericProvider("custom", token="new")
+ pf.set_model_provider(replacement)
+ self.assertEqual(pf.list_model_providers(), ["custom"])
+ self.assertIs(llm._resolve_provider(), replacement)
+
+ pf.set_model_provider("other", pf.GenericProvider("other"))
+ pf.set_default_model_provider("custom")
+ named = pf.GenericProvider("different-factory")
+ pf.set_model_provider(name="custom", provider=named)
+ self.assertEqual(pf.list_model_providers(), ["custom", "other"])
+ self.assertIs(llm._resolve_provider(), named)
+
+ pf.set_model_provider(provider=original)
+ other = pf.GenericProvider("another-factory")
+ pf.set_model_provider("other", other)
+ self.assertEqual(pf.list_model_providers(), ["custom", "other"])
+ self.assertIs(llm._resolve_provider(), original)
+ self.assertIs(llm._resolve_provider("other"), other)
+
+ def test_selection_requires_choice_only_with_multiple_providers(self):
+ chat = pf.GenericProvider("custom", task="chat")
+ embed = pf.GenericProvider("custom", task="embed")
+ with self.assertRaisesRegex(ValueError, "No model provider"):
+ llm._resolve_provider()
+
+ pf.set_model_provider("chat", chat)
+ self.assertIs(llm._resolve_provider(), chat)
+ pf.set_model_provider("embed", embed)
+ with self.assertRaisesRegex(ValueError, "Multiple model providers"):
+ llm._resolve_provider()
+
+ pf.set_default_model_provider("chat")
+ pf.set_model_provider("other", pf.GenericProvider("other"))
+ self.assertIs(llm._resolve_provider(), chat)
+ self.assertIs(llm._resolve_provider("embed"), embed)
+
+ def test_invalid_selection_preserves_default(self):
+ provider = pf.GenericProvider("custom")
+ pf.set_model_provider("chat", provider)
+ pf.set_default_model_provider("chat")
+ for name, error in [("unknown", ValueError), ("", ValueError),
+ (" ", ValueError), (1, TypeError), (None,
TypeError)]:
+ with self.subTest(name=name):
+ with self.assertRaises(error):
+ pf.set_default_model_provider(name)
+ self.assertIs(llm._resolve_provider(), provider)
+ if name is not None:
+ with self.assertRaises(error):
+ llm._resolve_provider(name)
+
+ def test_registry_and_default_survive_environment_replacement(self):
+ previous_environment = pf.get_table_environment()
+ buffered = dict(pf.config._buffered)
+ self.addCleanup(pf.config._buffered.update, buffered)
+ self.addCleanup(pf.config._buffered.clear)
+ self.addCleanup(pf.set_table_environment, previous_environment)
+ pf.config._buffered.clear()
+ provider = pf.GenericProvider("custom")
+ pf.set_model_provider("chat", provider)
+ pf.set_model_provider("embed", pf.GenericProvider("custom"))
+ pf.set_default_model_provider("chat")
+
+ for environment in (object.__new__(TableEnvironment), None,
+ object.__new__(TableEnvironment)):
+ pf.set_table_environment(environment)
+ self.assertEqual(pf.list_model_providers(), ["chat", "embed"])
+ self.assertIs(llm._resolve_provider(), provider)
+
+
+if __name__ == '__main__':
+ unittest.main()
diff --git a/flink-python/pyflink/dataframe/validation.py
b/flink-python/pyflink/dataframe/validation.py
index 49cc95d03aa..de55e8cba49 100644
--- a/flink-python/pyflink/dataframe/validation.py
+++ b/flink-python/pyflink/dataframe/validation.py
@@ -25,23 +25,42 @@ subclass and ``take(True)`` is almost certainly a mistake.
"""
import math
-from typing import Any, Sequence
+from typing import Any, Optional, Sequence, Union
-def _require_int(value: Any, name: str, minimum: int) -> None:
+def _require_int(
+ value: Any, name: str, minimum: Optional[int] = None, *, maximum:
Optional[int] = None,
+ include_minimum: bool = True,
+) -> None:
+ """Require an integer, optionally constrained by lower and upper bounds."""
if isinstance(value, bool) or not isinstance(value, int):
raise TypeError(f"{name} must be an integer")
- if value < minimum:
- raise ValueError(f"{name} must be at least {minimum}, got {value}")
+ _require_range(value, name, minimum, maximum, include_minimum)
-def _require_number(value: Any, name: str, minimum: float) -> None:
+def _require_number(
+ value: Any, name: str, minimum: Optional[float] = None, *, maximum:
Optional[float] = None,
+ include_minimum: bool = True,
+) -> None:
+ """Require a finite number, optionally constrained by lower and upper
bounds."""
if isinstance(value, bool) or not isinstance(value, (int, float)):
raise TypeError(f"{name} must be a number")
if not math.isfinite(value):
raise ValueError(f"{name} must be finite, got {value}")
- if value < minimum:
- raise ValueError(f"{name} must be at least {minimum}, got {value}")
+ _require_range(value, name, minimum, maximum, include_minimum)
+
+
+def _require_range(
+ value: Union[int, float], name: str, minimum: Optional[float], maximum:
Optional[float],
+ include_minimum: bool,
+) -> None:
+ if minimum is not None:
+ if include_minimum and value < minimum:
+ raise ValueError(f"{name} must be at least {minimum}, got {value}")
+ if not include_minimum and value <= minimum:
+ raise ValueError(f"{name} must be greater than {minimum}, got
{value}")
+ if maximum is not None and value > maximum:
+ raise ValueError(f"{name} must be at most {maximum}, got {value}")
def _require_non_empty_str(value: Any, name: str) -> None:
