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

opik/decorator/arguments_helpers.py

@@ -3,7 +3,7 @@ from typing import Any, Callable, Dict, List, Optional, Union
 
 from .. import datetime_helpers, llm_usage
 from ..api_objects import helpers, span
-from ..types import ErrorInfoDict, SpanType
+from ..types import ErrorInfoDict, SpanType, DistributedTraceHeadersDict
 
 
 @dataclasses.dataclass
@@ -34,6 +34,7 @@ class EndSpanParameters(BaseArguments):
     model: Optional[str] = None
     provider: Optional[str] = None
     error_info: Optional[ErrorInfoDict] = None
+    total_cost: Optional[float] = None
 
 
 @dataclasses.dataclass
@@ -51,6 +52,7 @@ class StartSpanParameters(BaseArguments):
     project_name: Optional[str] = None
     model: Optional[str] = None
     provider: Optional[str] = None
+    thread_id: Optional[str] = None  # used for traces only
 
 
 @dataclasses.dataclass
@@ -91,3 +93,9 @@ def create_span_data(
         provider=start_span_arguments.provider,
     )
     return span_data
+
+
+def extract_distributed_trace_headers(
+    kwargs: Dict[str, Any],
+) -> Optional[DistributedTraceHeadersDict]:
+    return kwargs.pop("opik_distributed_trace_headers", None)

opik/decorator/base_track_decorator.py

@@ -11,9 +11,10 @@ from typing import (
     Set,
     Tuple,
     Union,
+    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 (
@@ -21,8 +22,8 @@ from . import (
     error_info_collector,
     generator_wrappers,
     inspect_helpers,
+    opik_args,
     span_creation_handler,
-    tracing_runtime_config,
 )
 
 LOGGER = logging.getLogger(__name__)
@@ -30,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.
@@ -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,6 +183,54 @@ 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:
@@ -179,31 +238,19 @@ class BaseTrackDecorator(abc.ABC):
         def wrapper(*args, **kwargs) -> Any:  # type: ignore
             if not tracing_runtime_config.is_tracing_active():
                 return func(*args, **kwargs)
-            try:
-                opik_distributed_trace_headers: Optional[
-                    DistributedTraceHeadersDict
-                ] = kwargs.pop("opik_distributed_trace_headers", None)
 
-                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,
-                    inspect_helpers.get_function_name(func),
-                    (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,
+            )
 
             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,
                 )
@@ -229,31 +276,19 @@ class BaseTrackDecorator(abc.ABC):
         def wrapper(*args, **kwargs) -> Any:  # type: ignore
             if not tracing_runtime_config.is_tracing_active():
                 return func(*args, **kwargs)
-            try:
-                opik_distributed_trace_headers: Optional[
-                    DistributedTraceHeadersDict
-                ] = kwargs.pop("opik_distributed_trace_headers", None)
 
-                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,
-                    inspect_helpers.get_function_name(func),
-                    (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,
+            )
 
             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,
                 )
@@ -301,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
 
@@ -355,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
@@ -408,61 +441,19 @@ class BaseTrackDecorator(abc.ABC):
         args: Tuple,
         kwargs: Dict[str, Any],
     ) -> None:
-        opik_distributed_trace_headers: Optional[DistributedTraceHeadersDict] = None
-
-        try:
-            opik_distributed_trace_headers = kwargs.pop(
-                "opik_distributed_trace_headers", None
-            )
-
-            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,
-                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,
-            )
-
-        created_trace_data, created_span_data = (
-            span_creation_handler.create_span_respecting_context(
-                start_span_arguments=start_span_arguments,
-                distributed_trace_headers=opik_distributed_trace_headers,
-            )
+        track_start_options = self._prepare_tracking_start_options(
+            func=func,
+            track_options=track_options,
+            args=args,
+            kwargs=kwargs,
         )
-        client = opik_client.get_client_cached()
-
-        if (
-            client.config.log_start_trace_span
-            and tracing_runtime_config.is_tracing_active()
-        ):
-            client.span(**created_span_data.as_start_parameters)
 
-        if created_trace_data is not None:
-            context_storage.set_trace_data(created_trace_data)
-            TRACES_CREATED_BY_DECORATOR.add(created_trace_data.id)
-
-            if (
-                client.config.log_start_trace_span
-                and tracing_runtime_config.is_tracing_active()
-            ):
-                client.trace(**created_trace_data.as_start_parameters)
-
-        context_storage.add_span_data(created_span_data)
+        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,
@@ -619,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)

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/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",
+]