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

opik/configurator/configure.py

@@ -1,5 +1,6 @@
 import getpass
 import logging
+import os
 from typing import Final, Optional
 
 import httpx
@@ -12,7 +13,7 @@ from opik.config import (
 from opik.configurator.interactive_helpers import ask_user_for_approval, is_interactive
 from opik.configurator import opik_rest_helpers
 from opik.exceptions import ConfigurationError
-from opik import url_helpers
+import opik.url_helpers as url_helpers
 from opik.api_key import opik_api_key
 
 
@@ -46,9 +47,12 @@ class OpikConfigurator:
         # This URL set here might not be the final one.
         # It's possible that the URL will be extracted from the smart api key on the later stage.
         # In that case `self.base_url` field will be updated.
-        self.base_url = (
-            OPIK_BASE_URL_CLOUD if url is None else url_helpers.get_base_url(url)
-        )
+        if url is None:
+            self.base_url = (
+                OPIK_BASE_URL_LOCAL if self.use_local else OPIK_BASE_URL_CLOUD
+            )
+        else:
+            self.base_url = url_helpers.get_base_url(url)
 
     def configure(self) -> None:
         """
@@ -78,7 +82,7 @@ class OpikConfigurator:
     def _configure_cloud(self) -> None:
         """
         Configure the non-local Opik instance by handling API key and workspace settings.
-        non-local means both cloud and onprem.
+        Non-local means both cloud and onprem.
         """
         # Handle API key: get or prompt for one if needed
         update_config_with_api_key = self._set_api_key()
@@ -91,11 +95,16 @@ class OpikConfigurator:
             self._update_config()
         else:
             self._update_config(save_to_file=False)
+            _set_environment_variables_for_integrations(
+                api_key=self.api_key, workspace=self.workspace
+            )
             LOGGER.info(
                 "Opik is already configured. You can check the settings by viewing the config file at %s",
                 self.current_config.config_file_fullpath,
             )
 
+        self._log_project_configuration_message()
+
     def _configure_local(self) -> None:
         """
         Configure the local Opik instance by setting the local URL and workspace.
@@ -112,6 +121,7 @@ class OpikConfigurator:
         # Step 1: If the URL is provided and active, update the configuration
         if url_was_provided and opik_rest_helpers.is_instance_active(self.base_url):
             self._update_config(save_to_file=self.force)
+            self._log_project_configuration_message()
             return
 
         # Step 2: Check if the default local instance is active
@@ -123,6 +133,7 @@ class OpikConfigurator:
                 LOGGER.info(
                     f"Opik is already configured to local instance at {OPIK_BASE_URL_LOCAL}."
                 )
+                self._log_project_configuration_message()
                 return
 
             # Step 3: Ask user if they want to use the found local instance
@@ -142,6 +153,7 @@ class OpikConfigurator:
             if use_url:
                 self.base_url = OPIK_BASE_URL_LOCAL
                 self._update_config()
+                self._log_project_configuration_message()
                 return
 
         # Step 4: Ask user for URL if no valid local instance is found or approved
@@ -151,6 +163,7 @@ class OpikConfigurator:
             )
         self._ask_for_url()
         self._update_config()
+        self._log_project_configuration_message()
 
     def _set_api_key(self) -> bool:
         """
@@ -268,7 +281,7 @@ class OpikConfigurator:
                 )
             return True if self.force else False
 
-        # Case 2: Use workspace from current configuration if not forced to change
+        # Case 2: Use workspace from the current configuration if not forced to change
         if (
             "workspace" in self.current_config.model_fields_set
             and self.current_config.workspace != OPIK_WORKSPACE_DEFAULT_NAME
@@ -398,6 +411,17 @@ class OpikConfigurator:
             LOGGER.error(f"Failed to update config: {str(e)}")
             raise ConfigurationError("Failed to update configuration.")
 
+    def _log_project_configuration_message(self) -> None:
+        """
+        Log an informative message about project configuration after successful setup.
+        """
+        project_name = self.current_config.project_name
+
+        LOGGER.info(
+            f"Configuration completed successfully. Traces will be logged to '{project_name}' project. "
+            "To change the destination project, see: https://www.comet.com/docs/opik/tracing/log_traces#configuring-the-project-name"
+        )
+
     def _ask_for_url(self) -> None:
         """
         Prompt the user for an Opik instance URL and check if it is accessible.
