[FR] Support completions for `pyfunc` predict in OpenAI

[santiagxf]

Willingness to contribute

Yes. I can contribute this feature independently.

Proposal Summary

OpenAI flavor add PyFunc into the model for the tasks chat.completions and embeddings. I would like to extend this support to completions.

Motivation

What is the use case for this feature?

OpenAI for NLU supports 3 tasks: embeddings, chat.completions, and completions. Not all of them were implemented in MLflow flavor.

Why is this use case valuable to support for MLflow users in general?

Users willing to use completions can't use pyfunc implementation in MLflow.

Why is this use case valuable to support for your project(s) or organization?

We have customers looking to use the full set of features of OpenAI.

Why is it currently difficult to achieve this use case?

They are not supported as indicated at:

https://github.com/mlflow/mlflow/blob/50a4b658968db3a57f106e215b7ebb6b6f27ea78/mlflow/openai/__init__.py#L75C2-L75C24

Details

I would like to expand the support of MLflow pyfunc to include all the NLU tasks supported by OpenAI which should help users to adopt MLflow model format for these tasks. Today, this limitation imposes restrictions.

What component(s) does this bug affect?

• [ ] area/artifacts: Artifact stores and artifact logging
• [ ] area/build: Build and test infrastructure for MLflow
• [ ] area/docs: MLflow documentation pages
• [ ] area/examples: Example code
• [ ] area/gateway: AI Gateway service, Gateway client APIs, third-party Gateway integrations
• [ ] area/model-registry: Model Registry service, APIs, and the fluent client calls for Model Registry
• [X] area/models: MLmodel format, model serialization/deserialization, flavors
• [ ] area/recipes: Recipes, Recipe APIs, Recipe configs, Recipe Templates
• [ ] area/projects: MLproject format, project running backends
• [ ] area/scoring: MLflow Model server, model deployment tools, Spark UDFs
• [ ] area/server-infra: MLflow Tracking server backend
• [ ] area/tracking: Tracking Service, tracking client APIs, autologging

What interface(s) does this bug affect?

• [ ] area/uiux: Front-end, user experience, plotting, JavaScript, JavaScript dev server
• [ ] area/docker: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
• [ ] area/sqlalchemy: Use of SQLAlchemy in the Tracking Service or Model Registry
• [ ] area/windows: Windows support

What language(s) does this bug affect?

• [ ] language/r: R APIs and clients
• [ ] language/java: Java APIs and clients
• [ ] language/new: Proposals for new client languages

What integration(s) does this bug affect?

• [ ] integrations/azure: Azure and Azure ML integrations
• [ ] integrations/sagemaker: SageMaker integrations
• [ ] integrations/databricks: Databricks integrations

[BenWilson2]

Sounds great! Let us know when it's ready :)

[mlflow-automation]

@mlflow/mlflow-team Please assign a maintainer and start triaging this issue.