Add model config to prompt (backend and docs change) by chenmoneygithub · Pull Request #19174 · mlflow/mlflow

chenmoneygithub · 2025-12-02T19:51:54Z

🛠 DevTools 🛠

Install mlflow from this PR

# mlflow
pip install git+https://github.com/mlflow/mlflow.git@refs/pull/19174/merge
# mlflow-skinny
pip install git+https://github.com/mlflow/mlflow.git@refs/pull/19174/merge#subdirectory=libs/skinny

For Databricks, use the following command:

%sh curl -LsSf https://raw.githubusercontent.com/mlflow/mlflow/HEAD/dev/install-skinny.sh | sh -s pull/19174/merge

Related Issues/PRs

#xxx

What changes are proposed in this pull request?

Prompt should have model configs, because models use different prompts.

How is this PR tested?

Existing unit/integration tests
New unit/integration tests
Manual tests

Does this PR require documentation update?

Release Notes

Is this a user-facing change?

No. You can skip the rest of this section.
Yes. Give a description of this change to be included in the release notes for MLflow users.

What component(s), interfaces, languages, and integrations does this PR affect?

Components

How should the PR be classified in the release notes? Choose one:

rn/none - No description will be included. The PR will be mentioned only by the PR number in the "Small Bugfixes and Documentation Updates" section
rn/breaking-change - The PR will be mentioned in the "Breaking Changes" section
rn/feature - A new user-facing feature worth mentioning in the release notes
rn/bug-fix - A user-facing bug fix worth mentioning in the release notes
rn/documentation - A user-facing documentation change worth mentioning in the release notes

Should this PR be included in the next patch release?

Yes should be selected for bug fixes, documentation updates, and other small changes. No should be selected for new features and larger changes. If you're unsure about the release classification of this PR, leave this unchecked to let the maintainers decide.

What is a minor/patch release?

Minor release: a release that increments the second part of the version number (e.g., 1.2.0 -> 1.3.0).
Bug fixes, doc updates and new features usually go into minor releases.
Patch release: a release that increments the third part of the version number (e.g., 1.2.0 -> 1.2.1).
Bug fixes and doc updates usually go into patch releases.

Yes (this PR will be cherry-picked and included in the next patch release)
No (this PR will be included in the next minor release)

github-actions · 2025-12-02T19:52:10Z

@chenmoneygithub Thank you for the contribution! Could you fix the following issue(s)?

⚠ DCO check

The DCO check failed. Please sign off your commit(s) by following the instructions here. See https://github.com/mlflow/mlflow/blob/master/CONTRIBUTING.md#sign-your-work for more details.

harupy · 2025-12-03T00:23:51Z

/review

🚀 Review workflow started

✅ Review completed.

Review Output

Review Complete

I've reviewed PR #19174 "Add model config to prompt (backend and docs change)" and found 3 issues:

Type annotation style (line 91-98 in prompt_version.py): Using str | None union syntax instead of Optional[str] for consistency with the rest of the codebase
Comment inconsistency (line 1516 in rest_store.py): Comment says "Convert ModelConfig" but should say "Convert PromptModelConfig" to match the class name
Documentation typo (line 580 in client.py): Extra space in "model-specific configuration"

All comments have been posted to the PR with suggested fixes. The overall code quality is good, with comprehensive tests, proper validation, and clear documentation. The implementation correctly handles model configuration as a mutable property stored in tags.

mlflow/store/_unity_catalog/registry/rest_store.py

mlflow/tracking/client.py

Copilot

Pull request overview

This PR adds model configuration support to the MLflow Prompt Registry, allowing users to store model-specific settings (e.g., model name, temperature, max_tokens) alongside prompt templates. This ensures reproducibility by versioning both prompts and their associated model parameters together.

Key Changes:

Introduces a new PromptModelConfig dataclass with validation for common LLM parameters (temperature, max_tokens, top_p, etc.)
Adds model_config parameter to prompt registration and management APIs
Implements set_prompt_model_config() and delete_prompt_model_config() for mutable model configuration management

Reviewed changes

Copilot reviewed 11 out of 11 changed files in this pull request and generated 6 comments.

Show a summary per file

