Skip to content

Commit cd384f4

Browse files
cosminachocosminacho
authored
Capture LLMGW headers (#33)
1 parent a5ffd58 commit cd384f4

25 files changed

Lines changed: 784 additions & 135 deletions

File tree

‎.github/workflows/cd-langchain.yml‎

Lines changed: 37 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -6,37 +6,10 @@ on:
66
- main
77
paths:
88
- 'packages/uipath_langchain_client/src/uipath_langchain_client/__version__.py'
9-
workflow_run:
10-
workflows: ["cd"]
11-
types:
12-
- completed
139
workflow_dispatch:
1410

1511
jobs:
1612
build:
17-
if: >
18-
github.event_name == 'workflow_dispatch' ||
19-
20-
(
21-
github.event_name == 'push' &&
22-
!contains(
23-
join(github.event.commits.*.modified, ' '),
24-
'src/uipath_llm_client/__version__.py'
25-
)
26-
) ||
27-
28-
(
29-
github.event_name == 'workflow_run' &&
30-
github.event.workflow_run.conclusion == 'success' &&
31-
contains(
32-
join(github.event.workflow_run.head_commit.modified, ' '),
33-
'packages/uipath_langchain_client/src/uipath_langchain_client/__version__.py'
34-
) &&
35-
contains(
36-
join(github.event.workflow_run.head_commit.modified, ' '),
37-
'src/uipath_llm_client/__version__.py'
38-
)
39-
)
4013
name: Build package
4114
runs-on: ubuntu-latest
4215
environment: pypi
@@ -46,10 +19,44 @@ jobs:
4619

4720
steps:
4821
- uses: actions/checkout@v6
22+
with:
23+
fetch-depth: 2
4924

50-
- name: Wait for package indexing
51-
if: github.event_name == 'workflow_run'
52-
run: sleep 120
25+
- name: Wait for uipath-llm-client publish
26+
if: github.event_name == 'push'
27+
env:
28+
GH_TOKEN: ${{ github.token }}
29+
run: |
30+
# Check if the core package version was also modified in this push
31+
if ! git diff HEAD~1 --name-only | grep -q '^src/uipath_llm_client/__version__.py$'; then
32+
echo "Core package version unchanged — no need to wait."
33+
exit 0
34+
fi
35+
36+
echo "Core package version also changed — waiting for 'cd' workflow to succeed..."
37+
38+
# Poll the cd workflow for up to 15 minutes
39+
for i in $(seq 1 30); do
40+
STATUS=$(gh run list --workflow=cd.yml --branch=main --commit=${{ github.sha }} --json status,conclusion --jq '.[0] | .status + ":" + .conclusion')
41+
echo " Attempt $i/30 — cd workflow: $STATUS"
42+
43+
case "$STATUS" in
44+
completed:success)
45+
echo "cd workflow succeeded — waiting 120s for PyPI indexing..."
46+
sleep 120
47+
exit 0
48+
;;
49+
completed:*)
50+
echo "::error::cd workflow finished with: $STATUS"
51+
exit 1
52+
;;
53+
esac
54+
55+
sleep 30
56+
done
57+
58+
echo "::error::Timed out waiting for cd workflow (15 min)"
59+
exit 1
5360
5461
- name: Setup uv
5562
uses: astral-sh/setup-uv@v7

‎CHANGELOG.md‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,11 @@
22

33
All notable changes to `uipath_llm_client` (core package) will be documented in this file.
44

5+
## [1.2.3] - 2026-02-25
6+
7+
### Feature
8+
- Capture LLMGW headers
9+
510
## [1.2.2] - 2026-02-23
611

712
### Fix

‎packages/uipath_langchain_client/CHANGELOG.md‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,11 @@
22

33
All notable changes to `uipath_langchain_client` will be documented in this file.
44

5+
## [1.2.3] - 2026-02-25
6+
7+
### Feature
8+
- Capture headers and inject them in response_metadata.
9+
510
## [1.2.2] - 2026-02-23
611

712
### Fix

