mlflow server --backend-store-uri sqlite:///mlflow.db
cd mlflow/server/js & yarn install & yarn start

from datasets import load_dataset


def create_aime_evaluation_dataset(experiment_id: str) -> str:
    from mlflow.genai.datasets import create_dataset

    mlflow_dataset_name = "aime_1983_2024_tailored"

    print("Loading AIME dataset from HuggingFace...")
    hf_dataset = load_dataset("gneubig/aime-1983-2024", split="train")

    print(f"Loaded {len(hf_dataset)} samples")

    # Transform to MLflow format
    # - inputs: the question/problem to solve
    # - expectations: the expected answer (ground truth)
    records = []
    for item in hf_dataset:
        record = {
            "inputs": {
                "question": item["Question"],
            },
            "expectations": {
                "expected_response": item["Answer"],
            },
        }
        records.append(record)

    records = records[:20]

    print(f"Transformed {len(records)} records to MLflow format")

    # Create the dataset
    dataset_tags = {
        "source": "huggingface",
        "hf_dataset": "gneubig/aime-1983-2024",
        "purpose": "prompt_optimization",
    }

    print(f"Creating MLflow EvaluationDataset '{mlflow_dataset_name}'...")
    dataset = create_dataset(
        name=mlflow_dataset_name,
        experiment_id=experiment_id,
        tags=dataset_tags,
    )

    print(f"Created dataset with ID: {dataset.dataset_id}")

    # Add records to the dataset
    print(f"Merging {len(records)} records into dataset...")
    dataset.merge_records(records)

    print("Dataset created successfully!")
    print(f"  - Name: {dataset.name}")
    print(f"  - ID: {dataset.dataset_id}")
    print(f"  - Records: {len(records)}")

    return dataset.dataset_id


def main():
    import mlflow

    mlflow.set_tracking_uri("http://127.0.0.1:5000")
    experiment = mlflow.set_experiment("optimization_backend")
    experiment_id = experiment.experiment_id

    dataset_id = create_aime_evaluation_dataset(experiment_id)
    print(f"Dataset created with ID: {dataset_id}")


if __name__ == "__main__":
    main()

import json
import time

import requests

# Configuration
MLFLOW_SERVER_URL = "http://127.0.0.1:5000"
API_VERSION = 3  # API version from proto (prompt optimization APIs are v3)


def create_optimization_job(
    experiment_id: str,
    prompt_uri: str,
    optimizer_type: str = "gepa",
    scorers: list[str] | None = None,
    dataset_id: str | None = None,
    optimizer_config: dict | None = None,
    tags: list[dict] | None = None,
) -> dict:
    url = f"{MLFLOW_SERVER_URL}/ajax-api/{API_VERSION}.0/mlflow/prompt-optimization/jobs"

    # Convert string optimizer_type to proto enum value
    optimizer_type_to_enum = {
        "gepa": 1,  # OPTIMIZER_TYPE_GEPA
        "metaprompt": 2,  # OPTIMIZER_TYPE_METAPROMPT
    }
    optimizer_type_enum = optimizer_type_to_enum.get(optimizer_type.lower(), 0)

    # Build config - dataset_id and scorers are now separate fields
    config = {
        "target_prompt_uri": prompt_uri,
        "optimizer_type": optimizer_type_enum,
        "scorers": scorers or [],  # Empty list if None
    }

    # Add dataset_id if provided (optional for zero-shot metaprompting)
    if dataset_id:
        config["dataset_id"] = dataset_id

    # Add optimizer_config_json if provided
    if optimizer_config:
        config["optimizer_config_json"] = json.dumps(optimizer_config)

    payload = {
        "experiment_id": experiment_id,
        "config": config,
        "tags": tags or [],
    }

    print("Creating optimization job...")
    print(f"URL: {url}")
    print(f"Payload: {json.dumps(payload, indent=2)}")

    response = requests.post(url, json=payload)

    if response.status_code != 200:
        print(f"Error: {response.status_code}")
        print(f"Response: {response.text}")
        return None

    result = response.json()
    print(f"Response: {json.dumps(result, indent=2)}")
    return result


def cancel_optimization_job(job_id: str) -> dict:
    url = f"{MLFLOW_SERVER_URL}/ajax-api/{API_VERSION}.0/mlflow/prompt-optimization/jobs/{job_id}/cancel"

    print(f"Cancelling job: {job_id}")

    response = requests.post(url)

    if response.status_code != 200:
        print(f"Error: {response.status_code}")
        print(f"Response: {response.text}")
        return None

    result = response.json()
    print(f"Response: {json.dumps(result, indent=2)}")
    return result


def main():
    import mlflow

    dataset_id = "d-ecaeda79a412460bb6f0560af6bb7321"
    prompt_uri = "prompts:/aime_solver/1"

    mlflow.set_tracking_uri("http://127.0.0.1:5000")
    mlflow.set_experiment("optimization_backend")

    # Scorers - use built-in scorer names (e.g., "Correctness", "Safety")
    # or registered scorer names from the experiment
    scorers = ["Correctness"]  # Built-in scorer for checking expected_response

    # Optimizer-specific config
    optimizer_config = {
        "reflection_model": "openai:/gpt-4o",
        # "max_metric_calls": 100,
    }

    experiment = mlflow.get_experiment_by_name("optimization_backend")
    experiment_id = experiment.experiment_id

    result = create_optimization_job(
        experiment_id=experiment_id,
        prompt_uri=prompt_uri,
        dataset_id=None,
        scorers=[],
        optimizer_type="metaprompt",
        optimizer_config=optimizer_config,
        tags=[{"key": "test", "value": "true"}],
    )

    if result and "job" in result:
        job_id = result["job"].get("job_id")
        print(f"\nCreated job with ID: {job_id}")

        import pdb

        pdb.set_trace()

        cancel_optimization_job(job_id)
    else:
        print("Failed to create optimization job")

    print("\n" + "=" * 60)
    print("Test complete!")
    print("=" * 60)


if __name__ == "__main__":
    main()