Skip to content

Commit 523733d

Browse files
RunnanJiaRunnanJia
committed
feat: caller-owned trace_context_extra_baggage (v0.13.19)
1 parent b5bc6c5 commit 523733d

8 files changed

Lines changed: 226 additions & 44 deletions

File tree

‎pyproject.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
[project]
22
name = "uipath-langchain"
3-
version = "0.13.18"
3+
version = "0.13.19"
44
description = "Python SDK that enables developers to build and deploy LangGraph agents to the UiPath Cloud Platform"
55
readme = { file = "README.md", content-type = "text/markdown" }
66
requires-python = ">=3.11"

‎src/uipath_langchain/chat/_legacy/bedrock.py‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -72,12 +72,14 @@ def __init__(
7272
agenthub_config: Optional[str] = None,
7373
byo_connection_id: Optional[str] = None,
7474
header_capture: HeaderCapture | None = None,
75+
trace_context_extra_baggage: Optional[list[str]] = None,
7576
):
7677
self.model = model
7778
self.token = token
7879
self.api_flavor = api_flavor
7980
self.agenthub_config = agenthub_config
8081
self.byo_connection_id = byo_connection_id
82+
self._trace_context_extra_baggage = trace_context_extra_baggage
8183
self._vendor = "awsbedrock"
8284
self._url: Optional[str] = None
8385
self._is_override: bool = False
@@ -182,7 +184,11 @@ def _modify_request(self, request, **kwargs):
182184
)
183185
headers["X-UiPath-LlmGateway-ApiFlavor"] = self.api_flavor
184186
headers["X-UiPath-Streaming-Enabled"] = streaming
185-
headers.update(build_trace_context_headers(extra_baggage=["source=agents"]))
187+
headers.update(
188+
build_trace_context_headers(
189+
extra_baggage=self._trace_context_extra_baggage
190+
)
191+
)
186192

187193
request.headers.update(headers)
188194

@@ -204,6 +210,7 @@ def __init__(
204210
byo_connection_id: Optional[str] = None,
205211
retryer: Optional[Retrying] = None,
206212
aretryer: Optional[AsyncRetrying] = None,
213+
trace_context_extra_baggage: Optional[list[str]] = None,
207214
**kwargs,
208215
):
209216
org_id = org_id or os.getenv("UIPATH_ORGANIZATION_ID")
@@ -229,6 +236,7 @@ def __init__(
229236
api_flavor="converse",
230237
agenthub_config=agenthub_config,
231238
byo_connection_id=byo_connection_id,
239+
trace_context_extra_baggage=trace_context_extra_baggage,
232240
)
233241

234242
kwargs["client"] = passthrough_client.get_client()
@@ -266,6 +274,7 @@ def __init__(
266274
byo_connection_id: Optional[str] = None,
267275
retryer: Optional[Retrying] = None,
268276
aretryer: Optional[AsyncRetrying] = None,
277+
trace_context_extra_baggage: Optional[list[str]] = None,
269278
**kwargs,
270279
):
271280
org_id = org_id or os.getenv("UIPATH_ORGANIZATION_ID")
@@ -294,6 +303,7 @@ def __init__(
294303
agenthub_config=agenthub_config,
295304
byo_connection_id=byo_connection_id,
296305
header_capture=header_capture,
306+
trace_context_extra_baggage=trace_context_extra_baggage,
297307
)
298308

299309
kwargs["client"] = passthrough_client.get_client()

