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

opik/decorator/context_manager/span_context_manager.py

@@ -0,0 +1,123 @@
+import logging
+from contextlib import contextmanager
+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__)
+
+
+@contextmanager
+def start_as_current_span(
+    name: str,
+    type: SpanType = "general",
+    input: Optional[Dict[str, Any]] = None,
+    output: Optional[Dict[str, Any]] = None,
+    tags: Optional[List[str]] = None,
+    metadata: Optional[Dict[str, Any]] = None,
+    project_name: Optional[str] = None,
+    model: Optional[str] = None,
+    provider: Optional[str] = None,
+    flush: bool = False,
+    **kwargs: Dict[str, Any],
+) -> Generator[span.SpanData, Any, None]:
+    """
+    A context manager for starting and managing a span and parent trace.
+
+    This function creates a span and parent trace (if missing) with input parameters, processes outputs,
+    handles errors, and ensures the span/trace data is saved and flushed at the end of its lifecycle.
+    It integrates distributed tracing headers and allows additional metadata, tags, and other
+    contextual information to be provided.
+
+    Args:
+        name: The name of the span to create.
+        type: The type of the span. Defaults to "general".
+        input: A dictionary representing the input data associated with the span.
+        output: A dictionary for providing the output associated with the span.
+        tags: A list of tags to associate with the span.
+        metadata: A dictionary of additional metadata to attach to the span or trace.
+        project_name: The name of the project associated with this span.
+        model: The model name related to the span or trace.
+        provider: The provider responsible for the span or trace.
+        flush: Whether to flush the client data after the span is created and processed.
+        **kwargs (Dict[str, Any]): Additional parameters that may be passed to the
+            context manager.
+
+    Yields:
+        An iterator that provides the span data within the context of the span manager lifecycle.
+    """
+    start_span_parameters = arguments_helpers.StartSpanParameters(
+        name=name,
+        input=input,
+        type=type,
+        project_name=project_name,
+        model=model,
+        provider=provider,
+    )
+    distributed_headers = arguments_helpers.extract_distributed_trace_headers(kwargs)
+
+    # create span/trace with input parameters
+    span_creation_result = base_track_decorator.add_start_candidates(
+        start_span_parameters=start_span_parameters,
+        opik_distributed_trace_headers=distributed_headers,
+        opik_args_data=None,
+        tracing_active=True,
+    )
+
+    end_arguments = arguments_helpers.EndSpanParameters(
+        input=span_creation_result.span_data.input or input,
+        output=span_creation_result.span_data.output or output,
+        tags=span_creation_result.span_data.tags or tags,
+        metadata=span_creation_result.span_data.metadata or metadata,
+        provider=span_creation_result.span_data.provider or provider,
+        model=span_creation_result.span_data.model or model,
+    )
+    try:
+        yield span_creation_result.span_data
+
+        # update end arguments
+        end_arguments.input = span_creation_result.span_data.input or input
+        end_arguments.output = span_creation_result.span_data.output or output
+        end_arguments.tags = span_creation_result.span_data.tags or tags
+        end_arguments.metadata = span_creation_result.span_data.metadata or metadata
+        end_arguments.provider = span_creation_result.span_data.provider or provider
+        end_arguments.model = span_creation_result.span_data.model or model
+    except Exception as exception:
+        LOGGER.error(
+            "Error in user's script while executing span context manager: %s",
+            str(exception),
+            exc_info=True,
+        )
+
+        # collect error info
+        end_arguments.error_info = error_info_collector.collect(exception)
+        end_arguments.output = None
+        raise
+    finally:
+        # save span/trace data at the end of the context manager
+        client = opik_client.get_client_cached()
+
+        span_creation_result.span_data.init_end_time().update(
+            **end_arguments.to_kwargs(),
+        )
+        client.span(**span_creation_result.span_data.as_parameters)
+
+        if span_creation_result.trace_data is not None:
+            span_creation_result.trace_data.init_end_time().update(
+                **end_arguments.to_kwargs(ignore_keys=["usage", "model", "provider"]),
+            )
+            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

