opik 1.9.5__py3-none-any.whl → 1.9.39__py3-none-any.whl

opik/cli/usage_report/cli.py

@@ -7,11 +7,15 @@ import sys
 import traceback
 import webbrowser
 from pathlib import Path
-from typing import Optional
+from typing import List, Optional
 
 import click
 from rich.console import Console
 
+import opik.config as config
+import opik.url_helpers as url_helpers
+import opik.httpx_client as httpx_client
+
 from .extraction import extract_project_data
 from .pdf import create_pdf_report
 from .statistics import calculate_statistics
@@ -20,7 +24,7 @@ console = Console()
 
 
 @click.command(name="usage-report")
-@click.argument("workspace", type=str)
+@click.argument("workspaces", nargs=-1, type=str)
 @click.option(
     "--start-date",
     type=str,
@@ -50,35 +54,52 @@ console = Console()
     is_flag=True,
     help="Automatically open the generated PDF report in the default viewer.",
 )
+@click.option(
+    "--from-json",
+    "from_json",
+    is_flag=True,
+    help="Load data from existing JSON files instead of extracting from API. "
+    "JSON files should match the output pattern (e.g., opik_usage_report.json or opik_usage_report_{workspace}.json).",
+)
 @click.pass_context
 def usage_report(
     ctx: click.Context,
-    workspace: str,
+    workspaces: tuple,
     start_date: Optional[str],
     end_date: Optional[str],
     unit: str,
     output: str,
     open_pdf: bool,
+    from_json: bool,
 ) -> None:
     """
-    Extract Opik usage data for a workspace.
+    Extract Opik usage data for one or more workspaces.
 
-    This command extracts project-level metrics from Opik for a specific workspace:
-    - Loops through all projects in the workspace
+    This command extracts project-level metrics from Opik for specified workspace(s):
+    - Loops through all projects in each workspace
     - Gets trace count, cost, and token count
     - Gets experiment and dataset counts (workspace-level)
     - Aggregates data by the specified time unit (month, week, day, or hour)
     - Saves data to a JSON file
     - Generates a PDF report with charts and statistics
 
-    WORKSPACE: Workspace name to extract data from.
+    WORKSPACES: Zero or more workspace names to extract data from.
+                If no workspaces are provided, all workspaces will be processed.
 
     Examples:
 
-      Extract data with auto-detected date range, aggregated by month (default):
+      Extract data for a single workspace with auto-detected date range:
 
         opik usage-report my-workspace
 
+      Extract data for multiple workspaces:
+
+        opik usage-report workspace1 workspace2 workspace3
+
+      Extract data for all workspaces (no workspace specified):
+
+        opik usage-report
+
       Extract data aggregated by week:
 
         opik usage-report my-workspace --unit week
@@ -90,6 +111,10 @@ def usage_report(
       Extract data and automatically open the PDF report:
 
         opik usage-report my-workspace --open
+
+      Generate PDF from existing JSON file (skip data extraction):
+
+        opik usage-report my-workspace --from-json
     """
     try:
         # Get API key from context (set by main CLI)
