################################################################################
# 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.
################################################################################
"""
Model provider classes for the DataFrame LLM API.
Example::
import pyflink.dataframe as pf
# Typed provider with IDE auto-complete and validation
provider = pf.OpenAICompatProvider(
task="chat/completions",
temperature=0.7,
)
pf.set_model_provider(provider)
# Generic provider for unknown/custom providers
pf.set_model_provider("my-custom-provider", endpoint="https://...", api_key="sk-...")
"""
from abc import ABC, abstractmethod
from collections.abc import Mapping as MappingABC
import json
from typing import AbstractSet, Any, Dict, List, Literal, Mapping, Optional, Tuple, Union
_TritonDefaultValue = Union[
str, int, float, bool, List[Any], Tuple[Any, ...], Mapping[str, Any]
]
_CONTEXT_OVERFLOW_ACTIONS = frozenset((
"truncated-tail",
"truncated-tail-log",
"truncated-head",
"truncated-head-log",
"skipped",
"skipped-log",
))
_ContentType = Literal[
"TEXT",
"IMAGE_URL",
"MULTI_IMAGE_URLS",
]
_ContentTypeInput = Union[_ContentType, List[_ContentType], Tuple[_ContentType, ...]]
_CONTENT_TYPES = frozenset((
"TEXT",
"IMAGE_URL",
"MULTI_IMAGE_URLS",
))
_SINGLE_CONTENT_TYPE_OPTION_TYPES = frozenset((
"TEXT",
"IMAGE_URL",
))
_SINGLE_INPUT_CONTENT_TYPE_TASKS = frozenset((
"embeddings",
"multimodal-embedding",
))
_TEXT_CONTENT_TYPES = frozenset(("TEXT",))
_DASHSCOPE_MULTIMODAL_EMBEDDING_CONTENT_TYPES = frozenset((
"IMAGE_URL",
))
_MULTI_COLUMN_CONTENT_TYPE_TASKS = frozenset(("chat/completions",))
_MODEL_TASK_ENDPOINT_SUFFIXES = {
"chat/completions": "chat/completions",
"embeddings": "embeddings",
"multimodal-embedding":
"services/embeddings/multimodal-embedding/multimodal-embedding",
}
_OPENAI_COMPAT_TASK_CONTENT_TYPES = {
"chat/completions": _CONTENT_TYPES,
"embeddings": _TEXT_CONTENT_TYPES,
}
_DASHSCOPE_TASK_CONTENT_TYPES = {
"chat/completions": _CONTENT_TYPES,
"embeddings": _TEXT_CONTENT_TYPES,
"multimodal-embedding": _DASHSCOPE_MULTIMODAL_EMBEDDING_CONTENT_TYPES,
}
def _merge_options(
options: Mapping[str, Any],
overrides: Optional[Mapping[str, Any]] = None,
) -> Dict[str, Any]:
merged = dict(options)
if overrides:
for key, value in overrides.items():
if value is None:
merged.pop(key, None)
else:
merged[key] = value
return merged
def _merge_raw_options(
params: Mapping[str, Any],
extra_options: Mapping[str, Any],
constructor_param_names: AbstractSet[str],
overrides: Optional[Mapping[str, Any]] = None,
) -> Tuple[Dict[str, Any], Dict[str, Any]]:
constructor_overrides = {}
extra_overrides = {}
if overrides:
for key, value in overrides.items():
if key in constructor_param_names:
constructor_overrides[key] = value
else:
extra_overrides[key] = value
return (
_merge_options(params, constructor_overrides),
_merge_options(extra_options, extra_overrides),
)
[docs]class ModelProvider(ABC):
"""Base class for model providers.
A provider owns the Python-side configuration for a Java model provider
factory. Subclasses translate Python-style constructor parameters to the
Java-side option keys expected by Flink's ``ModelDescriptor``.
``to_options()`` is called in two modes:
- Without ``input_columns``, it serializes the provider's static
configuration.
- With ``input_columns``, it may additionally validate or adapt options
against the model input schema before the descriptor is created.
Custom provider subclasses should accept the optional ``input_columns``
argument even if they do not need schema-aware option handling.
"""
@abstractmethod
def provider_identifier(self) -> str:
"""Return the provider identifier recognized by Flink's Java runtime.
This is distinct from the Python-side lookup name used by
``set_model_provider(name, provider)``. Multiple registered providers
may share the same Java provider identifier.
"""
@abstractmethod
def to_options(
self,
input_columns: Optional[List[Tuple[str, Any]]] = None,
overrides: Optional[Mapping[str, Any]] = None,
) -> Dict[str, str]:
"""Return Java-side options for ``ModelDescriptor``.
Args:
input_columns: Optional list of ``(name, data_type)`` pairs for the
model input schema. DataFrame AI calls pass this when creating
a model descriptor. ``None`` means only static provider
configuration should be serialized.
overrides: Optional per-call provider constructor parameter
overrides. For typed providers, keys use Python constructor
parameter names. For ``GenericProvider``, keys are Java-side
option names.
Returns:
A dictionary whose keys and values are Java-side option strings.
"""
def model_option_key(self) -> str:
"""Return the Java-side option key used for a per-call model name.
Most providers use ``model``. Providers whose Java option differs, such
as Triton, override this method.
"""
return "model"
[docs]class OpenAICompatProvider(ModelProvider):
"""Model provider for all OpenAI-compatible endpoints (openai-compat).
Covers endpoints that implement the OpenAI chat/completions or embeddings
API.
.. note::
We recommend using Flink AI service. With Flink AI service, you do not
need to configure ``endpoint`` or ``api_key``.
Args:
task: The model task. Supported values are ``"chat/completions"``,
and ``"embeddings"``.
Required when using Flink AI Model Service.
endpoint: The endpoint to connect to. Required with ``api_key``.
api_key: The key used to authorize the access to the endpoint.
Required when using BYOK models.
model: The version of the model to use.
system_prompt: The system message of a chat. Can be disabled by
setting to empty string.
Defaults to ``"You are a helpful assistant."``.
user_prompt: The prompt of a chat, passed to the model service
through user's role. Can be disabled by setting to empty string.
temperature: Controls the randomness or "creativity" of the output.
Typical values are between 0.0 and 1.0.
top_p: The probability cutoff for token selection. Usually either
temperature or top_p are specified, but not both.
max_tokens: The maximum number of tokens that can be generated
in the chat completion.
stop: A CSV list of strings to pass as stop sequences to the model.
presence_penalty: Number between -2.0 and 2.0. Positive values
penalize new tokens based on whether they appear in the text
so far, increasing the model's likelihood to talk about new
topics.
n: How many chat completion choices to generate for each input
message. Keep n as 1 to minimize costs.
seed: If specified, the model platform will make a best effort to
sample deterministically. Determinism is not guaranteed.
content_type: Content type of the input column(s). For a single input
column or multiple input columns that all use the same content
type, pass one string, for example ``"TEXT"``. For multiple input
columns with different content types, pass a list or tuple of
strings; the count and order must match the input columns.
Supported values are ``"TEXT"`` (default), ``"IMAGE_URL"``,
and ``"MULTI_IMAGE_URLS"``. Task-specific limitations are
enforced by the model provider.
response_format: The format of the response (``"text"`` or
``"json_object"``).
dimension: The size of the embedding result array.
max_context_size: Max number of tokens for context.
``context_overflow_action`` is triggered if this threshold
is exceeded.
context_overflow_action: Action to handle context overflows. One
of ``"truncated-tail"``, ``"truncated-tail-log"``,
``"truncated-head"``, ``"truncated-head-log"``,
``"skipped"``, or ``"skipped-log"`` (case-insensitive).
Defaults to ``"truncated-tail"``.
error_handling_strategy: Strategy for handling errors during model
requests. ``"RETRY"`` retries the request (limited by
retry_num, retry_fallback_strategy, etc.); ``"FAILOVER"``
throws exceptions and fails the job; ``"IGNORE"`` skips the
error input and continues. Defaults to ``"RETRY"``.
retry_num: Number of retries for client requests.
Defaults to ``100``.
retry_backoff_strategy: The strategy to use for retry backoff.
``"FIXED"`` or ``"EXPONENTIAL"``. Defaults to ``"FIXED"``.
retry_backoff_base_interval: The base interval for retry backoff,
used as the initial delay before the first retry and as the
base for calculating subsequent retry delays.
Defaults to ``"1s"``.
retry_fallback_strategy: Fallback strategy to employ if the retry
attempts are exhausted. ``"FAILOVER"`` or ``"IGNORE"``.
Defaults to ``"FAILOVER"``.
extra_header: Additional headers for the requests. Should be a
JSON-format string whose values are strings or list of strings.
extra_body: Additional parameters to pass through the requests'
body. Should be a JSON-format string.
**extra_options: Additional options passed through as-is (keys are
not translated).
Note:
When a configured provider is used, DataFrame AI function calls can
override the provider's constructor parameters for that call.
LLM function ``config`` supports these runtime prediction options with
this provider:
- ``async``: Execution-mode hint. The provider supports async
prediction only, so omit this key or set it to ``"true"``. Setting
``"false"`` is not supported.
- ``output-mode``: Output ordering for async prediction. Allowed
values are ``"ORDERED"`` and ``"ALLOW_UNORDERED"``. Default: table
config ``table.exec.async-ml-predict.output-mode`` (``"ORDERED"`` by
default).
- ``max-concurrent-operations``: Maximum number of concurrent async
prediction calls. Allowed values are positive integer strings.
Default: table config
``table.exec.async-ml-predict.max-concurrent-operations`` (``"10"``
by default).
- ``timeout``: Timeout for async prediction calls. Allowed values are
duration strings such as ``"30s"`` or ``"3 min"``. Default: table
config ``table.exec.async-ml-predict.timeout`` (``"3 min"`` by
default).
Example::
>>> provider = OpenAICompatProvider(task="chat/completions")
>>> pf.set_model_provider(provider)
>>> df = pf.from_dict({
... "review_text": ["The shoes fit well but the sole arrived scratched."],
... "product_image_url": ["https://example.com/product.jpg"],
... })
>>> predictions = df.llm.predict(
... "review_text",
... "product_image_url",
... model="qwen3.6-plus",
... content_type=["TEXT", "IMAGE_URL"],
... system_prompt="Respond to the product review and image together.",
... temperature=0.3,
... config={
... "max-concurrent-operations": "20",
... "timeout": "30s",
... },
... )
>>> classified = df.llm.ai_classify(
... "review_text",
... ["positive", "negative"],
... model="qwen3.6-plus",
... temperature=0.1,
... )
"""
# Explicit mapping from Python parameter names to Java-side option keys.
_OPTION_MAP = {
"endpoint": "endpoint",
"api_key": "api-key",
"task": "task",
"model": "model",
"system_prompt": "system-prompt",
"user_prompt": "user-prompt",
"temperature": "temperature",
"top_p": "top-p",
"max_tokens": "max-tokens",
"stop": "stop",
"presence_penalty": "presence-penalty",
"n": "n",
"seed": "seed",
"content_type": "content-type",
"content_types": "content-types",
"response_format": "response-format",
"dimension": "dimension",
"max_context_size": "max-context-size",
"context_overflow_action": "context-overflow-action",
"error_handling_strategy": "error-handling-strategy",
"retry_num": "retry-num",
"retry_backoff_strategy": "retry-backoff-strategy",
"retry_backoff_base_interval": "retry-backoff-base-interval",
"retry_fallback_strategy": "retry-fallback-strategy",
"extra_header": "extra-header",
"extra_body": "extra-body",
}
_CONSTRUCTOR_PARAM_NAMES = frozenset((
"endpoint",
"api_key",
"task",
"model",
"system_prompt",
"user_prompt",
"temperature",
"top_p",
"max_tokens",
"stop",
"presence_penalty",
"n",
"seed",
"content_type",
"response_format",
"dimension",
"max_context_size",
"context_overflow_action",
"error_handling_strategy",
"retry_num",
"retry_backoff_strategy",
"retry_backoff_base_interval",
"retry_fallback_strategy",
"extra_header",
"extra_body",
))
_SUPPORTED_TASK_CONTENT_TYPES = _OPENAI_COMPAT_TASK_CONTENT_TYPES
def __init__(
self,
*,
task: Optional[str] = None,
endpoint: Optional[str] = None,
api_key: Optional[str] = None,
model: Optional[str] = None,
system_prompt: str = "You are a helpful assistant.",
user_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,
content_type: _ContentTypeInput = "TEXT",
response_format: Optional[str] = None,
dimension: Optional[int] = None,
max_context_size: Optional[int] = None,
context_overflow_action: str = "truncated-tail",
error_handling_strategy: str = "RETRY",
retry_num: int = 100,
retry_backoff_strategy: str = "FIXED",
retry_backoff_base_interval: str = "1s",
retry_fallback_strategy: str = "FAILOVER",
extra_header: Optional[str] = None,
extra_body: Optional[str] = None,
**extra_options: Any,
):
local_vars = {
"endpoint": endpoint,
"api_key": api_key,
"task": task,
"model": model,
"system_prompt": system_prompt,
"user_prompt": user_prompt,
"temperature": temperature,
"top_p": top_p,
"max_tokens": max_tokens,
"stop": stop,
"presence_penalty": presence_penalty,
"n": n,
"seed": seed,
"content_type": content_type,
"response_format": response_format,
"dimension": dimension,
"max_context_size": max_context_size,
"context_overflow_action": context_overflow_action,
"error_handling_strategy": error_handling_strategy,
"retry_num": retry_num,
"retry_backoff_strategy": retry_backoff_strategy,
"retry_backoff_base_interval": retry_backoff_base_interval,
"retry_fallback_strategy": retry_fallback_strategy,
"extra_header": extra_header,
"extra_body": extra_body,
}
self._validate_python_params(local_vars)
self._params: Dict[str, Any] = {
key: value for key, value in local_vars.items() if value is not None
}
self._extra_options = extra_options
def provider_identifier(self) -> str:
return "openai-compat"
@classmethod
def _validate_service_options(
cls,
endpoint: Optional[str],
api_key: Optional[str],
task: Optional[str],
) -> None:
if api_key is not None:
if endpoint is None or not endpoint.strip():
raise ValueError("'api_key' requires 'endpoint'.")
if task is not None:
raise ValueError("'api_key' cannot be combined with 'task'.")
if task is not None and task not in cls._SUPPORTED_TASK_CONTENT_TYPES:
supported = ", ".join(sorted(cls._SUPPORTED_TASK_CONTENT_TYPES))
raise ValueError(
f"Unsupported task {task!r}. Supported values: {supported}.")
@classmethod
def _resolve_task(cls, endpoint: Optional[str], task: Optional[str]) -> Optional[str]:
if task is not None:
return task
if endpoint is None or not endpoint.strip():
return None
normalized = endpoint.rstrip("/").lower()
for task_name, endpoint_suffix in _MODEL_TASK_ENDPOINT_SUFFIXES.items():
if normalized.endswith(endpoint_suffix):
if task_name not in cls._SUPPORTED_TASK_CONTENT_TYPES:
supported = ", ".join(
sorted(cls._SUPPORTED_TASK_CONTENT_TYPES))
raise ValueError(
f"Unsupported task {task_name!r}. "
f"Supported values: {supported}.")
return task_name
return None
@staticmethod
def _normalize_content_type(
content_type: _ContentTypeInput,
input_count: Optional[int] = None,
) -> Tuple[str, ...]:
if isinstance(content_type, (list, tuple)):
if not content_type:
raise ValueError("content_type must not be empty.")
content_types = list(content_type)
else:
content_types = [content_type]
supported = ", ".join(sorted(_CONTENT_TYPES))
normalized_content_types = []
for content_type_value in content_types:
if not isinstance(content_type_value, str):
raise ValueError(
f"Unsupported content_type {content_type_value!r}. "
f"Supported values: {supported}.")
normalized_content_type = content_type_value.upper()
if normalized_content_type not in _CONTENT_TYPES:
raise ValueError(
f"Unsupported content_type {content_type_value!r}. "
f"Supported values: {supported}.")
normalized_content_types.append(normalized_content_type)
if input_count is not None and input_count > 1:
if len(normalized_content_types) == 1:
normalized_content_types *= input_count
return tuple(normalized_content_types)
@staticmethod
def _content_type_param_name(content_types: Tuple[str, ...]) -> str:
if (len(content_types) == 1
and content_types[0] in _SINGLE_CONTENT_TYPE_OPTION_TYPES):
return "content_type"
return "content_types"
@classmethod
def _validate_content_type(
cls, content_types: Tuple[str, ...], task: Optional[str]) -> None:
if task is None:
return
supported_content_types = cls._SUPPORTED_TASK_CONTENT_TYPES[task]
if (len(content_types) > 1
and task not in _MULTI_COLUMN_CONTENT_TYPE_TASKS):
raise ValueError(
f"task {task!r} supports exactly one content_type, "
f"but got {len(content_types)}: {';'.join(content_types)}.")
for content_type_value in content_types:
if content_type_value not in supported_content_types:
supported = ", ".join(sorted(supported_content_types))
raise ValueError(
f"Unsupported content_type {content_type_value!r} "
f"for task {task!r}. Supported values: {supported}.")
@staticmethod
def _validate_context_overflow_action(value: str) -> None:
if value not in _CONTEXT_OVERFLOW_ACTIONS:
supported = ", ".join(sorted(_CONTEXT_OVERFLOW_ACTIONS))
raise ValueError(
f"Unsupported context_overflow_action {value!r}. "
f"Supported values: {supported}.")
@classmethod
def _validate_python_params(cls, params: Mapping[str, Any]) -> None:
cls._validate_service_options(
params.get("endpoint"), params.get("api_key"), params.get("task"))
context_overflow_action = params.get("context_overflow_action")
if context_overflow_action is not None:
cls._validate_context_overflow_action(context_overflow_action)
if "content_type" in params:
content_types = cls._normalize_content_type(params["content_type"])
resolved_task = cls._resolve_task(
params.get("endpoint"), params.get("task"))
cls._validate_content_type(content_types, resolved_task)
def _to_java_options(
self,
params: Mapping[str, Any],
extra_options: Mapping[str, Any],
input_count: Optional[int] = None,
) -> Dict[str, str]:
options: Dict[str, str] = {}
for py_key, value in params.items():
if py_key == "content_type":
content_types = self._normalize_content_type(value, input_count)
java_key = self._OPTION_MAP[
self._content_type_param_name(content_types)]
options[java_key] = ";".join(content_types)
continue
options[self._OPTION_MAP[py_key]] = str(value)
for key, value in extra_options.items():
options[str(key)] = str(value)
return options
def to_options(
self,
input_columns: Optional[List[Tuple[str, Any]]] = None,
overrides: Optional[Mapping[str, Any]] = None,
) -> Dict[str, str]:
params, extra_options = _merge_raw_options(
self._params,
self._extra_options,
self._CONSTRUCTOR_PARAM_NAMES,
overrides)
self._validate_python_params(params)
input_count = len(input_columns) if input_columns is not None else None
options = self._to_java_options(params, extra_options, input_count)
if input_columns is None:
return dict(options)
return self._align_content_type_input_count(options, input_columns)
def _align_content_type_input_count(
self,
provider_options: Mapping[str, str],
input_columns: List[Tuple[str, Any]]) -> Dict[str, str]:
"""Align content type options with the model input schema."""
aligned_options = dict(provider_options)
input_column_names = [name for name, _ in input_columns]
# Embedding-style tasks consume exactly one model input, so reject
# multi-column calls before default TEXT handling expands the option.
task = self._resolve_task(
provider_options.get("endpoint"), provider_options.get("task"))
if (task in _SINGLE_INPUT_CONTENT_TYPE_TASKS
and len(input_column_names) > 1):
raise ValueError(
f"task {task!r} supports exactly one input column, but got "
f"{len(input_column_names)}: {input_column_names}.")
scalar_option_value = provider_options.get("content-type")
plural_option_value = provider_options.get("content-types")
if scalar_option_value is not None and plural_option_value is not None:
raise ValueError(
"content-type and content-types are mutually exclusive for "
f"provider {self.provider_identifier()!r}. Use the Python "
"content_type parameter instead of raw Java options to set "
"input content types.")
if plural_option_value is not None:
resolved_content_types = plural_option_value.split(";")
reported_option_value = plural_option_value
elif scalar_option_value is not None:
resolved_content_types = scalar_option_value.split(";")
reported_option_value = scalar_option_value
else:
return aligned_options
if len(resolved_content_types) != len(input_column_names):
raise ValueError(
f"Number of content types ({len(resolved_content_types)}) must "
f"match number of input columns ({len(input_column_names)}) "
f"for provider {self.provider_identifier()!r}. Content type: "
f"{reported_option_value}, input columns: {input_column_names}.")
# The Java provider uses scalar content-type for single TEXT/IMAGE_URL
# inputs, and content-types for positional multi-input or array content.
if (len(resolved_content_types) == 1
and resolved_content_types[0] in _SINGLE_CONTENT_TYPE_OPTION_TYPES):
aligned_options.pop("content-types", None)
aligned_options["content-type"] = reported_option_value
else:
aligned_options.pop("content-type", None)
aligned_options["content-types"] = reported_option_value
return aligned_options
[docs]class DashScopeProvider(OpenAICompatProvider):
"""Model provider for Alibaba Cloud DashScope (dashscope).
DashScope reuses the OpenAI-compatible Java provider for normal chat and
embedding requests, and adds DashScope-specific multi-modal embedding support.
.. note::
We recommend using Flink AI service. With Flink AI service, you do not
need to configure ``endpoint`` or ``api_key``.
Args:
task: The model task. Supported values are ``"chat/completions"``,
``"embeddings"``, and ``"multimodal-embedding"``.
Required when using Flink AI Model Service.
endpoint: The endpoint to connect to. Required with ``api_key``.
api_key: The key used to authorize the access to the endpoint.
Required when using BYOK models.
model: The version of the model to use.
system_prompt: The system message of a chat. Can be disabled by
setting to empty string.
Defaults to ``"You are a helpful assistant."``.
user_prompt: The prompt of a chat, passed to the model service
through user's role. Can be disabled by setting to empty string.
temperature: Controls the randomness or "creativity" of the output.
Typical values are between 0.0 and 1.0.
top_p: The probability cutoff for token selection. Usually either
temperature or top_p are specified, but not both.
max_tokens: The maximum number of tokens that can be generated
in the chat completion.
stop: A CSV list of strings to pass as stop sequences to the model.
presence_penalty: Number between -2.0 and 2.0. Positive values
penalize new tokens based on whether they appear in the text
so far, increasing the model's likelihood to talk about new
topics.
n: How many chat completion choices to generate for each input
message. Keep n as 1 to minimize costs.
seed: If specified, the model platform will make a best effort to
sample deterministically. Determinism is not guaranteed.
content_type: Content type of the input column(s). For a single input
column or multiple input columns that all use the same content
type, pass one string, for example ``"TEXT"``. For multiple input
columns with different content types, pass a list or tuple of
strings; the count and order must match the input columns.
Supported values are ``"TEXT"`` (default), ``"IMAGE_URL"``,
and ``"MULTI_IMAGE_URLS"``. Task-specific limitations are
enforced by the model provider.
response_format: The format of the response (``"text"`` or
``"json_object"``).
dimension: The size of the embedding result array.
max_context_size: Max number of tokens for context.
``context_overflow_action`` is triggered if this threshold
is exceeded.
context_overflow_action: Action to handle context overflows. One
of ``"truncated-tail"``, ``"truncated-tail-log"``,
``"truncated-head"``, ``"truncated-head-log"``,
``"skipped"``, or ``"skipped-log"`` (case-insensitive).
Defaults to ``"truncated-tail"``.
error_handling_strategy: Strategy for handling errors during model
requests. ``"RETRY"`` retries the request (limited by
retry_num, retry_fallback_strategy, etc.); ``"FAILOVER"``
throws exceptions and fails the job; ``"IGNORE"`` skips the
error input and continues. Defaults to ``"RETRY"``.
retry_num: Number of retries for client requests.
Defaults to ``100``.
retry_backoff_strategy: The strategy to use for retry backoff.
``"FIXED"`` or ``"EXPONENTIAL"``. Defaults to ``"FIXED"``.
retry_backoff_base_interval: The base interval for retry backoff,
used as the initial delay before the first retry and as the
base for calculating subsequent retry delays.
Defaults to ``"1s"``.
retry_fallback_strategy: Fallback strategy to employ if the retry
attempts are exhausted. ``"FAILOVER"`` or ``"IGNORE"``.
Defaults to ``"FAILOVER"``.
extra_header: Additional headers for the requests. Should be a
JSON-format string whose values are strings or list of strings.
extra_body: Additional parameters to pass through the requests'
body. Should be a JSON-format string. Multi-modal
embedding tasks do not support this option.
**options: Additional options passed through as-is (keys are not
translated).
Note:
When a configured provider is used, DataFrame AI function calls can
override the provider's constructor parameters for that call.
LLM function ``config`` supports these runtime prediction options with
this provider:
- ``async``: Execution-mode hint. The provider supports async
prediction only, so omit this key or set it to ``"true"``. Setting
``"false"`` is not supported.
- ``output-mode``: Output ordering for async prediction. Allowed
values are ``"ORDERED"`` and ``"ALLOW_UNORDERED"``. Default: table
config ``table.exec.async-ml-predict.output-mode`` (``"ORDERED"`` by
default).
- ``max-concurrent-operations``: Maximum number of concurrent async
prediction calls. Allowed values are positive integer strings.
Default: table config
``table.exec.async-ml-predict.max-concurrent-operations`` (``"10"``
by default).
- ``timeout``: Timeout for async prediction calls. Allowed values are
duration strings such as ``"30s"`` or ``"3 min"``. Default: table
config ``table.exec.async-ml-predict.timeout`` (``"3 min"`` by
default).
Example::
>>> provider = DashScopeProvider(task="multimodal-embedding")
>>> pf.set_model_provider(provider)
>>> df = pf.from_dict({"image_url": [
... "https://dashscope.oss-cn-beijing.aliyuncs.com/images/tiger.png"
... ]})
>>> embeddings = df.llm.ai_embed(
... "image_url",
... model="qwen3-vl-embedding",
... dimension=512,
... content_type="IMAGE_URL",
... config={
... "max-concurrent-operations": "20",
... "timeout": "30s",
... },
... )
"""
_SUPPORTED_TASK_CONTENT_TYPES = _DASHSCOPE_TASK_CONTENT_TYPES
def provider_identifier(self) -> str:
return "dashscope"
[docs]class TritonProvider(ModelProvider):
"""Model provider for NVIDIA Triton Inference Server (triton).
Args:
endpoint: Full URL of the Triton Inference Server endpoint.
model_name: Name of the model to invoke. This can also be provided
through ``df.llm.predict(..., model="...")``, which maps to the
Java-side ``model-name`` option for Triton.
model_version: Version of the model to use.
Defaults to ``"latest"``.
timeout: HTTP request timeout, for example ``"10s"`` or
``"30000ms"``. Defaults to ``"30s"``.
flatten_batch_dim: Whether to flatten the leading batch dimension for
array inputs. For ``ARRAY<T>`` inputs, the default shape is
``[1, N]``, where ``N`` is the array length. Set this to ``True``
when the Triton model expects ``[N]`` instead.
Defaults to ``False``.
priority: Triton request priority level.
sequence_id: Sequence ID for stateful models.
sequence_start: Whether this request starts a stateful sequence.
Defaults to ``False``.
sequence_end: Whether this request ends a stateful sequence.
Defaults to ``False``.
compression: Compression algorithm for the request body. Currently
Triton provider supports ``"gzip"``.
auth_token: Authentication token for secured Triton servers. The Java
provider sends it as a Bearer token.
custom_headers: Custom HTTP headers as a Flink map string with
comma-separated ``key:value`` pairs (e.g.
``"X-Trace-Id:abc,X-Other:val"``).
max_retries: Maximum number of retries for failed inference requests.
Defaults to ``0``.
retry_initial_backoff: Initial backoff duration between retry
attempts. Defaults to ``"100ms"``.
retry_max_backoff: Upper bound on the delay between retry attempts.
Defaults to ``"30s"``.
default_value: Fallback value to return when inference fails after
retries or with a non-retryable error:
- If not specified, inference failures are propagated as
exceptions.
- For ``STRING`` outputs, pass plain text such as ``"FAILED"``.
- For numeric outputs, pass the numeric value or its string
representation, such as ``-1`` or ``"-1"``.
- For ``ARRAY`` or structured outputs, pass a JSON string or the
corresponding Python list, tuple, or mapping; Python containers
are serialized as JSON.
- To emit SQL ``NULL``, pass the lower-case literal ``"null"``.
For string outputs, ``"null"`` is therefore not usable as a
literal string sentinel; use values such as ``"NULL"``,
``"FAILED"``, or ``"<null>"`` instead.
health_check_enabled: Whether to enable periodic health checks for the
Triton server. Defaults to ``False``.
health_check_interval: Interval between health check requests.
Defaults to ``"30s"``.
circuit_breaker_enabled: Whether to enable circuit breaker protection
for Triton inference requests. Defaults to ``False``.
circuit_breaker_failure_threshold: Failure rate threshold that opens
the circuit breaker. Must be in ``(0.0, 1.0]``. Defaults to
``0.5``.
circuit_breaker_timeout: Duration to keep the circuit breaker open
before probing recovery. Defaults to ``"60s"``.
circuit_breaker_half_open_requests: Number of successful half-open
probe requests required to close the circuit. Defaults to ``3``.
**extra_options: Additional options passed through as-is.
Note:
LLM function ``config`` supports these runtime prediction options with
this provider:
- ``async``: Execution-mode hint. The provider supports async
prediction only, so omit this key or set it to ``"true"``. Setting
``"false"`` is not supported.
- ``output-mode``: Output ordering for async prediction. Allowed
values are ``"ORDERED"`` and ``"ALLOW_UNORDERED"``. Default: table
config ``table.exec.async-ml-predict.output-mode`` (``"ORDERED"`` by
default).
- ``max-concurrent-operations``: Maximum number of concurrent async
prediction calls. Allowed values are positive integer strings.
Default: table config
``table.exec.async-ml-predict.max-concurrent-operations`` (``"10"``
by default).
- ``timeout``: Timeout for async prediction calls. Allowed values are
duration strings such as ``"30s"`` or ``"3 min"``. Default: table
config ``table.exec.async-ml-predict.timeout`` (``"3 min"`` by
default). This is separate from the Triton HTTP request timeout
constructor option.
Examples::
>>> import pyflink.dataframe as pf
>>>
>>> # Classifier with ARRAY<FLOAT> features and BIGINT class output.
>>> provider = TritonProvider(
... endpoint="<Your Triton endpoint>",
... auth_token="<Your authentication token>",
... model_name="classifier",
... compression="gzip",
... )
>>> pf.set_model_provider(provider)
>>> df = pf.from_records(
... [([5.1, 3.5, 1.4, 0.2],), ...],
... schema=["features"])
>>> result = df.llm.predict(
... "features",
... output_type=pf.DataType.struct({
... "class_id": pf.DataType.int64(),
... }))
>>>
>>> # Stateful conversation model with a fixed Triton sequence.
>>> provider = TritonProvider(
... endpoint="<Your Triton endpoint>",
... auth_token="<Your authentication token>",
... model_name="chatbot_lstm",
... sequence_id="conv-001",
... sequence_start=True,
... sequence_end=False,
... )
>>> pf.set_model_provider(provider)
>>> chat_messages = pf.from_records(
... [("hello",), ...],
... schema=["message_text"])
>>> result = chat_messages.llm.predict(
... "message_text",
... output_type=pf.DataType.struct({
... "bot_response": pf.DataType.string(),
... }))
>>>
>>> # Vector transform model with ARRAY<FLOAT> input and output.
>>> provider = TritonProvider(
... endpoint="<Your Triton endpoint>",
... auth_token="<Your authentication token>",
... model_name="vector-transform",
... flatten_batch_dim=True, # Used when Triton model expects one-dimensional input
... )
>>> pf.set_model_provider(provider)
>>> vector_input = pf.from_records(
... [([0.1, 0.2, 0.3, ...],), ...],
... schema=["features"])
>>> result = vector_input.llm.predict(
... "features",
... output_type=pf.DataType.struct({
... "output_vector": pf.DataType.list(pf.DataType.float32()),
... }))
"""
_OPTION_MAP = {
"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": "compression",
"auth_token": "auth-token",
"custom_headers": "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"
),
}
_CONSTRUCTOR_PARAM_NAMES = frozenset(_OPTION_MAP)
def __init__(
self,
*,
endpoint: str,
model_name: Optional[str] = None,
model_version: str = "latest",
timeout: str = "30s",
flatten_batch_dim: bool = False,
priority: Optional[int] = None,
sequence_id: Optional[str] = None,
sequence_start: bool = False,
sequence_end: bool = False,
compression: Optional[Literal["gzip"]] = None,
auth_token: Optional[str] = None,
custom_headers: Optional[str] = None,
max_retries: int = 0,
retry_initial_backoff: str = "100ms",
retry_max_backoff: str = "30s",
default_value: Optional[_TritonDefaultValue] = None,
health_check_enabled: bool = False,
health_check_interval: str = "30s",
circuit_breaker_enabled: bool = False,
circuit_breaker_failure_threshold: float = 0.5,
circuit_breaker_timeout: str = "60s",
circuit_breaker_half_open_requests: int = 3,
**extra_options: Any,
):
local_vars = {
"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": compression,
"auth_token": auth_token,
"custom_headers": 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
),
}
self._validate_python_params(local_vars)
self._params: Dict[str, Any] = {
key: value for key, value in local_vars.items() if value is not None
}
self._extra_options = extra_options
def provider_identifier(self) -> str:
return "triton"
def model_option_key(self) -> str:
return "model-name"
def to_options(
self,
input_columns: Optional[List[Tuple[str, Any]]] = None,
overrides: Optional[Mapping[str, Any]] = None,
) -> Dict[str, str]:
params, extra_options = _merge_raw_options(
self._params,
self._extra_options,
self._CONSTRUCTOR_PARAM_NAMES,
overrides)
self._validate_python_params(params)
options: Dict[str, str] = {}
for py_key, value in params.items():
java_key = self._OPTION_MAP[py_key]
if py_key == "default_value":
value = self._format_default_value(value)
options[java_key] = self._stringify(value)
for key, value in extra_options.items():
options[str(key)] = self._stringify(value)
return options
@classmethod
def _validate_python_params(cls, params: Mapping[str, Any]) -> None:
endpoint = params.get("endpoint")
if not isinstance(endpoint, str):
raise TypeError(f"endpoint must be a string, got {type(endpoint)}")
compression = params.get("compression")
if compression is not None and compression != "gzip":
raise ValueError(
"compression must be 'gzip' when provided, got "
f"{compression!r}.")
priority = params.get("priority")
if priority is not None and not 0 <= priority <= 255:
raise ValueError("priority must be in range [0, 255].")
max_retries = params.get("max_retries")
if max_retries is not None and max_retries < 0:
raise ValueError("max_retries must be >= 0.")
failure_threshold = params.get("circuit_breaker_failure_threshold")
if (failure_threshold is not None
and not 0.0 < failure_threshold <= 1.0):
raise ValueError(
"circuit_breaker_failure_threshold must be in range "
"(0.0, 1.0].")
half_open_requests = params.get("circuit_breaker_half_open_requests")
if half_open_requests is not None and half_open_requests <= 0:
raise ValueError(
"circuit_breaker_half_open_requests must be positive.")
custom_headers = params.get("custom_headers")
if custom_headers is not None:
if not isinstance(custom_headers, str):
raise TypeError(
f"custom_headers must be a string, got "
f"{type(custom_headers)}")
if "\n" in custom_headers or "\r" in custom_headers:
raise ValueError("custom_headers cannot contain line breaks.")
if "default_value" in params:
cls._format_default_value(params["default_value"])
@staticmethod
def _format_default_value(
default_value: Optional[_TritonDefaultValue],
) -> Optional[str]:
if default_value is None:
return None
if isinstance(default_value, bool):
return "true" if default_value else "false"
if isinstance(default_value, (list, tuple, MappingABC)):
payload = (
dict(default_value)
if isinstance(default_value, MappingABC)
else default_value
)
try:
return json.dumps(payload, allow_nan=False)
except (TypeError, ValueError) as exc:
raise ValueError(
"default_value contains a non-JSON-serializable "
f"element: {exc}") from exc
return str(default_value)
@staticmethod
def _stringify(value: Any) -> str:
if isinstance(value, bool):
return str(value).lower()
return str(value)
[docs]class GenericProvider(ModelProvider):
"""Generic provider for unknown or custom model providers.
Options are passed through to the Java side as-is without any
key name translation.
Args:
name: Model provider identifier recognized by Flink's Java runtime.
**options: Arbitrary key-value options.
Example::
>>> provider = GenericProvider("my-provider", endpoint="https://...",
... **{"api-key": "sk-..."})
"""
def __init__(self, name: str, **options: Any):
self._name = name
self._options = options
def provider_identifier(self) -> str:
return self._name
def to_options(
self,
input_columns: Optional[List[Tuple[str, Any]]] = None,
overrides: Optional[Mapping[str, Any]] = None,
) -> Dict[str, str]:
options = _merge_options(self._options, overrides)
return {str(key): str(value) for key, value in options.items()}
_KNOWN_PROVIDER_CLASSES = {
"dashscope": DashScopeProvider,
"openai-compat": OpenAICompatProvider,
"triton": TritonProvider,
}
def create_provider(name: str, **options: Any) -> ModelProvider:
provider_cls = _KNOWN_PROVIDER_CLASSES.get(name)
if provider_cls is None:
return GenericProvider(name, **options)
return provider_cls(**options)