diff --git a/pyproject.toml b/pyproject.toml index 833d19fbb..0c4cb62a7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "uipath-langchain" -version = "0.13.18" +version = "0.13.19" description = "Python SDK that enables developers to build and deploy LangGraph agents to the UiPath Cloud Platform" readme = { file = "README.md", content-type = "text/markdown" } requires-python = ">=3.11" diff --git a/src/uipath_langchain/chat/_legacy/bedrock.py b/src/uipath_langchain/chat/_legacy/bedrock.py index dc99e3757..5f09c66d9 100644 --- a/src/uipath_langchain/chat/_legacy/bedrock.py +++ b/src/uipath_langchain/chat/_legacy/bedrock.py @@ -72,12 +72,14 @@ def __init__( agenthub_config: Optional[str] = None, byo_connection_id: Optional[str] = None, header_capture: HeaderCapture | None = None, + trace_context_extra_baggage: Optional[list[str]] = None, ): self.model = model self.token = token self.api_flavor = api_flavor self.agenthub_config = agenthub_config self.byo_connection_id = byo_connection_id + self._trace_context_extra_baggage = trace_context_extra_baggage self._vendor = "awsbedrock" self._url: Optional[str] = None self._is_override: bool = False @@ -182,7 +184,9 @@ def _modify_request(self, request, **kwargs): ) headers["X-UiPath-LlmGateway-ApiFlavor"] = self.api_flavor headers["X-UiPath-Streaming-Enabled"] = streaming - headers.update(build_trace_context_headers(extra_baggage=["source=agents"])) + headers.update( + build_trace_context_headers(extra_baggage=self._trace_context_extra_baggage) + ) request.headers.update(headers) @@ -204,6 +208,7 @@ def __init__( byo_connection_id: Optional[str] = None, retryer: Optional[Retrying] = None, aretryer: Optional[AsyncRetrying] = None, + trace_context_extra_baggage: Optional[list[str]] = None, **kwargs, ): org_id = org_id or os.getenv("UIPATH_ORGANIZATION_ID") @@ -229,6 +234,7 @@ def __init__( api_flavor="converse", agenthub_config=agenthub_config, byo_connection_id=byo_connection_id, + trace_context_extra_baggage=trace_context_extra_baggage, ) kwargs["client"] = passthrough_client.get_client() @@ -266,6 +272,7 @@ def __init__( byo_connection_id: Optional[str] = None, retryer: Optional[Retrying] = None, aretryer: Optional[AsyncRetrying] = None, + trace_context_extra_baggage: Optional[list[str]] = None, **kwargs, ): org_id = org_id or os.getenv("UIPATH_ORGANIZATION_ID") @@ -294,6 +301,7 @@ def __init__( agenthub_config=agenthub_config, byo_connection_id=byo_connection_id, header_capture=header_capture, + trace_context_extra_baggage=trace_context_extra_baggage, ) kwargs["client"] = passthrough_client.get_client() diff --git a/src/uipath_langchain/chat/_legacy/chat_model_factory.py b/src/uipath_langchain/chat/_legacy/chat_model_factory.py index 940cfcf75..e36c6eafc 100644 --- a/src/uipath_langchain/chat/_legacy/chat_model_factory.py +++ b/src/uipath_langchain/chat/_legacy/chat_model_factory.py @@ -43,6 +43,7 @@ def _create_openai_llm( max_tokens: int | None, agenthub_config: str, byo_connection_id: str | None = None, + trace_context_extra_baggage: list[str] | None = None, **kwargs: Any, ) -> BaseChatModel: """Create UiPathChatOpenAI for OpenAI models via LLMGateway.""" @@ -63,6 +64,7 @@ def _create_openai_llm( api_version=azure_open_ai_latest_api_version, agenthub_config=agenthub_config, byo_connection_id=byo_connection_id, + trace_context_extra_baggage=trace_context_extra_baggage, output_version="v1", **sampling_kwargs, **kwargs, @@ -75,6 +77,7 @@ def _create_openai_llm( api_version=azure_open_ai_latest_api_version, agenthub_config=agenthub_config, byo_connection_id=byo_connection_id, + trace_context_extra_baggage=trace_context_extra_baggage, output_version="v1", **sampling_kwargs, **kwargs, @@ -90,6 +93,7 @@ def _create_bedrock_llm( max_tokens: int | None, agenthub_config: str, byo_connection_id: str | None = None, + trace_context_extra_baggage: list[str] | None = None, **kwargs: Any, ) -> BaseChatModel: """Create UiPathChatBedrockConverse for Claude models via LLMGateway.""" @@ -109,6 +113,7 @@ def _create_bedrock_llm( max_tokens=max_tokens, agenthub_config=agenthub_config, byo_connection_id=byo_connection_id, + trace_context_extra_baggage=trace_context_extra_baggage, output_version="v1", **sampling_kwargs, **kwargs, @@ -119,6 +124,7 @@ def _create_bedrock_llm( max_tokens=max_tokens, agenthub_config=agenthub_config, byo_connection_id=byo_connection_id, + trace_context_extra_baggage=trace_context_extra_baggage, output_version="v1", **sampling_kwargs, **kwargs, @@ -134,6 +140,7 @@ def _create_vertex_llm( max_tokens: int | None, agenthub_config: str, byo_connection_id: str | None = None, + trace_context_extra_baggage: list[str] | None = None, **kwargs: Any, ) -> BaseChatModel: """Create UiPathChatVertex for Gemini models via LLMGateway.""" @@ -150,6 +157,7 @@ def _create_vertex_llm( max_tokens=max_tokens, agenthub_config=agenthub_config, byo_connection_id=byo_connection_id, + trace_context_extra_baggage=trace_context_extra_baggage, output_version="v1", **sampling_kwargs, **kwargs, @@ -248,6 +256,7 @@ def get_chat_model( max_tokens: int | None, agenthub_config: str, byo_connection_id: str | None = None, + trace_context_extra_baggage: list[str] | None = None, **kwargs: Any, ) -> BaseChatModel: """Create and configure LLM instance using LLMGateway API. @@ -272,7 +281,8 @@ def get_chat_model( effective_temperature, max_tokens, agenthub_config, - byo_connection_id, + byo_connection_id=byo_connection_id, + trace_context_extra_baggage=trace_context_extra_baggage, **kwargs, ) case LLMProvider.BEDROCK: @@ -282,7 +292,8 @@ def get_chat_model( effective_temperature, max_tokens, agenthub_config, - byo_connection_id, + byo_connection_id=byo_connection_id, + trace_context_extra_baggage=trace_context_extra_baggage, **kwargs, ) case LLMProvider.VERTEX: @@ -292,6 +303,7 @@ def get_chat_model( effective_temperature, max_tokens, agenthub_config, - byo_connection_id, + byo_connection_id=byo_connection_id, + trace_context_extra_baggage=trace_context_extra_baggage, **kwargs, ) diff --git a/src/uipath_langchain/chat/_legacy/openai.py b/src/uipath_langchain/chat/_legacy/openai.py index 93cf89b56..5c65bd198 100644 --- a/src/uipath_langchain/chat/_legacy/openai.py +++ b/src/uipath_langchain/chat/_legacy/openai.py @@ -54,38 +54,53 @@ def _inject_license_ref_id(request: httpx.Request) -> None: request.headers["X-UiPath-License-RefId"] = license_ref_id -def _inject_trace_context_headers(request: httpx.Request) -> None: +def _inject_trace_context_headers( + request: httpx.Request, + trace_context_extra_baggage: list[str] | None = None, +) -> None: """Inject trace context headers per-request from the active OTEL span.""" for key, value in build_trace_context_headers( - extra_baggage=["source=agents"] + extra_baggage=trace_context_extra_baggage ).items(): request.headers[key] = value class UiPathURLRewriteTransport(httpx.AsyncHTTPTransport): - def __init__(self, verify: bool = True, **kwargs): + def __init__( + self, + verify: bool = True, + trace_context_extra_baggage: list[str] | None = None, + **kwargs, + ): super().__init__(verify=verify, **kwargs) + self._trace_context_extra_baggage = trace_context_extra_baggage async def handle_async_request(self, request: httpx.Request) -> httpx.Response: new_url = _rewrite_openai_url(str(request.url), request.url.params) if new_url: request.url = new_url _inject_license_ref_id(request) - _inject_trace_context_headers(request) + _inject_trace_context_headers(request, self._trace_context_extra_baggage) return await super().handle_async_request(request) class UiPathSyncURLRewriteTransport(httpx.HTTPTransport): - def __init__(self, verify: bool = True, **kwargs): + def __init__( + self, + verify: bool = True, + trace_context_extra_baggage: list[str] | None = None, + **kwargs, + ): super().__init__(verify=verify, **kwargs) + self._trace_context_extra_baggage = trace_context_extra_baggage def handle_request(self, request: httpx.Request) -> httpx.Response: new_url = _rewrite_openai_url(str(request.url), request.url.params) if new_url: request.url = new_url _inject_license_ref_id(request) - _inject_trace_context_headers(request) + _inject_trace_context_headers(request, self._trace_context_extra_baggage) return super().handle_request(request) @@ -112,6 +127,7 @@ def __init__( agenthub_config: Optional[str] = None, extra_headers: Optional[dict[str, str]] = None, byo_connection_id: Optional[str] = None, + trace_context_extra_baggage: Optional[list[str]] = None, **kwargs, ): org_id = org_id or os.getenv("UIPATH_ORGANIZATION_ID") @@ -155,11 +171,17 @@ def __init__( model_name=model_name, default_headers=self._build_headers(token, inject_routing=is_override), http_async_client=httpx.AsyncClient( - transport=UiPathURLRewriteTransport(verify=verify), + transport=UiPathURLRewriteTransport( + verify=verify, + trace_context_extra_baggage=trace_context_extra_baggage, + ), **client_kwargs, ), http_client=httpx.Client( - transport=UiPathSyncURLRewriteTransport(verify=verify), + transport=UiPathSyncURLRewriteTransport( + verify=verify, + trace_context_extra_baggage=trace_context_extra_baggage, + ), **client_kwargs, ), api_key=token, diff --git a/src/uipath_langchain/chat/_legacy/vertex.py b/src/uipath_langchain/chat/_legacy/vertex.py index bef142b4b..251f572ad 100644 --- a/src/uipath_langchain/chat/_legacy/vertex.py +++ b/src/uipath_langchain/chat/_legacy/vertex.py @@ -79,10 +79,12 @@ def __init__( gateway_url: str, verify: bool = True, header_capture: HeaderCapture | None = None, + trace_context_extra_baggage: list[str] | None = None, ): super().__init__(verify=verify) self.gateway_url = gateway_url self.header_capture = header_capture + self._trace_context_extra_baggage = trace_context_extra_baggage def handle_request(self, request: httpx.Request) -> httpx.Response: original_url = str(request.url) @@ -98,7 +100,7 @@ def handle_request(self, request: httpx.Request) -> httpx.Response: request.url = new_url for key, value in build_trace_context_headers( - extra_baggage=["source=agents"] + extra_baggage=self._trace_context_extra_baggage ).items(): request.headers[key] = value @@ -117,10 +119,12 @@ def __init__( gateway_url: str, verify: bool = True, header_capture: HeaderCapture | None = None, + trace_context_extra_baggage: list[str] | None = None, ): super().__init__(verify=verify) self.gateway_url = gateway_url self.header_capture = header_capture + self._trace_context_extra_baggage = trace_context_extra_baggage async def handle_async_request(self, request: httpx.Request) -> httpx.Response: original_url = str(request.url) @@ -136,7 +140,7 @@ async def handle_async_request(self, request: httpx.Request) -> httpx.Response: request.url = new_url for key, value in build_trace_context_headers( - extra_baggage=["source=agents"] + extra_baggage=self._trace_context_extra_baggage ).items(): request.headers[key] = value @@ -176,6 +180,7 @@ def __init__( byo_connection_id: Optional[str] = None, retryer: Optional[Retrying] = None, aretryer: Optional[AsyncRetrying] = None, + trace_context_extra_baggage: Optional[list[str]] = None, **kwargs: Any, ): org_id = org_id or os.getenv("UIPATH_ORGANIZATION_ID") @@ -212,14 +217,20 @@ def __init__( http_options = genai_types.HttpOptions( httpx_client=httpx.Client( transport=_UrlRewriteTransport( - uipath_url, verify=verify, header_capture=header_capture + uipath_url, + verify=verify, + header_capture=header_capture, + trace_context_extra_baggage=trace_context_extra_baggage, ), headers=merged_headers, **client_kwargs, ), httpx_async_client=httpx.AsyncClient( transport=_AsyncUrlRewriteTransport( - uipath_url, verify=verify, header_capture=header_capture + uipath_url, + verify=verify, + header_capture=header_capture, + trace_context_extra_baggage=trace_context_extra_baggage, ), headers=merged_headers, **client_kwargs, diff --git a/src/uipath_langchain/chat/chat_model_factory.py b/src/uipath_langchain/chat/chat_model_factory.py index c22e16f76..0334f9bd9 100644 --- a/src/uipath_langchain/chat/chat_model_factory.py +++ b/src/uipath_langchain/chat/chat_model_factory.py @@ -39,9 +39,15 @@ class _TraceContextHeadersCallback(BaseCallbackHandler): run_inline: bool = True + def __init__(self, trace_context_extra_baggage: list[str] | None = None) -> None: + super().__init__() + self._trace_context_extra_baggage = trace_context_extra_baggage + def _merge_headers(self) -> None: existing = get_dynamic_request_headers() - existing.update(build_trace_context_headers(extra_baggage=["source=agents"])) + existing.update( + build_trace_context_headers(extra_baggage=self._trace_context_extra_baggage) + ) set_dynamic_request_headers(existing) def on_chat_model_start( @@ -78,6 +84,7 @@ def get_chat_model( callbacks: Callbacks = _UNSET, agenthub_config: str | None = None, use_new_llm_clients: bool = True, + trace_context_extra_baggage: list[str] | None = None, **kwargs: Any, ) -> BaseChatModel: """Create and configure a chat model, dispatching legacy vs new clients. @@ -109,6 +116,9 @@ def get_chat_model( use_new_llm_clients: Routes to the new ``uipath_langchain_client`` factory when True (default). When False, routes to the legacy in-repo clients. + trace_context_extra_baggage: Optional caller-owned trace baggage + additions, for example ``source=``. This shared package does + not default product source attribution. **kwargs: Forwarded to the underlying factory. The legacy factory accepts ``disable_streaming``; the new factory forwards extras as model kwargs to the LangChain constructor. @@ -120,7 +130,10 @@ def get_chat_model( # callback. For the new path the UiPathHttpxClient reads the ContextVar # set by the callback; for the legacy path the callback is a no-op but # keeps the wiring consistent. - callbacks = _ensure_trace_context_callback(callbacks) + callbacks = _ensure_trace_context_callback( + callbacks, + trace_context_extra_baggage=trace_context_extra_baggage, + ) if not use_new_llm_clients: return _legacy_chat_model( @@ -129,6 +142,7 @@ def get_chat_model( max_tokens=max_tokens, agenthub_config=agenthub_config, byo_connection_id=byo_connection_id, + trace_context_extra_baggage=trace_context_extra_baggage, **kwargs, ) @@ -158,14 +172,31 @@ def get_chat_model( ) -def _ensure_trace_context_callback(callbacks: Callbacks) -> list[BaseCallbackHandler]: - """Append a ``_TraceContextHeadersCallback`` if one is not already present.""" +def _ensure_trace_context_callback( + callbacks: Callbacks, + *, + trace_context_extra_baggage: list[str] | None = None, +) -> list[BaseCallbackHandler]: + """Append or update a ``_TraceContextHeadersCallback`` for trace headers.""" if callbacks is _UNSET or callbacks is None: cb_list: list[BaseCallbackHandler] = [] else: cb_list = list(callbacks) # type: ignore[arg-type] - if not any(isinstance(cb, _TraceContextHeadersCallback) for cb in cb_list): - cb_list.append(_TraceContextHeadersCallback()) + + existing_callback = next( + (cb for cb in cb_list if isinstance(cb, _TraceContextHeadersCallback)), + None, + ) + if existing_callback is not None: + if trace_context_extra_baggage is not None: + existing_callback._trace_context_extra_baggage = trace_context_extra_baggage + return cb_list + + cb_list.append( + _TraceContextHeadersCallback( + trace_context_extra_baggage=trace_context_extra_baggage + ) + ) return cb_list @@ -176,6 +207,7 @@ def _legacy_chat_model( max_tokens: int | None, agenthub_config: str | None, byo_connection_id: str | None, + trace_context_extra_baggage: list[str] | None, **kwargs: Any, ) -> BaseChatModel: if agenthub_config is None: @@ -191,5 +223,6 @@ def _legacy_chat_model( max_tokens, agenthub_config, byo_connection_id, + trace_context_extra_baggage=trace_context_extra_baggage, **kwargs, ) diff --git a/tests/chat/test_trace_context_callback.py b/tests/chat/test_trace_context_callback.py index 5b82e0f2e..3e733cd89 100644 --- a/tests/chat/test_trace_context_callback.py +++ b/tests/chat/test_trace_context_callback.py @@ -68,7 +68,7 @@ def test_returns_headers_with_active_span(self) -> None: ctx.span_id = 0x1122334455667788 mock_span.get_span_context.return_value = ctx - cb = _TraceContextHeadersCallback() + cb = _TraceContextHeadersCallback(trace_context_extra_baggage=["source=agents"]) with ( patch( "uipath.core.tracing.span_utils.UiPathSpanUtils" @@ -87,6 +87,29 @@ def test_returns_headers_with_active_span(self) -> None: assert "x-uipath-tracebaggage" in headers assert "source=agents" in headers["x-uipath-tracebaggage"] + def test_uses_no_extra_baggage_by_default(self) -> None: + cb = _TraceContextHeadersCallback() + with patch( + "uipath_langchain.chat.chat_model_factory.build_trace_context_headers", + return_value={}, + ) as mock_build_trace_context_headers: + cb._merge_headers() + + mock_build_trace_context_headers.assert_called_once_with(extra_baggage=None) + + def test_forwards_extra_baggage(self) -> None: + extra_baggage = ["source=agents", "executionType=1", "jobKey=job-123"] + cb = _TraceContextHeadersCallback(trace_context_extra_baggage=extra_baggage) + with patch( + "uipath_langchain.chat.chat_model_factory.build_trace_context_headers", + return_value={}, + ) as mock_build_trace_context_headers: + cb._merge_headers() + + mock_build_trace_context_headers.assert_called_once_with( + extra_baggage=extra_baggage + ) + def test_returns_empty_when_flag_disabled(self) -> None: FeatureFlags.configure_flags({"EnableTraceContextHeaders": False}) cb = _TraceContextHeadersCallback() @@ -131,12 +154,36 @@ def test_does_not_duplicate(self) -> None: count = sum(1 for cb in result if isinstance(cb, _TraceContextHeadersCallback)) assert count == 1 + def test_updates_existing_callback_with_extra_baggage(self) -> None: + callback = _TraceContextHeadersCallback() + extra_baggage = ["source=agents", "executionType=1"] + + result = _ensure_trace_context_callback( + [callback], + trace_context_extra_baggage=extra_baggage, + ) + + assert result == [callback] + assert callback._trace_context_extra_baggage == extra_baggage + def test_preserves_existing_callbacks(self) -> None: sentinel = MagicMock() result = _ensure_trace_context_callback([sentinel]) assert sentinel in result assert any(isinstance(cb, _TraceContextHeadersCallback) for cb in result) + def test_adds_callback_with_extra_baggage(self) -> None: + extra_baggage = ["source=agents", "executionType=1"] + result = _ensure_trace_context_callback( + None, + trace_context_extra_baggage=extra_baggage, + ) + + callback = next( + cb for cb in result if isinstance(cb, _TraceContextHeadersCallback) + ) + assert callback._trace_context_extra_baggage == extra_baggage + # --------------------------------------------------------------------------- # Legacy OpenAI transport — per-request trace context injection @@ -148,24 +195,36 @@ class TestOpenAITransportTraceContextHeaders: def test_inject_trace_context_headers_adds_headers(self) -> None: request = httpx.Request("POST", "https://example.com/completions") - with _patch_build_trace_context_headers( + extra_baggage = ["source=agents", "executionType=1", "jobKey=job-123"] + with patch( "uipath_langchain.chat._legacy.openai.build_trace_context_headers", - ): - _inject_trace_context_headers(request) + return_value=dict(_MOCK_TRACE_HEADERS), + ) as mock_build_trace_context_headers: + _inject_trace_context_headers( + request, + trace_context_extra_baggage=extra_baggage, + ) for key, value in _MOCK_TRACE_HEADERS.items(): assert request.headers[key] == value + mock_build_trace_context_headers.assert_called_once_with( + extra_baggage=extra_baggage + ) def test_sync_transport_calls_inject_trace_context(self) -> None: from uipath_langchain.chat._legacy.openai import UiPathSyncURLRewriteTransport - transport = UiPathSyncURLRewriteTransport() + extra_baggage = ["source=agents", "executionType=1"] + transport = UiPathSyncURLRewriteTransport( + trace_context_extra_baggage=extra_baggage + ) request = httpx.Request("POST", "https://example.com/completions") with ( - _patch_build_trace_context_headers( + patch( "uipath_langchain.chat._legacy.openai.build_trace_context_headers", - ), + return_value=dict(_MOCK_TRACE_HEADERS), + ) as mock_build_trace_context_headers, patch.object( httpx.HTTPTransport, "handle_request", @@ -176,17 +235,22 @@ def test_sync_transport_calls_inject_trace_context(self) -> None: for key, value in _MOCK_TRACE_HEADERS.items(): assert request.headers[key] == value + mock_build_trace_context_headers.assert_called_once_with( + extra_baggage=extra_baggage + ) async def test_async_transport_calls_inject_trace_context(self) -> None: from uipath_langchain.chat._legacy.openai import UiPathURLRewriteTransport - transport = UiPathURLRewriteTransport() + extra_baggage = ["source=agents", "executionType=1"] + transport = UiPathURLRewriteTransport(trace_context_extra_baggage=extra_baggage) request = httpx.Request("POST", "https://example.com/completions") with ( - _patch_build_trace_context_headers( + patch( "uipath_langchain.chat._legacy.openai.build_trace_context_headers", - ), + return_value=dict(_MOCK_TRACE_HEADERS), + ) as mock_build_trace_context_headers, patch.object( httpx.AsyncHTTPTransport, "handle_async_request", @@ -197,6 +261,9 @@ async def test_async_transport_calls_inject_trace_context(self) -> None: for key, value in _MOCK_TRACE_HEADERS.items(): assert request.headers[key] == value + mock_build_trace_context_headers.assert_called_once_with( + extra_baggage=extra_baggage + ) def test_inject_trace_context_headers_noop_when_empty(self) -> None: """When build_trace_context_headers returns {}, no headers are added.""" @@ -204,10 +271,11 @@ def test_inject_trace_context_headers_noop_when_empty(self) -> None: with patch( "uipath_langchain.chat._legacy.openai.build_trace_context_headers", return_value={}, - ): + ) as mock_build_trace_context_headers: _inject_trace_context_headers(request) assert "x-uipath-traceparent-id" not in request.headers + mock_build_trace_context_headers.assert_called_once_with(extra_baggage=None) # --------------------------------------------------------------------------- @@ -222,8 +290,10 @@ def test_sync_transport_injects_trace_headers(self) -> None: pytest.importorskip("google.genai") from uipath_langchain.chat._legacy.vertex import _UrlRewriteTransport + extra_baggage = ["source=agents", "executionType=1"] transport = _UrlRewriteTransport( - gateway_url="https://gateway.example.com/completions" + gateway_url="https://gateway.example.com/completions", + trace_context_extra_baggage=extra_baggage, ) request = httpx.Request( "POST", @@ -231,9 +301,10 @@ def test_sync_transport_injects_trace_headers(self) -> None: ) with ( - _patch_build_trace_context_headers( + patch( "uipath_langchain.chat._legacy.vertex.build_trace_context_headers", - ), + return_value=dict(_MOCK_TRACE_HEADERS), + ) as mock_build_trace_context_headers, patch.object( httpx.HTTPTransport, "handle_request", return_value=httpx.Response(200) ), @@ -242,13 +313,18 @@ def test_sync_transport_injects_trace_headers(self) -> None: for key, value in _MOCK_TRACE_HEADERS.items(): assert request.headers[key] == value + mock_build_trace_context_headers.assert_called_once_with( + extra_baggage=extra_baggage + ) async def test_async_transport_injects_trace_headers(self) -> None: pytest.importorskip("google.genai") from uipath_langchain.chat._legacy.vertex import _AsyncUrlRewriteTransport + extra_baggage = ["source=agents", "executionType=1"] transport = _AsyncUrlRewriteTransport( - gateway_url="https://gateway.example.com/completions" + gateway_url="https://gateway.example.com/completions", + trace_context_extra_baggage=extra_baggage, ) request = httpx.Request( "POST", @@ -256,9 +332,10 @@ async def test_async_transport_injects_trace_headers(self) -> None: ) with ( - _patch_build_trace_context_headers( + patch( "uipath_langchain.chat._legacy.vertex.build_trace_context_headers", - ), + return_value=dict(_MOCK_TRACE_HEADERS), + ) as mock_build_trace_context_headers, patch.object( httpx.AsyncHTTPTransport, "handle_async_request", @@ -269,6 +346,9 @@ async def test_async_transport_injects_trace_headers(self) -> None: for key, value in _MOCK_TRACE_HEADERS.items(): assert request.headers[key] == value + mock_build_trace_context_headers.assert_called_once_with( + extra_baggage=extra_baggage + ) # --------------------------------------------------------------------------- @@ -285,10 +365,12 @@ def test_modify_request_injects_trace_headers(self) -> None: AwsBedrockCompletionsPassthroughClient, ) + extra_baggage = ["source=agents", "executionType=1", "jobKey=job-123"] passthrough = AwsBedrockCompletionsPassthroughClient( model="anthropic.claude-haiku-4-5-20251001", token="test-token", api_flavor="converse", + trace_context_extra_baggage=extra_baggage, ) class _Req: @@ -297,9 +379,10 @@ class _Req: request = _Req() with ( - _patch_build_trace_context_headers( + patch( "uipath_langchain.chat._legacy.bedrock.build_trace_context_headers", - ), + return_value=dict(_MOCK_TRACE_HEADERS), + ) as mock_build_trace_context_headers, patch.object( passthrough, "_resolve_url", return_value=("https://gateway/x", False) ), @@ -308,6 +391,9 @@ class _Req: for key, value in _MOCK_TRACE_HEADERS.items(): assert request.headers[key] == value + mock_build_trace_context_headers.assert_called_once_with( + extra_baggage=extra_baggage + ) # --------------------------------------------------------------------------- diff --git a/uv.lock b/uv.lock index 208a31def..6948ba535 100644 --- a/uv.lock +++ b/uv.lock @@ -9,7 +9,7 @@ resolution-markers = [ ] [options] -exclude-newer = "0001-01-01T00:00:00Z" # This has no effect and is included for backwards compatibility when using relative exclude-newer values. +exclude-newer = "2026-06-13T19:20:23.757555Z" exclude-newer-span = "P2D" [options.exclude-newer-package] @@ -4439,7 +4439,7 @@ wheels = [ [[package]] name = "uipath-langchain" -version = "0.13.18" +version = "0.13.19" source = { editable = "." } dependencies = [ { name = "a2a-sdk" },