freeharness 0.4.1

package/canonical/skills/agent-evaluation/scripts/create_dataset_template.py

@@ -0,0 +1,382 @@
+"""
+Generate a template script for creating MLflow evaluation datasets.
+
+This script creates a customized Python script for dataset creation,
+handling both OSS MLflow and Databricks Unity Catalog scenarios.
+
+Usage:
+    # Minimum required
+    python create_dataset_template.py --test-cases-file test_cases.txt
+
+    # With custom dataset name
+    python create_dataset_template.py --test-cases-file test_cases.txt --dataset-name my-eval
+
+    # For Databricks Unity Catalog
+    python create_dataset_template.py --test-cases-file test_cases.txt \\
+        --catalog main --schema ml --table eval_v1
+"""
+
+import argparse
+import datetime
+import os
+import subprocess
+import sys
+
+from utils import check_databricks_config, validate_env_vars
+
+
+def parse_arguments():
+    """Parse command-line arguments."""
+    parser = argparse.ArgumentParser(description="Generate dataset creation script")
+    parser.add_argument(
+        "--test-cases-file",
+        required=True,
+        help="File with test cases (one per line, minimum 10)",
+    )
+    parser.add_argument(
+        "--dataset-name",
+        help="Dataset name (auto-generated if not provided)",
+    )
+    parser.add_argument("--catalog", help="UC catalog name (for Databricks)")
+    parser.add_argument("--schema", help="UC schema name (for Databricks)")
+    parser.add_argument("--table", help="UC table name (for Databricks)")
+    parser.add_argument("--output", default="create_evaluation_dataset.py", help="Output file name")
+    return parser.parse_args()
+
+
+def load_test_cases_from_file(file_path: str) -> list[str]:
+    """Load test cases from file (one per line)."""
+    try:
+        with open(file_path) as f:
+            test_cases = [line.strip() for line in f if line.strip()]
+
+        if len(test_cases) < 10:
+            print(f"✗ File has only {len(test_cases)} test cases (minimum: 10)")
+            print("  Please add more test cases to the file")
+            sys.exit(1)
+
+        print(f"✓ Loaded {len(test_cases)} test cases from {file_path}")
+        return test_cases
+
+    except FileNotFoundError:
+        print(f"✗ File not found: {file_path}")
+        sys.exit(1)
+    except Exception as e:
+        print(f"✗ Error reading file: {e}")
+        sys.exit(1)
+
+
+def get_uc_catalogs():
+    """Get available Unity Catalog catalogs."""
+    try:
+        code = """
+from databricks import sdk
+w = sdk.WorkspaceClient()
+catalogs = w.catalogs.list()
+for catalog in catalogs:
+    print(catalog.name)
+"""
+        result = subprocess.run(["python", "-c", code], capture_output=True, text=True, check=True)
+        return [line.strip() for line in result.stdout.strip().split("\n") if line.strip()]
+    except Exception:
+        return []
+
+
+def get_uc_schemas(catalog: str):
+    """Get available schemas in a catalog."""
+    try:
+        code = f"""
+from databricks import sdk
+w = sdk.WorkspaceClient()
+schemas = w.schemas.list(catalog_name='{catalog}')
+for schema in schemas:
+    print(schema.name)
+"""
+        result = subprocess.run(["python", "-c", code], capture_output=True, text=True, check=True)
+        return [line.strip() for line in result.stdout.strip().split("\n") if line.strip()]
+    except Exception:
+        return []
+
+
+def generate_dataset_creation_code(
+    tracking_uri: str,
+    experiment_id: str,
+    dataset_name: str,
+    test_cases: list[str],
+    catalog: str = None,
+    schema: str = None,
+    table: str = None,
+) -> str:
+    """Generate Python code for creating the dataset."""
+
+    # Escape test cases for Python code
+    test_cases_repr = repr(test_cases)
+
+    if catalog and schema and table:
+        # Databricks Unity Catalog version
+        uc_name = f"{catalog}.{schema}.{table}"
+        return f'''#!/usr/bin/env python3
+"""
+Create MLflow evaluation dataset.
+
+Generated by create_dataset_template.py
+"""
+
+import mlflow
+from mlflow.genai.datasets import create_dataset
+
+# Set tracking URI
+mlflow.set_tracking_uri("{tracking_uri}")
+
+# Configuration
+DATASET_NAME = "{uc_name}"
+EXPERIMENT_ID = "{experiment_id}"
+
+# Test cases
+TEST_CASES = {test_cases_repr}
+
+print("=" * 60)
+print("Creating MLflow Evaluation Dataset")
+print("=" * 60)
+print(f"Dataset: {{DATASET_NAME}}")
+print(f"Test cases: {{len(TEST_CASES)}}")
+print()
+
+# Create dataset
+print("Creating dataset...")
+try:
+    dataset = create_dataset(
+        name=DATASET_NAME,
+        source={{
+            "inputs": [
+                {{"query": query}}
+                for query in TEST_CASES
+            ]
+        }},
+        targets=[
+            {{"expected_output": "TODO: Add expected output"}}
+            for _ in TEST_CASES
+        ],
+        experiment_id=[EXPERIMENT_ID]
+    )
+
+    print(f"✓ Dataset created: {{DATASET_NAME}}")
+    print(f"  Location: Unity Catalog table")
+    print(f"  Queries: {{len(TEST_CASES)}}")
+    print()
+    print("=" * 60)
+    print("Next Steps")
+    print("=" * 60)
+    print()
+    print("1. Verify dataset in Databricks Unity Catalog")
+    print(f"2. Use in evaluation: python scripts/run_evaluation_template.py --dataset-name {uc_name}")
+    print()
+
+except Exception as e:
+    print(f"✗ Error creating dataset: {{e}}")
+    import traceback
+    traceback.print_exc()
+    exit(1)
+'''
+    else:
+        # OSS MLflow version (non-UC)
+        return f'''#!/usr/bin/env python3
+"""
+Create MLflow evaluation dataset.
+
+Generated by create_dataset_template.py
+"""
+
+import mlflow
+from mlflow.genai.datasets import create_dataset
+
+# Set tracking URI
+mlflow.set_tracking_uri("{tracking_uri}")
+
+# Configuration
+DATASET_NAME = "{dataset_name}"
+EXPERIMENT_ID = "{experiment_id}"
+
+# Test cases
+TEST_CASES = {test_cases_repr}
+
+print("=" * 60)
+print("Creating MLflow Evaluation Dataset")
+print("=" * 60)
+print(f"Dataset: {{DATASET_NAME}}")
+print(f"Test cases: {{len(TEST_CASES)}}")
+print()
+
+# Create dataset
+print("Creating dataset...")
+try:
+    dataset = create_dataset(
+        name=DATASET_NAME,
+        source={{
+            "inputs": [
+                {{"query": query}}
+                for query in TEST_CASES
+            ]
+        }},
+        targets=[
+            {{"expected_output": "TODO: Add expected output"}}
+            for _ in TEST_CASES
+        ],
+        experiment_id=[EXPERIMENT_ID]
+    )
+
+    print(f"✓ Dataset created: {{DATASET_NAME}}")
+    print(f"  Queries: {{len(TEST_CASES)}}")
+    print()
+    print("=" * 60)
+    print("Next Steps")
+    print("=" * 60)
+    print()
+    print("1. Verify dataset: python scripts/list_datasets.py")
+    print(f"2. Use in evaluation: python scripts/run_evaluation_template.py --dataset-name {{DATASET_NAME}}")
+    print()
+
+except Exception as e:
+    print(f"✗ Error creating dataset: {{e}}")
+    import traceback
+    traceback.print_exc()
+    exit(1)
+'''
+
+
+def main():
+    """Main workflow."""
+    args = parse_arguments()
+
+    print("=" * 60)
+    print("MLflow Dataset Creation Template Generator")
+    print("=" * 60)
+    print()
+
+    # Check environment
+    errors = validate_env_vars()
+    if errors:
+        print("✗ Environment validation failed:")
+        for error in errors:
+            print(f"  - {error}")
+        print("\nRun scripts/setup_mlflow.py first")
+        sys.exit(1)
+
+    tracking_uri = os.getenv("MLFLOW_TRACKING_URI")
+    experiment_id = os.getenv("MLFLOW_EXPERIMENT_ID")
+
+    print(f"Tracking URI: {tracking_uri}")
+    print(f"Experiment ID: {experiment_id}")
+    print()
+
+    # Load test cases
+    print("Loading test cases...")
+    test_cases = load_test_cases_from_file(args.test_cases_file)
+    print()
+
+    # Check if Databricks
+    is_databricks, profile = check_databricks_config()
+
+    # Handle dataset configuration
+    if is_databricks:
+        print("✓ Detected Databricks configuration")
+        print()
+
+        # Check if UC args provided
+        if args.catalog and args.schema and args.table:
+            catalog = args.catalog
+            schema = args.schema
+            table = args.table
+            print(f"Using Unity Catalog: {catalog}.{schema}.{table}")
+        elif args.catalog or args.schema or args.table:
+            print("✗ For Databricks UC, all three args are required:")
+            print("  --catalog, --schema, --table")
+            sys.exit(1)
+        else:
+            # Try to get catalogs
+            print("Fetching available catalogs...")
+            catalogs = get_uc_catalogs()
+
+            if not catalogs:
+                print("✗ Could not fetch catalogs")
+                print("  Please specify manually:")
+                print("    --catalog <catalog> --schema <schema> --table <table>")
+                sys.exit(1)
+
+            # Auto-select first catalog
+            catalog = catalogs[0]
+            print(f"✓ Found {len(catalogs)} catalog(s), using: {catalog}")
+
+            # Get schemas
+            print(f"Fetching schemas in {catalog}...")
+            schemas = get_uc_schemas(catalog)
+
+            if not schemas:
+                print("✗ Could not fetch schemas")
+                print("  Please specify manually:")
+                print("    --catalog <catalog> --schema <schema> --table <table>")
+                sys.exit(1)
+
+            # Auto-select first schema
+            schema = schemas[0]
+            print(f"✓ Found {len(schemas)} schema(s), using: {schema}")
+
+            # Auto-generate table name
+            timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+            table = f"eval_{timestamp}"
+            print(f"✓ Auto-generated table name: {table}")
+
+        print()
+        dataset_name = None  # UC uses catalog.schema.table
+    else:
+        # Non-Databricks: use simple dataset name
+        catalog = schema = table = None
+
+        if args.dataset_name:
+            dataset_name = args.dataset_name
+            print(f"Using dataset name: {dataset_name}")
+        else:
+            # Auto-generate dataset name
+            timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+            dataset_name = f"agent-eval-{timestamp}"
+            print(f"✓ Auto-generated dataset name: {dataset_name}")
+        print()
+
+    # Generate code
+    print("=" * 60)
+    print("Generating Dataset Creation Script")
+    print("=" * 60)
+
+    code = generate_dataset_creation_code(
+        tracking_uri, experiment_id, dataset_name, test_cases, catalog, schema, table
+    )
+
+    # Write to file
+    output_file = args.output
+    with open(output_file, "w") as f:
+        f.write(code)
+
+    print(f"\n✓ Script generated: {output_file}")
+    print()
+
+    # Make executable
+    try:
+        os.chmod(output_file, 0o755)
+        print(f"✓ Made executable: chmod +x {output_file}")
+    except Exception:
+        pass
+
+    print()
+    print("=" * 60)
+    print("Next Steps")
+    print("=" * 60)
+    print()
+    print(f"1. Review the generated script: {output_file}")
+    print(f"2. Execute it: python {output_file}")
+    print("3. Verify dataset was created: python scripts/list_datasets.py")
+    print()
+    print("=" * 60)
+
+
+if __name__ == "__main__":
+    main()

