opik 1.8.39__py3-none-any.whl → 1.9.71__py3-none-any.whl

opik/cli/exports/dataset.py

@@ -0,0 +1,278 @@
+"""Dataset export functionality."""
+
+import sys
+from datetime import datetime
+from pathlib import Path
+from typing import Optional
+
+import click
+from rich.console import Console
+
+import opik
+from .utils import (
+    debug_print,
+    dataset_to_csv_rows,
+    should_skip_file,
+    write_csv_data,
+    write_json_data,
+    print_export_summary,
+)
+
+console = Console()
+
+
+def export_single_dataset(
+    dataset: opik.Dataset,
+    output_dir: Path,
+    max_results: Optional[int],
+    force: bool,
+    debug: bool,
+    format: str,
+) -> int:
+    """Export a single dataset."""
+    try:
+        # Check if already exists and force is not set
+        if format.lower() == "csv":
+            dataset_file = output_dir / f"dataset_{dataset.name}.csv"
+        else:
+            dataset_file = output_dir / f"dataset_{dataset.name}.json"
+
+        if should_skip_file(dataset_file, force):
+            if debug:
+                debug_print(f"Skipping {dataset.name} (already exists)", debug)
+            return 0
+
+        # Get dataset items
+        if debug:
+            debug_print(f"Getting items for dataset: {dataset.name}", debug)
+        dataset_items = dataset.get_items()
+
+        # Format items for export
+        # Use all fields from each item (datasets can have any user-defined keys/values)
+        formatted_items = []
+        for item in dataset_items:
+            # Create a copy of the item, excluding the 'id' field if present
+            # (id is internal and not part of the dataset item content)
+            formatted_item = {k: v for k, v in item.items() if k != "id"}
+            formatted_items.append(formatted_item)
+
+        # Create dataset data structure
+        dataset_data = {
+            "name": dataset.name,
+            "description": dataset.description,
+            "items": formatted_items,
+            "downloaded_at": datetime.now().isoformat(),
+        }
+
+        # Save to file using the appropriate format
+        if format.lower() == "csv":
+            write_csv_data(dataset_data, dataset_file, dataset_to_csv_rows)
+        else:
+            write_json_data(dataset_data, dataset_file)
+
+        if debug:
+            debug_print(f"Exported dataset: {dataset.name}", debug)
+        return 1
+
+    except Exception as e:
+        console.print(f"[red]Error exporting dataset {dataset.name}: {e}[/red]")
+        return 0
+
+
+def export_dataset_by_name(
+    name: str,
+    workspace: str,
+    output_path: str,
+    max_results: Optional[int],
+    force: bool,
+    debug: bool,
+    format: str,
+    api_key: Optional[str] = None,
+) -> None:
+    """Export a dataset by exact name."""
+    try:
+        if debug:
+            debug_print(f"Exporting dataset: {name}", debug)
+
+        # Initialize client
+        if api_key:
+            client = opik.Opik(api_key=api_key, workspace=workspace)
+        else:
+            client = opik.Opik(workspace=workspace)
+
+        # Create output directory
+        output_dir = Path(output_path) / workspace / "datasets"
+        output_dir.mkdir(parents=True, exist_ok=True)
+
+        if debug:
+            debug_print(f"Target directory: {output_dir}", debug)
+
+        # Try to get dataset by exact name
+        try:
+            dataset = client.get_dataset(name)
+            if debug:
+                debug_print(f"Found dataset by direct lookup: {dataset.name}", debug)
+        except Exception as e:
+            console.print(f"[red]Dataset '{name}' not found: {e}[/red]")
+            sys.exit(1)
+
+        # Export the dataset
+        exported_count = export_single_dataset(
+            dataset, output_dir, max_results, force, debug, format
+        )
+
+        # Collect statistics for summary
+        stats = {
+            "datasets": 1 if exported_count > 0 else 0,
+            "datasets_skipped": 0 if exported_count > 0 else 1,
+        }
+
+        # Show export summary
+        print_export_summary(stats, format)
+
+        if exported_count > 0:
+            console.print(
+                f"[green]Successfully exported dataset '{name}' to {output_dir}[/green]"
+            )
+        else:
+            console.print(
+                f"[yellow]Dataset '{name}' already exists (use --force to re-download)[/yellow]"
+            )
+
+    except Exception as e:
+        console.print(f"[red]Error exporting dataset: {e}[/red]")
+        sys.exit(1)
+
+
+def export_experiment_datasets(
+    client: opik.Opik,
+    datasets_to_export: set[str],
+    datasets_dir: Path,
+    format: str,
+    debug: bool,
+    force: bool,
+) -> tuple[int, int]:
+    """Export datasets related to an experiment.
+
+    Args:
+        client: Opik client instance
+        datasets_to_export: Set of dataset names to export
+        datasets_dir: Directory to save datasets
+        format: Export format ('json' or 'csv')
+        debug: Enable debug output
+        force: Re-download datasets even if they already exist locally
+
+    Returns:
+        Tuple of (exported_count, skipped_count)
+    """
+    exported_count = 0
+    skipped_count = 0
+
+    for dataset_name in datasets_to_export:
+        try:
+            # Use format parameter to determine file extension
+            if format.lower() == "csv":
+                dataset_file = datasets_dir / f"dataset_{dataset_name}.csv"
+            else:
+                dataset_file = datasets_dir / f"dataset_{dataset_name}.json"
+            datasets_dir.mkdir(parents=True, exist_ok=True)
+
+            # Check if file already exists and should be skipped
+            if should_skip_file(dataset_file, force):
+                if debug:
+                    debug_print(
+                        f"Skipping dataset {dataset_name} (already exists)", debug
+                    )
+                else:
+                    console.print(
+                        f"[yellow]Skipping dataset: {dataset_name} (already exists)[/yellow]"
+                    )
+                skipped_count += 1
+                continue
+
+            dataset_obj = opik.Dataset(
+                name=dataset_name,
+                description=None,  # Description not available from experiment
+                rest_client=client.rest_client,
+            )
+            dataset_items = dataset_obj.get_items()
+
+            dataset_data = {
+                "dataset": {
+                    "name": dataset_name,
+                    "id": getattr(dataset_obj, "id", None),
+                },
+                # Use all fields from each item, excluding 'id' (internal field)
+                "items": [
+                    {k: v for k, v in item.items() if k != "id"}
+                    for item in dataset_items
+                ],
+                "downloaded_at": datetime.now().isoformat(),
+            }
+
+            # Save to file using the appropriate format
+            if format.lower() == "csv":
+                write_csv_data(dataset_data, dataset_file, dataset_to_csv_rows)
+            else:
+                write_json_data(dataset_data, dataset_file)
+
+            console.print(f"[green]Exported dataset: {dataset_name}[/green]")
+            exported_count += 1
+        except Exception as e:
+            if debug:
+                console.print(
+                    f"[yellow]Warning: Could not export dataset {dataset_name}: {e}[/yellow]"
+                )
+            else:
+                console.print(f"[red]Error exporting dataset {dataset_name}: {e}[/red]")
+
+    return exported_count, skipped_count
+
+
+@click.command(name="dataset")
+@click.argument("name", type=str)
+@click.option(
+    "--max-results",
+    type=int,
+    help="Maximum number of datasets to export. Limits the total number of datasets downloaded.",
+)
+@click.option(
+    "--path",
+    "-p",
+    type=click.Path(file_okay=False, dir_okay=True, writable=True),
+    default="opik_exports",
+    help="Directory to save exported data. Defaults to opik_exports.",
+)
+@click.option(
+    "--force",
+    is_flag=True,
+    help="Re-download items even if they already exist locally.",
+)
+@click.option(
+    "--debug",
+    is_flag=True,
+    help="Enable debug output to show detailed information about the export process.",
+)
+@click.option(
+    "--format",
+    type=click.Choice(["json", "csv"], case_sensitive=False),
+    default="json",
+    help="Format for exporting data. Defaults to json.",
+)
+@click.pass_context
+def export_dataset_command(
+    ctx: click.Context,
+    name: str,
+    max_results: Optional[int],
+    path: str,
+    force: bool,
+    debug: bool,
+    format: str,
+) -> None:
+    """Export a dataset by exact name to workspace/datasets."""
+    # Get workspace and API key from context
+    workspace = ctx.obj["workspace"]
+    api_key = ctx.obj.get("api_key") if ctx.obj else None
+    export_dataset_by_name(
+        name, workspace, path, max_results, force, debug, format, api_key
+    )