refactor: remove use_metadata_tags config option from MLflow provider

The use_metadata_tags configuration option has been removed to simplify
the MLflow provider API and ensure optimal performance. Metadata tags
are now always used for efficient prompt filtering and discovery.

Changes:
- Remove use_metadata_tags field from MLflowPromptsConfig
- Update PromptIDMapper to always generate metadata tags
- Remove conditional logic in list_prompts() - always uses tag filtering
- Update unit and integration tests
- Update documentation to remove references to use_metadata_tags
- Add GitHub Actions workflow for MLflow integration tests

Signed-off-by: William Caban <[email protected]>

Author: williamcaban

.github/workflows/integration-mlflow-tests.yml

@@ -0,0 +1,125 @@
+name: MLflow Prompts Integration Tests
+
+run-name: Run the integration test suite with MLflow Prompt Registry provider
+
+on:
+  push:
+    branches:
+      - main
+      - 'release-[0-9]+.[0-9]+.x'
+  pull_request:
+    branches:
+      - main
+      - 'release-[0-9]+.[0-9]+.x'
+    paths:
+      - 'src/llama_stack/providers/remote/prompts/mlflow/**'
+      - 'tests/integration/providers/remote/prompts/mlflow/**'
+      - 'tests/unit/providers/remote/prompts/mlflow/**'
+      - 'uv.lock'
+      - 'pyproject.toml'
+      - 'requirements.txt'
+      - '.github/workflows/integration-mlflow-tests.yml' # This workflow
+  schedule:
+    - cron: '0 0 * * *'  # Daily at 12 AM UTC
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref == 'refs/heads/main' && github.run_id || github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  test-mlflow:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ${{ github.event.schedule == '0 0 * * *' && fromJSON('["3.12", "3.13"]') || fromJSON('["3.12"]') }}
+      fail-fast: false
+
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3 # v6.0.0
+
+      - name: Install dependencies
+        uses: ./.github/actions/setup-runner
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Setup MLflow Server
+        run: |
+          docker run --rm -d --pull always \
+            --name mlflow \
+            -p 5555:5555 \
+            ghcr.io/mlflow/mlflow:latest \
+            mlflow server \
+              --host 0.0.0.0 \
+              --port 5555 \
+              --backend-store-uri sqlite:///mlflow.db \
+              --default-artifact-root ./mlruns
+
+      - name: Wait for MLflow to be ready
+        run: |
+          echo "Waiting for MLflow to be ready..."
+          for i in {1..60}; do
+            if curl -s http://localhost:5555/health | grep -q '"status": "OK"'; then
+              echo "MLflow is ready!"
+              exit 0
+            fi
+            echo "Not ready yet... ($i/60)"
+            sleep 2
+          done
+          echo "MLflow failed to start"
+          docker logs mlflow
+          exit 1
+
+      - name: Verify MLflow API
+        run: |
+          echo "Testing MLflow API..."
+          curl -X GET http://localhost:5555/api/2.0/mlflow/experiments/list
+          echo ""
+          echo "MLflow API is responding!"
+
+      - name: Build Llama Stack
+        run: |
+          uv run --no-sync llama stack list-deps ci-tests | xargs -L1 uv pip install
+
+      - name: Install MLflow Python client
+        run: |
+          uv pip install 'mlflow>=3.4.0'
+
+      - name: Check Storage and Memory Available Before Tests
+        if: ${{ always() }}
+        run: |
+          free -h
+          df -h
+
+      - name: Run MLflow Integration Tests
+        env:
+          MLFLOW_TRACKING_URI: http://localhost:5555
+        run: |
+          uv run --no-sync \
+            pytest -sv \
+            tests/integration/providers/remote/prompts/mlflow/
+
+      - name: Check Storage and Memory Available After Tests
+        if: ${{ always() }}
+        run: |
+          free -h
+          df -h
+
+      - name: Write MLflow logs to file
+        if: ${{ always() }}
+        run: |
+          docker logs mlflow > mlflow.log 2>&1 || true
+
+      - name: Upload all logs to artifacts
+        if: ${{ always() }}
+        uses: actions/upload-artifact@330a01c490aca151604b8cf639adc76d48f6c5d4 # v5.0.0
+        with:
+          name: mlflow-logs-${{ github.run_id }}-${{ github.run_attempt }}-${{ matrix.python-version }}
+          path: |
+            *.log
+          retention-days: 1
+
+      - name: Stop MLflow container
+        if: ${{ always() }}
+        run: |
+          docker stop mlflow || true