File	Description
`mlflow/entities/model_registry/prompt_version.py`	Adds `PromptModelConfig` dataclass with validation and `model_config` property to `PromptVersion`
`mlflow/entities/model_registry/__init__.py`	Exports `PromptModelConfig` for public API
`mlflow/genai/prompts/__init__.py`	Adds `model_config` parameter to `register_prompt()` and new functions `set_prompt_model_config()` and `delete_prompt_model_config()`
`mlflow/genai/__init__.py`	Exports new model config management functions
`mlflow/tracking/client.py`	Adds `model_config` parameter support in client's `register_prompt()` method
`mlflow/tracking/_model_registry/fluent.py`	Adds `model_config` parameter support in fluent API
`mlflow/store/model_registry/abstract_store.py`	Implements `model_config` handling in `create_prompt_version()`
`mlflow/store/_unity_catalog/registry/rest_store.py`	Implements `model_config` handling for Unity Catalog backend
`tests/genai/prompts/test_prompts.py`	Comprehensive tests for model config registration, loading, updating, and deletion
`tests/entities/test_prompt.py`	Unit tests for `PromptModelConfig` validation, serialization, and integration
`docs/docs/genai/prompt-registry/index.mdx`	Extensive documentation covering model config usage, validation, and provider-specific parameters

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

mlflow/store/_unity_catalog/registry/rest_store.py

mlflow/entities/model_registry/prompt_version.py

mlflow/genai/prompts/__init__.py

mlflow/tracking/client.py

mlflow/store/model_registry/abstract_store.py

Copilot

Pull request overview

Copilot reviewed 11 out of 11 changed files in this pull request and generated 1 comment.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

mlflow/tracking/client.py

harupy · 2025-12-03T06:17:28Z

/review rewrite PromptModelConfig using pydantic. Flag duplicate imports. Suggest using pytest.mark.parametrize.

🚀 Review workflow started

✅ Review completed.

Review Output

I've completed the PR review for PR #19174. Here's a summary of the issues found and comments added:

Review Summary

5 comments added:

PromptModelConfig implementation (mlflow/entities/model_registry/prompt_version.py:31-197)
- Suggested rewriting PromptModelConfig using Pydantic instead of dataclass for better validation, automatic type coercion, and consistency with the codebase (the file already imports pydantic)
- Pydantic provides cleaner validation with Field constraints and better error messages
Duplicate import Fix example file path #1 (mlflow/store/_unity_catalog/registry/rest_store.py:19-22)
- PromptModelConfig is imported but seems redundant - flagged for review
Duplicate import models.rst typo #2 (mlflow/store/model_registry/abstract_store.py:796-811)
- MODEL_CONFIG_TAG_KEY and PromptModelConfig imported inside function but already at module level (line 16)
- Provided suggestion to remove redundant import
Duplicate import TypeError: __init__() got an unexpected keyword argument 'file' #3 (mlflow/tracking/client.py:658-671)
- Same issue: MODEL_CONFIG_TAG_KEY and PromptModelConfig imported inside function but already at module level (line 49)
- Provided suggestion to remove redundant import
pytest.mark.parametrize suggestion (tests/entities/test_prompt.py:316-350)
- The validation tests check many different error cases separately
- Suggested consolidating into a single parametrized test for better maintainability and clearer failure diagnostics

All comments follow the requested focus areas and include the "🤖 Generated with Claude Code" signature.

mlflow/entities/model_registry/prompt_version.py

mlflow/store/_unity_catalog/registry/rest_store.py

tests/entities/test_prompt.py

mlflow/entities/model_registry/prompt_version.py

harupy · 2025-12-04T12:39:35Z

mlflow/genai/prompts/__init__.py

        registry_api._load_prompt_cached.cache_clear()
+
+
+@experimental(version="3.5.0")


can you fix the version?

I am using 3.8.0, which is the next minor release, let me know if that's the correct approach!

mlflow/store/model_registry/abstract_store.py

tests/genai/prompts/test_prompts.py

harupy · 2025-12-04T12:46:44Z

mlflow/entities/model_registry/prompt_version.py

 PromptVersionTag = ModelVersionTag


+MODEL_CONFIG_TAG_KEY = "_mlflow_prompt_model_config"


Lets rename the key to follow mlflow/utils/mlflow_tags.py, and move it in this file

harupy · 2025-12-04T14:48:24Z

mlflow/entities/model_registry/prompt_version.py

+from pydantic import BaseModel, Field
+
+
+class PromptModelConfig(BaseModel):


Do we really need this class?

Functionality-wise a wildcard dict can do the job, but I don't like that approach, which has no validation and format constraint. Also it's hard to document it in the docstring: model_config: dict[str, Any] is confusing.