@@ -104,52 +129,146 @@ def usage_report(
         if end_date:
             end_date_obj = datetime.datetime.strptime(end_date, "%Y-%m-%d")
 
-        console.print("[green]Starting Opik data extraction...[/green]\n")
+        # Determine which workspaces to process
+        workspaces_list: List[str] = list(workspaces) if workspaces else []
 
-        data = extract_project_data(
-            workspace, api_key, start_date_obj, end_date_obj, unit
-        )
+        # If no workspaces provided, fetch all workspaces
+        if not workspaces_list:
+            console.print(
+                "[blue]No workspaces specified. Fetching all workspaces...[/blue]"
+            )
+            cfg = config.OpikConfig()
+            # Use API key from context if available, otherwise from config
+            api_key_to_use = api_key or cfg.api_key
+            with httpx_client.get(
+                workspace=None,  # No workspace needed when fetching workspace list
+                api_key=api_key_to_use,
+                check_tls_certificate=cfg.check_tls_certificate,
+                compress_json_requests=cfg.enable_json_request_compression,
+            ) as client:
+                base_url = url_helpers.get_base_url(cfg.url_override)
+                workspace_list_url = url_helpers.get_workspace_list_url(base_url)
+                response = client.get(workspace_list_url)
+                workspaces_list = response.json().get("workspaceNames", [])
 
-        # Calculate and add summary statistics to the data
-        console.print("[blue]Calculating summary statistics...[/blue]")
-        stats = calculate_statistics(data)
-        data["statistics"] = stats
+            if not workspaces_list:
+                console.print("[yellow]No workspaces found.[/yellow]")
+                return
 
-        # Save to JSON file
-        console.print(f"\n[cyan]{'='*60}[/cyan]")
-        console.print(f"[blue]Saving data to {output}...[/blue]")
-        with open(output, "w") as f:
-            json.dump(data, f, indent=2, default=str)
+            console.print(f"[green]Found {len(workspaces_list)} workspace(s)[/green]\n")
 
-        console.print(
-            f"[green]Data extraction complete! Results saved to {output}[/green]"
-        )
+        # Process each workspace
+        for idx, workspace in enumerate(workspaces_list, 1):
+            console.print(f"\n[cyan]{'='*60}[/cyan]")
+            console.print(
+                f"[blue]Processing workspace {idx}/{len(workspaces_list)}: {workspace}[/blue]"
+            )
+            console.print(f"[cyan]{'='*60}[/cyan]\n")
 
-        # Generate PDF report
-        console.print(f"\n[cyan]{'='*60}[/cyan]")
-        console.print("[blue]Generating PDF report...[/blue]")
-        try:
-            output_path = Path(output)
-            output_dir = output_path.parent if output_path.parent != Path(".") else "."
-            pdf_filename = create_pdf_report(data, output_dir=str(output_dir))
-            console.print(f"[green]PDF report saved to {pdf_filename}[/green]")
-
-            # Open PDF if --open flag is set
-            if open_pdf:
-                pdf_path = os.path.abspath(pdf_filename)
-                if os.path.exists(pdf_path):
-                    webbrowser.open(f"file://{pdf_path}")
-                    console.print("[green]Opened PDF in default viewer[/green]")
+            try:
+                # Generate output filename for this workspace
+                if len(workspaces_list) == 1:
+                    # Single workspace: use the provided output filename
+                    workspace_output = output
                 else:
+                    # Multiple workspaces: append workspace name to output filename
+                    output_path = Path(output)
+                    workspace_output = str(
+                        output_path.parent
+                        / f"{output_path.stem}_{workspace}{output_path.suffix}"
+                    )
+
+                # Load from JSON or extract from API
+                if from_json:
                     console.print(
-                        f"[yellow]Warning: PDF file not found: {pdf_path}[/yellow]"
+                        f"[green]Loading data from {workspace_output}...[/green]\n"
                     )
-        except Exception as e:
-            console.print(
-                f"[yellow]Warning: Could not generate PDF report: {e}[/yellow]"
-            )
-            traceback.print_exc()
+                    if not os.path.exists(workspace_output):
+                        console.print(
+                            f"[red]Error: JSON file not found: {workspace_output}[/red]"
+                        )
+                        console.print(
+                            f"[yellow]Expected file: {workspace_output}[/yellow]"
+                        )
+                        continue
+
+                    with open(workspace_output, "r") as f:
+                        data = json.load(f)
+
+                    # Verify workspace matches
+                    if data.get("workspace") != workspace:
+                        console.print(
+                            f"[yellow]Warning: JSON file workspace '{data.get('workspace')}' "
+                            f"does not match expected workspace '{workspace}'[/yellow]"
+                        )
+
+                    # Calculate statistics if not present in JSON
+                    if "statistics" not in data:
+                        console.print("[blue]Calculating summary statistics...[/blue]")
+                        stats = calculate_statistics(data)
+                        data["statistics"] = stats
+
+                    console.print(f"[green]Loaded data from {workspace_output}[/green]")
+                else:
+                    console.print("[green]Starting Opik data extraction...[/green]\n")
+
+                    data = extract_project_data(
+                        workspace, api_key, start_date_obj, end_date_obj, unit
+                    )
+
+                    # Calculate and add summary statistics to the data
+                    console.print("[blue]Calculating summary statistics...[/blue]")
+                    stats = calculate_statistics(data)
+                    data["statistics"] = stats
+
+                    # Save to JSON file
+                    console.print(f"\n[cyan]{'='*60}[/cyan]")
+                    console.print(f"[blue]Saving data to {workspace_output}...[/blue]")
+                    with open(workspace_output, "w") as f:
+                        json.dump(data, f, indent=2, default=str)
+
+                    console.print(
+                        f"[green]Data extraction complete! Results saved to {workspace_output}[/green]"
+                    )
+
+                # Generate PDF report
+                console.print(f"\n[cyan]{'='*60}[/cyan]")
+                console.print("[blue]Generating PDF report...[/blue]")
+                try:
+                    output_path = Path(workspace_output)
+                    output_dir = (
+                        output_path.parent if output_path.parent != Path(".") else "."
+                    )
+                    pdf_filename = create_pdf_report(data, output_dir=str(output_dir))
+                    console.print(f"[green]PDF report saved to {pdf_filename}[/green]")
+
+                    # Open PDF if --open flag is set (only for the last workspace)
+                    if open_pdf and idx == len(workspaces_list):
+                        pdf_path = os.path.abspath(pdf_filename)
+                        if os.path.exists(pdf_path):
+                            webbrowser.open(f"file://{pdf_path}")
+                            console.print("[green]Opened PDF in default viewer[/green]")
+                        else:
+                            console.print(
+                                f"[yellow]Warning: PDF file not found: {pdf_path}[/yellow]"
+                            )
+                except Exception as e:
+                    console.print(
+                        f"[yellow]Warning: Could not generate PDF report: {e}[/yellow]"
+                    )
+                    traceback.print_exc()
+
+            except Exception as e:
+                console.print(f"[red]Error processing workspace {workspace}: {e}[/red]")
+                traceback.print_exc()
+                # Continue with next workspace instead of exiting
+                continue
+
+        console.print(
+            f"\n[green]Completed processing {len(workspaces_list)} workspace(s)[/green]"
+        )
 
     except Exception as e:
         console.print(f"[red]Error: {e}[/red]")
+        traceback.print_exc()
         sys.exit(1)

opik/cli/usage_report/pdf.py

@@ -184,7 +184,20 @@ def create_pdf_report(data: Dict[str, Any], output_dir: str = ".") -> str:
                             )
                             continue
 