‎packages/uipath_langchain_client/pyproject.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ readme = "README.md"
66
requires-python = ">=3.11"
77
dependencies = [
88
"langchain>=1.2.7",
9-
"uipath-llm-client>=1.2.2",
9+
"uipath-llm-client>=1.2.3",
1010
]
1111

1212
[project.optional-dependencies]

‎packages/uipath_langchain_client/src/uipath_langchain_client/__version__.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,3 @@
11
__title__ = "UiPath LangChain Client"
22
__description__ = "A Python client for interacting with UiPath's LLM services via LangChain."
3-
__version__ = "1.2.2"
3+
__version__ = "1.2.3"

‎packages/uipath_langchain_client/src/uipath_langchain_client/base_client.py‎

Lines changed: 102 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -25,21 +25,30 @@
2525

2626
import logging
2727
from abc import ABC
28-
from collections.abc import AsyncIterator, Iterator, Mapping
28+
from collections.abc import AsyncIterator, Iterator, Mapping, Sequence
2929
from functools import cached_property
3030
from typing import Any, Literal
3131

3232
from httpx import URL, Response
3333
from langchain_core.embeddings import Embeddings
3434
from langchain_core.language_models.chat_models import BaseChatModel
35+
from langchain_core.messages import BaseMessage
36+
from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
3537
from pydantic import AliasChoices, BaseModel, ConfigDict, Field
3638

3739
from uipath_langchain_client.settings import (
3840
UiPathAPIConfig,
3941
UiPathBaseSettings,
4042
get_default_client_settings,
4143
)
42-
from uipath_llm_client.httpx_client import UiPathHttpxAsyncClient, UiPathHttpxClient
44+
from uipath_llm_client.httpx_client import (
45+
UiPathHttpxAsyncClient,
46+
UiPathHttpxClient,
47+
)
48+
from uipath_llm_client.utils.headers import (
49+
get_captured_response_headers,
50+
set_captured_response_headers,
51+
)
4352
from uipath_llm_client.utils.retry import RetryConfig
4453

4554

