diff --git a/packages/uipath-platform/pyproject.toml b/packages/uipath-platform/pyproject.toml index 989952eb3..36bc271c1 100644 --- a/packages/uipath-platform/pyproject.toml +++ b/packages/uipath-platform/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "uipath-platform" -version = "0.1.43" +version = "0.1.44" description = "HTTP client library for programmatic access to UiPath Platform" readme = { file = "README.md", content-type = "text/markdown" } requires-python = ">=3.11" diff --git a/packages/uipath-platform/src/uipath/platform/chat/_llm_gateway_service.py b/packages/uipath-platform/src/uipath/platform/chat/_llm_gateway_service.py index d7c093d0d..ffe0bff99 100644 --- a/packages/uipath-platform/src/uipath/platform/chat/_llm_gateway_service.py +++ b/packages/uipath-platform/src/uipath/platform/chat/_llm_gateway_service.py @@ -35,6 +35,7 @@ ToolDefinition, ) from .llm_throttle import get_llm_semaphore +from .llm_trace_context import build_trace_context_headers # Common constants API_VERSION = "2024-10-21" # Standard API version for OpenAI-compatible endpoints @@ -224,7 +225,7 @@ async def embeddings( endpoint, json={"input": input}, params={"api-version": API_VERSION}, - headers=self._llm_headers, + headers={**self._llm_headers, **build_trace_context_headers()}, ) return TextEmbedding.model_validate(response.json()) @@ -355,7 +356,7 @@ class Country(BaseModel): endpoint, json=request_body, params={"api-version": API_VERSION}, - headers=self._llm_headers, + headers={**self._llm_headers, **build_trace_context_headers()}, ) return ChatCompletion.model_validate(response.json()) @@ -599,6 +600,7 @@ class Country(BaseModel): headers = { **self._llm_headers, + **build_trace_context_headers(), "X-UiPath-LlmGateway-NormalizedApi-ModelName": model, "X-UiPath-LLMGateway-AllowFull4xxResponse": "true", } diff --git a/packages/uipath-platform/src/uipath/platform/chat/llm_trace_context.py b/packages/uipath-platform/src/uipath/platform/chat/llm_trace_context.py new file mode 100644 index 000000000..4c6dd6062 --- /dev/null +++ b/packages/uipath-platform/src/uipath/platform/chat/llm_trace_context.py @@ -0,0 +1,42 @@ +"""W3C-style trace context headers for LLM Gateway requests.""" + +from opentelemetry import trace +from uipath.core.feature_flags import FeatureFlags + +from ..common._config import UiPathConfig + + +def build_trace_context_headers( + extra_baggage: list[str] | None = None, +) -> dict[str, str]: + """Build W3C-style trace context headers from the current OpenTelemetry span. + + Args: + extra_baggage: Additional baggage entries (e.g. ``["source=agents"]``) + that callers can inject alongside the platform-level entries. + + Returns an empty dict when the ``EnableTraceContextHeaders`` feature flag + is not enabled, or when no active span is present. + """ + if not FeatureFlags.is_flag_enabled("EnableTraceContextHeaders"): + return {} + + headers: dict[str, str] = {} + span = trace.get_current_span() + ctx = span.get_span_context() + if ctx and ctx.trace_id and ctx.span_id: + trace_id = format(ctx.trace_id, "032x") + span_id = format(ctx.span_id, "016x") + headers["x-uipath-traceparent-id"] = f"00-{trace_id}-{span_id}" + + baggage_parts: list[str] = list(extra_baggage) if extra_baggage else [] + if folder_key := UiPathConfig.folder_key: + baggage_parts.append(f"folderKey={folder_key}") + if agent_id := UiPathConfig.process_uuid: + baggage_parts.append(f"agentId={agent_id}") + if process_key := UiPathConfig.process_key: + baggage_parts.append(f"processKey={process_key}") + if baggage_parts: + headers["x-uipath-tracebaggage"] = ",".join(baggage_parts) + + return headers diff --git a/packages/uipath-platform/src/uipath/platform/common/_config.py b/packages/uipath-platform/src/uipath/platform/common/_config.py index cdb5eb383..40db82214 100644 --- a/packages/uipath-platform/src/uipath/platform/common/_config.py +++ b/packages/uipath-platform/src/uipath/platform/common/_config.py @@ -121,6 +121,12 @@ def folder_path(self) -> str | None: return os.getenv(ENV_FOLDER_PATH, None) + @property + def process_key(self) -> str | None: + from uipath.platform.common.constants import ENV_PROCESS_KEY + + return os.getenv(ENV_PROCESS_KEY, None) + @property def process_uuid(self) -> str | None: from uipath.platform.common.constants import ENV_UIPATH_PROCESS_UUID diff --git a/packages/uipath-platform/tests/services/test_llm_trace_context.py b/packages/uipath-platform/tests/services/test_llm_trace_context.py new file mode 100644 index 000000000..83bde3957 --- /dev/null +++ b/packages/uipath-platform/tests/services/test_llm_trace_context.py @@ -0,0 +1,161 @@ +"""Tests for build_trace_context_headers.""" + +import os +from unittest.mock import patch + +from opentelemetry import trace +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.trace import NonRecordingSpan, SpanContext, TraceFlags +from uipath.core.feature_flags import FeatureFlags + +from uipath.platform.chat.llm_trace_context import build_trace_context_headers + +FEATURE_FLAG = "EnableTraceContextHeaders" + + +class TestFeatureFlagDisabled: + """When the feature flag is off, no headers are returned.""" + + def setup_method(self) -> None: + FeatureFlags.reset_flags() + + def test_returns_empty_dict_by_default(self) -> None: + assert build_trace_context_headers() == {} + + def test_returns_empty_dict_when_explicitly_disabled(self) -> None: + FeatureFlags.configure_flags({FEATURE_FLAG: False}) + assert build_trace_context_headers() == {} + + +class TestTraceparentHeader: + """When enabled, x-uipath-traceparent-id is populated from the active span.""" + + def setup_method(self) -> None: + FeatureFlags.reset_flags() + FeatureFlags.configure_flags({FEATURE_FLAG: True}) + + def test_traceparent_from_active_span(self) -> None: + provider = TracerProvider() + tracer = provider.get_tracer("test") + with tracer.start_as_current_span("test-span") as span: + ctx = span.get_span_context() + expected_trace_id = format(ctx.trace_id, "032x") + expected_span_id = format(ctx.span_id, "016x") + + headers = build_trace_context_headers() + + assert "x-uipath-traceparent-id" in headers + value = headers["x-uipath-traceparent-id"] + assert value == f"00-{expected_trace_id}-{expected_span_id}" + # Verify format: version (2) + dash + trace_id (32) + dash + span_id (16) + parts = value.split("-") + assert len(parts) == 3 + assert parts[0] == "00" + assert len(parts[1]) == 32 + assert len(parts[2]) == 16 + + def test_no_traceparent_without_active_span(self) -> None: + # INVALID_SPAN has trace_id=0 and span_id=0 + from opentelemetry.context import attach, detach + + ctx = SpanContext( + trace_id=0, + span_id=0, + is_remote=False, + trace_flags=TraceFlags(0), + ) + non_recording = NonRecordingSpan(ctx) + token = attach(trace.set_span_in_context(non_recording)) + try: + headers = build_trace_context_headers() + finally: + detach(token) + + assert "x-uipath-traceparent-id" not in headers + + +class TestBaggageHeader: + """When enabled, x-uipath-tracebaggage is populated from UiPathConfig.""" + + def setup_method(self) -> None: + FeatureFlags.reset_flags() + FeatureFlags.configure_flags({FEATURE_FLAG: True}) + + def test_all_env_vars_present(self) -> None: + env = { + "UIPATH_FOLDER_KEY": "folder-abc", + "UIPATH_PROCESS_UUID": "agent-123", + "UIPATH_PROCESS_KEY": "process-789", + } + with patch.dict(os.environ, env, clear=True): + headers = build_trace_context_headers() + + baggage = headers["x-uipath-tracebaggage"] + assert "folderKey=folder-abc" in baggage + assert "agentId=agent-123" in baggage + assert "processKey=process-789" in baggage + + def test_partial_env_vars(self) -> None: + env = {"UIPATH_FOLDER_KEY": "folder-only"} + with patch.dict(os.environ, env, clear=True): + headers = build_trace_context_headers() + + baggage = headers["x-uipath-tracebaggage"] + assert "folderKey=folder-only" in baggage + + def test_no_baggage_without_env_vars(self) -> None: + with patch.dict(os.environ, {}, clear=True): + headers = build_trace_context_headers() + + assert "x-uipath-tracebaggage" not in headers + + def test_baggage_comma_separated(self) -> None: + env = { + "UIPATH_FOLDER_KEY": "f1", + "UIPATH_PROCESS_UUID": "a1", + } + with patch.dict(os.environ, env, clear=True): + headers = build_trace_context_headers() + + baggage = headers["x-uipath-tracebaggage"] + parts = baggage.split(",") + assert len(parts) == 2 # folderKey + agentId + + def test_extra_baggage_included(self) -> None: + env = {"UIPATH_FOLDER_KEY": "f1"} + with patch.dict(os.environ, env, clear=True): + headers = build_trace_context_headers(extra_baggage=["source=agents"]) + + baggage = headers["x-uipath-tracebaggage"] + assert "source=agents" in baggage + assert "folderKey=f1" in baggage + + def test_extra_baggage_only(self) -> None: + with patch.dict(os.environ, {}, clear=True): + headers = build_trace_context_headers( + extra_baggage=["source=agents", "custom=value"] + ) + + baggage = headers["x-uipath-tracebaggage"] + assert baggage == "source=agents,custom=value" + + +class TestBothHeaders: + """When enabled with an active span and env vars, both headers are present.""" + + def setup_method(self) -> None: + FeatureFlags.reset_flags() + FeatureFlags.configure_flags({FEATURE_FLAG: True}) + + def test_both_headers_present(self) -> None: + provider = TracerProvider() + tracer = provider.get_tracer("test") + env = {"UIPATH_FOLDER_KEY": "folder-abc"} + with ( + tracer.start_as_current_span("test-span"), + patch.dict(os.environ, env, clear=True), + ): + headers = build_trace_context_headers() + + assert "x-uipath-traceparent-id" in headers + assert "x-uipath-tracebaggage" in headers diff --git a/packages/uipath-platform/uv.lock b/packages/uipath-platform/uv.lock index 66c8eb62a..435bc613b 100644 --- a/packages/uipath-platform/uv.lock +++ b/packages/uipath-platform/uv.lock @@ -1088,7 +1088,7 @@ dev = [ [[package]] name = "uipath-platform" -version = "0.1.43" +version = "0.1.44" source = { editable = "." } dependencies = [ { name = "httpx" }, diff --git a/packages/uipath/uv.lock b/packages/uipath/uv.lock index 2eacfed8a..dc6b7286a 100644 --- a/packages/uipath/uv.lock +++ b/packages/uipath/uv.lock @@ -2682,7 +2682,7 @@ dev = [ [[package]] name = "uipath-platform" -version = "0.1.43" +version = "0.1.44" source = { editable = "../uipath-platform" } dependencies = [ { name = "httpx" },