Commit 1e68fa8
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- docs/docs/providers/prompts
- scripts
- src/llama_stack
- core/prompts
- providers
- inline/prompts
- reference
- registry
- remote/prompts
- mlflow
- tests
- integration/providers/remote/prompts
- mlflow
- unit/providers/remote/prompts
- mlflow
24 files changed
+4071
-229
lines changed| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
| 1 | + | |
| 2 | + | |
| 3 | + | |
| 4 | + | |
| 5 | + | |
| 6 | + | |
| 7 | + | |
| 8 | + | |
| 9 | + | |
| 10 | + | |
| 11 | + | |
| 12 | + | |
| 13 | + | |
| 14 | + | |
| 15 | + | |
| 16 | + | |
| 17 | + | |
| 18 | + | |
| 19 | + | |
| 20 | + | |
| 21 | + | |
| 22 | + | |
| 23 | + | |
| 24 | + | |
| 25 | + | |
| 26 | + | |
| 27 | + | |
| 28 | + | |
| 29 | + | |
| 30 | + | |
| 31 | + | |
| 32 | + | |
| 33 | + | |
| 34 | + | |
| 35 | + | |
| 36 | + | |
| 37 | + | |
| 38 | + | |
| 39 | + | |
| 40 | + | |
| 41 | + | |
| 42 | + | |
| 43 | + | |
| 44 | + | |
| 45 | + | |
| 46 | + | |
| 47 | + | |
| 48 | + | |
| 49 | + | |
| 50 | + | |
| 51 | + | |
| 52 | + | |
| 53 | + | |
| 54 | + | |
| 55 | + | |
| 56 | + | |
| 57 | + | |
| 58 | + | |
| 59 | + | |
| 60 | + | |
| 61 | + | |
| 62 | + | |
| 63 | + | |
| 64 | + | |
| 65 | + | |
| 66 | + | |
| 67 | + | |
| 68 | + | |
| 69 | + | |
| 70 | + | |
| 71 | + | |
| 72 | + | |
| 73 | + | |
| 74 | + | |
| 75 | + | |
| 76 | + | |
| 77 | + | |
| 78 | + | |
| 79 | + | |
| 80 | + | |
| 81 | + | |
| 82 | + | |
| 83 | + | |
| 84 | + | |
| 85 | + | |
| 86 | + | |
| 87 | + | |
| 88 | + | |
| 89 | + | |
| 90 | + | |
| 91 | + | |
| 92 | + | |
| 93 | + | |
0 commit comments