Skip to content

Commit 4cea46b

Browse files
gcuipgcuip
authored
feat: send jobkey to ECS [ECS-1819] (#1687)
1 parent 3ff1078 commit 4cea46b

8 files changed

Lines changed: 206 additions & 12 deletions

File tree

‎packages/uipath-platform/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-platform"
3-
version = "0.1.58"
3+
version = "0.1.59"
44
description = "HTTP client library for programmatic access to UiPath Platform"
55
readme = { file = "README.md", content-type = "text/markdown" }
66
requires-python = ">=3.11"

‎packages/uipath-platform/src/uipath/platform/common/_job_context.py‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
from ._config import UiPathConfig
2+
from .constants import HEADER_JOB_KEY
3+
4+
5+
def header_job_key() -> dict[str, str]:
6+
"""Return the X-UiPath-JobKey header when the orchestrator job key is set.
7+
8+
Returns an empty dict when ``UiPathConfig.job_key`` is unset or empty.
9+
"""
10+
job_key = UiPathConfig.job_key
11+
if not job_key:
12+
return {}
13+
return {HEADER_JOB_KEY: job_key}

‎packages/uipath-platform/src/uipath/platform/context_grounding/_context_grounding_service.py‎

Lines changed: 28 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
from ..common._execution_context import UiPathExecutionContext
1313
from ..common._folder_context import FolderContext, header_folder
1414
from ..common._http_config import get_httpx_client_kwargs
15+
from ..common._job_context import header_job_key
1516
from ..common._models import Endpoint, RequestSpec
1617
from ..common.constants import (
1718
ORCHESTRATOR_STORAGE_BUCKET_DATA_SOURCE,
@@ -221,6 +222,7 @@ def retrieve_across_folders(
221222
spec.method,
222223
spec.endpoint,
223224
params=spec.params,
225+
headers=spec.headers,
224226
).json()
225227

226228
return [
@@ -250,6 +252,7 @@ async def retrieve_across_folders_async(
250252
spec.method,
251253
spec.endpoint,
252254
params=spec.params,
255+
headers=spec.headers,
253256
)
254257
).json()
255258

@@ -268,6 +271,7 @@ def _retrieve_system_indexes(
268271
spec.method,
269272
spec.endpoint,
270273
params=spec.params,
274+
headers=spec.headers,
271275
).json()
272276

273277
return [
@@ -286,6 +290,7 @@ async def _retrieve_system_indexes_async(
286290
spec.method,
287291
spec.endpoint,
288292
params=spec.params,
293+
headers=spec.headers,
289294
)
290295
).json()
291296

@@ -455,7 +460,7 @@ def list(
455460
"GET",
456461
Endpoint("/ecs_/v2/indexes"),
457462
params={"$expand": "dataSource"},
458-
headers={**header_folder(folder_key, None)},
463+
headers={**header_folder(folder_key, None), **header_job_key()},
459464
).json()
460465
return [
461466
ContextGroundingIndex.model_validate(item)
@@ -483,7 +488,7 @@ async def list_async(
483488
"GET",
484489
Endpoint("/ecs_/v2/indexes"),
485490
params={"$expand": "dataSource"},
486-
headers={**header_folder(folder_key, None)},
491+
headers={**header_folder(folder_key, None), **header_job_key()},
487492
)
488493
).json()
489494
return [
@@ -521,6 +526,7 @@ def retrieve_by_id(
521526
spec.method,
522527
spec.endpoint,
523528
params=spec.params,
529+
headers=spec.headers,
524530
).json()
525531

526532
@traced(name="contextgrounding_retrieve_by_id", run_type="uipath")
@@ -553,6 +559,7 @@ async def retrieve_by_id_async(
553559
spec.method,
554560
spec.endpoint,
555561
params=spec.params,
562+
headers=spec.headers,
556563
)
557564

558565
return response.json()
@@ -1997,6 +2004,7 @@ def _ingest_spec(
19972004
endpoint=Endpoint(f"/ecs_/v2/indexes/{key}/ingest"),
19982005
headers={
19992006
**header_folder(folder_key, None),
2007+
**header_job_key(),
20002008
},
20012009
)
20022010

@@ -2024,6 +2032,7 @@ def _retrieve_across_folders_spec(
20242032
method="GET",
20252033
endpoint=Endpoint("/ecs_/v2/indexes/allacrossfolders"),
20262034
params=params,
2035+
headers={**header_job_key()},
20272036
)
20282037

20292038
def _retrieve_system_indexes_spec(
@@ -2040,6 +2049,7 @@ def _retrieve_system_indexes_spec(
20402049
method="GET",
20412050
endpoint=Endpoint("/ecs_/v2/indexes/allsystemindexes"),
20422051
params=params,
2052+
headers={**header_job_key()},
20432053
)
20442054

20452055
def _list_spec(
@@ -2054,6 +2064,7 @@ def _list_spec(
20542064
},
20552065
headers={
20562066
**header_folder(folder_key, None),
2067+
**header_job_key(),
20572068
},
20582069
)
20592070

@@ -2074,6 +2085,7 @@ def _retrieve_spec(
20742085
},
20752086
headers={
20762087
**header_folder(folder_key, None),
2088+
**header_job_key(),
20772089
},
20782090
)
20792091

@@ -2126,6 +2138,7 @@ def _create_spec(
21262138
json=payload.model_dump(by_alias=True, exclude_none=True),
21272139
headers={
21282140
**header_folder(folder_key, None),
2141+
**header_job_key(),
21292142
},
21302143
)
21312144

@@ -2156,7 +2169,7 @@ def _create_ephemeral_spec(
21562169
method="POST",
21572170
endpoint=Endpoint("/ecs_/v2/indexes/createephemeral"),
21582171
json=payload.model_dump(by_alias=True, exclude_none=True),
2159-
headers={**header_folder(folder_key, None)},
2172+
headers={**header_folder(folder_key, None), **header_job_key()},
21602173
)
21612174

21622175
def _build_data_source(self, source: SourceConfig) -> Dict[str, Any]:
@@ -2256,6 +2269,7 @@ def _retrieve_by_id_spec(
22562269
endpoint=Endpoint(f"/ecs_/v2/indexes/{id}"),
22572270
headers={
22582271
**header_folder(folder_key, None),
2272+
**header_job_key(),
22592273
},
22602274
)
22612275

@@ -2272,6 +2286,7 @@ def _delete_by_id_spec(
22722286
endpoint=Endpoint(f"/ecs_/v2/indexes/{id}"),
22732287
headers={
22742288
**header_folder(folder_key, None),
2289+
**header_job_key(),
22752290
},
22762291
)
22772292

@@ -2299,6 +2314,7 @@ def _search_spec(
22992314
},
23002315
headers={
23012316
**header_folder(folder_key, None),
2317+
**header_job_key(),
23022318
},
23032319
)
23042320

@@ -2335,6 +2351,7 @@ def _unified_search_spec(
23352351
json=json_body,
23362352
headers={
23372353
**header_folder(folder_key, None),
2354+
**header_job_key(),
23382355
},
23392356
)
23402357

@@ -2364,6 +2381,7 @@ def _deep_rag_creation_spec(
23642381
},
23652382
headers={
23662383
**header_folder(folder_key, None),
2384+
**header_job_key(),
23672385
},
23682386
)
23692387

@@ -2387,7 +2405,7 @@ def _deep_rag_ephemeral_creation_spec(
23872405
params={
23882406
"$select": "id,lastDeepRagStatus,createdDate",
23892407
},
2390-
headers={},
2408+
headers={**header_job_key()},
23912409
)
23922410

23932411
def _batch_transform_creation_spec(
@@ -2437,6 +2455,7 @@ def _batch_transform_creation_spec(
24372455
},
24382456
headers={
24392457
**header_folder(folder_key, None),
2458+
**header_job_key(),
24402459
},
24412460
)
24422461

@@ -2463,7 +2482,7 @@ def _batch_transform_ephemeral_creation_spec(
24632482
column.model_dump(by_alias=True) for column in output_columns
24642483
],
24652484
},
2466-
headers={},
2485+
headers={**header_job_key()},
24672486
)
24682487

