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

opik/cli/exports/experiment.py

@@ -0,0 +1,784 @@
+"""Experiment export functionality."""
+
+import sys
+from concurrent.futures import Future, ThreadPoolExecutor, as_completed
+from datetime import datetime
+from pathlib import Path
+from typing import Optional, List, Dict, Tuple
+
+import click
+from rich.console import Console
+from rich.progress import (
+    Progress,
+    SpinnerColumn,
+    TextColumn,
+    BarColumn,
+    TaskProgressColumn,
+)
+
+import opik
+from opik import exceptions
+from .utils import (
+    create_experiment_data_structure,
+    debug_print,
+    write_json_data,
+    write_csv_data,
+    print_export_summary,
+    should_skip_file,
+    trace_to_csv_rows,
+)
+from .dataset import export_experiment_datasets
+from .prompt import (
+    export_related_prompts_by_name,
+    export_prompts_by_ids,
+)
+
+console = Console()
+
+# Batch size for parallel trace fetching
+BATCH_SIZE = 100
+# Maximum number of concurrent workers for parallel execution
+MAX_WORKERS = 20
+
+
+def _fetch_trace_data(
+    client: opik.Opik,
+    trace_id: str,
+    project_name_cache: dict[str, str],
+    debug: bool,
+) -> Optional[Tuple[str, dict, str]]:
+    """Fetch trace and span data for a single trace ID.
+
+    Returns:
+        Tuple of (trace_id, trace_data_dict, project_name) or None if failed.
+    """
+    try:
+        # Get trace by ID
+        trace = client.get_trace_content(trace_id)
+
+        # Get project name for this trace
+        if not trace.project_id:
+            return None
+
+        # Get project name (use cache if available)
+        if trace.project_id not in project_name_cache:
+            try:
+                project = client.get_project(trace.project_id)
+                project_name_cache[trace.project_id] = project.name
+            except Exception as e:
+                if debug:
+                    debug_print(
+                        f"Warning: Could not get project for trace {trace_id}: {e}",
+                        debug,
+                    )
+                return None
+
+        project_name = project_name_cache[trace.project_id]
+
+        # Get spans for this trace
+        spans = client.search_spans(
+            trace_id=trace_id,
+            max_results=1000,
+            truncate=False,
+        )
+
+        # Create trace data structure
+        trace_data = {
+            "trace": trace.model_dump(),
+            "spans": [span.model_dump() for span in spans],
+            "downloaded_at": datetime.now().isoformat(),
+            "project_name": project_name,
+        }
+
+        return (trace_id, trace_data, project_name)
+    except Exception as e:
+        if debug:
+            import traceback
+
+            debug_print(
+                f"Error fetching trace {trace_id}: {e}\n{traceback.format_exc()}", debug
+            )
+        return None
+
+
+def _write_trace_file(
+    trace_id: str,
+    trace_data: dict,
+    project_name: str,
+    workspace_root: Path,
+    format: str,
+    force: bool,
+    debug: bool,
+) -> bool:
+    """Write a single trace to file. Returns True if exported, False if skipped."""
+    try:
+        # Save trace in projects/PROJECT_NAME/ directory
+        project_dir = workspace_root / "projects" / project_name
+        project_dir.mkdir(parents=True, exist_ok=True)
+
+        # Determine file path based on format
+        if format.lower() == "csv":
+            file_path = project_dir / f"trace_{trace_id}.csv"
+        else:
+            file_path = project_dir / f"trace_{trace_id}.json"
+
+        # Check if file already exists and should be skipped
+        if should_skip_file(file_path, force):
+            if debug:
+                debug_print(f"Skipping trace {trace_id} (already exists)", debug)
+            return False
+
+        # Save to file using the appropriate format
+        if format.lower() == "csv":
+            write_csv_data(trace_data, file_path, trace_to_csv_rows)
+            if debug:
+                debug_print(f"Wrote CSV file: {file_path}", debug)
+        else:
+            write_json_data(trace_data, file_path)
+            if debug:
+                debug_print(f"Wrote JSON file: {file_path}", debug)
+
+        return True
+    except Exception as e:
+        console.print(f"[red]Error writing trace {trace_id} to file: {e}[/red]")
+        if debug:
+            import traceback
+
+            debug_print(f"Traceback: {traceback.format_exc()}", debug)
+        return False
+
+
+def export_traces_by_ids(
+    client: opik.Opik,
+    trace_ids: List[str],
+    workspace_root: Path,
+    max_traces: Optional[int],
+    format: str,
+    debug: bool,
+    force: bool,
+) -> tuple[int, int]:
+    """Export traces by their IDs using parallel batch processing.
+
+    Traces are saved in projects/PROJECT_NAME/ directory based on each trace's project.
+    Uses parallel execution to fetch traces/spans and write files concurrently.
+    """
+    exported_count = 0
+    skipped_count = 0
+
+    if max_traces:
+        trace_ids = trace_ids[:max_traces]
+
+    if not trace_ids:
+        return 0, 0
+
+    if debug:
+        debug_print(
+            f"Exporting {len(trace_ids)} trace(s) in batches of {BATCH_SIZE}", debug
+        )
+
+    # Cache project names to avoid repeated API calls (shared across threads)
+    project_name_cache: dict[str, str] = {}
+
+    # Use progress bar for trace export
+    with Progress(
+        SpinnerColumn(),
+        TextColumn("[progress.description]{task.description}"),
+        BarColumn(),
+        TaskProgressColumn(),
+        console=console,
+    ) as progress:
+        task = progress.add_task(
+            f"Exporting {len(trace_ids)} traces...", total=len(trace_ids)
+        )
+
+        # Process traces in batches
+        for batch_start in range(0, len(trace_ids), BATCH_SIZE):
+            batch_end = min(batch_start + BATCH_SIZE, len(trace_ids))
+            batch_trace_ids = trace_ids[batch_start:batch_end]
+
+            if debug:
+                debug_print(
+                    f"Batch {batch_start // BATCH_SIZE + 1}: traces {batch_start + 1}-{batch_end}",
+                    debug,
+                )
+
+            # Fetch trace data in parallel
+            fetched_traces: dict[str, Tuple[dict, str]] = {}
+
+            with ThreadPoolExecutor(max_workers=MAX_WORKERS) as fetch_executor:
+                # Submit all trace fetch tasks and track trace_id for each future
+                fetch_futures: Dict[Future[Optional[Tuple[str, dict, str]]], str] = {}
+                for trace_id in batch_trace_ids:
+                    fetch_future: Future[Optional[Tuple[str, dict, str]]] = (
+                        fetch_executor.submit(
+                            _fetch_trace_data,
+                            client,
+                            trace_id,
+                            project_name_cache,
+                            debug,
+                        )
+                    )
+                    fetch_futures[fetch_future] = trace_id
+
+                # Collect completed fetches
+                for fetch_future in as_completed(fetch_futures):
+                    trace_id = fetch_futures[fetch_future]
+                    try:
+                        result = fetch_future.result()
+                        if result is not None:
+                            fetched_trace_id, trace_data, project_name = result
+                            fetched_traces[fetched_trace_id] = (
+                                trace_data,
+                                project_name,
+                            )
+                    except Exception as e:
+                        if debug:
+                            console.print(
+                                f"[red]Error fetching trace {trace_id}: {e}[/red]"
+                            )
+
+            # Write files in parallel
+            with ThreadPoolExecutor(max_workers=MAX_WORKERS) as write_executor:
+                # Submit all write tasks and track trace_id for each future
+                write_futures: Dict[Future[bool], str] = {}
+                for trace_id, (trace_data, project_name) in fetched_traces.items():
+                    write_future: Future[bool] = write_executor.submit(
+                        _write_trace_file,
+                        trace_id,
+                        trace_data,
+                        project_name,
+                        workspace_root,
+                        format,
+                        force,
+                        debug,
+                    )
+                    write_futures[write_future] = trace_id
+
+                # Process completed writes
+                for write_future in as_completed(write_futures):
+                    trace_id = write_futures[write_future]
+                    try:
+                        if write_future.result():
+                            exported_count += 1
+                        else:
+                            skipped_count += 1
+                    except Exception as e:
+                        if debug:
+                            console.print(
+                                f"[red]Error writing trace {trace_id}: {e}[/red]"
+                            )
+                    finally:
+                        progress.update(
+                            task,
+                            advance=1,
+                            description=f"Exported {exported_count}/{len(trace_ids)} traces",
+                        )
+
+                # Update progress for traces that failed to fetch
+                for trace_id in batch_trace_ids:
+                    if trace_id not in fetched_traces:
+                        progress.update(task, advance=1)
+
+    return exported_count, skipped_count
+
+
+def export_experiment_by_id(
+    client: opik.Opik,
+    output_dir: Path,
+    experiment_id: str,
+    max_traces: Optional[int],
+    force: bool,
+    debug: bool,
+    format: str,
+    trace_ids_collector: Optional[set[str]] = None,
+) -> tuple[Dict[str, int], int]:
+    """Export a specific experiment by ID, including related datasets and traces.
+
+    Returns:
+        Tuple of (stats dictionary, file_written flag) where:
+        - stats: Dictionary with keys "datasets", "prompts", "traces" and their counts
+        - file_written: 1 if experiment file was written, 0 if skipped or error
+    """
+    try:
+        console.print(f"[blue]Fetching experiment by ID: {experiment_id}[/blue]")
+
+        # Get the specific experiment by ID
+        experiment = client.get_experiment_by_id(experiment_id)
+        if not experiment:
+            console.print(f"[red]Experiment '{experiment_id}' not found[/red]")
+            # Return empty stats and 0 for file written when not found
+            return ({"datasets": 0, "prompts": 0, "traces": 0}, 0)
+
+        debug_print(f"Found experiment: {experiment.name}", debug)
+
+        # Get experiment items first (this can be slow for large experiments)
+        console.print("[blue]Fetching experiment items...[/blue]")
+        with Progress(
+            SpinnerColumn(),
+            TextColumn("[progress.description]{task.description}"),
+            console=console,
+        ) as progress:
+            task = progress.add_task("Getting experiment items...", total=None)
+            experiment_items = experiment.get_items()
+            progress.update(task, description="Got experiment items")
+
+        # Create experiment data structure
+        experiment_data = create_experiment_data_structure(experiment, experiment_items)
+
+        # Save experiment data
+        # Include experiment ID in filename to handle multiple experiments with same name
+        experiment_file = (
+            output_dir / f"experiment_{experiment.name}_{experiment.id}.json"
+        )
+        file_already_exists = experiment_file.exists()
+        experiment_file_written = False
+
+        if not file_already_exists or force:
+            write_json_data(experiment_data, experiment_file)
+            experiment_file_written = True
+            debug_print(
+                f"Exported experiment: {experiment.name} (ID: {experiment.id})", debug
+            )
+        else:
+            debug_print(
+                f"Skipping experiment {experiment.name} (ID: {experiment.id}) (already exists)",
+                debug,
+            )
+
+        # Related prompts and traces are handled at the batch level
+        # Only export related prompts by name (this is experiment-specific and can't be easily deduplicated)
+        stats = {
+            "datasets": 0,
+            "datasets_skipped": 0,
+            "prompts": 0,
+            "prompts_skipped": 0,
+            "traces": 0,
+            "traces_skipped": 0,
+        }
+        stats["prompts"] = export_related_prompts_by_name(
+            client, experiment, output_dir, force, debug, format
+        )
+
+        # Collect trace IDs from experiment items (for batch export later)
+        trace_ids = [item.trace_id for item in experiment_items if item.trace_id]
+        if trace_ids_collector is not None:
+            trace_ids_collector.update(trace_ids)
+
+        # Traces are exported at batch level, so we don't export them here
+        stats["traces"] = 0
+        stats["traces_skipped"] = 0
+
+        if debug:
+            console.print(
+                f"[green]Experiment {experiment.name} exported with stats: {stats}[/green]"
+            )
+
+        # Return stats dictionary and whether file was written
+        return (stats, 1 if experiment_file_written else 0)
+
+    except Exception as e:
+        console.print(f"[red]Error exporting experiment {experiment_id}: {e}[/red]")
+        # Return empty stats and 0 for file written on error
+        return ({"datasets": 0, "prompts": 0, "traces": 0}, 0)
+
+
+def export_experiment_by_name(
+    name: str,
+    workspace: str,
+    output_path: str,
+    dataset: Optional[str],
+    max_traces: Optional[int],
+    force: bool,
+    debug: bool,
+    format: str,
+    api_key: Optional[str] = None,
+) -> None:
+    """Export an experiment by exact name."""
+    try:
+        if debug:
+            debug_print(f"Exporting experiment: {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 / "experiments"
+        output_dir.mkdir(parents=True, exist_ok=True)
+        datasets_dir = Path(output_path) / workspace / "datasets"
+        datasets_dir.mkdir(parents=True, exist_ok=True)
+
+        if debug:
+            debug_print(f"Target directory: {output_dir}", debug)
+
+        # Try to get experiments by exact name
+        try:
+            experiments = client.get_experiments_by_name(name)
+            if not experiments:
+                console.print(f"[red]Experiment '{name}' not found[/red]")
+                return
+
+            if debug:
+                debug_print(
+                    f"Found {len(experiments)} experiment(s) with name '{name}'", debug
+                )
+
+            if len(experiments) > 1:
+                console.print(
+                    f"[blue]Found {len(experiments)} experiments with name '{name}', exporting all of them[/blue]"
+                )
+        except Exception as e:
+            console.print(f"[red]Experiment '{name}' not found: {e}[/red]")
+            return
+
+        # Filter experiments by dataset if specified (client-side filtering)
+        if dataset:
+            experiments = [exp for exp in experiments if exp.dataset_name == dataset]
+            if not experiments:
+                console.print(
+                    f"[yellow]No experiments found with name '{name}' using dataset '{dataset}'[/yellow]"
+                )
+                return
+            if debug:
+                debug_print(
+                    f"Filtered to {len(experiments)} experiment(s) using dataset '{dataset}'",
+                    debug,
+                )
+
+        # Collect all unique resources from all experiments first
+        unique_datasets = set()
+        unique_prompt_ids: set[str] = set()
+
+        # First pass: collect datasets and prompt IDs (these are available without fetching items)
+        for experiment in experiments:
+            if experiment.dataset_name:
+                unique_datasets.add(experiment.dataset_name)
+            # Get experiment data to access prompt_versions
+            experiment_data = experiment.get_experiment_data()
+            if experiment_data.prompt_versions:
+                for prompt_version in experiment_data.prompt_versions:
+                    if prompt_version.prompt_id:
+                        unique_prompt_ids.add(prompt_version.prompt_id)
+
+        # Export all unique datasets once before processing experiments
+        datasets_exported = 0
+        datasets_skipped = 0
+        if unique_datasets:
+            if len(unique_datasets) > 1:
+                console.print(
+                    f"[blue]Exporting {len(unique_datasets)} unique dataset(s) used by these experiments...[/blue]"
+                )
+            datasets_exported, datasets_skipped = export_experiment_datasets(
+                client, unique_datasets, datasets_dir, format, debug, force
+            )
+
+        # Export all unique prompts once before processing experiments
+        prompts_dir = output_dir.parent / "prompts"
+        prompts_dir.mkdir(parents=True, exist_ok=True)
+        prompts_exported = 0
+        prompts_skipped = 0
+        if unique_prompt_ids:
+            if len(unique_prompt_ids) > 1:
+                console.print(
+                    f"[blue]Exporting {len(unique_prompt_ids)} unique prompt(s) used by these experiments...[/blue]"
+                )
+            prompts_exported, prompts_skipped = export_prompts_by_ids(
+                client, unique_prompt_ids, prompts_dir, format, debug, force
+            )
+
+        # Collect all unique trace IDs from all experiments as we process them
+        # We'll collect them during the first pass, then export once
+        all_trace_ids: set[str] = set()
+
+        # Export all matching experiments
+        exported_count = 0
+        skipped_count = 0
+
+        # Aggregate stats from all experiments (prompts and traces already exported at batch level)
+        aggregated_stats = {
+            "prompts": 0,
+            "prompts_skipped": 0,
+        }
+
+        for experiment in experiments:
+            if debug:
+                debug_print(
+                    f"Exporting experiment: {experiment.name} (ID: {experiment.id})",
+                    debug,
+                )
+
+            result = export_experiment_by_id(
+                client,
+                output_dir,
+                experiment.id,
+                max_traces,
+                force,
+                debug,
+                format,
+                all_trace_ids,
+            )
+
+            # result is a tuple: (stats_dict, file_written_flag)
+            exp_stats, file_written = result
+            # Aggregate stats (only related prompts, traces already handled)
+            aggregated_stats["prompts"] += exp_stats.get("prompts", 0)
+            aggregated_stats["prompts_skipped"] += exp_stats.get("prompts_skipped", 0)
+
+            if file_written > 0:
+                exported_count += 1
+            else:
+                skipped_count += 1
+
+        # Export all unique traces once after collecting them from all experiments
+        workspace_root = output_dir.parent
+        traces_exported = 0
+        traces_skipped = 0
+        if all_trace_ids:
+            trace_ids_list = list(all_trace_ids)
+            if max_traces:
+                trace_ids_list = trace_ids_list[:max_traces]
+            if len(trace_ids_list) > 0:
+                if len(all_trace_ids) > 1:
+                    console.print(
+                        f"[blue]Exporting {len(trace_ids_list)} unique trace(s) from these experiments...[/blue]"
+                    )
+                traces_exported, traces_skipped = export_traces_by_ids(
+                    client, trace_ids_list, workspace_root, None, format, debug, force
+                )
+
+        # Collect statistics for summary
+        stats = {
+            "experiments": exported_count,
+            "experiments_skipped": skipped_count,
+            "datasets": datasets_exported,
+            "datasets_skipped": datasets_skipped,
+            "prompts": prompts_exported + aggregated_stats["prompts"],
+            "prompts_skipped": prompts_skipped + aggregated_stats["prompts_skipped"],
+            "traces": traces_exported,
+            "traces_skipped": traces_skipped,
+        }
+
+        # Show export summary
+        print_export_summary(stats, format)
+
+        if exported_count > 0:
+            if len(experiments) > 1:
+                console.print(
+                    f"[green]Successfully exported {exported_count} experiment(s) with name '{name}' to {output_dir}[/green]"
+                )
+            else:
+                console.print(
+                    f"[green]Successfully exported experiment '{name}' to {output_dir}[/green]"
+                )
+        else:
+            console.print(
+                f"[yellow]All {len(experiments)} experiment(s) with name '{name}' already exist (use --force to re-download)[/yellow]"
+            )
+
+    except Exception as e:
+        console.print(f"[red]Error exporting experiment: {e}[/red]")
+        sys.exit(1)
+
+
+def export_experiment_by_name_or_id(
+    name_or_id: str,
+    workspace: str,
+    output_path: str,
+    dataset: Optional[str],
+    max_traces: Optional[int],
+    force: bool,
+    debug: bool,
+    format: str,
+    api_key: Optional[str] = None,
+) -> None:
+    """Export an experiment by name or ID.
+
+    First tries to get the experiment by ID. If not found, tries by name.
+    """
+    try:
+        if debug:
+            debug_print(f"Attempting to export experiment: {name_or_id}", 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 / "experiments"
+        output_dir.mkdir(parents=True, exist_ok=True)
+        datasets_dir = Path(output_path) / workspace / "datasets"
+        datasets_dir.mkdir(parents=True, exist_ok=True)
+
+        # Try to get experiment by ID first
+        try:
+            if debug:
+                debug_print(f"Trying to get experiment by ID: {name_or_id}", debug)
+            experiment = client.get_experiment_by_id(name_or_id)
+
+            # Successfully found by ID, export it
+            if debug:
+                debug_print(
+                    f"Found experiment by ID: {experiment.name} (ID: {experiment.id})",
+                    debug,
+                )
+
+            # Collect trace IDs as we export
+            trace_ids_collector: set[str] = set()
+
+            # Use the ID-based export function
+            result = export_experiment_by_id(
+                client,
+                output_dir,
+                name_or_id,
+                max_traces,
+                force,
+                debug,
+                format,
+                trace_ids_collector,
+            )
+
+            exp_stats, file_written = result
+
+            # Export related datasets
+            unique_datasets = set()
+            if experiment.dataset_name:
+                unique_datasets.add(experiment.dataset_name)
+
+            datasets_exported = 0
+            datasets_skipped = 0
+            if unique_datasets:
+                datasets_exported, datasets_skipped = export_experiment_datasets(
+                    client, unique_datasets, datasets_dir, format, debug, force
+                )
+
+            # Export traces collected from experiment items
+            workspace_root = output_dir.parent
+            traces_exported = 0
+            traces_skipped = 0
+            if trace_ids_collector:
+                trace_ids_list = list(trace_ids_collector)
+                if max_traces:
+                    trace_ids_list = trace_ids_list[:max_traces]
+                if len(trace_ids_list) > 0:
+                    traces_exported, traces_skipped = export_traces_by_ids(
+                        client,
+                        trace_ids_list,
+                        workspace_root,
+                        None,
+                        format,
+                        debug,
+                        force,
+                    )
+
+            # Collect statistics for summary
+            stats = {
+                "experiments": 1 if file_written > 0 else 0,
+                "experiments_skipped": 0 if file_written > 0 else 1,
+                "datasets": datasets_exported,
+                "datasets_skipped": datasets_skipped,
+                "prompts": exp_stats.get("prompts", 0),
+                "prompts_skipped": exp_stats.get("prompts_skipped", 0),
+                "traces": traces_exported,
+                "traces_skipped": traces_skipped,
+            }
+
+            # Show export summary
+            print_export_summary(stats, format)
+
+            if file_written > 0:
+                console.print(
+                    f"[green]Successfully exported experiment '{experiment.name}' (ID: {experiment.id}) to {output_dir}[/green]"
+                )
+            else:
+                console.print(
+                    f"[yellow]Experiment '{experiment.name}' (ID: {experiment.id}) already exists (use --force to re-download)[/yellow]"
+                )
+            return
+
+        except exceptions.ExperimentNotFound:
+            # Not found by ID, try by name
+            if debug:
+                debug_print(
+                    f"Experiment not found by ID, trying by name: {name_or_id}", debug
+                )
+            # Fall through to name-based export
+            pass
+
+        # Try by name (either because ID lookup failed or we're explicitly trying name)
+        export_experiment_by_name(
+            name_or_id,
+            workspace,
+            output_path,
+            dataset,
+            max_traces,
+            force,
+            debug,
+            format,
+            api_key,
+        )
+
+    except Exception as e:
+        console.print(f"[red]Error exporting experiment: {e}[/red]")
+        sys.exit(1)
+
+
+@click.command(name="experiment")
+@click.argument("name_or_id", type=str)
+@click.option(
+    "--dataset",
+    type=str,
+    help="Filter experiments by dataset name. Only experiments using this dataset will be exported.",
+)
+@click.option(
+    "--max-traces",
+    type=int,
+    help="Maximum number of traces to export per experiment. Limits the total number of traces 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_experiment_command(
+    ctx: click.Context,
+    name_or_id: str,
+    dataset: Optional[str],
+    max_traces: Optional[int],
+    path: str,
+    force: bool,
+    debug: bool,
+    format: str,
+) -> None:
+    """Export an experiment by exact name to workspace/experiments.
+
+    The command will first try to find the experiment by ID. If not found, it will try by name.
+    """
+    # Get workspace and API key from context
+    workspace = ctx.obj["workspace"]
+    api_key = ctx.obj.get("api_key") if ctx.obj else None
+    export_experiment_by_name_or_id(
+        name_or_id, workspace, path, dataset, max_traces, force, debug, format, api_key
+    )