"""Import command for Opik CLI.""" import sys from pathlib import Path from typing import Dict, Optional import click from rich.console import Console import opik from ..migration_manifest import MigrationManifest from .dataset import import_datasets_from_directory from .experiment import import_experiments_from_directory from .project import import_projects_from_directory from .prompt import import_prompts_from_directory from .utils import print_import_summary, debug_print console = Console() IMPORT_CONTEXT_SETTINGS = {"help_option_names": ["-h", "--help"]} def _import_by_type( import_type: str, path: str, workspace: str, dry_run: bool, name_pattern: Optional[str], debug: bool, recreate_experiments: bool = False, api_key: Optional[str] = None, force: bool = False, ) -> None: """ Import data by type (dataset, project, experiment) with pattern matching. Args: import_type: Type of data to import ("dataset", "project", "experiment") path: Base directory containing the exported data workspace: Target workspace name dry_run: Whether to show what would be imported without importing name_pattern: Optional string pattern to filter items by name debug: Enable debug output recreate_experiments: Whether to recreate experiments after importing force: Discard any existing manifest and restart from scratch """ try: debug_print(f"DEBUG: Starting {import_type} import from {path}", debug) # Initialize Opik client # _use_batching=True speeds up bulk ingestion. It is safe here because # the import flow only creates traces/spans (no update calls) and # explicitly calls client.flush() before exiting. if api_key: client = opik.Opik(api_key=api_key, workspace=workspace, _use_batching=True) else: client = opik.Opik(workspace=workspace, _use_batching=True) base_path = Path(path) # ------------------------------------------------------------------ # Manifest lifecycle management # MigrationManifest is only constructed for real (non-dry-run) imports # so that --dry-run never creates the manifest DB on disk. # ------------------------------------------------------------------ manifest: Optional[MigrationManifest] = None if not dry_run: manifest = MigrationManifest(base_path) if force: if MigrationManifest.exists(base_path): manifest.reset() console.print( "[yellow]--force: discarding existing manifest, starting fresh[/yellow]" ) else: if manifest.is_completed: console.print( "[green]Import already completed. Use --force to re-import.[/green]" ) return elif manifest.is_in_progress: console.print( f"[blue]Resuming interrupted import: " f"{manifest.completed_count()} file(s) already completed[/blue]" ) manifest.start() # ------------------------------------------------------------------ # Determine source directory # ------------------------------------------------------------------ if import_type == "dataset": source_dir = base_path / "datasets" elif import_type == "project": source_dir = base_path / "projects" elif import_type == "experiment": source_dir = base_path / "experiments" elif import_type == "prompt": source_dir = base_path / "prompts" else: console.print(f"[red]Unknown import type: {import_type}[/red]") return if not source_dir.exists(): console.print(f"[red]Source directory {source_dir} does not exist[/red]") sys.exit(1) debug_print(f"Source directory: {source_dir}", debug) stats: Dict[str, int] = {} if import_type == "dataset": stats = import_datasets_from_directory( client, source_dir, dry_run, name_pattern, debug, manifest=manifest ) elif import_type == "project": stats = import_projects_from_directory( client, source_dir, dry_run, name_pattern, debug, recreate_experiments, manifest=manifest, ) elif import_type == "experiment": stats = import_experiments_from_directory( client, source_dir, dry_run, name_pattern, debug, manifest=manifest ) elif import_type == "prompt": stats = import_prompts_from_directory( client, source_dir, dry_run, name_pattern, debug, manifest=manifest ) # Flush the async ingestion queue before returning so that all # enqueued traces/spans are persisted on the server. Without this, # the process can exit while the background worker is still sending # data, silently dropping items (especially under rate-limiting). if not dry_run: flushed = client.flush() if not flushed: console.print( "[yellow]Warning: flush timed out — some traces/spans may not have been ingested. " "Re-run the import to retry.[/yellow]" ) sys.exit(1) # FileUploadManager.flush() returns True even when individual # uploads fail (only False on timeout), so check failed_uploads # separately to catch silent upload failures. failed = client.__internal_api__failed_uploads__(timeout=None) if failed > 0: console.print( f"[yellow]Warning: {failed} file upload(s) failed during import. " "Re-run the import to retry.[/yellow]" ) sys.exit(1) # ------------------------------------------------------------------ # Mark manifest complete (only on full success, not dry-run) # ------------------------------------------------------------------ assert manifest is not None # guaranteed: manifest is set when not dry_run manifest.complete() # Display summary print_import_summary(stats) # Map import_type to the key used in stats dictionary type_key_map = { "dataset": "datasets", "prompt": "prompts", "project": "projects", "experiment": "experiments", } stats_key = type_key_map.get(import_type, import_type + "s") imported_count = stats.get(stats_key, 0) errors = stats.get(stats_key + "_errors", 0) if dry_run: console.print( f"[blue]Dry run complete: Would import {imported_count} {import_type}s[/blue]" ) else: if errors > 0: console.print( f"[yellow]Import completed with {errors} error(s) while importing {import_type}s[/yellow]" ) elif imported_count == 0: console.print(f"[yellow]No {import_type}s were imported[/yellow]") else: console.print( f"[green]Successfully imported {imported_count} {import_type}s[/green]" ) except Exception as e: console.print(f"[red]Error importing {import_type}s: {e}[/red]") sys.exit(1) @click.group(name="import", context_settings=IMPORT_CONTEXT_SETTINGS) @click.argument("workspace", type=str) @click.option( "--api-key", type=str, help="Opik API key. If not provided, will use OPIK_API_KEY environment variable or configuration.", ) @click.pass_context def import_group(ctx: click.Context, workspace: str, api_key: Optional[str]) -> None: """Import data to Opik workspace. This command allows you to import previously exported data back into an Opik workspace. Supported data types include projects, datasets, experiments, and prompts. A migration_manifest.db file is automatically maintained in the import directory. If an import is interrupted, re-running the same command will resume from where it left off without creating duplicates. Use --force to discard the manifest and restart. \b General Usage: opik import WORKSPACE TYPE NAME [OPTIONS] \b Data Types: project Import projects from path/projects/ (default: opik_exports) dataset Import datasets from path/datasets/ (default: opik_exports) experiment Import experiments from path/experiments/ (default: opik_exports) prompt Import prompts from path/prompts/ (default: opik_exports) \b Common Options: --path, -p Directory containing exported data (default: opik_exports) --dry-run Preview what would be imported without actually importing --force Discard manifest and restart from scratch --debug Show detailed information about the import process \b Examples: # Preview an experiment that would be imported opik import my-workspace experiment "my-experiment" --dry-run # Import a specific project opik import my-workspace project "my-project" # Import a specific dataset opik import my-workspace dataset "my-dataset" # Import from a custom path opik import my-workspace project "my-project" --path ./custom-exports/ # Resume an interrupted import (automatic — just re-run the same command) opik import my-workspace project "my-project" # Discard progress and restart from scratch opik import my-workspace project "my-project" --force """ ctx.ensure_object(dict) ctx.obj["workspace"] = workspace # Use API key from this command or from parent context ctx.obj["api_key"] = api_key or ( ctx.parent.obj.get("api_key") if ctx.parent and ctx.parent.obj else None ) # Set subcommand metavar to ITEM instead of COMMAND import_group.subcommand_metavar = "ITEM [ARGS]..." def format_commands( self: click.Group, ctx: click.Context, formatter: click.HelpFormatter ) -> None: """Override to change 'Commands' heading to 'Items'.""" commands = [] for subcommand in self.list_commands(ctx): cmd = self.get_command(ctx, subcommand) if cmd is None or cmd.hidden: continue commands.append((subcommand, cmd)) if len(commands): limit = formatter.width - 6 - max(len(cmd[0]) for cmd in commands) rows = [] for subcommand, cmd in commands: help = cmd.get_short_help_str(limit) rows.append((subcommand, help)) if rows: with formatter.section("Items"): formatter.write_dl(rows) # Override format_commands method setattr( import_group, "format_commands", format_commands.__get__(import_group, type(import_group)), ) @import_group.command(name="dataset") @click.argument("name", type=str) @click.option( "--path", "-p", type=click.Path(file_okay=False, dir_okay=True, readable=True), default="opik_exports", help="Directory containing exported data. Defaults to opik_exports.", ) @click.option( "--dry-run", is_flag=True, help="Show what would be imported without actually importing.", ) @click.option( "--force", is_flag=True, help="Discard the migration manifest and re-import everything from scratch.", ) @click.option( "--debug", is_flag=True, help="Enable debug output to show detailed information about the import process.", ) @click.pass_context def import_dataset( ctx: click.Context, name: str, path: str, dry_run: bool, force: bool, debug: bool, ) -> None: """Import datasets from workspace/datasets directory. This command imports datasets matching the specified name from the path/datasets/ directory. The name is matched using case-insensitive substring matching. \b Examples: \b # Preview a dataset that would be imported opik import my-workspace dataset "my-dataset" --dry-run \b # Import a specific dataset opik import my-workspace dataset "my-dataset" \b # Import datasets containing "training" in the name opik import my-workspace dataset "training" \b # Import from a custom path opik import my-workspace dataset "my-dataset" --path ./custom-exports/ """ workspace = ctx.obj["workspace"] api_key = ctx.obj.get("api_key") if ctx.obj else None _import_by_type( "dataset", path, workspace, dry_run, name, debug, api_key=api_key, force=force ) @import_group.command(name="project") @click.argument("name", type=str) @click.option( "--path", "-p", type=click.Path(file_okay=False, dir_okay=True, readable=True), default="opik_exports", help="Directory containing exported data. Defaults to opik_exports.", ) @click.option( "--dry-run", is_flag=True, help="Show what would be imported without actually importing.", ) @click.option( "--force", is_flag=True, help="Discard the migration manifest and re-import everything from scratch.", ) @click.option( "--debug", is_flag=True, help="Enable debug output to show detailed information about the import process.", ) @click.pass_context def import_project( ctx: click.Context, name: str, path: str, dry_run: bool, force: bool, debug: bool, ) -> None: """Import projects from workspace/projects directory. This command imports projects matching the specified name from the path/projects/ directory. The name is matched using case-insensitive substring matching. If a previous import was interrupted, re-running the same command automatically resumes from where it left off. Use --force to start over from scratch. \b Examples: \b # Preview a project that would be imported opik import my-workspace project "my-project" --dry-run \b # Import a specific project opik import my-workspace project "my-project" \b # Import projects containing "test" in the name opik import my-workspace project "test" \b # Import projects with debug output opik import my-workspace project "my-project" --debug \b # Import from a custom path opik import my-workspace project "my-project" --path ./custom-exports/ \b # Discard progress and restart opik import my-workspace project "my-project" --force """ workspace = ctx.obj["workspace"] api_key = ctx.obj.get("api_key") if ctx.obj else None _import_by_type( "project", path, workspace, dry_run, name, debug, True, # Always recreate experiments when importing projects api_key=api_key, force=force, ) @import_group.command(name="experiment") @click.argument("name", type=str) @click.option( "--path", "-p", type=click.Path(file_okay=False, dir_okay=True, readable=True), default="opik_exports", help="Directory containing exported data. Defaults to opik_exports.", ) @click.option( "--dry-run", is_flag=True, help="Show what would be imported without actually importing.", ) @click.option( "--force", is_flag=True, help="Discard the migration manifest and re-import everything from scratch.", ) @click.option( "--debug", is_flag=True, help="Enable debug output to show detailed information about the import process.", ) @click.pass_context def import_experiment( ctx: click.Context, name: str, path: str, dry_run: bool, force: bool, debug: bool, ) -> None: """Import experiments from workspace/experiments directory. This command imports experiments matching the specified name from the path/experiments/ directory. The name is matched using case-insensitive substring matching. If a previous import was interrupted, re-running the same command automatically resumes from where it left off. Use --force to start over from scratch. \b Examples: \b # Preview an experiment that would be imported opik import my-workspace experiment "my-experiment" --dry-run \b # Import a specific experiment opik import my-workspace experiment "my-experiment" \b # Import from a custom path opik import my-workspace experiment "my-experiment" --path ./custom-exports/ """ workspace = ctx.obj["workspace"] api_key = ctx.obj.get("api_key") if ctx.obj else None _import_by_type( "experiment", path, workspace, dry_run, name, debug, True, api_key=api_key, force=force, ) @import_group.command(name="prompt") @click.argument("name", type=str) @click.option( "--path", "-p", type=click.Path(file_okay=False, dir_okay=True, readable=True), default="opik_exports", help="Directory containing exported data. Defaults to opik_exports.", ) @click.option( "--dry-run", is_flag=True, help="Show what would be imported without actually importing.", ) @click.option( "--force", is_flag=True, help="Discard the migration manifest and re-import everything from scratch.", ) @click.option( "--debug", is_flag=True, help="Enable debug output to show detailed information about the import process.", ) @click.pass_context def import_prompt( ctx: click.Context, name: str, path: str, dry_run: bool, force: bool, debug: bool, ) -> None: """Import prompts from workspace/prompts directory. This command imports prompts matching the specified name from the path/prompts/ directory. The name is matched using case-insensitive substring matching. \b Examples: \b # Preview a prompt that would be imported opik import my-workspace prompt "my-prompt" --dry-run \b # Import a specific prompt opik import my-workspace prompt "my-prompt" \b # Import from a custom path opik import my-workspace prompt "my-prompt" --path ./custom-exports/ """ workspace = ctx.obj["workspace"] api_key = ctx.obj.get("api_key") if ctx.obj else None _import_by_type( "prompt", path, workspace, dry_run, name, debug, api_key=api_key, force=force )