-                        img = Image(chart_path, width=7 * inch, height=4.5 * inch)
+                        # All charts are exactly 14x8 inches (4200x2400 pixels at 300 DPI)
+                        # Scale to fit page with margins
+                        # Aspect ratio: 14/8 = 1.75 (always wider than tall)
+                        max_width = 7.5 * inch  # Leave margin
+                        chart_aspect_ratio = 14.0 / 8.0  # 1.75
+
+                        # Charts are always wider than tall, so always scale by width
+                        display_width = max_width
+                        display_height = max_width / chart_aspect_ratio
+
+                        # All charts use the same dimensions, so use fixed scaling
+                        img = Image(
+                            chart_path, width=display_width, height=display_height
+                        )
                         story.append(img)
                         story.append(Spacer(1, 0.1 * inch))
                         story.append(PageBreak())

opik/config.py

@@ -7,7 +7,6 @@ import pathlib
 import urllib.parse
 from typing import Any, Dict, Final, List, Literal, Optional, Tuple, Type, Union
 
-import opik.decorator.tracing_runtime_config as tracing_runtime_config
 import pydantic
 import pydantic_settings
 from pydantic_settings import BaseSettings, InitSettingsSource
@@ -257,10 +256,6 @@ class OpikConfig(pydantic_settings.BaseSettings):
     def guardrails_backend_host(self) -> str:
         return url_helpers.get_base_url(self.url_override) + "guardrails/"
 
-    @property
-    def runtime(self) -> tracing_runtime_config.TracingRuntimeConfig:
-        return tracing_runtime_config.runtime_cfg
-
     @pydantic.model_validator(mode="after")
     def _set_url_override_from_api_key(self) -> "OpikConfig":
         url_was_not_provided = (

opik/decorator/base_track_decorator.py

@@ -14,7 +14,7 @@ from typing import (
     NamedTuple,
 )
 
-from .. import context_storage, logging_messages
+from .. import context_storage, logging_messages, tracing_runtime_config
 from ..api_objects import opik_client, span, trace
 from ..types import DistributedTraceHeadersDict, ErrorInfoDict, SpanType
 from . import (
@@ -24,7 +24,6 @@ from . import (
     inspect_helpers,
     opik_args,
     span_creation_handler,
-    tracing_runtime_config,
 )
 
 LOGGER = logging.getLogger(__name__)
@@ -337,25 +336,24 @@ class BaseTrackDecorator(abc.ABC):
                 )
                 error_info = error_info_collector.collect(exception)
                 func_exception = exception
-            finally:
-                stream_or_stream_manager = self._streams_handler(
-                    result,
-                    track_options.capture_output,
-                    track_options.generations_aggregator,
-                )
-                if stream_or_stream_manager is not None:
-                    return stream_or_stream_manager
-
-                self._after_call(
-                    output=result,
-                    error_info=error_info,
-                    capture_output=track_options.capture_output,
-                    flush=track_options.flush,
-                )
-                if func_exception is not None:
-                    raise func_exception
-                else:
-                    return result
+
+            stream_or_stream_manager = self._streams_handler(
+                result,
+                track_options.capture_output,
+                track_options.generations_aggregator,
+            )
+            if stream_or_stream_manager is not None:
+                return stream_or_stream_manager
+
+            self._after_call(
+                output=result,
+                error_info=error_info,
+                capture_output=track_options.capture_output,
+                flush=track_options.flush,
+            )
+            if func_exception is not None:
+                raise func_exception
+            return result
 
         wrapper.opik_tracked = True  # type: ignore
 
@@ -391,25 +389,24 @@ class BaseTrackDecorator(abc.ABC):
                 )
                 error_info = error_info_collector.collect(exception)
                 func_exception = exception