@@ -99,6 +108,13 @@ class UiPathBaseLLMClient(BaseModel, ABC):
99108
},
100109
description="Default request headers to include in requests",
101110
)
111+
captured_headers: tuple[str, ...] = Field(
112+
default=("x-uipath-",),
113+
description="Case-insensitive response header prefixes to capture from LLM Gateway responses. "
114+
"Captured headers appear in response_metadata under the 'uipath_llmgateway_headers' key. "
115+
"Set to an empty tuple to disable.",
116+
)
117+
102118
request_timeout: float | None = Field(
103119
alias="timeout",
104120
validation_alias=AliasChoices("timeout", "request_timeout", "default_request_timeout"),
@@ -113,6 +129,7 @@ class UiPathBaseLLMClient(BaseModel, ABC):
113129
default=None,
114130
description="Retry configuration for failed requests",
115131
)
132+
116133
logger: logging.Logger | None = Field(
117134
default=None,
118135
description="Logger for request/response logging",
@@ -135,6 +152,7 @@ def uipath_sync_client(self) -> UiPathHttpxClient:
135152
model_name=self.model_name, api_config=self.api_config
136153
),
137154
},
155+
captured_headers=self.captured_headers,
138156
timeout=self.request_timeout,
139157
max_retries=self.max_retries,
140158
retry_config=self.retry_config,
@@ -158,6 +176,7 @@ def uipath_async_client(self) -> UiPathHttpxAsyncClient:
158176
model_name=self.model_name, api_config=self.api_config
159177
),
160178
},
179+
captured_headers=self.captured_headers,
161180
timeout=self.request_timeout,
162181
max_retries=self.max_retries,
163182
retry_config=self.retry_config,
@@ -283,7 +302,87 @@ async def uipath_astream(
283302

284303

285304
class UiPathBaseChatModel(UiPathBaseLLMClient, BaseChatModel):
286-
pass
305+
"""Base chat model that captures LLM Gateway response headers into response_metadata.
306+
307+
Wraps _generate/_agenerate/_stream/_astream to automatically read captured headers
308+
from the ContextVar (populated by the httpx client's send()) and inject them into
309+
the AIMessage's response_metadata under the 'uipath_llmgateway_headers' key.
310+
311+
Passthrough clients that delegate to vendor SDKs should inherit from this class
312+
so that headers are captured transparently.
313+
"""
314+
315+
def _generate(
316+
self,
317+
messages: list[BaseMessage],
318+
*args: Any,
319+
**kwargs: Any,
320+
) -> ChatResult:
321+
set_captured_response_headers({})
322+
try:
323+
result = super()._generate(messages, *args, **kwargs)
324+
self._inject_gateway_headers(result.generations)
325+
return result
326+
finally:
327+
set_captured_response_headers({})
328+
329+
async def _agenerate(
330+
self,
331+
messages: list[BaseMessage],
332+
*args: Any,
333+
**kwargs: Any,
334+
) -> ChatResult:
335+
set_captured_response_headers({})
336+
try:
337+
result = await super()._agenerate(messages, *args, **kwargs)
338+
self._inject_gateway_headers(result.generations)
339+
return result
340+
finally:
341+
set_captured_response_headers({})
342+
343+
def _stream(
344+
self,
345+
messages: list[BaseMessage],
346+
*args: Any,
347+
**kwargs: Any,
348+
) -> Iterator[ChatGenerationChunk]:
349+
set_captured_response_headers({})
350+
try:
351+
first = True
352+
for chunk in super()._stream(messages, *args, **kwargs):
353+
if first:
354+
self._inject_gateway_headers([chunk])
355+
first = False
356+
yield chunk
357+
finally:
358+
set_captured_response_headers({})
359+
360+
async def _astream(
361+
self,
362+
messages: list[BaseMessage],
363+
*args: Any,
364+
**kwargs: Any,
365+
) -> AsyncIterator[ChatGenerationChunk]:
366+
set_captured_response_headers({})
367+
try:
368+
first = True
369+
async for chunk in super()._astream(messages, *args, **kwargs):
370+
if first:
371+
self._inject_gateway_headers([chunk])
372+
first = False
373+
yield chunk
374+
finally:
375+
set_captured_response_headers({})
376+
377+
def _inject_gateway_headers(self, generations: Sequence[ChatGeneration]) -> None:
378+
"""Inject captured gateway headers into each generation's response_metadata."""
379+
if not self.captured_headers:
380+
return
381+
headers = get_captured_response_headers()
382+
if not headers:
383+
return
384+
for generation in generations:
385+
generation.message.response_metadata["uipath_llmgateway_headers"] = headers
287386

288387

289388
class UiPathBaseEmbeddings(UiPathBaseLLMClient, Embeddings):

‎packages/uipath_langchain_client/src/uipath_langchain_client/clients/anthropic/chat_models.py‎

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
from pydantic import Field, model_validator
55
from typing_extensions import override
66

7-
from uipath_langchain_client.base_client import UiPathBaseLLMClient
7+
from uipath_langchain_client.base_client import UiPathBaseChatModel
88
from uipath_langchain_client.settings import UiPathAPIConfig
99

1010
try:
@@ -26,7 +26,7 @@
2626
) from e
2727

2828

