Updated trace url link (#366)

comet-ml · Oct 11, 2024 · 8a29238 · 8a29238
1 parent 29ca4fe
commit 8a29238
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 24 deletions.
diff --git a/sdks/python/src/opik/api_objects/opik_client.py b/sdks/python/src/opik/api_objects/opik_client.py
@@ -19,7 +19,7 @@
 from ..rest_api import client as rest_api_client
 from ..rest_api.types import dataset_public, trace_public, span_public, project_public
 from ..rest_api.core.api_error import ApiError
-from .. import datetime_helpers, config, httpx_client
+from .. import datetime_helpers, config, httpx_client, url_helpers
 
 
 LOGGER = logging.getLogger(__name__)
@@ -51,6 +51,7 @@ def __init__(
         self._workspace: str = config_.workspace
         self._project_name: str = config_.project_name
         self._flush_timeout: Optional[int] = config_.default_flush_timeout
+        self._project_name_most_recent_trace: Optional[str] = None
 
         self._initialize_streamer(
             base_url=config_.url_override,
@@ -79,6 +80,18 @@ def _initialize_streamer(
             use_batching=use_batching,
         )
 
+    def _display_trace_url(self, workspace: str, project_name: str) -> None:
+        projects_url = url_helpers.get_projects_url(workspace=workspace)
+
+        if (
+            self._project_name_most_recent_trace is None
+            or self._project_name_most_recent_trace != project_name
+        ):
+            LOGGER.info(
+                f'Started logging traces to the "{project_name}" project at {projects_url}.'
+            )
+            self._project_name_most_recent_trace = project_name
+
     def trace(
         self,
         id: Optional[str] = None,
@@ -126,6 +139,9 @@ def trace(
             tags=tags,
         )
         self._streamer.put(create_trace_message)
+        self._display_trace_url(
+            workspace=self._workspace, project_name=project_name or self._project_name
+        )
 
         if feedback_scores is not None:
             for feedback_score in feedback_scores:

diff --git a/sdks/python/src/opik/message_processing/streamer.py b/sdks/python/src/opik/message_processing/streamer.py
@@ -7,8 +7,6 @@
 from .. import synchronization
 from .batching import batch_manager
 
-from .. import url_helpers
-
 LOGGER = logging.getLogger(__name__)
 
 
@@ -31,9 +29,6 @@ def __init__(
         if self._batch_manager is not None:
             self._batch_manager.start()
 
-        # Used to know when to display the project URL
-        self._project_name_most_recent_trace: Optional[str] = None
-
     def put(self, message: messages.BaseMessage) -> None:
         with self._lock:
             if self._drain:
@@ -47,19 +42,6 @@ def put(self, message: messages.BaseMessage) -> None:
             else:
                 self._message_queue.put(message)
 
-        # Display message in console
-        if isinstance(message, messages.CreateTraceMessage):
-            projects_url = url_helpers.get_projects_url()
-            project_name = message.project_name
-            if (
-                self._project_name_most_recent_trace is None
-                or self._project_name_most_recent_trace != project_name
-            ):
-                LOGGER.info(
-                    f'Started logging traces to the "{project_name}" project at {projects_url}.'
-                )
-                self._project_name_most_recent_trace = project_name
-
     def close(self, timeout: Optional[int]) -> bool:
         """
         Stops data sending threads

diff --git a/sdks/python/src/opik/url_helpers.py b/sdks/python/src/opik/url_helpers.py
@@ -1,5 +1,4 @@
 import opik.config
-import opik.api_objects.opik_client
 
 
 def get_ui_url() -> str:
@@ -24,8 +23,6 @@ def get_experiment_url(dataset_name: str, experiment_id: str) -> str:
     return f"{ui_url}/{config.workspace}/experiments/{dataset_id}/compare?experiments=%5B%22{experiment_id}%22%5D"
 
 
-def get_projects_url() -> str:
-    config = opik.config.OpikConfig()
+def get_projects_url(workspace: str) -> str:
     ui_url = get_ui_url()
-
-    return f"{ui_url}/{config.workspace}/projects"
+    return f"{ui_url}/{workspace}/projects"