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

opik/cli/exports/prompt.py

@@ -0,0 +1,578 @@
+"""Prompt export functionality."""
+
+import json
+import sys
+from datetime import datetime
+from pathlib import Path
+from typing import Any, List, Optional, Union
+
+import click
+from rich.console import Console
+
+import opik
+from opik.api_objects.prompt import Prompt, ChatPrompt
+from .utils import (
+    debug_print,
+    prompt_to_csv_rows,
+    should_skip_file,
+    write_csv_data,
+    write_json_data,
+    print_export_summary,
+)
+
+console = Console()
+
+
+def _get_prompt_content(prompt: Any) -> Any:
+    """Extract prompt content based on prompt type.
+
+    Args:
+        prompt: A Prompt or ChatPrompt instance
+
+    Returns:
+        For Prompt: the prompt string
+        For ChatPrompt: the template (list of message dicts)
+        Otherwise: None
+    """
+    if isinstance(prompt, Prompt):
+        return prompt.prompt
+    elif isinstance(prompt, ChatPrompt):
+        return prompt.template
+    return None
+
+
+def _get_template_structure(prompt: Any) -> str:
+    """Get template_structure based on prompt type.
+
+    Args:
+        prompt: A Prompt or ChatPrompt instance
+
+    Returns:
+        "text" for Prompt, "chat" for ChatPrompt, "text" as default
+    """
+    if isinstance(prompt, ChatPrompt):
+        return "chat"
+    elif isinstance(prompt, Prompt):
+        return "text"
+    return "text"
+
+
+def _get_prompt_type_string(prompt: Any) -> Optional[str]:
+    """Get prompt type as uppercase string.
+
+    Args:
+        prompt: A Prompt or ChatPrompt instance
+
+    Returns:
+        Uppercase type string (e.g., "JINJA2", "MUSTACHE") or None
+    """
+    prompt_type = getattr(prompt, "type", None)
+    if prompt_type is None:
+        return None
+
+    # If it's an enum, get the value and convert to uppercase
+    if hasattr(prompt_type, "value"):
+        return prompt_type.value.upper()
+    # If it's already a string, convert to uppercase
+    if isinstance(prompt_type, str):
+        return prompt_type.upper()
+    # Otherwise, convert to string and uppercase
+    return str(prompt_type).upper()
+
+
+def export_single_prompt(
+    client: opik.Opik,
+    prompt: Union[Prompt, ChatPrompt],
+    output_dir: Path,
+    max_results: Optional[int],
+    force: bool,
+    debug: bool,
+    format: str,
+) -> int:
+    """Export a single prompt."""
+    try:
+        # Check if already exists and force is not set
+        if format.lower() == "csv":
+            prompt_file = output_dir / f"prompts_{prompt.name.replace('/', '_')}.csv"
+        else:
+            prompt_file = output_dir / f"prompt_{prompt.name.replace('/', '_')}.json"
+
+        if should_skip_file(prompt_file, force):
+            if debug:
+                debug_print(f"Skipping {prompt.name} (already exists)", debug)
+            return 0
+
+        # Get prompt history - use appropriate method based on prompt type
+        prompt_history: List[Union[Prompt, ChatPrompt]]
+        if isinstance(prompt, ChatPrompt):
+            prompt_history = list(client.get_chat_prompt_history(prompt.name))
+        else:
+            prompt_history = list(client.get_prompt_history(prompt.name))
+
+        # Create prompt data structure
+        prompt_data = {
+            "name": prompt.name,
+            "current_version": {
+                "prompt": _get_prompt_content(prompt),
+                "metadata": getattr(prompt, "metadata", None),
+                "type": _get_prompt_type_string(prompt),
+                "commit": getattr(prompt, "commit", None),
+                "template_structure": _get_template_structure(prompt),
+            },
+            "history": [
+                {
+                    "prompt": _get_prompt_content(version),
+                    "metadata": getattr(version, "metadata", None),
+                    "type": _get_prompt_type_string(version),
+                    "commit": getattr(version, "commit", None),
+                    "template_structure": _get_template_structure(version),
+                }
+                for version in prompt_history
+            ],
+            "downloaded_at": datetime.now().isoformat(),
+        }
+
+        # Save to file using the appropriate format
+        if format.lower() == "csv":
+            write_csv_data(prompt_data, prompt_file, prompt_to_csv_rows)
+        else:
+            write_json_data(prompt_data, prompt_file)
+
+        if debug:
+            debug_print(f"Exported prompt: {prompt.name}", debug)
+        return 1
+
+    except Exception as e:
+        console.print(f"[red]Error exporting prompt {prompt.name}: {e}[/red]")
+        return 0
+
+
+def export_prompt_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 prompt by exact name."""
+    try:
+        if debug:
+            debug_print(f"Exporting prompt: {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 / "prompts"
+        output_dir.mkdir(parents=True, exist_ok=True)
+
+        if debug:
+            debug_print(f"Target directory: {output_dir}", debug)
+
+        # Try to get prompt by exact name
+        # Try ChatPrompt first, then regular Prompt
+        prompt: Optional[Union[Prompt, ChatPrompt]] = None
+        try:
+            prompt = client.get_chat_prompt(name)
+            if debug and prompt:
+                debug_print(f"Found ChatPrompt by direct lookup: {prompt.name}", debug)
+        except Exception:
+            # Not a ChatPrompt, try regular Prompt
+            try:
+                prompt = client.get_prompt(name)
+                if not prompt:
+                    console.print(f"[red]Prompt '{name}' not found[/red]")
+                    return
+                if debug:
+                    debug_print(f"Found Prompt by direct lookup: {prompt.name}", debug)
+            except Exception as e:
+                console.print(f"[red]Prompt '{name}' not found: {e}[/red]")
+                return
+
+        if prompt is None:
+            console.print(f"[red]Prompt '{name}' not found[/red]")
+            return
+
+        # Export the prompt
+        exported_count = export_single_prompt(
+            client, prompt, output_dir, max_results, force, debug, format
+        )
+
+        # Collect statistics for summary
+        stats = {
+            "prompts": 1 if exported_count > 0 else 0,
+            "prompts_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 prompt '{name}' to {output_dir}[/green]"
+            )
+        else:
+            console.print(
+                f"[yellow]Prompt '{name}' already exists (use --force to re-download)[/yellow]"
+            )
+
+    except Exception as e:
+        console.print(f"[red]Error exporting prompt: {e}[/red]")
+        sys.exit(1)
+
+
+def export_prompts_by_ids(
+    client: opik.Opik,
+    prompt_ids: set[str],
+    prompts_dir: Path,
+    format: str,
+    debug: bool,
+    force: bool,
+) -> tuple[int, int]:
+    """Export prompts by their IDs.
+
+    Args:
+        client: Opik client instance
+        prompt_ids: Set of prompt IDs to export
+        prompts_dir: Directory to save prompts
+        format: Export format ('json' or 'csv')
+        debug: Enable debug output
+        force: Re-download prompts even if they already exist locally
+
+    Returns:
+        Tuple of (exported_count, skipped_count)
+    """
+    exported_count = 0
+    skipped_count = 0
+
+    for prompt_id in prompt_ids:
+        try:
+            # Get the prompt - try ChatPrompt first, then regular Prompt
+            prompt: Optional[Union[Prompt, ChatPrompt]] = None
+            try:
+                prompt = client.get_chat_prompt(prompt_id)
+            except Exception:
+                # Not a ChatPrompt, try regular Prompt
+                prompt = client.get_prompt(prompt_id)
+
+            if not prompt:
+                if debug:
+                    console.print(
+                        f"[yellow]Warning: Prompt {prompt_id} not found[/yellow]"
+                    )
+                continue
+
+            # Determine file path
+            if format.lower() == "csv":
+                prompt_file = (
+                    prompts_dir
+                    / f"prompts_{prompt.name or getattr(prompt, 'id', 'unknown')}.csv"
+                )
+            else:
+                prompt_file = (
+                    prompts_dir
+                    / f"prompt_{prompt.name or getattr(prompt, 'id', 'unknown')}.json"
+                )
+
+            # Check if file already exists and should be skipped
+            if should_skip_file(prompt_file, force):
+                if debug:
+                    debug_print(
+                        f"Skipping prompt {prompt.name or prompt_id} (already exists)",
+                        debug,
+                    )
+                else:
+                    console.print(
+                        f"[yellow]Skipping prompt: {prompt.name or prompt_id} (already exists)[/yellow]"
+                    )
+                skipped_count += 1
+                continue
+
+            # Get prompt history - use appropriate method based on prompt type
+            prompt_history: List[Union[Prompt, ChatPrompt]]
+            if isinstance(prompt, ChatPrompt):
+                prompt_history = list(client.get_chat_prompt_history(prompt.name))
+            else:
+                prompt_history = list(client.get_prompt_history(prompt_id))
+
+            # Create prompt data structure
+            prompt_data = {
+                "prompt": {
+                    "id": getattr(prompt, "id", None),
+                    "name": prompt.name,
+                    "description": getattr(prompt, "description", None),
+                    "created_at": (
+                        created_at.isoformat()
+                        if (created_at := getattr(prompt, "created_at", None))
+                        else None
+                    ),
+                    "last_updated_at": (
+                        last_updated_at.isoformat()
+                        if (last_updated_at := getattr(prompt, "last_updated_at", None))
+                        else None
+                    ),
+                },
+                "current_version": {
+                    "prompt": _get_prompt_content(prompt),
+                    "metadata": getattr(prompt, "metadata", None),
+                    "type": _get_prompt_type_string(prompt),
+                    "commit": getattr(prompt, "commit", None),
+                    "template_structure": _get_template_structure(prompt),
+                },
+                "history": [
+                    {
+                        "prompt": _get_prompt_content(version),
+                        "metadata": getattr(version, "metadata", None),
+                        "type": _get_prompt_type_string(version),
+                        "commit": getattr(version, "commit", None),
+                        "template_structure": _get_template_structure(version),
+                    }
+                    for version in prompt_history
+                ],
+                "downloaded_at": datetime.now().isoformat(),
+            }
+
+            # Save prompt data using the appropriate format
+            if format.lower() == "csv":
+                write_csv_data(prompt_data, prompt_file, prompt_to_csv_rows)
+            else:
+                write_json_data(prompt_data, prompt_file)
+
+            console.print(f"[green]Exported prompt: {prompt.name or prompt_id}[/green]")
+            exported_count += 1
+
+        except Exception as e:
+            if debug:
+                console.print(
+                    f"[yellow]Warning: Could not export prompt {prompt_id}: {e}[/yellow]"
+                )
+            else:
+                console.print(f"[red]Error exporting prompt {prompt_id}: {e}[/red]")
+            continue
+
+    return exported_count, skipped_count
+
+
+def export_related_prompts_by_name(
+    client: opik.Opik,
+    experiment: Any,
+    output_dir: Path,
+    force: bool,
+    debug: bool,
+    format: str = "json",
+) -> int:
+    """Export prompts explicitly related to the experiment from experiment metadata."""
+    try:
+        prompts_dir = output_dir.parent / "prompts"
+        prompts_dir.mkdir(parents=True, exist_ok=True)
+
+        # Get experiment data to access metadata
+        experiment_data = experiment.get_experiment_data()
+        if not experiment_data:
+            debug_print("Could not get experiment data", debug)
+            return 0
+
+        # Extract prompt names from experiment metadata
+        prompt_names = []
+        metadata = experiment_data.metadata
+
+        if metadata:
+            # Metadata can be a dict, list, or string (JsonListStringPublic)
+            # Parse if it's a string, otherwise use directly
+            if isinstance(metadata, str):
+                try:
+                    metadata = json.loads(metadata)
+                except (json.JSONDecodeError, Exception) as e:
+                    if debug:
+                        debug_print(f"Could not parse metadata as JSON: {e}", debug)
+                    metadata = None
+
+            # Check if metadata is a dict and has "prompts" key
+            if isinstance(metadata, dict) and "prompts" in metadata:
+                prompts_dict = metadata["prompts"]
+                if isinstance(prompts_dict, dict):
+                    # Prompts are stored as a dict with prompt names as keys
+                    prompt_names = list(prompts_dict.keys())
+                    if debug:
+                        debug_print(
+                            f"Found {len(prompt_names)} prompt(s) in experiment metadata: {prompt_names}",
+                            debug,
+                        )
+                else:
+                    if debug:
+                        debug_print(
+                            f"Metadata 'prompts' is not a dict, got: {type(prompts_dict)}",
+                            debug,
+                        )
+            else:
+                if debug:
+                    debug_print("No 'prompts' key found in experiment metadata", debug)
+
+        if not prompt_names:
+            debug_print("No prompts found in experiment metadata", debug)
+            return 0
+
+        console.print(
+            f"[blue]Exporting {len(prompt_names)} prompt(s) from experiment metadata...[/blue]"
+        )
+
+        exported_count = 0
+        # Export each prompt by name from metadata
+        for prompt_name in prompt_names:
+            try:
+                debug_print(f"Exporting prompt: {prompt_name}", debug)
+
+                # Try to get the prompt - try ChatPrompt first, then regular Prompt
+                prompt: Optional[Union[Prompt, ChatPrompt]] = None
+                try:
+                    prompt = client.get_chat_prompt(prompt_name)
+                except Exception:
+                    # Not a ChatPrompt, try regular Prompt
+                    try:
+                        prompt = client.get_prompt(prompt_name)
+                    except Exception as e:
+                        if debug:
+                            console.print(
+                                f"[yellow]Warning: Could not get prompt '{prompt_name}': {e}[/yellow]"
+                            )
+                        continue
+
+                if not prompt:
+                    if debug:
+                        console.print(
+                            f"[yellow]Warning: Prompt '{prompt_name}' not found[/yellow]"
+                        )
+                    continue
+
+                # Get prompt history - use appropriate method based on prompt type
+                prompt_history: List[Union[Prompt, ChatPrompt]]
+                if isinstance(prompt, ChatPrompt):
+                    prompt_history = list(client.get_chat_prompt_history(prompt.name))
+                else:
+                    prompt_history = list(client.get_prompt_history(prompt.name))
+
+                # Create prompt data structure
+                prompt_data = {
+                    "prompt": {
+                        "id": getattr(prompt, "__internal_api__prompt_id__", None),
+                        "name": prompt.name,
+                        "description": getattr(prompt, "description", None),
+                        "created_at": getattr(prompt, "created_at", None),
+                        "last_updated_at": getattr(prompt, "last_updated_at", None),
+                    },
+                    "current_version": {
+                        "prompt": _get_prompt_content(prompt),
+                        "metadata": getattr(prompt, "metadata", None),
+                        "type": _get_prompt_type_string(prompt),
+                        "commit": getattr(prompt, "commit", None),
+                        "template_structure": _get_template_structure(prompt),
+                    },
+                    "history": [
+                        {
+                            "prompt": _get_prompt_content(version),
+                            "metadata": getattr(version, "metadata", None),
+                            "type": _get_prompt_type_string(version),
+                            "commit": getattr(version, "commit", None),
+                            "template_structure": _get_template_structure(version),
+                        }
+                        for version in prompt_history
+                    ],
+                    "downloaded_at": datetime.now().isoformat(),
+                    "related_to_experiment": experiment.name or experiment.id,
+                }
+
+                # Save prompt data using the appropriate format
+                # Sanitize prompt name for filename (replace / with _)
+                sanitized_name = prompt.name.replace("/", "_")
+                if format.lower() == "csv":
+                    prompt_file = prompts_dir / f"prompts_{sanitized_name}.csv"
+                else:
+                    prompt_file = prompts_dir / f"prompt_{sanitized_name}.json"
+
+                # Check if file should be skipped using the standard utility
+                if should_skip_file(prompt_file, force):
+                    debug_print(
+                        f"Skipping prompt {prompt.name} (already exists)", debug
+                    )
+                    continue
+
+                # File doesn't exist or force is set, so export it
+                if format.lower() == "csv":
+                    write_csv_data(prompt_data, prompt_file, prompt_to_csv_rows)
+                else:
+                    write_json_data(prompt_data, prompt_file)
+
+                console.print(f"[green]Exported related prompt: {prompt.name}[/green]")
+                exported_count += 1
+
+            except Exception as e:
+                if debug:
+                    prompt_display_name = prompt.name if prompt else prompt_name
+                    console.print(
+                        f"[yellow]Warning: Could not export related prompt {prompt_display_name}: {e}[/yellow]"
+                    )
+                continue
+
+        return exported_count
+
+    except Exception as e:
+        if debug:
+            console.print(
+                f"[yellow]Warning: Could not export related prompts: {e}[/yellow]"
+            )
+        return 0
+
+
+@click.command(name="prompt")
+@click.argument("name", type=str)
+@click.option(
+    "--max-results",
+    type=int,
+    help="Maximum number of prompts to export. Limits the total number of prompts 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_prompt_command(
+    ctx: click.Context,
+    name: str,
+    max_results: Optional[int],
+    path: str,
+    force: bool,
+    debug: bool,
+    format: str,
+) -> None:
+    """Export a prompt by exact name to workspace/prompts."""
+    # Get workspace and API key from context
+    workspace = ctx.obj["workspace"]
+    api_key = ctx.obj.get("api_key") if ctx.obj else None
+    export_prompt_by_name(
+        name, workspace, path, max_results, force, debug, format, api_key
+    )