auroflow commented on code in PR #29454: URL: https://github.com/apache/flink/pull/29454#discussion_r4253974446
########## flink-python/pyflink/dataframe/llm.py: ########## @@ -0,0 +1,585 @@ +################################################################################ +# 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 _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``. + + 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__( Review Comment: Makes sense, I've added an optional `options` to both `TritonProvider` and `OpenAIProvider`. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