docs/docs/providers/prompts/remote_mlflow.mdx

@@ -155,7 +155,6 @@ prompts:
     config:
       mlflow_tracking_uri: https://mlflow.example.com
       experiment_name: team-prompts
-      use_metadata_tags: true
       timeout_seconds: 45
 ```
 
@@ -732,7 +731,6 @@ See [MLflow's documentation](https://mlflow.org/docs/latest/prompts.html) for mo
 | `mlflow_registry_uri` | `str \| None` | No | None | MLflow model registry URI (defaults to tracking URI if not set) |
 | `experiment_name` | `str` | No | llama-stack-prompts | MLflow experiment name for storing prompts |
 | `auth_credential` | `SecretStr \| None` | No | None | MLflow API token for authentication. Can be overridden via provider data header. |
-| `use_metadata_tags` | `bool` | No | True | Whether to store Llama Stack metadata as MLflow tags |
 | `timeout_seconds` | `int` | No | 30 | Timeout for MLflow API calls (1-300 seconds) |
 
 ## Sample Configuration
@@ -741,7 +739,6 @@ See [MLflow's documentation](https://mlflow.org/docs/latest/prompts.html) for mo
 ```yaml
 mlflow_tracking_uri: http://localhost:5555
 experiment_name: llama-stack-prompts
-use_metadata_tags: true
 timeout_seconds: 30
 ```
 
@@ -750,6 +747,5 @@ timeout_seconds: 30
 mlflow_tracking_uri: ${env.MLFLOW_TRACKING_URI:=http://localhost:5000}
 experiment_name: llama-stack-prompts
 auth_credential: ${env.MLFLOW_TRACKING_TOKEN:=}
-use_metadata_tags: true
 timeout_seconds: 30
 ```

src/llama_stack/providers/remote/prompts/mlflow/config.py

@@ -45,7 +45,6 @@ class MLflowPromptsConfig(BaseModel):
         mlflow_registry_uri: MLflow registry URI (optional, defaults to tracking_uri)
         experiment_name: MLflow experiment name for prompt storage
         auth_credential: MLflow API token for authentication (optional, can be overridden by provider data)
-        use_metadata_tags: Store Llama Stack metadata in MLflow tags (default: True)
         timeout_seconds: Timeout for MLflow API calls in seconds (default: 30)
     """
 
@@ -65,10 +64,6 @@ class MLflowPromptsConfig(BaseModel):
         default=None,
         description="MLflow API token for authentication. Can be overridden via provider data header.",
     )
-    use_metadata_tags: bool = Field(
-        default=True,
-        description="Store Llama Stack metadata (prompt_id, variables) in MLflow tags",
-    )
     timeout_seconds: int = Field(
         default=30,
         ge=1,

src/llama_stack/providers/remote/prompts/mlflow/mapping.py

@@ -31,13 +31,13 @@ class PromptIDMapper:
     # Prefix for MLflow prompt names managed by Llama Stack
     MLFLOW_NAME_PREFIX = "llama_prompt_"
 
-    def __init__(self, use_metadata: bool = True):
+    def __init__(self):
         """Initialize ID mapper.
 
-        Args:
-            use_metadata: Store Llama Stack ID in MLflow tags for discoverability
+        Llama Stack metadata is always stored in MLflow tags for efficient
+        filtering and prompt discovery.
         """
-        self.use_metadata = use_metadata
+        pass
 
     def to_mlflow_name(self, prompt_id: str) -> str:
         """Convert Llama Stack prompt_id to MLflow prompt name.
@@ -119,9 +119,6 @@ def get_metadata_tags(self, prompt_id: str, variables: list[str] | None = None)
             >>> tags
             {"llama_stack_id": "pmpt_abc123...", "llama_stack_managed": "true", "variables": "var1,var2"}
         """
-        if not self.use_metadata:
-            return {}
-
         tags = {
             "llama_stack_id": prompt_id,
             "llama_stack_managed": "true",

src/llama_stack/providers/remote/prompts/mlflow/mlflow.py

@@ -61,7 +61,7 @@ def __init__(self, config: MLflowPromptsConfig):
         """
         self.config = config
         self.mlflow_client: "MlflowClient | None" = None