@@ -449,6 +473,21 @@ class OpikConfigurator:
         self.base_url = extracted_base_url
 
 
+def _set_environment_variables_for_integrations(
+    api_key: Optional[str], workspace: Optional[str]
+) -> None:
+    """
+    Environment variables are set for use by some integrations (liteLLM, etc.) when both the API key and workspace name
+    are provided by the user. According to the current implementation logic, these values will not be
+    saved to the OPIK configuration file. As a result, some third-party integrations will not be able to use them.
+    This is a workaround for this issue: https://github.com/comet-ml/opik/issues/2118
+    """
+    if api_key is not None:
+        os.environ["OPIK_API_KEY"] = api_key
+    if workspace is not None:
+        os.environ["OPIK_WORKSPACE"] = workspace
+
+
 def _extract_base_url_from_api_key(api_key: str) -> Optional[str]:
     opik_api_key_ = opik_api_key.parse_api_key(api_key)
 

opik/configurator/opik_rest_helpers.py

@@ -4,7 +4,9 @@ from typing import Final, List, Optional
 import httpx
 
 from opik.exceptions import ConfigurationError
-from opik import url_helpers, config, httpx_client
+import opik.url_helpers as url_helpers
+import opik.config as config
+import opik.httpx_client as httpx_client
 
 LOGGER = logging.getLogger(__name__)
 
@@ -17,6 +19,7 @@ def _get_httpx_client(api_key: Optional[str] = None) -> httpx.Client:
         workspace=None,
         api_key=api_key,
         check_tls_certificate=config_.check_tls_certificate,
+        compress_json_requests=config_.enable_json_request_compression,
     )
 
     return client

opik/context_storage.py

@@ -1,82 +1,181 @@
 import contextvars
 import contextlib
 
-from typing import List, Optional, Generator
+from typing import List, Optional, Generator, Tuple
 from opik.api_objects import span, trace
 
-_current_trace_data_context: contextvars.ContextVar[Optional[trace.TraceData]] = (
-    contextvars.ContextVar("current_trace_data", default=None)
-)
-
-_spans_data_stack_context: contextvars.ContextVar[List[span.SpanData]] = (
-    contextvars.ContextVar("spans_data_stack", default=[])
-)
-
-# Read if you are going to change this module.
-#
-# If you want to add or pop the span from the stack
-# you must NOT do it like that: spans_stack_context.get().append()
-#
-# If you modify the default object in ContextVar (which is a MUTABLE list)
-# these changes will be reflected in any new context which breaks the whole idea
-# of ContextVar.
-#
-# You should append to it like that:
-# span_data_stack = span_data_stack_context.get().copy()
-# span_data_stack_context.set(span_data_stack + [new_element])
-#
-# And pop like that:
-# span_data_stack = spans_stack_context.get().copy()
-# span_data_stack_context.set(span_data_stack[:-1])
-#
-# The following functions provide an API to work with ContextVars this way
-
-
-def top_span_data() -> Optional[span.SpanData]:
-    if span_data_stack_empty():
-        return None
-
-    stack = _get_data_stack()
-    return stack[-1]
-
-
-def pop_span_data() -> Optional[span.SpanData]:
-    stack = _get_data_stack()
-    _spans_data_stack_context.set(stack[:-1])
-    return stack[-1]
-
-
-def add_span_data(span: span.SpanData) -> None:
-    stack = _get_data_stack()
-    _spans_data_stack_context.set(stack + [span])
-
-
-def span_data_stack_empty() -> bool:
-    return len(_get_data_stack()) == 0
 
