opik 1.9.41__py3-none-any.whl → 1.9.86__py3-none-any.whl

opik/api_objects/trace/trace_data.py

@@ -1,6 +1,5 @@
 import dataclasses
 import datetime
-import logging
 from typing import Any, Dict, List, Optional, Union
 
 import opik.api_objects.attachment as attachment
@@ -14,40 +13,19 @@ from opik.types import (
     LLMProvider,
     SpanType,
 )
-from .. import span, data_helpers
+from .. import span
+from ..observation_data import ObservationData
 
-LOGGER = logging.getLogger(__name__)
 
-
-# Engineer note:
-#
-# After moving to minimal python version 3.10, a lot of common content
-# from SpanData and TraceData can be moved to ObservationData parent dataclass.
-# Before that it's impossible because of the dataclasses limitation to have optional arguments
-# strictly after positional ones (including the attributes from the parent class).
-# In python 3.10 @dataclass(kw_only=True) should help.
 @dataclasses.dataclass
-class TraceData:
+class TraceData(ObservationData):
     """
     The TraceData object is returned when calling :func:`opik.opik_context.get_current_trace_data` from a tracked function.
     """
 
     id: str = dataclasses.field(default_factory=id_helpers.generate_id)
-    name: Optional[str] = None
-    start_time: Optional[datetime.datetime] = dataclasses.field(
-        default_factory=datetime_helpers.local_timestamp
-    )
-    end_time: Optional[datetime.datetime] = None
-    metadata: Optional[Dict[str, Any]] = None
-    input: Optional[Dict[str, Any]] = None
-    output: Optional[Dict[str, Any]] = None
-    tags: Optional[List[str]] = None
-    feedback_scores: Optional[List[FeedbackScoreDict]] = None
-    project_name: Optional[str] = None
     created_by: Optional[CreatedByType] = None
-    error_info: Optional[ErrorInfoDict] = None
     thread_id: Optional[str] = None
-    attachments: Optional[List[attachment.Attachment]] = None
 
     def create_child_span_data(
         self,
@@ -91,55 +69,6 @@ class TraceData:
             attachments=attachments,
         )
 
-    def update(self, **new_data: Any) -> "TraceData":
-        for key, value in new_data.items():
-            if value is None:
-                continue
-
-            if key not in self.__dict__ and key != "prompts":
-                LOGGER.debug(
-                    "An attempt to update span with parameter name it doesn't have: %s",
-                    key,
-                )
-                continue
-
-            if key == "metadata":
-                self.metadata = data_helpers.merge_metadata(
-                    self.metadata, new_metadata=value
-                )
-                continue
-            elif key == "output":
-                self.output = data_helpers.merge_outputs(self.output, new_outputs=value)
-                continue
-            elif key == "input":
-                self.input = data_helpers.merge_inputs(self.input, new_inputs=value)
-                continue
-            elif key == "attachments":
-                self._update_attachments(value)
-                continue
-            elif key == "tags":
-                self.tags = data_helpers.merge_tags(self.tags, new_tags=value)
-                continue
-            elif key == "prompts":
-                self.metadata = data_helpers.merge_metadata(
-                    self.metadata, new_metadata=new_data.get("metadata"), prompts=value
-                )
-                continue
-
-            self.__dict__[key] = value
-
-        return self
-
-    def init_end_time(self) -> "TraceData":
-        self.end_time = datetime_helpers.local_timestamp()
-        return self
-
-    def _update_attachments(self, attachments: List[attachment.Attachment]) -> None:
-        if self.attachments is None:
-            self.attachments = attachments
-        else:
-            self.attachments.extend(attachments)
-
     @property
     def as_start_parameters(self) -> Dict[str, Any]:
         """Returns parameters of this trace to be sent to the server when starting a new trace."""

opik/api_objects/validation_helpers.py

@@ -1,7 +1,7 @@
 import logging
 from typing import Any, Optional, cast, Union, Dict
 
-from ..types import FeedbackScoreDict
+from ..types import BatchFeedbackScoreDict
 from ..validation import feedback_score as feedback_score_validator
 from .. import logging_messages, llm_usage
 from opik.types import LLMProvider