package/canonical/skills/agent-evaluation/scripts/list_datasets.py

@@ -0,0 +1,273 @@
+"""
+List and compare MLflow evaluation datasets in an experiment.
+
+This script discovers existing datasets before prompting to create new ones,
+preventing duplicate work and helping users make informed choices.
+
+Features:
+- Diversity metrics (query length variability, unique vocabulary)
+- Timeout protection for large experiments
+- Multiple output formats (table, JSON, names-only)
+- Sample query preview
+
+Usage:
+    python scripts/list_datasets.py                      # Table format (default)
+    python scripts/list_datasets.py --format json        # JSON output
+    python scripts/list_datasets.py --format names-only  # Names only (for piping)
+    python scripts/list_datasets.py --detailed          # Include diversity analysis
+
+Environment variables required:
+    MLFLOW_TRACKING_URI
+    MLFLOW_EXPERIMENT_ID
+"""
+
+import argparse
+import json
+import os
+import signal
+import sys
+
+import numpy as np
+
+from mlflow import MlflowClient
+from mlflow.genai.datasets import get_dataset
+from utils import validate_env_vars
+
+
+class TimeoutError(Exception):
+    """Custom timeout exception."""
+
+
+def timeout_handler(signum, frame):
+    """Handle timeout signal."""
+    raise TimeoutError()
+
+
+def parse_arguments():
+    """Parse command-line arguments."""
+    parser = argparse.ArgumentParser(description="List and compare MLflow evaluation datasets")
+    parser.add_argument("--dataset-name", help="Specific dataset to display")
+    parser.add_argument(
+        "--show-samples", type=int, default=5, help="Number of sample queries to show (default: 5)"
+    )
+    parser.add_argument(
+        "--format",
+        choices=["table", "json", "names-only"],
+        default="table",
+        help="Output format (default: table)",
+    )
+    parser.add_argument(
+        "--timeout",
+        type=int,
+        default=30,
+        help="Timeout in seconds for dataset search (default: 30)",
+    )
+    parser.add_argument(
+        "--detailed", action="store_true", help="Include detailed diversity analysis (slower)"
+    )
+    return parser.parse_args()
+
+
+def calculate_diversity_metrics(queries):
+    """Calculate diversity metrics for a list of queries."""
+    if not queries:
+        return 0.0, 0.0, 0.0
+
+    # Query length statistics
+    lengths = [len(q) for q in queries]
+    avg_length = np.mean(lengths)
+    std_length = np.std(lengths)
+
+    # Unique word count (simple diversity measure)
+    all_words = set()
+    for query in queries:
+        words = query.lower().split()
+        all_words.update(words)
+
+    unique_word_ratio = len(all_words) / len(queries) if queries else 0
+
+    return avg_length, std_length, unique_word_ratio
+
+
+def classify_diversity(std_length, unique_word_ratio, query_count):
+    """Classify diversity as HIGH, MEDIUM, or LOW."""
+    # Heuristics based on variability and vocabulary
+    if query_count < 5:
+        return "LOW (too few queries)"
+
+    if std_length > 30 and unique_word_ratio > 5:
+        return "HIGH"
+    elif std_length > 15 and unique_word_ratio > 3:
+        return "MEDIUM"
+    else:
+        return "LOW"
+
+
+def get_datasets_with_timeout(client, experiment_ids, timeout_seconds):
+    """Get datasets with timeout protection."""
+    # Set alarm for timeout
+    signal.signal(signal.SIGALRM, timeout_handler)
+    signal.alarm(timeout_seconds)
+
+    try:
+        datasets = client.search_datasets(experiment_ids=experiment_ids)
+        signal.alarm(0)  # Cancel alarm
+        return datasets
+    except TimeoutError:
+        signal.alarm(0)
+        print(f"⚠ Dataset search timed out after {timeout_seconds}s")
+        print("  Try: --timeout <seconds> to increase timeout")
+        return []
+    except Exception as e:
+        signal.alarm(0)
+        print(f"✗ Error searching datasets: {str(e)[:100]}")
+        return []
+
+
+def print_table_format(dataset_info, args):
+    """Print datasets in table format."""
+    if not dataset_info:
+        print("\n✗ No datasets found in this experiment")
+        print("\nTo create a new dataset:")
+        print("  python scripts/create_dataset_template.py --test-cases-file test_cases.txt")
+        return
+
+    print(f"\n✓ Found {len(dataset_info)} dataset(s):")
+    print("=" * 80)
+
+    for i, info in enumerate(dataset_info, 1):
+        print(f"\n{i}. {info['name']}")
+        print(f"   Queries: {info.get('count', '?')}")
+
+        if args.detailed:
+            if "avg_length" in info:
+                print(f"   Avg length: {info['avg_length']:.1f} chars")
+                print(f"   Std length: {info['std_length']:.1f} chars")
+                print(f"   Unique words/query: {info['unique_word_ratio']:.1f}")
+                print(f"   Diversity: {info.get('diversity', 'N/A')}")
+
+            if "samples" in info:
+                print(f"\n   Sample queries:")
+                for j, sample in enumerate(info["samples"], 1):
+                    preview = sample[:60] + "..." if len(sample) > 60 else sample
+                    print(f"     {j}. {preview}")
+
+    print("\n" + "=" * 80)
+    print("\nTo use a dataset in evaluation:")
+    print('  python scripts/run_evaluation_template.py --dataset-name "dataset_name"')
+
+
+def print_json_format(dataset_info):
+    """Print datasets in JSON format."""
+    print(json.dumps(dataset_info, indent=2))
+
+
+def print_names_only(dataset_info):
+    """Print dataset names only (one per line)."""
+    for info in dataset_info:
+        print(info["name"])
+
+
+def main():
+    """Main workflow."""
+    args = parse_arguments()
+
+    print("=" * 80)
+    print("MLflow Evaluation Datasets")
+    print("=" * 80)
+
+    # Check environment using utility
+    errors = validate_env_vars()
+    if errors:
+        print("\n✗ Environment validation failed:")
+        for error in errors:
+            print(f"  - {error}")
+        print("\nRun scripts/setup_mlflow.py to configure environment")
+        sys.exit(1)
+
+    experiment_id = os.getenv("MLFLOW_EXPERIMENT_ID")
+    print(f"\nExperiment ID: {experiment_id}")
+
+    # Get datasets
+    print("\nSearching for datasets...")
+    client = MlflowClient()
+
+    try:
+        if args.dataset_name:
+            # Search for specific dataset
+            print(f"  Looking for: {args.dataset_name}")
+            datasets = get_datasets_with_timeout(client, [experiment_id], args.timeout)
+            datasets = [d for d in datasets if d.name == args.dataset_name]
+
+            if not datasets:
+                print(f"\n✗ Dataset '{args.dataset_name}' not found")
+                sys.exit(1)
+        else:
+            # Get all datasets
+            datasets = get_datasets_with_timeout(client, [experiment_id], args.timeout)
+
+    except Exception as e:
+        print(f"\n✗ Error: {str(e)[:200]}")
+        sys.exit(1)
+
+    # Process datasets
+    dataset_info = []
+
+    for dataset in datasets:
+        info = {"name": dataset.name}
+
+        # Try to load dataset for detailed info
+        if args.detailed or args.show_samples > 0:
+            try:
+                ds = get_dataset(dataset.name)
+                df = ds.to_df()
+
+                info["count"] = len(df)
+
+                # Extract queries (flexible extraction from various input formats)
+                queries = []
+                for _, row in df.iterrows():
+                    inputs = row.get("inputs", {})
+                    if isinstance(inputs, dict):
+                        # Try common keys first, then use first non-empty value
+                        query = (
+                            inputs.get("query")
+                            or inputs.get("question")
+                            or inputs.get("input")
+                            or inputs.get("prompt")
+                            or next((v for v in inputs.values() if v), str(inputs))
+                        )
+                        queries.append(str(query))
+                    else:
+                        # If inputs is not a dict, use it directly
+                        queries.append(str(inputs))
+
+                # Calculate diversity metrics
+                if queries and args.detailed:
+                    avg_len, std_len, unique_ratio = calculate_diversity_metrics(queries)
+                    info["avg_length"] = avg_len
+                    info["std_length"] = std_len
+                    info["unique_word_ratio"] = unique_ratio
+                    info["diversity"] = classify_diversity(std_len, unique_ratio, len(queries))
+
+                # Sample queries
+                if queries and args.show_samples > 0:
+                    info["samples"] = queries[: args.show_samples]
+
+            except Exception as e:
+                info["count"] = "?"
+                info["error"] = str(e)[:50]
+
+        dataset_info.append(info)
+
+    # Output in requested format
+    if args.format == "json":
+        print_json_format(dataset_info)
+    elif args.format == "names-only":
+        print_names_only(dataset_info)
+    else:  # table
+        print_table_format(dataset_info, args)
+
+
+if __name__ == "__main__":
+    main()