Skip to content

Commit 1e68fa8

Browse files
williamcabanclaude
andcommitted
feat: Add MLflow Prompt Registry provider
This commit adds a new remote provider for the Prompts API that integrates with MLflow's Prompt Registry (MLflow 3.4+), enabling centralized prompt management and versioning. Key features: - Full CRUD operations (create, read, update, list) - Immutable version control with default version management - Automatic variable extraction from {{ variable }} placeholders - Deterministic ID mapping between Llama Stack and MLflow - Comprehensive test coverage (41 unit + 14 integration tests) Provider implementation: - src/llama_stack/providers/remote/prompts/mlflow/ - config.py: Configuration schema with validation - mapping.py: ID mapping and metadata utilities - mlflow.py: Main adapter implementation (520 lines) - src/llama_stack/providers/registry/prompts.py: Provider registration Testing: - tests/unit/providers/remote/prompts/mlflow/ - test_config.py: 14 configuration tests - test_mapping.py: 27 ID mapping tests - tests/integration/providers/remote/prompts/mlflow/ - test_end_to_end.py: 14 end-to-end scenarios - conftest.py: Test fixtures and server availability checks - README.md: Testing documentation - scripts/test_mlflow_prompts_manual.py: Manual validation script Documentation: - docs/docs/providers/prompts/remote_mlflow.mdx: User guide - docs/docs/providers/prompts/index.mdx: Provider category index Dependencies: - Requires mlflow>=3.4.0 (added to provider pip_packages) Limitations: - No deletion support (MLflow constraint) - Sequential version numbers only (1, 2, 3...) Signed-off-by: William Caban <[email protected]> Co-Authored-By: Claude <[email protected]> feat: Add inline::reference provider and MLflow authentication This commit addresses the PR feedback to refactor prompts architecture and add authentication support for the MLflow provider. - Created new inline::reference provider in providers/inline/prompts/reference/ - Moved implementation from core/prompts/prompts.py to reference provider - Updated core/prompts/prompts.py to delegate to inline provider for backward compatibility - Follows Llama Stack pattern where core routes, providers implement - Registered inline::reference provider (InlineProviderSpec) - Zero dependencies, uses KVStore backend (SQLite, PostgreSQL, etc.) - Full CRUD support including deletion - Registered remote::mlflow provider (RemoteProviderSpec) - Requires mlflow>=3.4.0 - Added provider_data_validator for authentication support Added three-tier authentication support following inference provider pattern: 1. **Per-request provider data** (highest priority) - Via x-llamastack-provider-data header - Enables multi-tenant deployments with user-specific credentials 2. **Configuration auth_credential** (fallback) - Server-level authentication in config file - Supports environment variable substitution 3. **Environment variables** (lowest priority) - Standard MLflow behavior (MLFLOW_TRACKING_TOKEN, etc.) - Added MLflowProviderDataValidator for request header validation - Added auth_credential field (SecretStr) to MLflowPromptsConfig - MLflowPromptsAdapter inherits from NeedsRequestProviderData - Implemented _get_mlflow_token() method with proper precedence - Added @json_schema_type decorator and sample_run_config() method - Complete documentation for inline::reference provider - Configuration examples for SQLite and PostgreSQL - Migration guide from core implementation - Best practices and troubleshooting - Added comprehensive Authentication section - Documented all three authentication methods with examples - Added Databricks authentication guide - Added enterprise MLflow examples - Security best practices - Updated configuration table and sample configs - Added "Available Providers" section - Added "Choosing a Provider" decision guide - Quick start examples for both providers - Common features overview - Updated test_config.py with authentication tests - MLflowProviderDataValidator tests - auth_credential configuration tests - sample_run_config() tests - inline::reference unit tests: 12/12 - remote::mlflow unit tests: 49/49 - remote::mlflow integration tests: 14/14 All existing imports continue to work: - from llama_stack.core.prompts import PromptServiceConfig - from llama_stack.core.prompts import PromptServiceImpl - from llama_stack.core.prompts import get_provider_impl Prompts work without MLflow installed (uses inline::reference) Signed-off-by: William Caban <[email protected]> Co-Authored-By: Claude <[email protected]>
1 parent dabebdd commit 1e68fa8

File tree

24 files changed

+4071
-229
lines changed

24 files changed

+4071
-229
lines changed
Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
---
2+
sidebar_label: Prompts
3+
title: Prompts
4+
---
5+
6+
# Prompts
7+
8+
## Overview
9+
10+
This section contains documentation for all available providers for the **prompts** API.
11+
12+
The Prompts API enables centralized management of prompt templates with versioning, variable handling, and team collaboration capabilities.
13+
14+
## Available Providers
15+
16+
### Inline Providers
17+
18+
Inline providers run in the same process as the Llama Stack server and require no external dependencies:
19+
20+
- **[inline::reference](inline_reference.mdx)** - Reference implementation using KVStore backend (SQLite, PostgreSQL, etc.)
21+
- Zero external dependencies
22+
- Supports local SQLite or PostgreSQL storage
23+
- Full CRUD operations including deletion
24+
- Ideal for local development and single-server deployments
25+
26+
### Remote Providers
27+
28+
Remote providers connect to external services for centralized prompt management:
29+
30+
- **[remote::mlflow](remote_mlflow.mdx)** - MLflow Prompt Registry integration (requires MLflow 3.4+)
31+
- Centralized prompt management across teams
32+
- Built-in versioning and audit trail
33+
- Supports authentication (per-request, config, or environment variables)
34+
- Integrates with Databricks and enterprise MLflow deployments
35+
- Ideal for team collaboration and production environments
36+
37+
## Choosing a Provider
38+
39+
### Use `inline::reference` when:
40+
- Developing locally or deploying to a single server
41+
- You want zero external dependencies
42+
- SQLite or PostgreSQL storage is sufficient
43+
- You need full CRUD operations (including deletion)
44+
- You prefer simple configuration
45+
46+
### Use `remote::mlflow` when:
47+
- Working in a team environment with multiple users
48+
- You need centralized prompt management
49+
- Integration with existing MLflow infrastructure
50+
- You want to leverage Databricks or enterprise MLflow
51+
- You need authentication and multi-tenant support
52+
- Advanced versioning and audit trail capabilities are required
53+
54+
## Quick Start Examples
55+
56+
### Using inline::reference
57+
58+
```yaml
59+
prompts:
60+
- provider_id: local-prompts
61+
provider_type: inline::reference
62+
config:
63+
run_config:
64+
storage:
65+
stores:
66+
prompts:
67+
type: sqlite
68+
db_path: ./prompts.db
69+
```
70+
71+
### Using remote::mlflow
72+
73+
```yaml
74+
prompts:
75+
- provider_id: mlflow-prompts
76+
provider_type: remote::mlflow
77+
config:
78+
mlflow_tracking_uri: http://localhost:5555
79+
experiment_name: llama-stack-prompts
80+
auth_credential: ${env.MLFLOW_TRACKING_TOKEN}
81+
```
82+
83+
## Common Features
84+
85+
All prompt providers support:
86+
- Create and store prompts with version control
87+
- Retrieve prompts by ID and version
88+
- Update prompts (creates new versions)
89+
- List all prompts or versions of a specific prompt
90+
- Set default version for a prompt
91+
- Automatic variable extraction from `{{ variable }}` templates
92+
93+
For detailed documentation on each provider, see the individual provider pages linked above.

0 commit comments

Comments
 (0)