Prompt Registry in Databricks

What is MLflow Prompt Registry?

MLflow Prompt Registry is a powerful tool that streamlines prompt engineering and management in your Generative AI (GenAI) applications. It enables you to version, track, and reuse prompts across your organization, helping maintain consistency and improving collaboration in prompt development.

Key Features

• Reusability - Store and manage prompts in a centralized registry and reuse them across multiple applications.
• Version Control - Track the evolution of your prompts with Git-inspired commit-based versioning and side-by-side comparison of prompt versions with diff highlighting.
• Aliasing - Build robust yet flexible deployment pipelines for prompts using aliases, allowing you to isolate prompt versions from main application code and perform tasks such as A/B testing and roll-backs with ease.
• Lineage - Seamlessly integrate with MLflow's existing features such as model tracking and evaluation for end-to-end GenAI lifecycle management.
• Collaboration - Share prompts across your organization with a centralized registry, enabling teams to build upon each other's work.

note

The MLflow Prompt Registry integration with Databricks Unity Catalog is coming soon.

It will be a centralized system for managing, versioning, and governing prompt templates used in your GenAI applications. It will be deeply integrated with Databricks Unity Catalog, providing:

• Centralized Prompt Management: Store and organize prompt templates in a governed, searchable registry
• Version Control: Track changes to prompts over time with full lineage and rollback capabilities
• Access Control: Leverage Unity Catalog's permissions system to control who can view, edit, and use specific prompts
• Collaboration: Enable teams to share and collaborate on prompt development
• Governance: Apply data governance policies to prompt templates, including classification and compliance requirements

Integration with Unity Catalog

The Prompt Registry will leverage Unity Catalog's governance framework to provide:

• Fine-grained Access Control: Control access to prompts at the individual, team, or organization level
• Data Lineage: Track how prompts are used across different applications and experiments
• Audit Trails: Monitor who accessed or modified prompts and when
• Classification and Tagging: Organize prompts with metadata and apply governance policies
• Cross-workspace Sharing: Share prompts securely across different Databricks workspaces

Future Capabilities

When available, the Prompt Registry will support:

• Template Management: Create, edit, and version prompt templates with variable substitution
• Testing and Validation: Test prompts against evaluation datasets before deployment
• A/B Testing: Compare different prompt versions to optimize performance
• Integration with MLflow Tracking: Automatically link prompt versions to experiments and model runs
• API Access: Programmatic access to prompts for use in automated workflows

Stay Updated

For the latest updates on Prompt Registry availability and features, monitor the MLflow documentation and Databricks release notes.

Prompt Registry in OSS MLflow

1. Create a Prompt

UI

1. Run mlflow ui in your terminal to start the MLflow UI.
2. Navigate to the Prompts tab in the MLflow UI.
3. Click on the Create Prompt button.
4. Fill in the prompt details such as name, prompt template text, and commit message (optional).
5. Click Create to register the prompt.

note

Prompt template text can contain variables in {{variable}} format. These variables can be filled with dynamic content when using the prompt in your GenAI application. MLflow also provides the to_single_brace_format() API to convert templates into single brace format for frameworks like LangChain or LlamaIndex that require single brace interpolation.

Python

To create a new prompt using the Python API, use mlflow.genai.register_prompt() API:

import mlflow

# Use double curly braces for variables in the template
initial_template = """\
Summarize content you are provided with in {{ num_sentences }} sentences.

Sentences: {{ sentences }}
"""

# Register a new prompt
prompt = mlflow.genai.register_prompt(
    name="summarization-prompt",
    template=initial_template,
    # Optional: Provide a commit message to describe the changes
    commit_message="Initial commit",
    # Optional: Set tags applies to the prompt (across versions)
    tags={
        "author": "author@example.com",
        "task": "summarization",
        "language": "en",
    },
)

# The prompt object contains information about the registered prompt
print(f"Created prompt '{prompt.name}' (version {prompt.version})")

This creates a new prompt with the specified template text and metadata. The prompt is now available in the MLflow UI for further management.

2. Update the Prompt with a New Version

UI

1. The previous step leads to the created prompt page. (If you closed the page, navigate to the Prompts tab in the MLflow UI and click on the prompt name.)
2. Click on the Create prompt Version button.
3. The popup dialog is pre-filled with the existing prompt text. Modify the prompt as you wish.
4. Click Create to register the new version.

Python

To update an existing prompt with a new version, use the mlflow.genai.register_prompt() API with the existing prompt name:

import mlflow

new_template = """\
You are an expert summarizer. Condense the following content into exactly {{ num_sentences }} clear and informative sentences that capture the key points.

Sentences: {{ sentences }}

Your summary should:
- Contain exactly {{ num_sentences }} sentences
- Include only the most important information
- Be written in a neutral, objective tone
- Maintain the same level of formality as the original text
"""