+class OpikContextStorage:
+    """
+    Manages span and trace context using Python's contextvars.
 
-def _get_data_stack() -> List[span.SpanData]:
-    return _spans_data_stack_context.get().copy()
+    ## IMPORTANT: Working with ContextVars
 
+    This class uses ContextVars to maintain isolated stacks across different
+    execution contexts (like threads or async tasks). To ensure proper isolation and safety,
+    this implementation uses immutable tuples for stack storage.
 
-def get_trace_data() -> Optional[trace.TraceData]:
-    trace = _current_trace_data_context.get()
-    return trace
+    ### DO use immutable data structures and create-new-set pattern:
 
+    For adding elements:
+    ```python
+    # Get current tuple and create a new one with added element
+    stack = spans_stack_context.get()
+    spans_stack_context.set(stack + (new_element,))
+    ```
 
-def pop_trace_data() -> Optional[trace.TraceData]:
-    trace = _current_trace_data_context.get()
-    set_trace_data(None)
-    return trace
+    For removing elements:
+    ```python
+    # Get current tuple and create a new one without the last element
+    stack = spans_stack_context.get()
+    spans_stack_context.set(stack[:-1])
+    ```
 
+    The methods in this class follow these patterns and provide a safe API
+    for manipulating the context stacks.
+    """
 
