feat: Add MLflow Prompt Registry provider

[williamcaban]

MLflow Prompt Registry Provider

Summary

This PR adds a new remote MLflow provider for the Prompts API, enabling centralized prompt management and versioning using MLflow's Prompt Registry (MLflow 3.4+).

What's New

Remote Provider: remote::mlflow

A production-ready provider that integrates Llama Stack's Prompts API with MLflow's centralized prompt registry, supporting:

• Version Control: Immutable prompt versioning with full history
• Default Version Management: Easy version switching via aliases
• Auto Variable Extraction: Automatic detection of {{ variable }} placeholders
• Centralized Storage: Team collaboration via shared MLflow server
• Metadata Preservation: Llama Stack metadata stored as MLflow tags

Quick Start

1. Configure Llama Stack

Basic configuration with SQLite (default):

prompts:
  - provider_id: reference-prompts
    provider_type: inline::reference
    config:
      run_config:
        storage:
          stores:
            prompts:
              type: sqlite
              db_path: ./prompts.db

With PostgreSQL:

prompts:
  - provider_id: postgres-prompts
    provider_type: inline::reference
    config:
      run_config:
        storage:
          stores:
            prompts:
              type: postgres
              url: postgresql://user:pass@localhost/llama_stack

2. Use the Prompts API

from llama_stack_client import LlamaStackClient

client = LlamaStackClient(base_url="http://localhost:5000")

# Create a prompt
prompt = client.prompts.create(
    prompt="Summarize the following text in {{ num_sentences }} sentences:\n\n{{ text }}",
    variables=["num_sentences", "text"]
)
print(f"Created prompt: {prompt.prompt_id} (v{prompt.version})")

# Retrieve prompt
retrieved = client.prompts.get(prompt_id=prompt.prompt_id)
print(f"Retrieved: {retrieved.prompt}")

# Update prompt (creates version 2)
updated = client.prompts.update(
    prompt_id=prompt.prompt_id,
    prompt="Summarize in exactly {{ num_sentences }} sentences:\n\n{{ text }}",
    version=1,
    set_as_default=True
)
print(f"Updated to version: {updated.version}")

# List all prompts
prompts = client.prompts.list()
print(f"Found {len(prompts.data)} prompts")

# Delete prompt
client.prompts.delete(prompt_id=prompt.prompt_id)

[mattf]

@williamcaban this is the right direction

few things jump out for me -

• the existing impl needs to be moved to be an inline::reference impl (let the llm know it's in src/llama_stack/core/prompts/prompts.py)
• mlflow credential handling needs to be added, should follow the pattern in inference (provider-data backstopped by config)
• tests don't need to be standalone

[williamcaban]

PR is now ready for review and includes the following updates:

• Moved the previous prompts.py as an inline provider (see inline_reference.mdx for details)
• Defined a remote provider with MLflow supporting authentication (see remote_mlflow.mdx for details)
• Removed any dependencies on prompt caching

[franciscojavierarceo]

some of these tests could be shortened. i understand Claude often generates them easily but sometimes they're a little excessive.

[franciscojavierarceo]

some small nits but this is looking very exciting!

[franciscojavierarceo]

thank you for adding the prompts docs!!

we also have this in the UI, maybe we should mention it?

[williamcaban]

Are you referring to making it available in the UI? If so, could that be a follow-up PR? I would prefer to avoid adding more to this one.

[mattf]

why do we need to maintain a mapping from prompt id to mlflow prompt name?

[williamcaban]

@mattf

why do we need to maintain a mapping from prompt id to mlflow prompt name?

The idea is to distinguish Llama Stack-managed prompts from other prompts that might exist in the same MLflow registry.

[mattf]

why do we need to maintain a mapping from prompt id to mlflow prompt name?

The idea is to distinguish Llama Stack-managed prompts from other prompts that might exist in the same MLflow registry.

when would a deployer want to set use_metadata_tags=False? can we always use metadata and skip the id/name translations?

[williamcaban]

why do we need to maintain a mapping from prompt id to mlflow prompt name?

The idea is to distinguish Llama Stack-managed prompts from other prompts that might exist in the same MLflow registry.

when would a deployer want to set use_metadata_tags=False? can we always use metadata and skip the id/name translations?

We can remove that option because, in practice, setting it to false has significant downsides. Let me remove the option.

[mattf]

why do we need to maintain a mapping from prompt id to mlflow prompt name?

The idea is to distinguish Llama Stack-managed prompts from other prompts that might exist in the same MLflow registry.

when would a deployer want to set use_metadata_tags=False? can we always use metadata and skip the id/name translations?

We can remove that option because, in practice, setting it to false has significant downsides. Let me remove the option.

thanks. do we still need to have the id mapping?

[mattf]

@williamcaban why create a func for token extraction and a client but then not use them?

by requesting pr review you're asking others to read all this code (over 3k lines). please review it all before requesting.

[mattf]

because you use env.MLFLOW_TRACKING_TOKEN in the sample_run_config, the auth_credential will already be set from the MLFLOW_TRACKING_TOKEN.

is setting MLFLOW_TRACKING_TOKEN the only way to communicate the token to the client?

[williamcaban]

This can be either—an environment variable at runtime or a setting in the config file.

[mattf]

this is unused.

[williamcaban]

This is used in the edge case of managing prompts created outside Llama Stack. I'll update the code for clarity.