@@ -0,0 +1,84 @@
+import logging
+from contextlib import contextmanager
+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__)
+
+
+@contextmanager
+def start_as_current_trace(
+    name: str,
+    input: Optional[Dict[str, Any]] = None,
+    output: Optional[Dict[str, Any]] = None,
+    tags: Optional[List[str]] = None,
+    metadata: Optional[Dict[str, Any]] = None,
+    project_name: Optional[str] = None,
+    thread_id: Optional[str] = None,
+    flush: bool = False,
+) -> Generator[trace.TraceData, Any, None]:
+    """
+    Starts a trace context manager to collect and manage tracing data during the
+    execution of a code block. This function initializes a trace, allows for
+    modifications within the context, and finishes the trace, sending data to the
+    Opik tracing infrastructure.
+
+    Args:
+        name: The name of the trace for identification purposes.
+        input: Optional input data associated with the trace.
+        output: Optional output data expected or associated with the trace.
+        tags: Optional list of string tags for labeling or describing the trace.
+        metadata: Optional dictionary containing additional information about the
+            trace.
+        project_name: Optional name of the project under which the trace belongs.
+        thread_id: Optional thread identifier to associate the trace with a
+            specific thread.
+        flush: A boolean indicating whether to flush the trace data immediately
+            after finishing the trace context.
+
+    Yields:
+        Provides the initialized trace data for manipulation during execution in the context.
+    """
+    trace_data = trace.TraceData(
+        id=helpers.generate_id(),
+        start_time=datetime_helpers.local_timestamp(),
+        name=name,
+        input=input,
+        output=output,
+        metadata=metadata,
+        tags=tags,
+        project_name=project_name,
+        thread_id=thread_id,
+    )
+    base_track_decorator.add_start_trace_candidate(
+        trace_data=trace_data,
+        opik_args_data=None,
+        tracing_active=True,
+    )
+
+    try:
+        yield trace_data
+    except Exception as exception:
+        LOGGER.error(
+            "Error in user's script while executing trace context manager: %s",
+            str(exception),
+            exc_info=True,
+        )
+        trace_data.error_info = error_info_collector.collect(exception)
+        trace_data.output = None
+        raise
+    finally:
+        # save trace data at the end of the context manager
+        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/decorator/generator_wrappers.py

@@ -11,7 +11,8 @@ from typing import (
     Generic,
 )
 
-from opik import context_storage, logging_messages
+import opik.context_storage as context_storage
+import opik.logging_messages as logging_messages
 from opik.api_objects import span, trace
 from opik.types import DistributedTraceHeadersDict, ErrorInfoDict
 