-def set_trace_data(trace: Optional[trace.TraceData]) -> None:
-    _current_trace_data_context.set(trace)
+    def __init__(self) -> None:
+        self._current_trace_data_context: contextvars.ContextVar[
+            Optional[trace.TraceData]
+        ] = contextvars.ContextVar("current_trace_data", default=None)
+        default_span_stack: Tuple[span.SpanData, ...] = tuple()
+        self._spans_data_stack_context: contextvars.ContextVar[
+            Tuple[span.SpanData, ...]
+        ] = contextvars.ContextVar("spans_data_stack", default=default_span_stack)
+
+    def _has_span_id(self, span_id: str) -> bool:
+        return any(span.id == span_id for span in self._spans_data_stack_context.get())
+
+    def trim_span_data_stack_to_certain_span(self, span_id: str) -> None:
+        """
+        If span with the given id exists in the stack, eliminates the spans from the stack
+        until the span with the given id is at the top.
+
+        Intended to be used in the modules that perform unsafe manipulations with the
+        span data stack (when there is a risk of missing the pop operation, e.g. in callback-based integrations).
+
+        When the id of the span that SHOULD be on top is known, we can trim
+        the stack to remove hanged spans if there are any.
+
+        Args:
+            span_id: The id of the span to trim the stack to.
+        Returns:
+            None
+        """
+        if not self._has_span_id(span_id):
+            return
+
+        stack = self._spans_data_stack_context.get()
+        new_stack_list: List[span.SpanData] = []
+        for span_data in stack:
+            new_stack_list.append(span_data)
+            if span_data.id == span_id:
+                break
+
+        self._spans_data_stack_context.set(tuple(new_stack_list))
+
+    def top_span_data(self) -> Optional[span.SpanData]:
+        if self.span_data_stack_empty():
+            return None
+        stack = self._spans_data_stack_context.get()
+        return stack[-1]
+
+    def pop_span_data(
+        self,
+        ensure_id: Optional[str] = None,
+    ) -> Optional[span.SpanData]:
+        """
+        Pops the span from the stack.
+        Args:
+            ensure_id: If provided, it will pop the span only if it has the given id.
+                Intended to be used in the modules that perform unsafe manipulations with the
+                span data stack (when there is a risk of missing the add or pop operation,
+                e.g. in callback-based integrations), to make sure the correct span is popped.
+        Returns:
+            The span that was popped from the stack or None.
+        """
+        if self.span_data_stack_empty():
+            return None
+
+        if ensure_id is None:
+            stack = self._spans_data_stack_context.get()
+            self._spans_data_stack_context.set(stack[:-1])
+            return stack[-1]
+
+        if self.top_span_data().id == ensure_id:  # type: ignore
+            return self.pop_span_data()
+
+        STACK_IS_EMPTY_OR_THE_ID_DOES_NOT_MATCH = None
+        return STACK_IS_EMPTY_OR_THE_ID_DOES_NOT_MATCH
+
+    def add_span_data(self, span: span.SpanData) -> None:
+        stack = self._spans_data_stack_context.get()
+        self._spans_data_stack_context.set(stack + (span,))
+
+    def span_data_stack_empty(self) -> bool:
+        return len(self._spans_data_stack_context.get()) == 0
+
+    def span_data_stack_size(self) -> int:
+        return len(self._spans_data_stack_context.get())
+
+    def get_trace_data(self) -> Optional[trace.TraceData]:
+        trace_data = self._current_trace_data_context.get()
+        return trace_data
+
+    def pop_trace_data(
+        self, ensure_id: Optional[str] = None
+    ) -> Optional[trace.TraceData]:
+        """
+        Pops the trace from the context.
+        Args:
+            ensure_id: If provided, it will pop the trace only if it has the given id.
+                Intended to be used in the modules that perform unsafe manipulations with the
+                trace data (when there is a risk of missing the set operation,
+                e.g. in callback-based integrations), to make sure the correct trace is popped.
+        Returns:
+            The trace that was popped from the context or None.
+        """
+        trace_data = self._current_trace_data_context.get()
+
+        if trace_data is None:
+            return None
+
+        if ensure_id is not None and trace_data.id != ensure_id:
+            return None
+
+        self.set_trace_data(None)
+        return trace_data
+
+    def set_trace_data(self, trace: Optional[trace.TraceData]) -> None:
+        self._current_trace_data_context.set(trace)
+
+    def clear_spans(self) -> None:
+        self._spans_data_stack_context.set(tuple())
+
+    def clear_all(self) -> None:
+        self._current_trace_data_context.set(None)
+        self.clear_spans()
+
+
+_context_storage = OpikContextStorage()
+
+top_span_data = _context_storage.top_span_data
+pop_span_data = _context_storage.pop_span_data
+add_span_data = _context_storage.add_span_data
+span_data_stack_empty = _context_storage.span_data_stack_empty
+get_trace_data = _context_storage.get_trace_data
+pop_trace_data = _context_storage.pop_trace_data
+set_trace_data = _context_storage.set_trace_data
+clear_all = _context_storage.clear_all
+span_data_stack_size = _context_storage.span_data_stack_size
+trim_span_data_stack_to_certain_span = (
+    _context_storage.trim_span_data_stack_to_certain_span
+)
 
 
-def clear_all() -> None:
-    _current_trace_data_context.set(None)
-    _spans_data_stack_context.set([])
+def get_current_context_instance() -> OpikContextStorage:
+    return _context_storage
 
 
 @contextlib.contextmanager
@@ -84,7 +183,7 @@ def temporary_context(
     span_data: span.SpanData, trace_data: Optional[trace.TraceData]
 ) -> Generator[None, None, None]:
     """
-    Temporary adds span and trace data to the context.
+    Temporarily adds span and trace data to the context.
     If trace_data is None, it has no effect on the current trace in the context.
     """
     try:

opik/datetime_helpers.py

@@ -1,4 +1,5 @@
 import datetime
+from typing import Optional
 
 
 def local_timestamp() -> datetime.datetime:
@@ -8,3 +9,14 @@ def local_timestamp() -> datetime.datetime:
 
 def datetime_to_iso8601(value: datetime.datetime) -> str:
     return value.isoformat()
+
+
+def parse_iso_timestamp(timestamp_str: Optional[str]) -> Optional[datetime.datetime]:
+    """Parse an ISO 8601 timestamp string to datetime."""
+    if timestamp_str is None:
+        return None
+    try:
+        timestamp_str = timestamp_str.replace("Z", "+00:00")
+        return datetime.datetime.fromisoformat(timestamp_str)
+    except (ValueError, TypeError):
+        return None

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)