-        self.mapper = PromptIDMapper(use_metadata=config.use_metadata_tags)
+        self.mapper = PromptIDMapper()
         logger.info(
             f"MLflowPromptsAdapter initialized: tracking_uri={config.mlflow_tracking_uri}, "
             f"experiment={config.experiment_name}"
@@ -416,25 +416,14 @@ async def list_prompts(self) -> ListPromptsResponse:
             (those with llama_stack_managed=true tag)
         """
         try:
-            # Search for Llama Stack managed prompts
-            if self.config.use_metadata_tags:
-                prompts = mlflow.genai.search_prompts(filter_string="tag.llama_stack_managed='true'")
-            else:
-                # If not using metadata, search for prompts with llama_prompt_ prefix
-                # Note: MLflow search_prompts may not support name prefix filtering
-                # In this case, we need to list all and filter
-                prompts = mlflow.genai.search_prompts()
+            # Search for Llama Stack managed prompts using metadata tags
+            prompts = mlflow.genai.search_prompts(filter_string="tag.llama_stack_managed='true'")
         except Exception as e:
             logger.error(f"Failed to search prompts in MLflow: {e}")
             return ListPromptsResponse(data=[])
 
         llama_prompts = []
         for mlflow_prompt in prompts:
-            # Filter by name prefix if not using metadata
-            if not self.config.use_metadata_tags:
-                if not mlflow_prompt.name.startswith(self.mapper.MLFLOW_NAME_PREFIX):
-                    continue
-
             try:
                 # Convert MLflow name to Llama Stack ID
                 prompt_id = self.mapper.to_llama_id(mlflow_prompt.name)

tests/integration/providers/remote/prompts/mlflow/conftest.py

@@ -66,7 +66,6 @@ async def mlflow_config(mlflow_tracking_uri, mlflow_server_available):
     return MLflowPromptsConfig(
         mlflow_tracking_uri=mlflow_tracking_uri,
         experiment_name="test-llama-stack-prompts",
-        use_metadata_tags=True,
         timeout_seconds=30,
     )
 

tests/unit/providers/remote/prompts/mlflow/test_config.py

@@ -25,7 +25,6 @@ def test_default_config(self):
         assert config.mlflow_tracking_uri == "http://localhost:5000"
         assert config.mlflow_registry_uri is None
         assert config.experiment_name == "llama-stack-prompts"
-        assert config.use_metadata_tags is True
         assert config.timeout_seconds == 30
 
     def test_custom_config(self):
@@ -34,14 +33,12 @@ def test_custom_config(self):
             mlflow_tracking_uri="http://mlflow.example.com:8080",
             mlflow_registry_uri="http://registry.example.com:8080",
             experiment_name="my-prompts",
-            use_metadata_tags=False,
             timeout_seconds=60,
         )
 
         assert config.mlflow_tracking_uri == "http://mlflow.example.com:8080"
         assert config.mlflow_registry_uri == "http://registry.example.com:8080"
         assert config.experiment_name == "my-prompts"
-        assert config.use_metadata_tags is False
         assert config.timeout_seconds == 60
 
     def test_databricks_uri(self):
@@ -113,14 +110,6 @@ def test_registry_uri_defaults_to_none(self):
         config = MLflowPromptsConfig()
         assert config.mlflow_registry_uri is None
 
-    def test_use_metadata_tags_boolean(self):
-        """Test use_metadata_tags accepts boolean values."""
-        config_true = MLflowPromptsConfig(use_metadata_tags=True)
-        assert config_true.use_metadata_tags is True
-
-        config_false = MLflowPromptsConfig(use_metadata_tags=False)
-        assert config_false.use_metadata_tags is False
-
 
 class TestMLflowProviderDataValidator:
     """Tests for MLflowProviderDataValidator."""

tests/unit/providers/remote/prompts/mlflow/test_mapping.py

@@ -17,7 +17,7 @@ class TestPromptIDMapper:
     @pytest.fixture
     def mapper(self):
         """Create ID mapper instance."""
-        return PromptIDMapper(use_metadata=True)
+        return PromptIDMapper()
 
     def test_to_mlflow_name_valid_id(self, mapper):
         """Test converting valid prompt_id to MLflow name."""
@@ -148,16 +148,6 @@ def test_get_metadata_tags_empty_variables(self, mapper):
         assert tags["llama_stack_managed"] == "true"
         assert "variables" not in tags
 
-    def test_get_metadata_tags_disabled(self):
-        """Test metadata tags returns empty dict when disabled."""
-        mapper = PromptIDMapper(use_metadata=False)
-        prompt_id = "pmpt_" + "f" * 48
-        variables = ["var1"]
-
-        tags = mapper.get_metadata_tags(prompt_id, variables)
-
-        assert tags == {}
-
     def test_extract_variables_from_tags(self, mapper):
         """Test extracting variables from tags."""
         tags = {"variables": "var1,var2,var3"}