set user agent

microsoft · Jun 17, 2024 · 3d045c7 · 3d045c7
1 parent 003ffc0
commit 3d045c7
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/src/promptflow-devkit/promptflow/_sdk/_utilities/tracing_utils.py b/src/promptflow-devkit/promptflow/_sdk/_utilities/tracing_utils.py
@@ -31,9 +31,10 @@
 )
 from promptflow._sdk._constants import HOME_PROMPT_FLOW_DIR, AzureMLWorkspaceTriad
 from promptflow._sdk._telemetry.telemetry import get_telemetry_logger
+from promptflow._sdk._user_agent import USER_AGENT
 from promptflow._sdk.entities._trace import Span
 from promptflow._utils.logger_utils import get_cli_sdk_logger
-from promptflow._utils.user_agent_utils import ClientUserAgentUtil
+from promptflow._utils.user_agent_utils import ClientUserAgentUtil, setup_user_agent_to_operation_context
 from promptflow.core._errors import MissingRequiredPackage
 
 from .general_utils import convert_time_unix_nano_to_timestamp, json_load
@@ -346,6 +347,7 @@ class TraceTelemetryHelper:
     CUSTOM_DIMENSIONS_TRACE_COUNT = "trace_count"
 
     def __init__(self):
+        setup_user_agent_to_operation_context(USER_AGENT)
         self._user_agent = ClientUserAgentUtil.get_user_agent()
         self._telemetry_logger = get_telemetry_logger()
         self._lock = multiprocessing.Lock()

diff --git a/src/promptflow-devkit/tests/sdk_cli_test/unittests/test_trace.py b/src/promptflow-devkit/tests/sdk_cli_test/unittests/test_trace.py
@@ -341,6 +341,7 @@ def mock_info(*args, **kwargs):
             extra: dict = kwargs.get("extra")
             custom_dimensions: dict = extra.get("custom_dimensions")
             assert "user_agent" in custom_dimensions.keys()
+            assert "promptflow-sdk/" in custom_dimensions["user_agent"]
 
         telemetry_helper = TraceTelemetryHelper()
         telemetry_helper._telemetry_logger.info = mock_info