feat: add MLflow Prompt Registry export/import support (#5)

* feat: add MLflow Prompt Registry export/import support

* fix: add newline

* fix: add small fixes

* fix: types

* fix: replace union

* fix: typo

* fix: types

* fix: types

* fix: add import

* fix: typo

---------

Co-authored-by: odorfer <odorfer@amazon.de>

Author: VShturma

docs/bulk.md

@@ -24,7 +24,8 @@ Notes:
 | | [import-models](#Import-registered-models) | [code](mlflow_migration/bulk/import_models.py) | Imports registered models from a directory. |
 | Experiment | [export-experiments](#Export-experiments) | [code](mlflow_migration/bulk/export_experiments.py) | Export several (or all) experiments to a directory. |
 | | [import-experiments](#Import-experiments) | [code](mlflow_migration/bulk/import_experiments.py) | Imports experiments from a directory. |
-
+| Prompt        | [export-prompts](#Export-prompts)             | [code](mlflow_migration/bulk/export_prompts.py)              | Export prompts from the MLflow Prompt Registry (MLflow 2.21.0+).                                                           |
+|               | [import-prompts](#Import-prompts)             | [code](mlflow_migration/bulk/import_prompts.py)              | Imports prompts to the MLflow Prompt Registry. |
 
 ## All MLflow Objects Tools
 
@@ -387,3 +388,87 @@ cat experiment-names.csv
 /Users/me@mycompany.com,/Users/you@mycompany.com
 /Users/foo@mycompany.com,/Users/bar@mycompany.com
 ```
+ 
+## Prompts
+ 
+Export/import prompts from the MLflow Prompt Registry (MLflow 2.21.0+).
+ 
+**Notes:** 
+* Prompt Registry support requires MLflow 2.21.0 or higher. The export/import will be skipped with a warning message if the MLflow version doesn't support prompts.
+* **Version handling**: MLflow 3.0+ uses efficient pagination to export all prompt versions. MLflow 2.21-2.x uses iterative discovery with early stopping (stops after 3 consecutive missing versions).
+ 
+### Export prompts
+ 
+Export prompts from the MLflow Prompt Registry to a directory.
+ 
+Source: [export_prompts.py](mlflow_migration/bulk/export_prompts.py).
+ 
+#### Usage
+ 
+```
+export-prompts --help
+ 
+Options:
+  --output-dir TEXT        Output directory.  [required]
+  --prompts TEXT           Prompt names: 'all' for all prompts, comma-delimited
+                          list (e.g., 'prompt1,prompt2'), or file path ending
+                          with '.txt' containing prompt names (one per line).
+                          [required]
+  --use-threads BOOLEAN    Use multithreading for export.  [default: False]
+```
+ 
+#### Examples
+ 
+##### Export all prompts
+```
+export-prompts \
+  --output-dir out/prompts \
+  --prompts all
+```
+ 
+##### Export specific prompts
+```
+export-prompts \
+  --output-dir out/prompts \
+  --prompts my-prompt-1,my-prompt-2
+```
+
+##### Export prompts from filename
+```
+export-prompts \
+  --output-dir out/prompts \
+  --prompts my-prompts.txt
+```
+
+where `my-prompts.txt` is:
+```
+my-prompt-1
+my-prompt-2
+```
+ 
+### Import prompts
+ 
+Import prompts to the MLflow Prompt Registry from a directory.
+ 
+Source: [import_prompts.py](mlflow_migration/bulk/import_prompts.py).
+ 
+#### Usage
+ 
+```
+import-prompts --help
+ 
+Options:
+  --input-dir TEXT         Input directory containing exported prompts.  [required]
+  --delete-prompt BOOLEAN  Delete existing prompt before importing.  [default: False]
+  --use-threads BOOLEAN    Use multithreading for import.  [default: False]
+```
+ 
+#### Examples
+ 
+```
+import-prompts --input-dir out/prompts
+```
+ 
+**Notes:** 
+* Prompts are imported with their original names and version numbers are preserved.
+* All versions of each prompt are exported and imported to maintain complete version history.

docs/single.md

@@ -23,6 +23,8 @@ See sample JSON export files [here](export_format.md#sample-export-json-files).
 | Run | [export-run](#export-run) | [code](mlflow_migration/run/export_run.py) |
 |  | [import-run](#import-run) | [code](mlflow_migration/run/import_run.py) |
 |  | [copy-run](copy.md#copy-run) | [code](mlflow_migration/copy/copy_run.py) |
+| Prompt | [export-prompt](#export-prompt) | [code](mlflow_migration/prompt/export_prompt.py) |
+|  | [import-prompt](#import-prompt) | [code](mlflow_migration/prompt/import_prompt.py) |
 
 ## Experiment Tools
 
@@ -493,3 +495,66 @@ Options:
                                   metadata (description and tags).  [default:
                                   False]
 ```
+
+## Prompt Tools
+
+Export and import prompts from the MLflow Prompt Registry (MLflow 2.21.0+).
+
+**Note:** Prompt Registry support requires MLflow 2.21.0 or higher.
+
+### Export Prompt
+
+Export a single prompt version to a directory.
+
+Source: [export_prompt.py](mlflow_migration/prompt/export_prompt.py).
+
+#### Usage
+
+```
+export-prompt --help
+
+Options:
+  --prompt-name TEXT     Name of the prompt to export.  [required]
+  --prompt-version TEXT  Version of the prompt to export.  [required]
+  --output-dir TEXT      Output directory.  [required]
+```
+
+#### Example
+
+```
+export-prompt \
+  --prompt-name my-greeting-prompt \
+  --prompt-version 1 \
+  --output-dir out
+```
+
+### Import Prompt
+
+Import a prompt from an exported directory.
+
+Source: [import_prompt.py](mlflow_migration/prompt/import_prompt.py).
+
+#### Usage
+
+```
+import-prompt --help
+
+Options:
+  --input-dir TEXT    Input directory containing exported prompt.  [required]
+  --prompt-name TEXT  Optional new name for the imported prompt. If not
+                     specified, uses original name.
+```
+
+#### Examples
+
+##### Import with original name
+```
+import-prompt --input-dir out
+```
+
+##### Import with new name
+```
+import-prompt \
+  --input-dir out \
+  --prompt-name my-new-prompt-name
+```

mlflow_migration/bulk/export_all.py

@@ -23,6 +23,7 @@
 from mlflow_migration.client.client_utils import create_mlflow_client
 from mlflow_migration.bulk.export_models import export_models
 from mlflow_migration.bulk.export_experiments import export_experiments
+from mlflow_migration.bulk.export_prompts import export_prompts
 
 ALL_STAGES = "Production,Staging,Archived,None"
 
@@ -73,6 +74,26 @@ def export_all(
         notebook_formats=notebook_formats,
         use_threads=use_threads,
     )
+
+    # Export prompts (returns dict with status)
+    res_prompts = None
+    try:
+        _logger.info("Exporting prompts...")
+        res_prompts = export_prompts(
+            output_dir=os.path.join(output_dir, "prompts"),
+            prompt_names=None,  # Export all prompts
+            use_threads=use_threads,
+            mlflow_client=mlflow_client
+        )
+        # Log if unsupported but don't fail
+        if res_prompts and "unsupported" in res_prompts:
+            _logger.warning(f"Prompts not supported in MLflow {res_prompts.get('mlflow_version')}")
+        elif res_prompts and "error" in res_prompts:
+            _logger.warning(f"Failed to export prompts: {res_prompts['error']}")
+    except Exception as e:
+        _logger.warning(f"Failed to export prompts: {e}")
+        res_prompts = {"error": str(e)}
+    
     duration = round(time.time() - start_time, 1)
     info_attr = {
         "options": {
@@ -83,7 +104,12 @@ def export_all(
             "use_threads": use_threads,
             "output_dir": output_dir,
         },
-        "status": {"duration": duration, "models": res_models, "experiments": res_exps},
+        "status": {
+            "duration": duration,
+            "models": res_models,
+            "experiments": res_exps,
+            "prompts": res_prompts
+        }
     }
     io_utils.write_export_file(output_dir, "manifest.json", __file__, {}, info_attr)
     _logger.info(f"Duration for entire tracking server export: {duration} seconds")