harupy · 2025-12-05T03:18:59Z

tests/entities/test_prompt.py

+    # Test valid values
+    PromptModelConfig(temperature=0.0, max_tokens=1000, top_p=0.9, top_k=50)


Suggested change

# Test valid values

PromptModelConfig(temperature=0.0, max_tokens=1000, top_p=0.9, top_k=50)

def test_...

# Test valid values

PromptModelConfig(temperature=0.0, max_tokens=1000, top_p=0.9, top_k=50)

can we separate this test?

this line is actually redundant, and has been covered by other test cases, removed!

harupy · 2025-12-05T03:19:20Z

mlflow/utils/mlflow_tags.py

 LATEST_CHECKPOINT_ARTIFACT_TAG_KEY = "mlflow.latest_checkpoint_artifact"

+# For prompt model configuration
+MLFLOW_PROMPT_MODEL_CONFIG = "_mlflow_prompt_model_config"


can we use the same format as other tags?

oops, missed this one, thanks for catching!

harupy · 2025-12-05T03:53:47Z

mlflow/entities/model_registry/prompt_version.py

+    frequency_penalty: float | None = None
+    presence_penalty: float | None = None
+    stop_sequences: list[str] | None = None
+    extra_params: dict[str, Any] = Field(default_factory=dict)


can extra_params contain the known fields?

yes good question - it's allowed and would work fine, but not the good practice because it can break the UI rendering or so.

harupy

Left a few more comments/questions, otherwise LGTM :)

harupy

Left a few more comments, otherwise LGTM!

harupy · 2025-12-05T06:52:33Z

mlflow/entities/model_registry/prompt_version.py

+        known_fields = {
+            "model_name",
+            "temperature",
+            "max_tokens",
+            "top_p",
+            "top_k",
+            "frequency_penalty",
+            "presence_penalty",
+            "stop_sequences",
+        }


Any pydantic model field we can leverage to de-hardcode these?

yes good call!

github-actions · 2025-12-08T18:32:52Z

Documentation preview for cc4fc4f is available at:

https://pr-19174--mlflow-docs-preview.netlify.app/docs/latest/

Changed Pages (1)

genai/prompt-registry (modified)

More info

Ignore this comment if this PR does not change the documentation.
The preview is updated when a new commit is pushed to this PR.
This comment was created by this workflow run.
The documentation was built by this workflow run.

TomeHirata · 2025-12-10T09:04:06Z

mlflow/entities/model_registry/prompt_version.py