29-
class UiPathChatAnthropic(UiPathBaseLLMClient, ChatAnthropic):
29+
class UiPathChatAnthropic(UiPathBaseChatModel, ChatAnthropic):
3030
api_config: UiPathAPIConfig = UiPathAPIConfig(
3131
api_type="completions",
3232
client_type="passthrough",
@@ -65,7 +65,7 @@ def _anthropic_client(
6565
api_key="PLACEHOLDER",
6666
base_url=str(self.uipath_sync_client.base_url),
6767
default_headers=dict(self.uipath_sync_client.headers),
68-
max_retries=0, # handled by the UiPathBaseLLMClient
68+
max_retries=0, # handled by the UiPathBaseChatModel
6969
http_client=self.uipath_sync_client,
7070
)
7171
case "vertexai":
@@ -75,7 +75,7 @@ def _anthropic_client(
7575
access_token="PLACEHOLDER",
7676
base_url=str(self.uipath_sync_client.base_url),
7777
default_headers=dict(self.uipath_sync_client.headers),
78-
max_retries=0, # handled by the UiPathBaseLLMClient
78+
max_retries=0, # handled by the UiPathBaseChatModel
7979
http_client=self.uipath_sync_client,
8080
)
8181
case "awsbedrock":
@@ -85,15 +85,15 @@ def _anthropic_client(
8585
aws_region="PLACEHOLDER",
8686
base_url=str(self.uipath_sync_client.base_url),
8787
default_headers=dict(self.uipath_sync_client.headers),
88-
max_retries=0, # handled by the UiPathBaseLLMClient
88+
max_retries=0, # handled by the UiPathBaseChatModel
8989
http_client=self.uipath_sync_client,
9090
)
9191
case "anthropic":
9292
return Anthropic(
9393
api_key="PLACEHOLDER",
9494
base_url=str(self.uipath_sync_client.base_url),
9595
default_headers=dict(self.uipath_sync_client.headers),
96-
max_retries=0, # handled by the UiPathBaseLLMClient
96+
max_retries=0, # handled by the UiPathBaseChatModel
9797
http_client=self.uipath_sync_client,
9898
)
9999

@@ -107,7 +107,7 @@ def _async_anthropic_client(
107107
api_key="PLACEHOLDER",
108108
base_url=str(self.uipath_async_client.base_url),
109109
default_headers=dict(self.uipath_async_client.headers),
110-
max_retries=0, # handled by the UiPathBaseLLMClient
110+
max_retries=0, # handled by the UiPathBaseChatModel
111111
http_client=self.uipath_async_client,
112112
)
113113
case "vertexai":
@@ -117,7 +117,7 @@ def _async_anthropic_client(
117117
access_token="PLACEHOLDER",
118118
base_url=str(self.uipath_async_client.base_url),
119119
default_headers=dict(self.uipath_async_client.headers),
120-
max_retries=0, # handled by the UiPathBaseLLMClient
120+
max_retries=0, # handled by the UiPathBaseChatModel
121121
http_client=self.uipath_async_client,
122122
)
123123
case "awsbedrock":
@@ -127,15 +127,15 @@ def _async_anthropic_client(
127127
aws_region="PLACEHOLDER",
128128
base_url=str(self.uipath_async_client.base_url),
129129
default_headers=dict(self.uipath_async_client.headers),
130-
max_retries=0, # handled by the UiPathBaseLLMClient
130+
max_retries=0, # handled by the UiPathBaseChatModel
131131
http_client=self.uipath_async_client,
132132
)
133133
case _:
134134
return AsyncAnthropic(
135135
api_key="PLACEHOLDER",
136136
base_url=str(self.uipath_async_client.base_url),
137137
default_headers=dict(self.uipath_async_client.headers),
138-
max_retries=0, # handled by the UiPathBaseLLMClient
138+
max_retries=0, # handled by the UiPathBaseChatModel
139139
http_client=self.uipath_async_client,
140140
)
141141

‎packages/uipath_langchain_client/src/uipath_langchain_client/clients/azure/chat_models.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22

33
from pydantic import model_validator
44

5-
from uipath_langchain_client.base_client import UiPathBaseLLMClient
5+
from uipath_langchain_client.base_client import UiPathBaseChatModel
66
from uipath_langchain_client.settings import UiPathAPIConfig
77

88
try:
@@ -17,7 +17,7 @@
1717
) from e
1818

1919

20-
class UiPathAzureAIChatCompletionsModel(UiPathBaseLLMClient, AzureAIChatCompletionsModel): # type: ignore[override]
20+
class UiPathAzureAIChatCompletionsModel(UiPathBaseChatModel, AzureAIChatCompletionsModel): # type: ignore[override]
2121
api_config: UiPathAPIConfig = UiPathAPIConfig(
2222
api_type="completions",
2323
client_type="passthrough",

0 commit comments

Comments
 (0)