################################################################################
# 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(
endpoint="https://dashscope.aliyuncs.com/compatible-mode/v1/chat/completions",
api_key="sk-...",
temperature=0.7,
)
pf.set_provider(provider)
# Generic provider for unknown/custom providers
pf.set_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 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,
}
[docs]class Provider(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.
"""
[docs] @abstractmethod
def provider_name(self) -> str:
"""Return the provider identifier recognized by Flink's Java runtime."""
[docs] @abstractmethod
def to_options(
self,
input_columns: Optional[List[Tuple[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.
Returns:
A dictionary whose keys and values are Java-side option strings.
"""
[docs] 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(Provider):
"""Provider for all OpenAI-compatible endpoints (openai-compat).
Covers OpenAI, DeepSeek, Bailian, and any other service that
implements the OpenAI chat/completions or embeddings API.
Args:
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.
task: The model task. Supported values are ``"chat/completions"``,
and ``"embeddings"``.
Required when using Flink AI Model Service.
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 columns. Pass a string
for a single input column, or a list/tuple of strings for
multiple input columns. Multiple content types are positional:
their 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).
Example::
>>> provider = OpenAICompatProvider(
... endpoint="https://dashscope.aliyuncs.com/compatible-mode/v1/chat/completions",
... api_key="sk-...",
... temperature=0.7,
... )
"""
# 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",
}
_SUPPORTED_TASK_CONTENT_TYPES = _OPENAI_COMPAT_TASK_CONTENT_TYPES
def __init__(
self,
*,
endpoint: Optional[str] = None,
api_key: Optional[str] = None,
task: 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,
):
self._validate_service_options(endpoint, api_key, task)
self._validate_context_overflow_action(context_overflow_action)
content_types = self._normalize_content_type(content_type)
resolved_task = self._resolve_task(endpoint, task)
self._validate_content_type(content_types, resolved_task)
content_type_param = self._content_type_param_name(content_types)
content_type_value = ";".join(content_types)
skip_default_content_type = (
content_type_param == "content_type"
and content_type_value == "TEXT"
and "content-types" in extra_options
)
self._params: Dict[str, Any] = {}
# Collect all explicitly named parameters that are not None.
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,
"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,
}
if not skip_default_content_type:
local_vars[content_type_param] = content_type_value
for key, value in local_vars.items():
if value is not None:
self._params[key] = value
self._extra_options = extra_options
[docs] def provider_name(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) -> 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)
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}.")
[docs] def to_options(
self,
input_columns: Optional[List[Tuple[str, Any]]] = None
) -> Dict[str, str]:
options: Dict[str, str] = {}
for py_key, value in self._params.items():
java_key = self._OPTION_MAP[py_key]
options[java_key] = str(value)
# Extra options are passed through as-is.
for key, value in self._extra_options.items():
options[key] = str(value)
if input_columns is None:
return 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_name()!r}.")
# The Java provider uses scalar content-type for single TEXT/IMAGE_URL
# inputs, and content-types for positional multi-input or array content.
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
if len(resolved_content_types) > 1:
aligned_options.pop("content-type", None)
aligned_options["content-types"] = scalar_option_value
else:
resolved_content_types = ["TEXT"]
reported_option_value = "TEXT"
aligned_options["content-type"] = "TEXT"
# An omitted content_type means every input column is TEXT, not that the
# model has only one input column.
if resolved_content_types == ["TEXT"] and len(input_column_names) > 1:
resolved_content_types = ["TEXT"] * len(input_column_names)
reported_option_value = ";".join(resolved_content_types)
aligned_options.pop("content-type", None)
aligned_options["content-types"] = reported_option_value
# Explicit content types are positional and must match predict(*cols).
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_name()!r}. Content type: "
f"{reported_option_value}, input columns: {input_column_names}.")
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):
"""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.
Args:
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.
task: The model task. Supported values are ``"chat/completions"``,
``"embeddings"``, and ``"multimodal-embedding"``.
Required when using Flink AI Model Service.
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 columns. Pass a string
for a single input column, or a list/tuple of strings for
multiple input columns. Multiple content types are positional:
their 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).
Examples:
>>> provider = DashScopeProvider(
... endpoint="https://dashscope.aliyuncs.com/compatible-mode/v1/chat/completions",
... api_key="sk-...",
... model="qwen-plus",
... temperature=0.7,
... )
>>> provider = DashScopeProvider(
... endpoint="https://dashscope.aliyuncs.com/api/v1/services/"
... "embeddings/multimodal-embedding/multimodal-embedding",
... api_key="sk-...",
... model="tongyi-embedding-vision-plus",
... content_type="IMAGE_URL",
... )
>>> pf.set_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", dimension=512)
"""
_SUPPORTED_TASK_CONTENT_TYPES = _DASHSCOPE_TASK_CONTENT_TYPES
[docs] def provider_name(self) -> str:
return "dashscope"
[docs]class TritonProvider(Provider):
"""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.
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_provider(provider)
>>> df = pf.from_records(
... [([5.1, 3.5, 1.4, 0.2],), ...],
... schema=["features"])
>>> result = df.llm.predict(
... "features",
... output_type={"class_id": "BIGINT"})
>>>
>>> # 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_provider(provider)
>>> chat_messages = pf.from_records(
... [("hello",), ...],
... schema=["message_text"])
>>> result = chat_messages.llm.predict(
... "message_text",
... output_type={"bot_response": "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_provider(provider)
>>> vector_input = pf.from_records(
... [([0.1, 0.2, 0.3, ...],), ...],
... schema=["features"])
>>> result = vector_input.llm.predict(
... "features",
... output_type={"output_vector": "ARRAY<FLOAT>"})
"""
_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"
),
}
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,
):
if not isinstance(endpoint, str):
raise TypeError(f"endpoint must be a string, got {type(endpoint)}")
if compression is not None and compression != "gzip":
raise ValueError(
"compression must be 'gzip' when provided, got "
f"{compression!r}.")
if priority is not None and not 0 <= priority <= 255:
raise ValueError("priority must be in range [0, 255].")
if max_retries < 0:
raise ValueError("max_retries must be >= 0.")
if not 0.0 < circuit_breaker_failure_threshold <= 1.0:
raise ValueError(
"circuit_breaker_failure_threshold must be in range "
"(0.0, 1.0].")
if circuit_breaker_half_open_requests <= 0:
raise ValueError(
"circuit_breaker_half_open_requests must be positive.")
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.")
default_value_value = self._format_default_value(default_value)
self._params: Dict[str, 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_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 key, value in local_vars.items():
if value is not None:
self._params[key] = value
self._extra_options = extra_options
[docs] def provider_name(self) -> str:
return "triton"
[docs] def model_option_key(self) -> str:
return "model-name"
[docs] def to_options(
self,
input_columns: Optional[List[Tuple[str, Any]]] = None
) -> Dict[str, str]:
options: Dict[str, str] = {}
for py_key, value in self._params.items():
java_key = self._OPTION_MAP[py_key]
options[java_key] = self._stringify(value)
for key, value in self._extra_options.items():
options[key] = self._stringify(value)
return options
@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(Provider):
"""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: 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
[docs] def provider_name(self) -> str:
return self._name
[docs] def to_options(
self,
input_columns: Optional[List[Tuple[str, Any]]] = None
) -> Dict[str, str]:
return {k: str(v) for k, v in self._options.items()}
_KNOWN_PROVIDER_CLASSES = {
"dashscope": DashScopeProvider,
"openai-compat": OpenAICompatProvider,
"triton": TritonProvider,
}
def create_provider(name: str, **options: Any) -> Provider:
provider_cls = _KNOWN_PROVIDER_CLASSES.get(name)
if provider_cls is None:
return GenericProvider(name, **options)
return provider_cls(**options)