‎src/uipath_langchain/chat/_legacy/chat_model_factory.py‎

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@ def _create_openai_llm(
4343
max_tokens: int | None,
4444
agenthub_config: str,
4545
byo_connection_id: str | None = None,
46+
trace_context_extra_baggage: list[str] | None = None,
4647
**kwargs: Any,
4748
) -> BaseChatModel:
4849
"""Create UiPathChatOpenAI for OpenAI models via LLMGateway."""
@@ -63,6 +64,7 @@ def _create_openai_llm(
6364
api_version=azure_open_ai_latest_api_version,
6465
agenthub_config=agenthub_config,
6566
byo_connection_id=byo_connection_id,
67+
trace_context_extra_baggage=trace_context_extra_baggage,
6668
output_version="v1",
6769
**sampling_kwargs,
6870
**kwargs,
@@ -75,6 +77,7 @@ def _create_openai_llm(
7577
api_version=azure_open_ai_latest_api_version,
7678
agenthub_config=agenthub_config,
7779
byo_connection_id=byo_connection_id,
80+
trace_context_extra_baggage=trace_context_extra_baggage,
7881
output_version="v1",
7982
**sampling_kwargs,
8083
**kwargs,
@@ -90,6 +93,7 @@ def _create_bedrock_llm(
9093
max_tokens: int | None,
9194
agenthub_config: str,
9295
byo_connection_id: str | None = None,
96+
trace_context_extra_baggage: list[str] | None = None,
9397
**kwargs: Any,
9498
) -> BaseChatModel:
9599
"""Create UiPathChatBedrockConverse for Claude models via LLMGateway."""
@@ -109,6 +113,7 @@ def _create_bedrock_llm(
109113
max_tokens=max_tokens,
110114
agenthub_config=agenthub_config,
111115
byo_connection_id=byo_connection_id,
116+
trace_context_extra_baggage=trace_context_extra_baggage,
112117
output_version="v1",
113118
**sampling_kwargs,
114119
**kwargs,
@@ -119,6 +124,7 @@ def _create_bedrock_llm(
119124
max_tokens=max_tokens,
120125
agenthub_config=agenthub_config,
121126
byo_connection_id=byo_connection_id,
127+
trace_context_extra_baggage=trace_context_extra_baggage,
122128
output_version="v1",
123129
**sampling_kwargs,
124130
**kwargs,
@@ -134,6 +140,7 @@ def _create_vertex_llm(
134140
max_tokens: int | None,
135141
agenthub_config: str,
136142
byo_connection_id: str | None = None,
143+
trace_context_extra_baggage: list[str] | None = None,
137144
**kwargs: Any,
138145
) -> BaseChatModel:
139146
"""Create UiPathChatVertex for Gemini models via LLMGateway."""
@@ -150,6 +157,7 @@ def _create_vertex_llm(
150157
max_tokens=max_tokens,
151158
agenthub_config=agenthub_config,
152159
byo_connection_id=byo_connection_id,
160+
trace_context_extra_baggage=trace_context_extra_baggage,
153161
output_version="v1",
154162
**sampling_kwargs,
155163
**kwargs,
@@ -248,6 +256,7 @@ def get_chat_model(
248256
max_tokens: int | None,
249257
agenthub_config: str,
250258
byo_connection_id: str | None = None,
259+
trace_context_extra_baggage: list[str] | None = None,
251260
**kwargs: Any,
252261
) -> BaseChatModel:
253262
"""Create and configure LLM instance using LLMGateway API.
@@ -272,7 +281,8 @@ def get_chat_model(
272281
effective_temperature,
273282
max_tokens,
274283
agenthub_config,
275-
byo_connection_id,
284+
byo_connection_id=byo_connection_id,
285+
trace_context_extra_baggage=trace_context_extra_baggage,
276286
**kwargs,
277287
)
278288
case LLMProvider.BEDROCK:
@@ -282,7 +292,8 @@ def get_chat_model(
282292
effective_temperature,
283293
max_tokens,
284294
agenthub_config,
285-
byo_connection_id,
295+
byo_connection_id=byo_connection_id,
296+
trace_context_extra_baggage=trace_context_extra_baggage,
286297
**kwargs,
287298
)
288299
case LLMProvider.VERTEX:
@@ -292,6 +303,7 @@ def get_chat_model(
292303
effective_temperature,
293304
max_tokens,
294305
agenthub_config,
295-
byo_connection_id,
306+
byo_connection_id=byo_connection_id,
307+
trace_context_extra_baggage=trace_context_extra_baggage,
296308
**kwargs,
297309
)

‎src/uipath_langchain/chat/_legacy/openai.py‎

Lines changed: 30 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -54,38 +54,53 @@ def _inject_license_ref_id(request: httpx.Request) -> None:
5454
request.headers["X-UiPath-License-RefId"] = license_ref_id
5555

5656

57-
def _inject_trace_context_headers(request: httpx.Request) -> None:
57+
def _inject_trace_context_headers(
58+
request: httpx.Request,
59+
trace_context_extra_baggage: list[str] | None = None,
60+
) -> None:
5861
"""Inject trace context headers per-request from the active OTEL span."""
5962
for key, value in build_trace_context_headers(
60-
extra_baggage=["source=agents"]
63+
extra_baggage=trace_context_extra_baggage
6164
).items():
6265
request.headers[key] = value
6366

6467

6568
class UiPathURLRewriteTransport(httpx.AsyncHTTPTransport):
66-
def __init__(self, verify: bool = True, **kwargs):
69+
def __init__(
70+
self,
71+
verify: bool = True,
72+
trace_context_extra_baggage: list[str] | None = None,
73+
**kwargs,
74+
):
6775
super().__init__(verify=verify, **kwargs)
76+
self._trace_context_extra_baggage = trace_context_extra_baggage
6877

6978
async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
7079
new_url = _rewrite_openai_url(str(request.url), request.url.params)
7180
if new_url:
7281
request.url = new_url
7382
_inject_license_ref_id(request)
74-
_inject_trace_context_headers(request)
83+
_inject_trace_context_headers(request, self._trace_context_extra_baggage)
7584

7685
return await super().handle_async_request(request)
7786

7887

7988
class UiPathSyncURLRewriteTransport(httpx.HTTPTransport):
80-
def __init__(self, verify: bool = True, **kwargs):
89+
def __init__(
90+
self,
91+
verify: bool = True,
92+
trace_context_extra_baggage: list[str] | None = None,
93+
**kwargs,
94+
):
8195
super().__init__(verify=verify, **kwargs)
96+
self._trace_context_extra_baggage = trace_context_extra_baggage
8297

8398
def handle_request(self, request: httpx.Request) -> httpx.Response:
8499
new_url = _rewrite_openai_url(str(request.url), request.url.params)
85100
if new_url:
86101
request.url = new_url
87102
_inject_license_ref_id(request)
88-
_inject_trace_context_headers(request)
103+
_inject_trace_context_headers(request, self._trace_context_extra_baggage)
89104

90105
return super().handle_request(request)
91106

@@ -112,6 +127,7 @@ def __init__(
112127
agenthub_config: Optional[str] = None,
113128
extra_headers: Optional[dict[str, str]] = None,
114129
byo_connection_id: Optional[str] = None,
130+
trace_context_extra_baggage: Optional[list[str]] = None,
115131
**kwargs,
116132
):
117133
org_id = org_id or os.getenv("UIPATH_ORGANIZATION_ID")
@@ -155,11 +171,17 @@ def __init__(
155171
model_name=model_name,
156172
default_headers=self._build_headers(token, inject_routing=is_override),
157173
http_async_client=httpx.AsyncClient(
158-
transport=UiPathURLRewriteTransport(verify=verify),
174+
transport=UiPathURLRewriteTransport(
175+
verify=verify,
176+
trace_context_extra_baggage=trace_context_extra_baggage,
177+
),
159178
**client_kwargs,
160179
),
161180
http_client=httpx.Client(
162-
transport=UiPathSyncURLRewriteTransport(verify=verify),
181+
transport=UiPathSyncURLRewriteTransport(
182+
verify=verify,
183+
trace_context_extra_baggage=trace_context_extra_baggage,
184+
),
163185
**client_kwargs,
164186
),
165187
api_key=token,

‎src/uipath_langchain/chat/_legacy/vertex.py‎

Lines changed: 15 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -79,10 +79,12 @@ def __init__(
7979
gateway_url: str,
8080
verify: bool = True,
8181
header_capture: HeaderCapture | None = None,
82+
trace_context_extra_baggage: list[str] | None = None,
8283
):
8384
super().__init__(verify=verify)
8485
self.gateway_url = gateway_url
8586
self.header_capture = header_capture
87+
self._trace_context_extra_baggage = trace_context_extra_baggage
8688

8789
def handle_request(self, request: httpx.Request) -> httpx.Response:
8890
original_url = str(request.url)
@@ -98,7 +100,7 @@ def handle_request(self, request: httpx.Request) -> httpx.Response:
98100
request.url = new_url
99101

100102
for key, value in build_trace_context_headers(
101-
extra_baggage=["source=agents"]
103+
extra_baggage=self._trace_context_extra_baggage
102104
).items():
103105
request.headers[key] = value
104106

@@ -117,10 +119,12 @@ def __init__(
117119
gateway_url: str,
118120
verify: bool = True,
119121
header_capture: HeaderCapture | None = None,
122+
trace_context_extra_baggage: list[str] | None = None,
120123
):
121124
super().__init__(verify=verify)
122125
self.gateway_url = gateway_url
123126
self.header_capture = header_capture
127+
self._trace_context_extra_baggage = trace_context_extra_baggage
124128

125129
async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
126130
original_url = str(request.url)
@@ -136,7 +140,7 @@ async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
136140
request.url = new_url
137141

138142
for key, value in build_trace_context_headers(
139-
extra_baggage=["source=agents"]
143+
extra_baggage=self._trace_context_extra_baggage
140144
).items():
141145
request.headers[key] = value
142146

@@ -176,6 +180,7 @@ def __init__(
176180
byo_connection_id: Optional[str] = None,
177181
retryer: Optional[Retrying] = None,
178182
aretryer: Optional[AsyncRetrying] = None,
183+
trace_context_extra_baggage: Optional[list[str]] = None,
179184
**kwargs: Any,
180185
):
181186
org_id = org_id or os.getenv("UIPATH_ORGANIZATION_ID")
@@ -212,14 +217,20 @@ def __init__(
212217
http_options = genai_types.HttpOptions(
213218
httpx_client=httpx.Client(
214219
transport=_UrlRewriteTransport(
215-
uipath_url, verify=verify, header_capture=header_capture
220+
uipath_url,
221+
verify=verify,
222+
header_capture=header_capture,
223+
trace_context_extra_baggage=trace_context_extra_baggage,
216224
),
217225
headers=merged_headers,
218226
**client_kwargs,
219227
),
220228
httpx_async_client=httpx.AsyncClient(
221229
transport=_AsyncUrlRewriteTransport(
222-
uipath_url, verify=verify, header_capture=header_capture
230+
uipath_url,
231+
verify=verify,
232+
header_capture=header_capture,
233+
trace_context_extra_baggage=trace_context_extra_baggage,
223234
),
224235
headers=merged_headers,
225236
**client_kwargs,

0 commit comments

Comments
 (0)