+        )
+    """
+
+    model_name: str | None = None


@chenmoneygithub Sorry for the late review! How do you feel about making model_name required? I think it's unlikely that users want to configure other parameters without configuring model name. Also as we discussed, can we add provider name?

TomeHirata · 2025-12-10T09:08:50Z

mlflow/utils/mlflow_tags.py

 LATEST_CHECKPOINT_ARTIFACT_TAG_KEY = "mlflow.latest_checkpoint_artifact"

+# For prompt model configuration
+MLFLOW_PROMPT_MODEL_CONFIG = "mlflow.prompt.modelConfig"


@chenmoneygithub Can we replace "." with "_" and move this config to mlflow/prompt/constants.py? Managed prompt registry does not allow to have "." in the tag key name.

chenmoneygithub added 3 commits December 1, 2025 19:40

init

9585b3a

Merge branch 'master' into prompt-model-config

39b5c87

init wrapping

8176050

chenmoneygithub changed the title ~~Add model config to prompt~~ [WIP] Add model config to prompt (backend and docs change) Dec 2, 2025

github-actions bot added area/prompts MLflow Prompt Registry and Optimization rn/feature Mention under Features in Changelogs. labels Dec 2, 2025

chenmoneygithub added 2 commits December 2, 2025 14:18

clean up various things

1e2f384

Merge branch 'master' into prompt-model-config

9e6b4a3

chenmoneygithub changed the title ~~[WIP] Add model config to prompt (backend and docs change)~~ Add model config to prompt (backend and docs change) Dec 2, 2025

chenmoneygithub requested review from B-Step62, TomeHirata and harupy December 2, 2025 23:57

harupy requested a review from Copilot December 3, 2025 00:23

Copilot started reviewing on behalf of harupy December 3, 2025 00:24 View session

github-actions bot reviewed Dec 3, 2025

View reviewed changes

mlflow/store/_unity_catalog/registry/rest_store.py Show resolved Hide resolved

github-actions bot reviewed Dec 3, 2025

View reviewed changes

mlflow/tracking/client.py Outdated Show resolved Hide resolved

mlflow deleted a comment from github-actions bot Dec 3, 2025

Copilot finished reviewing on behalf of harupy December 3, 2025 00:27

Copilot AI reviewed Dec 3, 2025

View reviewed changes

fix comments

c10ea5b

chenmoneygithub requested a review from Copilot December 3, 2025 01:49

Copilot started reviewing on behalf of chenmoneygithub December 3, 2025 01:50 View session

Copilot finished reviewing on behalf of chenmoneygithub December 3, 2025 01:51

Copilot AI reviewed Dec 3, 2025

View reviewed changes

mlflow/tracking/client.py Outdated Show resolved Hide resolved

fix doc build

e53502b

github-actions bot reviewed Dec 3, 2025

View reviewed changes

mlflow/entities/model_registry/prompt_version.py Outdated Show resolved Hide resolved

github-actions bot reviewed Dec 3, 2025

View reviewed changes

mlflow/store/_unity_catalog/registry/rest_store.py Show resolved Hide resolved

github-actions bot reviewed Dec 3, 2025

View reviewed changes

tests/entities/test_prompt.py Outdated Show resolved Hide resolved

chenmoneygithub added 2 commits December 3, 2025 11:30

move to pydantic

f37f238

fix docs

89a8292

harupy reviewed Dec 4, 2025

View reviewed changes

mlflow/entities/model_registry/prompt_version.py Show resolved Hide resolved

harupy reviewed Dec 4, 2025

View reviewed changes

mlflow/entities/model_registry/prompt_version.py Show resolved Hide resolved

harupy reviewed Dec 4, 2025

View reviewed changes

mlflow/store/model_registry/abstract_store.py Outdated Show resolved Hide resolved

harupy reviewed Dec 4, 2025

View reviewed changes

tests/genai/prompts/test_prompts.py Outdated Show resolved Hide resolved

harupy reviewed Dec 4, 2025

View reviewed changes

chenmoneygithub added 4 commits December 4, 2025 10:32

comments

4c1b056

merge master

9c4988e

nit

c2eb02e

fix tests

a079f51

harupy reviewed Dec 5, 2025

View reviewed changes

harupy approved these changes Dec 5, 2025

View reviewed changes

harupy reviewed Dec 5, 2025

View reviewed changes

chenmoneygithub added 2 commits December 5, 2025 11:15

comments

90cbdbc

fix docs build

cc4fc4f

chenmoneygithub added this pull request to the merge queue Dec 8, 2025

Merged via the queue into mlflow:master with commit 339fc03 Dec 8, 2025
47 checks passed

chenmoneygithub deleted the prompt-model-config branch December 8, 2025 19:53

chenmoneygithub mentioned this pull request Dec 9, 2025

Add prompt model config to the prompt registry UI #19279

Merged

29 tasks

TomeHirata reviewed Dec 10, 2025

View reviewed changes

		registry_api._load_prompt_cached.cache_clear()


		@experimental(version="3.5.0")

		PromptVersionTag = ModelVersionTag


		MODEL_CONFIG_TAG_KEY = "_mlflow_prompt_model_config"

		from pydantic import BaseModel, Field


		class PromptModelConfig(BaseModel):

		# Test valid values
		PromptModelConfig(temperature=0.0, max_tokens=1000, top_p=0.9, top_k=50)

Conversation

chenmoneygithub commented Dec 2, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Install mlflow from this PR

Related Issues/PRs

What changes are proposed in this pull request?

How is this PR tested?

Does this PR require documentation update?

Release Notes

Is this a user-facing change?

What component(s), interfaces, languages, and integrations does this PR affect?

How should the PR be classified in the release notes? Choose one:

Should this PR be included in the next patch release?

Uh oh!

github-actions bot commented Dec 2, 2025

⚠ DCO check

Uh oh!

harupy commented Dec 3, 2025 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Review Complete

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

harupy commented Dec 3, 2025 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Review Summary

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

harupy Dec 5, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

harupy left a comment

Choose a reason for hiding this comment

Uh oh!

harupy left a comment

chenmoneygithub commented Dec 2, 2025 •

edited

Loading

harupy commented Dec 3, 2025 •

edited by github-actions bot

Loading

harupy commented Dec 3, 2025 •

edited by github-actions bot

Loading

harupy Dec 5, 2025 •

edited

Loading