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

opik/decorator/base_track_decorator.py

@@ -11,9 +11,10 @@ from typing import (
     Set,
     Tuple,
     Union,
+    NamedTuple,
 )
 
-from .. import config, 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 (
@@ -21,6 +22,7 @@ from . import (
     error_info_collector,
     generator_wrappers,
     inspect_helpers,
+    opik_args,
     span_creation_handler,
 )
 
@@ -29,6 +31,12 @@ LOGGER = logging.getLogger(__name__)
 TRACES_CREATED_BY_DECORATOR: Set[str] = set()
 
 
+class TrackingStartOptions(NamedTuple):
+    start_span_parameters: arguments_helpers.StartSpanParameters
+    opik_args: Optional[opik_args.OpikArgs]
+    opik_distributed_trace_headers: Optional[DistributedTraceHeadersDict]
+
+
 class BaseTrackDecorator(abc.ABC):
     """
     For internal usage.
@@ -48,11 +56,6 @@ class BaseTrackDecorator(abc.ABC):
         self.provider: Optional[str] = None
         """ Name of the LLM provider. Used in subclasses in integrations track decorators. """
 
-    @functools.cached_property
-    def disabled(self) -> bool:
-        config_ = config.OpikConfig()
-        return config_.track_disable
-
     def track(
         self,
         name: Optional[Union[Callable, str]] = None,
@@ -94,6 +97,10 @@ class BaseTrackDecorator(abc.ABC):
             This decorator can be used to track both synchronous and asynchronous functions,
             and also synchronous and asynchronous generators.
             It automatically detects the function type and applies the appropriate tracking logic.
+
+            Tracing is checked only once at the start of the call; a call that
+            began while tracing was enabled will still be logged even if
+            tracing is disabled before it returns.
         """
         track_options = arguments_helpers.TrackOptions(
             name=None,
@@ -137,11 +144,11 @@ class BaseTrackDecorator(abc.ABC):
 
             * Regular sync and async functions/methods: start the span when the
         function is called, end the span when the function is finished. While the
-        function is working, span is kept in opik context, so it can be a parent for the
+        function is working, the span is kept in opik context, so it can be a parent for the
         spans created by nested tracked functions.
 
-            * Generators and async generators: start the span when generator started
-        yielding values, end the trace when generator finished yielding values.
+            * Generators and async generators: start the span when the generator started
+        yielding values, end the trace when the generator finished yielding values.
         Span is kept in the opik context only while __next__ or __anext__ method is working.
         It means that the span can be a parent only for spans created by tracked functions
         called inside __next__ or __anext__.
@@ -152,6 +159,10 @@ class BaseTrackDecorator(abc.ABC):
         So these spans can't be parents for other spans. This is usually the case LLM API calls
         with `stream=True`.
         """
+        # Idempotency: skip re-decoration if already tracked
+        if hasattr(func, "opik_tracked") and func.opik_tracked:  # type: ignore
+            return func
+
         if inspect.isgeneratorfunction(func):
             return self._tracked_sync_generator(func=func, track_options=track_options)
 
@@ -172,37 +183,74 @@ class BaseTrackDecorator(abc.ABC):
             track_options=track_options,
         )
 
+    def _prepare_tracking_start_options(
+        self,
+        func: Callable,
+        track_options: arguments_helpers.TrackOptions,
+        args: Tuple,
+        kwargs: Dict[str, Any],
+    ) -> TrackingStartOptions:
+        opik_distributed_trace_headers = (
+            arguments_helpers.extract_distributed_trace_headers(kwargs)
+        )
+
+        opik_args_ = None
+        try:
+            opik_args_ = opik_args.extract_opik_args(kwargs, func)
+
+            start_span_arguments = self._start_span_inputs_preprocessor(
+                func=func,
+                track_options=track_options,
+                args=args,
+                kwargs=kwargs,
+            )
+
+            # Apply opik_args to start span arguments
+            start_span_arguments = opik_args.apply_opik_args_to_start_span_params(
+                params=start_span_arguments,
+                opik_args=opik_args_,
+            )
+        except Exception as exception:
+            LOGGER.error(
+                logging_messages.UNEXPECTED_EXCEPTION_ON_SPAN_CREATION_FOR_TRACKED_FUNCTION,
+                inspect_helpers.get_function_name(func),
+                (args, kwargs),
+                str(exception),
+                exc_info=True,
+            )
+
+            start_span_arguments = arguments_helpers.StartSpanParameters(
+                name=inspect_helpers.get_function_name(func),
+                type=track_options.type,
+                tags=track_options.tags,
+                metadata=track_options.metadata,
+                project_name=track_options.project_name,
+            )
+
+        return TrackingStartOptions(
+            start_span_arguments, opik_args_, opik_distributed_trace_headers
+        )
+
     def _tracked_sync_generator(
         self, func: Callable, track_options: arguments_helpers.TrackOptions
     ) -> Callable:
         @functools.wraps(func)
         def wrapper(*args, **kwargs) -> Any:  # type: ignore
-            try:
-                opik_distributed_trace_headers: Optional[
-                    DistributedTraceHeadersDict
-                ] = kwargs.pop("opik_distributed_trace_headers", None)
+            if not tracing_runtime_config.is_tracing_active():
+                return func(*args, **kwargs)
 
-                start_span_arguments = self._start_span_inputs_preprocessor(
-                    func=func,
-                    track_options=track_options,
-                    args=args,
-                    kwargs=kwargs,
-                )
-            except Exception as exception:
-                LOGGER.error(
-                    logging_messages.UNEXPECTED_EXCEPTION_ON_SPAN_CREATION_FOR_TRACKED_FUNCTION,
-                    func.__name__,
-                    (args, kwargs),
-                    str(exception),
-                    exc_info=True,
-                )
+            track_start_options = self._prepare_tracking_start_options(
+                func=func,
+                track_options=track_options,
+                args=args,
+                kwargs=kwargs,
+            )
 
-            result = None
             try:
                 result = generator_wrappers.SyncTrackedGenerator(
                     func(*args, **kwargs),
-                    start_span_arguments=start_span_arguments,
-                    opik_distributed_trace_headers=opik_distributed_trace_headers,
+                    start_span_arguments=track_start_options.start_span_parameters,
+                    opik_distributed_trace_headers=track_start_options.opik_distributed_trace_headers,
                     track_options=track_options,
                     finally_callback=self._after_call,
                 )
@@ -210,7 +258,7 @@ class BaseTrackDecorator(abc.ABC):
             except Exception as exception:
                 LOGGER.debug(
                     logging_messages.EXCEPTION_RAISED_FROM_TRACKED_FUNCTION,
-                    func.__name__,
+                    inspect_helpers.get_function_name(func),
                     (args, kwargs),
                     str(exception),
                     exc_info=True,
@@ -226,32 +274,21 @@ class BaseTrackDecorator(abc.ABC):
     ) -> Callable:
         @functools.wraps(func)
         def wrapper(*args, **kwargs) -> Any:  # type: ignore
-            try:
-                opik_distributed_trace_headers: Optional[
-                    DistributedTraceHeadersDict
-                ] = kwargs.pop("opik_distributed_trace_headers", None)
+            if not tracing_runtime_config.is_tracing_active():
+                return func(*args, **kwargs)
 
-                start_span_arguments = self._start_span_inputs_preprocessor(
-                    func=func,
-                    track_options=track_options,
-                    args=args,
-                    kwargs=kwargs,
-                )
-            except Exception as exception:
-                LOGGER.error(
-                    logging_messages.UNEXPECTED_EXCEPTION_ON_SPAN_CREATION_FOR_TRACKED_FUNCTION,
-                    func.__name__,
-                    (args, kwargs),
-                    str(exception),
-                    exc_info=True,
-                )
+            track_start_options = self._prepare_tracking_start_options(
+                func=func,
+                track_options=track_options,
+                args=args,
+                kwargs=kwargs,
+            )
 
-            result = None
             try:
                 result = generator_wrappers.AsyncTrackedGenerator(
                     func(*args, **kwargs),
-                    start_span_arguments=start_span_arguments,
-                    opik_distributed_trace_headers=opik_distributed_trace_headers,
+                    start_span_arguments=track_start_options.start_span_parameters,
+                    opik_distributed_trace_headers=track_start_options.opik_distributed_trace_headers,
                     track_options=track_options,
                     finally_callback=self._after_call,
                 )
@@ -259,7 +296,7 @@ class BaseTrackDecorator(abc.ABC):
             except Exception as exception:
                 LOGGER.debug(
                     logging_messages.EXCEPTION_RAISED_FROM_TRACKED_FUNCTION,
-                    func.__name__,
+                    inspect_helpers.get_function_name(func),
                     (args, kwargs),
                     str(exception),
                     exc_info=True,
@@ -275,6 +312,8 @@ class BaseTrackDecorator(abc.ABC):
     ) -> Callable:
         @functools.wraps(func)
         def wrapper(*args, **kwargs) -> Any:  # type: ignore
+            if not tracing_runtime_config.is_tracing_active():
+                return func(*args, **kwargs)
             self._before_call(
                 func=func,
                 track_options=track_options,
@@ -284,35 +323,37 @@ class BaseTrackDecorator(abc.ABC):
 
             result = None
             error_info: Optional[ErrorInfoDict] = None
+            func_exception = None
             try:
                 result = func(*args, **kwargs)
             except Exception as exception:
                 LOGGER.debug(
                     logging_messages.EXCEPTION_RAISED_FROM_TRACKED_FUNCTION,
-                    func.__name__,
+                    inspect_helpers.get_function_name(func),
                     (args, kwargs),
                     str(exception),
                     exc_info=True,
                 )
                 error_info = error_info_collector.collect(exception)
-                raise 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 result is not None:
-                    return result
+                func_exception = exception
+
+            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
 
@@ -325,6 +366,8 @@ class BaseTrackDecorator(abc.ABC):
     ) -> Callable:
         @functools.wraps(func)
         async def wrapper(*args, **kwargs) -> Any:  # type: ignore
+            if not tracing_runtime_config.is_tracing_active():
+                return await func(*args, **kwargs)
             self._before_call(
                 func=func,
                 track_options=track_options,
@@ -333,35 +376,37 @@ class BaseTrackDecorator(abc.ABC):
             )
             result = None
             error_info: Optional[ErrorInfoDict] = None
+            func_exception = None
             try:
                 result = await func(*args, **kwargs)
             except Exception as exception:
                 LOGGER.debug(
                     logging_messages.EXCEPTION_RAISED_FROM_TRACKED_FUNCTION,
-                    func.__name__,
+                    inspect_helpers.get_function_name(func),
                     (args, kwargs),
                     str(exception),
                     exc_info=True,
                 )
                 error_info = error_info_collector.collect(exception)
-                raise 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 result is not None:
-                    return result
+                func_exception = exception
+
+            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
@@ -374,38 +419,42 @@ class BaseTrackDecorator(abc.ABC):
         kwargs: Dict[str, Any],
     ) -> None:
         try:
-            opik_distributed_trace_headers: Optional[DistributedTraceHeadersDict] = (
-                kwargs.pop("opik_distributed_trace_headers", None)
-            )
-
-            start_span_arguments = self._start_span_inputs_preprocessor(
+            self.__before_call_unsafe(
                 func=func,
                 track_options=track_options,
                 args=args,
                 kwargs=kwargs,
             )
-
-            created_trace_data, created_span_data = (
-                span_creation_handler.create_span_for_current_context(
-                    start_span_arguments=start_span_arguments,
-                    distributed_trace_headers=opik_distributed_trace_headers,
-                )
-            )
-            if created_trace_data is not None:
-                context_storage.set_trace_data(created_trace_data)
-                TRACES_CREATED_BY_DECORATOR.add(created_trace_data.id)
-
-            context_storage.add_span_data(created_span_data)
-
         except Exception as exception:
             LOGGER.error(
                 logging_messages.UNEXPECTED_EXCEPTION_ON_SPAN_CREATION_FOR_TRACKED_FUNCTION,
-                func.__name__,
+                inspect_helpers.get_function_name(func),
                 (args, kwargs),
                 str(exception),
                 exc_info=True,
             )
 
+    def __before_call_unsafe(
+        self,
+        func: Callable,
+        track_options: arguments_helpers.TrackOptions,
+        args: Tuple,
+        kwargs: Dict[str, Any],
+    ) -> None:
+        track_start_options = self._prepare_tracking_start_options(
+            func=func,
+            track_options=track_options,
+            args=args,
+            kwargs=kwargs,
+        )
+
+        add_start_candidates(
+            start_span_parameters=track_start_options.start_span_parameters,
+            opik_distributed_trace_headers=track_start_options.opik_distributed_trace_headers,
+            opik_args_data=track_start_options.opik_args,
+            tracing_active=tracing_runtime_config.is_tracing_active(),
+        )
+
     def _after_call(
         self,
         output: Optional[Any],
@@ -415,56 +464,78 @@ class BaseTrackDecorator(abc.ABC):
         generators_trace_to_end: Optional[trace.TraceData] = None,
         flush: bool = False,
     ) -> None:
-        if self.disabled:
-            return
-
         try:
-            if generators_span_to_end is None:
-                span_data_to_end, trace_data_to_end = pop_end_candidates()
-            else:
-                span_data_to_end, trace_data_to_end = (
-                    generators_span_to_end,
-                    generators_trace_to_end,
-                )
+            self.__after_call_unsafe(
+                output=output,
+                error_info=error_info,
+                capture_output=capture_output,
+                generators_span_to_end=generators_span_to_end,
+                generators_trace_to_end=generators_trace_to_end,
+                flush=flush,
+            )
+        except Exception as exception:
+            LOGGER.error(
+                logging_messages.UNEXPECTED_EXCEPTION_ON_SPAN_FINALIZATION_FOR_TRACKED_FUNCTION,
+                output,
+                str(exception),
+                exc_info=True,
+            )
 
-            if output is not None:
+    def __after_call_unsafe(
+        self,
+        output: Optional[Any],
+        error_info: Optional[ErrorInfoDict],
+        capture_output: bool,
+        generators_span_to_end: Optional[span.SpanData] = None,
+        generators_trace_to_end: Optional[trace.TraceData] = None,
+        flush: bool = False,
+    ) -> None:
+        if generators_span_to_end is None:
+            span_data_to_end, trace_data_to_end = pop_end_candidates()
+        else:
+            span_data_to_end, trace_data_to_end = (
+                generators_span_to_end,
+                generators_trace_to_end,
+            )
+
+        if output is not None:
+            try:
                 end_arguments = self._end_span_inputs_preprocessor(
                     output=output,
                     capture_output=capture_output,
                     current_span_data=span_data_to_end,
                 )
-            else:
-                end_arguments = arguments_helpers.EndSpanParameters(
-                    error_info=error_info
+            except Exception as e:
+                LOGGER.error(
+                    logging_messages.UNEXPECTED_EXCEPTION_ON_SPAN_FINALIZATION_FOR_TRACKED_FUNCTION,
+                    output,
+                    str(e),
+                    exc_info=True,
                 )
 
-            client = opik_client.get_client_cached()
+                end_arguments = arguments_helpers.EndSpanParameters(
+                    output={"output": output}
+                )
+        else:
+            end_arguments = arguments_helpers.EndSpanParameters(error_info=error_info)
 
-            span_data_to_end.init_end_time().update(
-                **end_arguments.to_kwargs(),
-            )
+        client = opik_client.get_client_cached()
 
-            client.span(**span_data_to_end.__dict__)
+        span_data_to_end.init_end_time().update(
+            **end_arguments.to_kwargs(),
+        )
 
-            if trace_data_to_end is not None:
-                trace_data_to_end.init_end_time().update(
-                    **end_arguments.to_kwargs(
-                        ignore_keys=["usage", "model", "provider"]
-                    ),
-                )
+        client.span(**span_data_to_end.as_parameters)
 
-                client.trace(**trace_data_to_end.__dict__)
+        if trace_data_to_end is not None:
+            trace_data_to_end.init_end_time().update(
+                **end_arguments.to_kwargs(ignore_keys=["usage", "model", "provider"]),
+            )
 
-            if flush:
-                client.flush()
+            client.trace(**trace_data_to_end.as_parameters)
 
-        except Exception as exception:
-            LOGGER.error(
-                logging_messages.UNEXPECTED_EXCEPTION_ON_SPAN_FINALIZATION_FOR_TRACKED_FUNCTION,
-                output,
-                str(exception),
-                exc_info=True,
-            )
+        if flush:
+            client.flush()
 
     @abc.abstractmethod
     def _streams_handler(
@@ -520,7 +591,7 @@ def pop_end_candidates() -> Tuple[span.SpanData, Optional[trace.TraceData]]:
     from the current context, returns popped objects.
 
     Decorator can't attach any child objects to the popped ones because
-    they are no longer in context stack.
+    they are no longer in the context stack.
     """
     span_data_to_end = context_storage.pop_span_data()
     assert (
@@ -539,3 +610,84 @@ def pop_end_candidates() -> Tuple[span.SpanData, Optional[trace.TraceData]]:
         TRACES_CREATED_BY_DECORATOR.discard(possible_trace_data_to_end.id)
 
     return span_data_to_end, trace_data_to_end
+
+
+def add_start_candidates(
+    start_span_parameters: arguments_helpers.StartSpanParameters,
+    opik_distributed_trace_headers: Optional[DistributedTraceHeadersDict],
+    opik_args_data: Optional[opik_args.OpikArgs],
+    tracing_active: bool,
+) -> span_creation_handler.SpanCreationResult:
+    """
+    Handles the creation and registration of a new start span and trace while respecting the
+    tracing context based on given parameters. It also applies relevant arguments
+    to the trace if it was created and handles client logging if the tracing is active.
+
+    Args:
+        start_span_parameters: The parameters used to start the span, including the
+            span name and other configurations.
+        opik_distributed_trace_headers: Optional headers for distributed tracing, which
+            are passed to the span creation process.
+        opik_args_data : Optional additional arguments that can be applied to the trace
+            data after the span is created.
+        tracing_active: A boolean indicating whether a tracing is active.
+
+    Returns:
+        The result of the span creation, including the span and trace data.
+    """
+    span_creation_result = span_creation_handler.create_span_respecting_context(
+        start_span_arguments=start_span_parameters,
+        distributed_trace_headers=opik_distributed_trace_headers,
+    )
+    context_storage.add_span_data(span_creation_result.span_data)
+
+    if tracing_active:
+        client = opik_client.get_client_cached()
+
+        if client.config.log_start_trace_span:
+            client.span(**span_creation_result.span_data.as_start_parameters)
+
+    if span_creation_result.trace_data is not None:
+        add_start_trace_candidate(
+            trace_data=span_creation_result.trace_data,
+            opik_args_data=opik_args_data,
+            tracing_active=tracing_active,
+        )
+
+    return span_creation_result
+
+
+def add_start_trace_candidate(
+    trace_data: trace.TraceData,
+    opik_args_data: Optional[opik_args.OpikArgs],
+    tracing_active: bool,
+) -> None:
+    """
+    Adds a start trace candidate to the current context storage and updates
+    it with the given Opik arguments if applicable.
+
+    This function initializes the trace data in the current context and
+    tracks its creation. It also applies provided Opik argument modifications
+    to the trace and logs the start trace span in the client if tracing is
+    active and logging is enabled.
+
+    Args:
+        trace_data: The trace data object to be added and initialized in the
+            current context storage. It contains details about the trace.
+        opik_args_data: Optional OpikArgs object containing additional data
+            to be applied to the trace. This may include configurations
+            that modify or enrich the trace data.
+        tracing_active: A boolean indicating whether a tracing is active.
+    """
+    context_storage.set_trace_data(trace_data)
+    TRACES_CREATED_BY_DECORATOR.add(trace_data.id)
+
+    # Handle thread_id and trace updates after span/trace creation
+    opik_args.apply_opik_args_to_trace(opik_args=opik_args_data, trace_data=trace_data)
+
+    if not tracing_active:
+        return
+
+    client = opik_client.get_client_cached()
+    if client.config.log_start_trace_span:
+        client.trace(**trace_data.as_start_parameters)