@@ -38,7 +38,7 @@ def validate_and_parse_usage(
 
 def validate_feedback_score(
     feedback_score: Any, logger: logging.Logger
-) -> Optional[FeedbackScoreDict]:
+) -> Optional[BatchFeedbackScoreDict]:
     feedback_score_validator_ = feedback_score_validator.FeedbackScoreValidator(
         feedback_score
     )
@@ -51,4 +51,4 @@ def validate_feedback_score(
         )
         return None
 
-    return cast(FeedbackScoreDict, feedback_score)
+    return cast(BatchFeedbackScoreDict, feedback_score)

opik/cli/exports/__init__.py

@@ -0,0 +1,131 @@
+"""Download command for Opik CLI."""
+
+from typing import Optional
+
+import click
+
+from .dataset import export_dataset_command
+from .experiment import export_experiment_command
+from .prompt import export_prompt_command
+from .project import export_project_command
+
+EXPORT_CONTEXT_SETTINGS = {"help_option_names": ["-h", "--help"]}
+
+
+@click.group(
+    name="export", context_settings=EXPORT_CONTEXT_SETTINGS, invoke_without_command=True
+)
+@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 export_group(ctx: click.Context, workspace: str, api_key: Optional[str]) -> None:
+    """Export data from Opik workspace.
+
+    This command allows you to export specific data from an Opik workspace to local files.
+    Supported data types include datasets, projects, experiments, and prompts.
+
+    \b
+    General Usage:
+        opik export WORKSPACE ITEM NAME [OPTIONS]
+
+    \b
+    Data Types (ITEM):
+        dataset      Export a dataset by exact name (exports dataset definition and items)
+        project      Export a project by name or ID (exports project traces and metadata)
+        experiment   Export an experiment by name or ID (exports experiment configuration and results)
+        prompt       Export a prompt by exact name (exports prompt templates and versions)
+
+    \b
+    Common Options:
+        --path, -p       Directory to save exported data (default: opik_exports)
+        --format         Export format: json or csv (default: json)
+        --max-results    Maximum number of items to export (varies by data type)
+        --force          Re-download items even if they already exist locally
+        --debug          Show detailed information about the export process
+
+    \b
+    Examples:
+        # Export a specific dataset
+        opik export my-workspace dataset "my-dataset"
+
+        # Export a project with OQL filter
+        opik export my-workspace project "my-project" --filter "status:completed"
+
+        # Export an experiment with dataset filter (by name or ID)
+        opik export my-workspace experiment "my-experiment" --dataset "my-dataset"
+        opik export my-workspace experiment "01234567-89ab-cdef-0123-456789abcdef" --dataset "my-dataset"
+
+        # Export in CSV format to a specific directory
+        opik export my-workspace prompt "my-template" --format csv --path ./custom-exports
+    """
+    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
+    )
+
+    # If no subcommand was invoked, show helpful error
+    if ctx.invoked_subcommand is None:
+        available_items = ", ".join(
+            sorted(["dataset", "experiment", "prompt", "project"])
+        )
+        click.echo(
+            f"Error: Missing ITEM.\n\n"
+            f"Available items: {available_items}\n\n"
+            f"Usage: opik export {workspace} ITEM NAME [OPTIONS]\n\n"
+            f"Examples:\n"
+            f'  opik export {workspace} dataset "my-dataset"\n'
+            f'  opik export {workspace} project "my-project"\n'
+            f'  opik export {workspace} experiment "my-experiment"\n'
+            f'  opik export {workspace} prompt "my-template"\n\n'
+            f"Run 'opik export {workspace} --help' for more information.",
+            err=True,
+        )
+        ctx.exit(2)
+
+
+# Set subcommand metavar to ITEM instead of COMMAND
+export_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(
+    export_group,
+    "format_commands",
+    format_commands.__get__(export_group, type(export_group)),
+)
+
+
+# Add the subcommands
+export_group.add_command(export_dataset_command)
+export_group.add_command(export_experiment_command)
+export_group.add_command(export_prompt_command)
+export_group.add_command(export_project_command)

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
+    )