-            finally:
-                stream_or_stream_manager = self._streams_handler(
-                    result,
-                    track_options.capture_output,
-                    track_options.generations_aggregator,
-                )
-                if stream_or_stream_manager is not None:
-                    return stream_or_stream_manager
-
-                self._after_call(
-                    output=result,
-                    error_info=error_info,
-                    capture_output=track_options.capture_output,
-                    flush=track_options.flush,
-                )
-                if func_exception is not None:
-                    raise func_exception
-                else:
-                    return result
+
+            stream_or_stream_manager = self._streams_handler(
+                result,
+                track_options.capture_output,
+                track_options.generations_aggregator,
+            )
+            if stream_or_stream_manager is not None:
+                return stream_or_stream_manager
+
+            self._after_call(
+                output=result,
+                error_info=error_info,
+                capture_output=track_options.capture_output,
+                flush=track_options.flush,
+            )
+            if func_exception is not None:
+                raise func_exception
+            return result
 
         wrapper.opik_tracked = True  # type: ignore
         return wrapper

opik/decorator/context_manager/span_context_manager.py

@@ -4,6 +4,7 @@ from typing import Optional, Dict, Any, List, Generator
 
 from opik.api_objects import span, opik_client
 from opik.types import SpanType
+from opik import context_storage
 from .. import arguments_helpers, base_track_decorator, error_info_collector
 
 LOGGER = logging.getLogger(__name__)
@@ -110,5 +111,13 @@ def start_as_current_span(
             )
             client.trace(**span_creation_result.trace_data.as_parameters)
 
+        # Clean up span and trace from context
+        opik_context_storage = context_storage.get_current_context_instance()
+        opik_context_storage.pop_span_data(ensure_id=span_creation_result.span_data.id)
+        if span_creation_result.trace_data is not None:
+            opik_context_storage.pop_trace_data(
+                ensure_id=span_creation_result.trace_data.id
+            )
+
         if flush:
             client.flush()

opik/decorator/context_manager/trace_context_manager.py

@@ -4,6 +4,7 @@ from typing import Any, Generator, Optional, Dict, List
 
 from opik import datetime_helpers
 from opik.api_objects import trace, opik_client, helpers
+from opik import context_storage
 from .. import base_track_decorator, error_info_collector
 
 LOGGER = logging.getLogger(__name__)
@@ -75,5 +76,9 @@ def start_as_current_trace(
         client = opik_client.get_client_cached()
         client.trace(**trace_data.init_end_time().as_parameters)
 
+        # Clean up trace from context
+        opik_context_storage = context_storage.get_current_context_instance()
+        opik_context_storage.pop_trace_data(ensure_id=trace_data.id)
+
         if flush:
             client.flush()

opik/dict_utils.py

@@ -1,6 +1,6 @@
 import copy
 import logging
-from typing import Any, Dict, Mapping, Optional, List, Tuple, TypeVar, Type
+from typing import Any, Dict, Optional, List, Tuple, TypeVar, Type
 
 from . import logging_messages
 
@@ -11,7 +11,7 @@ def flatten_dict(
     d: Dict[str, Any], parent_key: str, delim: str = "."
 ) -> Dict[str, Any]:
     """
-    Current implementation does not have max depth restrictions or cyclic references handling!
+    The current implementation does not have max depth restrictions or cyclic references handling!
     """
     items = []  # type: ignore
 
@@ -60,7 +60,7 @@ def deepmerge(
     return merged
 
 
-def remove_none_from_dict(original: Mapping[str, Optional[Any]]) -> Mapping[str, Any]:
+def remove_none_from_dict(original: Dict[str, Optional[Any]]) -> Dict[str, Any]:
     new: Dict[str, Any] = {}
 
     for key, value in original.items():

opik/evaluation/__init__.py

@@ -1,4 +1,15 @@
-from .evaluator import evaluate, evaluate_prompt, evaluate_experiment
+from .evaluator import (
+    evaluate,
+    evaluate_prompt,
+    evaluate_experiment,
+    evaluate_on_dict_items,
+)
 from .threads.evaluator import evaluate_threads
 
-__all__ = ["evaluate", "evaluate_prompt", "evaluate_experiment", "evaluate_threads"]
+__all__ = [
+    "evaluate",
+    "evaluate_prompt",
+    "evaluate_experiment",
+    "evaluate_on_dict_items",
+    "evaluate_threads",
+]