24692488
def _deep_rag_retrieve_spec(
@@ -2476,6 +2495,7 @@ def _deep_rag_retrieve_spec(
24762495
params={
24772496
"$expand": "content",
24782497
},
2498+
headers={**header_job_key()},
24792499
)
24802500

24812501
def _batch_transform_retrieve_spec(
@@ -2485,6 +2505,7 @@ def _batch_transform_retrieve_spec(
24852505
return RequestSpec(
24862506
method="GET",
24872507
endpoint=Endpoint(f"/ecs_/v2/batchRag/{id}"),
2508+
headers={**header_job_key()},
24882509
)
24892510

24902511
def _batch_transform_get_read_uri_spec(
@@ -2494,6 +2515,7 @@ def _batch_transform_get_read_uri_spec(
24942515
return RequestSpec(
24952516
method="GET",
24962517
endpoint=Endpoint(f"/ecs_/v2/batchRag/{id}/GetReadUri"),
2518+
headers={**header_job_key()},
24972519
)
24982520

24992521
def _batch_transform_download_blob_spec(
@@ -2503,6 +2525,7 @@ def _batch_transform_download_blob_spec(
25032525
return RequestSpec(
25042526
method="GET",
25052527
endpoint=Endpoint(f"/ecs_/v2/batchRag/{id}/DownloadBlob"),
2528+
headers={**header_job_key()},
25062529
)
25072530

25082531
def _resolve_folder_key(self, folder_key, folder_path):

‎packages/uipath-platform/tests/common/test_job_context.py‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
import pytest
2+
3+
from uipath.platform.common._job_context import header_job_key
4+
from uipath.platform.common.constants import ENV_JOB_KEY, HEADER_JOB_KEY
5+
6+
7+
def test_returns_header_when_env_var_set(monkeypatch: pytest.MonkeyPatch) -> None:
8+
monkeypatch.setenv(ENV_JOB_KEY, "test-job-key")
9+
10+
assert header_job_key() == {HEADER_JOB_KEY: "test-job-key"}
11+
12+
13+
def test_returns_empty_when_env_var_unset(monkeypatch: pytest.MonkeyPatch) -> None:
14+
monkeypatch.delenv(ENV_JOB_KEY, raising=False)
15+
16+
assert header_job_key() == {}
17+
18+
19+
def test_returns_empty_when_env_var_blank(monkeypatch: pytest.MonkeyPatch) -> None:
20+
monkeypatch.setenv(ENV_JOB_KEY, "")
21+
22+
assert header_job_key() == {}

0 commit comments

Comments
 (0)