@@ -58,7 +59,7 @@ class BaseTrackedGenerator(Generic[YieldType]):
             return
 
         self._created_trace_data, self._created_span_data = (
-            span_creation_handler.create_span_for_current_context(
+            span_creation_handler.create_span_respecting_context(
                 self._start_span_arguments, self._opik_distributed_trace_headers
             )
         )

opik/decorator/inspect_helpers.py

@@ -1,4 +1,5 @@
 import inspect
+import functools
 
 from typing import Callable, Tuple, Any, Dict
 
@@ -44,3 +45,13 @@ def is_async(func: Callable) -> bool:
         return True
 
     return False
+
+
+def get_function_name(func: Callable) -> str:
+    """Safely get the name of a function, handling functools.partial objects."""
+    if isinstance(func, functools.partial):
+        # For partial objects, get the name from the underlying function
+        return getattr(func.func, "__name__", "<unknown>")
+
+    # For regular functions and other callables
+    return getattr(func, "__name__", "<unknown>")

opik/decorator/opik_args/__init__.py

@@ -0,0 +1,13 @@
+from .api_classes import OpikArgs
+from .helpers import (
+    extract_opik_args,
+    apply_opik_args_to_start_span_params,
+    apply_opik_args_to_trace,
+)
+
+__all__ = [
+    "OpikArgs",
+    "extract_opik_args",
+    "apply_opik_args_to_start_span_params",
+    "apply_opik_args_to_trace",
+]

opik/decorator/opik_args/api_classes.py

@@ -0,0 +1,71 @@
+import logging
+from typing import Any, Dict, List, Optional
+
+import pydantic
+
+
+LOGGER = logging.getLogger(__name__)
+
+
+class OpikArgsSpan(pydantic.BaseModel):
+    """
+    Configuration for span updates passed via opik_args parameter.
+    """
+
+    tags: Optional[List[str]] = None
+    metadata: Optional[Dict[str, Any]] = None
+
+
+class OpikArgsTrace(pydantic.BaseModel):
+    """
+    Configuration for trace updates passed via opik_args parameter.
+    """
+
+    thread_id: Optional[str] = None
+    tags: Optional[List[str]] = None
+    metadata: Optional[Dict[str, Any]] = None
+
+
+class OpikArgs(pydantic.BaseModel):
+    """
+    Configuration structure for special opik arguments passed to tracked functions.
+
+    This allows users to specify span and trace updates without modifying function bodies.
+    Tags and metadata are merged (not overwritten) with existing values.
+    """
+
+    span_args: Optional[OpikArgsSpan] = None
+    trace_args: Optional[OpikArgsTrace] = None
+
+    @classmethod
+    def from_dict(cls, config_dict: Optional[Dict[str, Any]]) -> Optional["OpikArgs"]:
+        """Create OpikArgs from dictionary, with validation."""
+        if not isinstance(config_dict, dict):
+            LOGGER.warning("opik_args must be a dictionary, got %s", type(config_dict))
+            return None
+
+        span_args = None
+        trace_args = None
+
+        if "span" in config_dict:
+            span_data = config_dict["span"]
+            if isinstance(span_data, dict):
+                span_args = OpikArgsSpan(
+                    tags=span_data.get("tags"),
+                    metadata=span_data.get("metadata"),
+                )
+            else:
+                LOGGER.warning("opik_args['span'] must be a dictionary")
+
+        if "trace" in config_dict:
+            trace_data = config_dict["trace"]
+            if isinstance(trace_data, dict):
+                trace_args = OpikArgsTrace(
+                    thread_id=trace_data.get("thread_id"),
+                    tags=trace_data.get("tags"),
+                    metadata=trace_data.get("metadata"),
+                )
+            else:
+                LOGGER.warning("opik_args['trace'] must be a dictionary")
+
+        return cls(span_args=span_args, trace_args=trace_args)

opik/decorator/opik_args/helpers.py

@@ -0,0 +1,120 @@
+import inspect
+import logging
+from typing import Any, Dict, Optional, Callable
+
+from .. import arguments_helpers
+from ...api_objects import trace, data_helpers
+from . import api_classes
+
+
+LOGGER = logging.getLogger(__name__)
+
+
+def extract_opik_args(
+    kwargs: Dict[str, Any], func: Callable
+) -> Optional[api_classes.OpikArgs]:
+    """
+    Extracts opik_args from kwargs and returns the parsed OpikArgs or None.
+
+    Mutates kwargs in place by removing "opik_args" if the function does not explicitly
+    declare opik_args as a parameter.
+
+    Args:
+        kwargs: Function keyword arguments
+        func: The function being decorated
+
+    Returns:
+        opik_args if found, otherwise None.
+    """
+
+    # Check if a function has opik_args in its signature
+    try:
+        sig = inspect.signature(func)
+        has_opik_args_param = "opik_args" in sig.parameters
+    except (ValueError, TypeError):
+        # If we can't inspect the signature, assume no opik_args parameter
+        has_opik_args_param = False
+
+    # If function explicitly has opik_args parameter, just get it without popping
+    if has_opik_args_param:
+        opik_args_dict = kwargs.get("opik_args", None)
+    else:
+        # If function doesn't have opik_args parameter, pop it from kwargs
+        opik_args_dict = kwargs.pop("opik_args", None)
+
+    if opik_args_dict is None:
+        return None
+
+    return api_classes.OpikArgs.from_dict(opik_args_dict)
+
+
+def apply_opik_args_to_trace(
+    opik_args: Optional[api_classes.OpikArgs],
+    trace_data: Optional[trace.TraceData],
+) -> None:
+    """
+    Apply opik_args to the trace data, including thread_id, tags, and metadata.
+
+    Args:
+        opik_args: The configuration extracted from function arguments
+        trace_data: The current trace data to modify
+    """
+    if opik_args is None or opik_args.trace_args is None or trace_data is None:
+        return
+
+    # Handle thread_id
+    if opik_args.trace_args.thread_id is not None:
+        # Check for thread_id conflict
+        if (
+            trace_data.thread_id is not None
+            and trace_data.thread_id != opik_args.trace_args.thread_id
+        ):
+            LOGGER.warning(
+                "Trace already has thread_id='%s', but opik_args specifies thread_id='%s'. Keeping existing thread_id.",
+                trace_data.thread_id,
+                opik_args.trace_args.thread_id,
+            )
+        else:
+            # Apply thread_id to trace
+            trace_data.thread_id = opik_args.trace_args.thread_id
+
+    # Apply trace tags if specified
+    if opik_args.trace_args.tags:
+        existing_tags = trace_data.tags or []
+        merged_tags = data_helpers.merge_tags(existing_tags, opik_args.trace_args.tags)
+        trace_data.tags = merged_tags
+
+    # Apply trace metadata if specified
+    if opik_args.trace_args.metadata:
+        existing_metadata = trace_data.metadata or {}
+        merged_metadata = data_helpers.merge_metadata(
+            existing_metadata, opik_args.trace_args.metadata
+        )
+        trace_data.metadata = merged_metadata
+
+
+def apply_opik_args_to_start_span_params(
+    params: arguments_helpers.StartSpanParameters,
+    opik_args: Optional[api_classes.OpikArgs],
+) -> arguments_helpers.StartSpanParameters:
+    """Apply opik_args to StartSpanParameters, merging tags and metadata."""
+    if opik_args is None or opik_args.span_args is None:
+        return params
+
+    span_config = opik_args.span_args
+
+    # Merge tags and metadata
+    merged_tags = data_helpers.merge_tags(params.tags, span_config.tags)
+    merged_metadata = data_helpers.merge_metadata(params.metadata, span_config.metadata)
+
+    # Create updated parameters
+    return arguments_helpers.StartSpanParameters(
+        type=params.type,
+        name=params.name,
+        tags=merged_tags,
+        metadata=merged_metadata,
+        input=params.input,
+        project_name=params.project_name,
+        model=params.model,
+        provider=params.provider,
+    )

opik/decorator/span_creation_handler.py

@@ -1,24 +1,45 @@
 from typing import (
     Optional,
-    Tuple,
+    NamedTuple,
 )
 
-from opik import context_storage, datetime_helpers
+import opik.context_storage as context_storage
+import opik.datetime_helpers as datetime_helpers
 from opik.api_objects import helpers, span, trace
 from opik.types import DistributedTraceHeadersDict
 
 from . import arguments_helpers
 
 
-def create_span_for_current_context(
+class SpanCreationResult(NamedTuple):
+    """
+    Represents the result of a span creation process.
+
+    This class encapsulates the data resulting from the creation of a new
+    span, including trace information and span-specific details.
+
+    Attributes:
+        trace_data: Trace-related data associated
+            with the span if a new trace was created. Can be None if no new trace was created.
+        span_data : Data specific to the created span, containing
+            information such as span identifiers and timestamps.
+    """
+
+    trace_data: Optional[trace.TraceData]
+    span_data: span.SpanData
+
+
+def create_span_respecting_context(
     start_span_arguments: arguments_helpers.StartSpanParameters,
     distributed_trace_headers: Optional[DistributedTraceHeadersDict],
-) -> Tuple[Optional[trace.TraceData], span.SpanData]:
+    opik_context_storage: Optional[context_storage.OpikContextStorage] = None,
+) -> SpanCreationResult:
     """
     Handles different span creation flows.
     """
-    span_data: span.SpanData
-    trace_data: trace.TraceData
+
+    if opik_context_storage is None:
+        opik_context_storage = context_storage.get_current_context_instance()
 
     if distributed_trace_headers:
         span_data = arguments_helpers.create_span_data(
@@ -27,20 +48,26 @@ def create_span_for_current_context(
             trace_id=distributed_trace_headers["opik_trace_id"],
         )
 
-        return None, span_data
+        return SpanCreationResult(None, span_data)
 
-    current_span_data = context_storage.top_span_data()
-    current_trace_data = context_storage.get_trace_data()
+    current_span_data = opik_context_storage.top_span_data()
+    current_trace_data = opik_context_storage.get_trace_data()
 
     if current_span_data is not None:
-        # There is already at least one span in current context.
-        # Simply attach a new span to it.
-        assert current_trace_data is not None
+        # There is already at least one span in the current context - attach a new span to it.
+        #
+        # NOTE: We can have a situation when span data is in context, but there is no trace data
+        # because we are in a distributed environment and trace data was created in another thread.
+        # See: https://github.com/comet-ml/opik/pull/2244
+        if current_trace_data is None:
+            show_warning = False
+        else:
+            show_warning = current_trace_data.created_by != "evaluation"
 
         project_name = helpers.resolve_child_span_project_name(
             parent_project_name=current_span_data.project_name,
             child_project_name=start_span_arguments.project_name,
-            show_warning=current_trace_data.created_by != "evaluation",
+            show_warning=show_warning,
         )
 
         start_span_arguments.project_name = project_name
@@ -51,13 +78,13 @@ def create_span_for_current_context(
             trace_id=current_span_data.trace_id,
         )
 
-        return None, span_data
+        return SpanCreationResult(None, span_data)
 
     if current_trace_data is not None and current_span_data is None:
         # By default, we expect trace to be created with a span.
         # But there can be cases when trace was created and added
         # to context manually (not via decorator).
-        # In that case decorator should just create a span for the existing trace.
+        # In that case the decorator should just create a span for the existing trace.
 
         project_name = helpers.resolve_child_span_project_name(
             parent_project_name=current_trace_data.project_name,
@@ -73,12 +100,12 @@ def create_span_for_current_context(
             trace_id=current_trace_data.id,
         )
 
-        return None, span_data
+        return SpanCreationResult(None, span_data)
 
     if current_span_data is None and current_trace_data is None:
         # Create a trace and root span because it is
-        # the first decorated function run in current context.
-        trace_data = trace.TraceData(
+        # the first decorated function run in the current context.
+        current_trace_data = trace.TraceData(
             id=helpers.generate_id(),
             start_time=datetime_helpers.local_timestamp(),
             name=start_span_arguments.name,
@@ -86,12 +113,13 @@ def create_span_for_current_context(
             metadata=start_span_arguments.metadata,
             tags=start_span_arguments.tags,
             project_name=start_span_arguments.project_name,
+            thread_id=start_span_arguments.thread_id,
         )
 
-        span_data = arguments_helpers.create_span_data(
+        current_span_data = arguments_helpers.create_span_data(
             start_span_arguments=start_span_arguments,
             parent_span_id=None,
-            trace_id=trace_data.id,
+            trace_id=current_trace_data.id,
         )
 
-        return trace_data, span_data
+    return SpanCreationResult(current_trace_data, current_span_data)

opik/decorator/tracker.py

@@ -1,5 +1,6 @@
 import logging
 from typing import Any, Callable, Dict, List, Optional, Tuple
+from typing_extensions import override
 
 from ..api_objects import opik_client, span
 from . import arguments_helpers, base_track_decorator, inspect_helpers
@@ -12,6 +13,7 @@ class OpikTrackDecorator(base_track_decorator.BaseTrackDecorator):
     Default implementation of BaseTrackDecorator
     """
 
+    @override
     def _start_span_inputs_preprocessor(
         self,
         func: Callable,
@@ -29,7 +31,11 @@ class OpikTrackDecorator(base_track_decorator.BaseTrackDecorator):
             for argument in track_options.ignore_arguments:
                 input.pop(argument, None)
 
-        name = track_options.name if track_options.name is not None else func.__name__
+        name = (
+            track_options.name
+            if track_options.name is not None
+            else inspect_helpers.get_function_name(func)
+        )
 
         result = arguments_helpers.StartSpanParameters(
             name=name,
@@ -42,6 +48,7 @@ class OpikTrackDecorator(base_track_decorator.BaseTrackDecorator):
 
         return result
 
+    @override
     def _end_span_inputs_preprocessor(
         self,
         output: Any,
@@ -57,6 +64,7 @@ class OpikTrackDecorator(base_track_decorator.BaseTrackDecorator):
 
         return result
 
+    @override
     def _streams_handler(
         self,
         output: Any,

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/environment.py

@@ -7,9 +7,11 @@ import socket
 import sys
 from importlib import metadata
 from typing import Dict, Literal
+import tqdm
+from tqdm.utils import Comparable
 
 import opik.config
-from opik import url_helpers
+import opik.url_helpers as url_helpers
 
 LOGGER = logging.getLogger(__name__)
 
@@ -140,3 +142,13 @@ def get_installed_packages() -> Dict[str, str]:
         pkg.metadata["Name"]: pkg.version for pkg in metadata.distributions()
     }
     return installed_packages
+
+
+def get_tqdm_for_current_environment() -> Comparable:
+    """
+    Get a tqdm progress bar for your environment.
+    """
+    if in_jupyter() or in_colab():
+        return tqdm.tqdm_notebook
+    else:
+        return tqdm.tqdm