# Register a new version of an existing prompt
updated_prompt = mlflow.genai.register_prompt(
    name="summarization-prompt",  # Specify the existing prompt name
    template=new_template,
    commit_message="Improvement",
    tags={
        "author": "author@example.com",
    },
)

3. Compare the Prompt Versions

Once you have multiple versions of a prompt, you can compare them to understand the changes between versions. To compare prompt versions in the MLflow UI, click on the Compare tab in the prompt details page:

4. Load and Use the Prompt

To use a prompt in your GenAI application, you can load it with the mlflow.genai.load_prompt() API and fill in the variables using the mlflow.entities.Prompt.format() method of the prompt object:

import mlflow
import openai

target_text = """
MLflow is an open source platform for managing the end-to-end machine learning lifecycle.
It tackles four primary functions in the ML lifecycle: Tracking experiments, packaging ML
code for reuse, managing and deploying models, and providing a central model registry.
MLflow currently offers these functions as four components: MLflow Tracking,
MLflow Projects, MLflow Models, and MLflow Registry.
"""

# Load the prompt
prompt = mlflow.genai.load_prompt("prompts:/summarization-prompt/2")

# Use the prompt with an LLM
client = openai.OpenAI()
response = client.chat.completions.create(
    messages=[
        {
            "role": "user",
            "content": prompt.format(num_sentences=1, sentences=target_text),
        }
    ],
    model="gpt-4o-mini",
)

print(response.choices[0].message.content)

5. Search Prompts

You can discover prompts by name, tag or other registry fields:

import mlflow

# Fluent API: returns a flat list of all matching prompts
prompts = mlflow.genai.search_prompts(filter_string="task='summarization'")
print(f"Found {len(prompts)} prompts")

# For pagination control, use the client API:
from mlflow.tracking import MlflowClient

client = MlflowClient()
all_prompts = []
token = None
while True:
    page = client.search_prompts(
        filter_string="task='summarization'",
        max_results=50,
        page_token=token,
    )
    all_prompts.extend(page)
    token = page.token
    if not token:
        break
print(f"Total prompts across pages: {len(all_prompts)}")

Prompt Object

The Prompt object is the core entity in MLflow Prompt Registry. It represents a versioned template text that can contain variables for dynamic content.

Key attributes of a Prompt object:

• Name: A unique identifier for the prompt.
• Template: The text of the prompt, which can include variables in {{variable}} format.
• Version: A sequential number representing the revision of the prompt.
• Commit Message: A description of the changes made in the prompt version, similar to Git commit messages.
• Version Metadata: Optional key-value pairs for adding metadata to the prompt version. For example, you may use this for tracking the author of the prompt version.
• Tags: Optional key-value pairs assigned at the prompt level (across versions) for categorization and filtering. For example, you may add tags for project name, language, etc, which apply to all versions of the prompt.
• Alias: An mutable named reference to the prompt. For example, you can create an alias named production to refer to the version used in your production system. See Aliases for more details.

Log Prompts with Models

Prompts are often used as a part of GenAI applications. Managing the association between prompts and models is crucial for tracking the evolution of models and ensuring consistency across different environments. MLflow Prompt Registry is integrated with MLflow's model tracking capability, allowing you to track which prompts (and versions) are used by your models and applications.

Basic Usage

To log a model with associated prompts, use the prompts parameter in the log_model method. The prompts parameter accepts a list of prompt URLs or prompt objects that are associated with the model. The associated prompts are displayed in the MLflow UI for the model run.

import mlflow

with mlflow.start_run():
    mlflow.<flavor>.log_model(
        model,
        ...
        # Specify a list of prompt URLs or prompt objects.
        prompts=["prompts:/summarization-prompt/2"]
    )

warning

The prompts parameter for associating prompts with models is only supported for GenAI flavors such as OpenAI, LangChain, LlamaIndex, DSPy, etc.

Example 1: Logging Prompts with LangChain

1. Create a prompt

If you haven't already created a prompt, follow this step to create a new prompt.

2. Define a Chain using the registered prompts

from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI

# Load registered prompt
prompt = mlflow.genai.load_prompt("prompts:/summarization-prompt/2")

# Create LangChain prompt object
langchain_prompt = ChatPromptTemplate.from_messages(
    [
        (
            # IMPORTANT: Convert prompt template from double to single curly braces format
            "system",
            prompt.to_single_brace_format(),
        ),
        ("placeholder", "{messages}"),
    ]
)

# Define the LangChain chain
llm = ChatOpenAI()
chain = langchain_prompt | llm

# Invoke the chain
response = chain.invoke({"num_sentences": 1, "sentences": "This is a test sentence."})
print(response)

3. Log the Chain to MLflow

Then log the chain to MLflow and specify the prompt URL in the prompts parameter:

with mlflow.start_run(run_name="summarizer-model"):
    mlflow.langchain.log_model(
        chain, name="model", prompts=["prompts:/summarization-prompt/2"]
    )

Now you can view the associated prompts to the model in MLflow UI:

Moreover, you can view the list of models (runs) that use a specific prompt in the prompt details page:

Example 2: Automatic Prompt Logging with Models-from-Code

Models-from-Code is a feature that allows you to define and log models in code. Logging a model with code brings several benefits, such as portability, readability, avoiding serialization, and more.

Combining with MLflow Prompt Registry, the feature unlocks even more flexibility to manage prompt versions. Notably, if your model code uses a prompt from MLflow Prompt Registry, MLflow automatically logs it with the model for you.

In the following example, we use LangGraph to define a very simple chat bot using the registered prompt.

1. Create a prompt

import mlflow

# Register a new prompt
prompt = mlflow.genai.register_prompt(
    name="chat-prompt",
    template="You are an expert in programming. Please answer the user's question about programming.",
)

2. Define a Graph using the registered prompt

Create a Python script chatbot.py with the following content.

tip

If you are using Jupyter notebook, you can uncomment the %writefile magic command and run the following code in a cell to generate the script.

# %%writefile chatbot.py

import mlflow
from typing import Annotated
from typing_extensions import TypedDict

from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages


class State(TypedDict):
    messages: list


llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.1)
system_prompt = mlflow.genai.load_prompt("prompts:/chat-prompt/1")


def add_system_message(state: State):
    return {
        "messages": [
            {
                "role": "system",
                "content": system_prompt.to_single_brace_format(),
            },
            *state["messages"],
        ]
    }


def chatbot(state: State):
    return {"messages": [llm.invoke(state["messages"])]}


graph_builder = StateGraph(State)
graph_builder.add_node("add_system_message", add_system_message)
graph_builder.add_node("chatbot", chatbot)
graph_builder.add_edge(START, "add_system_message")
graph_builder.add_edge("add_system_message", "chatbot")
graph_builder.add_edge("chatbot", END)

graph = graph_builder.compile()

mlflow.models.set_model(graph)

3. Log the Graph to MLflow

Specify the file path to the script in the model parameter:

with mlflow.start_run():
    model_info = mlflow.langchain.log_model(
        lc_model="./chatbot.py",
        name="graph",
    )

We didn't specify the prompts parameter this time, but MLflow automatically logs the prompt loaded within the script to the logged model. Now you can view the associated prompt in MLflow UI:

4. Load the graph back and invoke

Finally, let's load the graph back and invoke it to see the chatbot in action.

# Enable MLflow tracing for LangChain to view the prompt passed to LLM.
mlflow.langchain.autolog()

# Load the graph
graph = mlflow.langchain.load_model(model_info.model_uri)

graph.invoke(
    {
        "messages": [
            {
                "role": "user",
                "content": "What is the difference between multi-threading and multi-processing?",
            }
        ]
    }
)

FAQ

Q: How do I delete a prompt version?

A: You can delete a prompt version using the MLflow UI or Python API:

import mlflow

# Delete a prompt version
client = mlflow.MlflowClient()
client.delete_prompt_version("summarization-prompt", version=2)

To avoid accidental deletion, you can only delete one version at a time via API. If you delete the all versions of a prompt, the prompt itself will be deleted.

Q: Can I update the prompt template of an existing prompt version?

A: No, prompt versions are immutable once created. To update a prompt, create a new version with the desired changes.

Q: Can I use prompt templates with frameworks like LangChain or LlamaIndex?

A: Yes, you can load prompts from MLflow and use them with any framework. For example, the following example demonstrates how to use a prompt registered in MLflow with LangChain. Also refer to Logging Prompts with LangChain for more details.

import mlflow
from langchain.prompts import PromptTemplate

# Load prompt from MLflow
prompt = mlflow.genai.load_prompt("question_answering")

# Convert the prompt to single brace format for LangChain (MLflow uses double braces),
# using the `to_single_brace_format` method.
langchain_prompt = PromptTemplate.from_template(prompt.to_single_brace_format())
print(langchain_prompt.input_variables)
# Output: ['num_sentences', 'sentences']

Q: Is Prompt Registry integrated with the Prompt Engineering UI?

A. Direct integration between the Prompt Registry and the Prompt Engineering UI is coming soon. In the meantime, you can iterate on prompt template in the Prompt Engineering UI and register the final version in the Prompt Registry by manually copying the prompt template.

What's Next

Manage prompts using other MLflow capabilities:

• Store prompt templates as parameters in MLflow experiments
• Version prompts alongside your application code in Git
• Use MLflow's version tracking features to link specific prompt versions to application versions
• Leverage MLflow's tagging system to organize and categorize prompts

See the Version Tracking section for current best practices on some of these other approaches.