auroflow commented on code in PR #29478:
URL: https://github.com/apache/flink/pull/29478#discussion_r4279789490
##########
flink-python/pyflink/dataframe/llm.py:
##########
@@ -610,3 +666,189 @@ def _resolve_provider(name: Optional[str] = None) ->
ModelProvider:
"Multiple model providers are registered; call
set_default_model_provider() "
"or select a provider explicitly"
)
+
+
+def _override_options(
+ options: Mapping[str, str], overrides: Mapping[str, object],
+) -> Dict[str, str]:
+ result = dict(options)
+ for key, value in overrides.items():
+ if not isinstance(key, str):
+ raise TypeError("Option keys must be strings")
+ if value is None:
+ result.pop(key, None)
+ elif isinstance(value, str):
+ result[key] = value
+ else:
+ raise TypeError(f"Option {key!r} must have a string value or None")
+ return result
+
+
+def _provider_options(
+ provider: ModelProvider, overrides: Mapping[str, object],
+) -> Dict[str, str]:
+ configured = provider.to_options()
+ if not isinstance(configured, MappingABC):
+ raise TypeError("ModelProvider.to_options() must return a mapping of
strings to strings")
+ options = _merge_options(configured, {})
+ if type(provider) not in TYPED_PROVIDERS:
+ return _override_options(options, overrides)
+ if not overrides:
+ return options
+
+ typed_provider = cast(Union[OpenAIProvider, TritonProvider], provider)
+ parameters = dict(overrides)
+ raw_supplied = "options" in parameters
+ raw = parameters.pop("options", None)
+ if raw is not None and not isinstance(raw, MappingABC):
+ raise TypeError("options must be a mapping of strings to strings or
None")
+
+ for name in parameters:
+ if name not in typed_provider._OPTION_MAP:
+ raise TypeError(f"Unexpected provider parameter {name!r}")
+ constructor: Callable[..., ModelProvider] = type(typed_provider)
+ validated = constructor(**{**typed_provider._parameters,
**parameters}).to_options()
+ typed = {typed_provider._OPTION_MAP[name]:
validated.get(typed_provider._OPTION_MAP[name])
+ for name in parameters}
+ if raw_supplied and raw is None:
+ options = dict(typed_provider._typed_options)
+ elif raw is not None:
+ # Validate raw values even if an equal typed override supplies the
same key.
+ _override_options({}, raw)
+ for key in typed.keys() & raw.keys():
+ if typed[key] != raw[key]:
+ raise ValueError(f"conflicting values for option {key!r}")
+
+ options = _override_options(options, typed)
+ if raw is not None:
+ options = _override_options(options, raw)
+ typed_provider._validate_required_options(options)
+ return options
+
+
+def _prediction_output_columns(
+ output_type: Optional[Union[str, DataType]],
+) -> List[Tuple[str, TableDataType]]:
+ if output_type is None:
+ return [("output", DataTypes.STRING())]
+ if isinstance(output_type, str):
+ return [("output",
DataType._from_sql(output_type)._to_table_data_type())]
+ if not isinstance(output_type, DataType):
+ raise TypeError("output_type must be a SQL type string or a
pyflink.dataframe.DataType")
+ table_type = output_type._to_table_data_type()
+ if not isinstance(table_type, RowType):
+ return [("output", table_type)]
+ if not table_type.fields:
+ raise ValueError("Struct output_type must contain at least one field")
+ return [(field.name, field.data_type) for field in table_type.fields]
+
+
+@PublicEvolving()
+class LLMAccessor:
+ """Model inference operations accessed through :attr:`DataFrame.llm`.
+
+ .. versionadded:: 2.4.0
+ """
+
+ def __init__(self, dataframe: DataFrame):
+ self._df = dataframe
+
+ @PublicEvolving()
+ def predict(
+ self, *input_cols: str, provider: Optional[str] = None, model:
Optional[str] = None,
+ output_type: Optional[Union[str, DataType]] = None,
+ config: Optional[Mapping[str, str]] = None, **kwargs: object,
+ ) -> DataFrame:
+ """
+ Predict each row and append the model's output columns.
+
+ The selected provider validates input and output types and the number
of columns.
+
+ :param input_cols: One or more existing column names, in model input
order.
+ :param provider: Registered provider name. If omitted, use the
selected default
+ or the sole registered provider.
+ :param model: Override the provider's model name for this call.
+ :param output_type: SQL type string or :class:`DataType` declaring the
provider's
+ native output schema. Defaults to ``STRING``. A
scalar or container
+ type appends one column named ``output``.
+ ``DataType.struct(...)`` appends its top-level
fields as
+ named columns, preserving nested types. Output
names must
+ be unique and must not occur in the input
DataFrame.
+ :param config: Runtime configuration as string keys and values.
+ :param kwargs: Provider constructor overrides for this call. Built-ins
accept
+ their typed Python parameter names and merge
``options`` per key.
+ Custom providers accept raw option names and string
values.
+ ``None`` removes an inherited option; ``options=None``
clears
+ inherited raw options. The registered provider stays
unchanged.
+ :return: A new DataFrame containing the input and prediction columns.
+
+ Examples::
+
+ >>> import pyflink.dataframe as pf
+ >>> pf.set_model_provider(pf.OpenAIProvider(
+ ... "https://api.openai.com/v1/chat/completions";,
api_key="key", retry_num=3))
+ >>> df = pf.from_records([{"question": "What is Apache Flink?"}])
+ >>> answer = df.llm.predict("question", model="gpt-6-luna")
+
+ >>> # Override the provider's retry_num parameter for this call.
+ >>> answer = df.llm.predict(
+ ... "question", model="gpt-6-luna", retry_num=0)
+
+ >>> # Name the prediction column "answer" instead of the default
"output".
+ >>> answer = df.llm.predict(
+ ... "question", model="gpt-6-luna",
+ ... output_type=pf.DataType.struct({"answer":
pf.DataType.string()}))
+
+ .. versionadded:: 2.4.0
+ """
+ if not input_cols:
+ raise ValueError("predict() requires at least one input column")
+ for name in input_cols:
+ _validate_name(name, "input column")
+ if len(set(input_cols)) != len(input_cols):
+ raise ValueError("predict() input columns must be unique")
+ if model is not None:
+ _validate_name(model, "model")
+ if config is not None and not isinstance(config, MappingABC):
+ raise TypeError("config must be a mapping of strings to strings")
+ runtime_options = _merge_options(config, {})
+ selected = _resolve_provider(provider)
+ identifier = _validate_name(selected.provider_identifier(), "provider
identifier")
+ options = _provider_options(selected, kwargs)
+ if model is not None:
+ key = _validate_name(selected.model_option_key(), "model option
key")
+ options[key] = model
+ output_columns = _prediction_output_columns(output_type)
+ output_names = [name for name, _ in output_columns]
+ for name in output_names:
+ _validate_name(name, "output field name")
+ if len(set(output_names)) != len(output_names):
+ raise ValueError("Struct output_type contains duplicate field
names")
+
+ resolved = self._df._table.get_resolved_schema()
+ names = list(resolved.get_column_names())
+ missing = [name for name in input_cols if name not in names]
+ if missing:
+ raise ValueError(f"Input columns not found in DataFrame:
{missing}")
+ collisions = sorted(set(names).intersection(output_names))
+ if collisions:
+ raise ValueError(
+ f"Prediction output columns conflict with existing columns:
{collisions}")
+
+ input_schema = Schema.new_builder()
+ column_indexes = {name: index for index, name in enumerate(names)}
+ j_column_types = resolved._j_resolved_schema.getColumnDataTypes()
+ for name in input_cols:
Review Comment:
Agreed. I have created a separate JIRA issue
[FLINK-41001](https://issues.apache.org/jira/browse/FLINK-41001), since this
affects SQL and Table API as well.
--
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]
