Capture LLMGW headers

.github/workflows/cd-langchain.yml

@@ -6,37 +6,10 @@ on:
       - main
     paths:
       - 'packages/uipath_langchain_client/src/uipath_langchain_client/__version__.py'
-  workflow_run:
-    workflows: ["cd"]
-    types:
-      - completed
   workflow_dispatch:
 
 jobs:
   build:
-    if: >
-      github.event_name == 'workflow_dispatch' ||
-
-      (
-        github.event_name == 'push' &&
-        !contains(
-          join(github.event.commits.*.modified, ' '),
-          'src/uipath_llm_client/__version__.py'
-        )
-      ) ||
-
-      (
-        github.event_name == 'workflow_run' &&
-        github.event.workflow_run.conclusion == 'success' &&
-        contains(
-          join(github.event.workflow_run.head_commit.modified, ' '),
-          'packages/uipath_langchain_client/src/uipath_langchain_client/__version__.py'
-        ) &&
-        contains(
-          join(github.event.workflow_run.head_commit.modified, ' '),
-          'src/uipath_llm_client/__version__.py'
-        )
-      )
     name: Build package
     runs-on: ubuntu-latest
     environment: pypi
@@ -46,10 +19,44 @@ jobs:
 
     steps:
       - uses: actions/checkout@v6
+        with:
+          fetch-depth: 2
 
-      - name: Wait for package indexing
-        if: github.event_name == 'workflow_run'
-        run: sleep 120
+      - name: Wait for uipath-llm-client publish
+        if: github.event_name == 'push'
+        env:
+          GH_TOKEN: ${{ github.token }}
+        run: |
+          # Check if the core package version was also modified in this push
+          if ! git diff HEAD~1 --name-only | grep -q '^src/uipath_llm_client/__version__.py$'; then
+            echo "Core package version unchanged — no need to wait."
+            exit 0
+          fi
+
+          echo "Core package version also changed — waiting for 'cd' workflow to succeed..."
+
+          # Poll the cd workflow for up to 15 minutes
+          for i in $(seq 1 30); do
+            STATUS=$(gh run list --workflow=cd.yml --branch=main --commit=${{ github.sha }} --json status,conclusion --jq '.[0] | .status + ":" + .conclusion')
+            echo "  Attempt $i/30 — cd workflow: $STATUS"
+
+            case "$STATUS" in
+              completed:success)
+                echo "cd workflow succeeded — waiting 120s for PyPI indexing..."
+                sleep 120
+                exit 0
+                ;;
+              completed:*)
+                echo "::error::cd workflow finished with: $STATUS"
+                exit 1
+                ;;
+            esac
+
+            sleep 30
+          done
+
+          echo "::error::Timed out waiting for cd workflow (15 min)"
+          exit 1
 
       - name: Setup uv
         uses: astral-sh/setup-uv@v7

CHANGELOG.md

@@ -2,6 +2,11 @@
 
 All notable changes to `uipath_llm_client` (core package) will be documented in this file.
 
+## [1.2.3] - 2026-02-25
+
+### Feature
+- Capture LLMGW headers
+
 ## [1.2.2] - 2026-02-23
 
 ### Fix

packages/uipath_langchain_client/CHANGELOG.md

@@ -2,6 +2,11 @@
 
 All notable changes to `uipath_langchain_client` will be documented in this file.
 
+## [1.2.3] - 2026-02-25
+
+### Feature
+- Capture headers and inject them in response_metadata.
+
 ## [1.2.2] - 2026-02-23
 
 ### Fix

packages/uipath_langchain_client/pyproject.toml

@@ -6,7 +6,7 @@ readme = "README.md"
 requires-python = ">=3.11"
 dependencies = [
     "langchain>=1.2.7",
-    "uipath-llm-client>=1.2.2",
+    "uipath-llm-client>=1.2.3",
 ]
 
 [project.optional-dependencies]

packages/uipath_langchain_client/src/uipath_langchain_client/__version__.py

@@ -1,3 +1,3 @@
 __title__ = "UiPath LangChain Client"
 __description__ = "A Python client for interacting with UiPath's LLM services via LangChain."
-__version__ = "1.2.2"
+__version__ = "1.2.3"

packages/uipath_langchain_client/src/uipath_langchain_client/base_client.py

@@ -25,21 +25,30 @@
 
 import logging
 from abc import ABC
-from collections.abc import AsyncIterator, Iterator, Mapping
+from collections.abc import AsyncIterator, Iterator, Mapping, Sequence
 from functools import cached_property
 from typing import Any, Literal
 
 from httpx import URL, Response
 from langchain_core.embeddings import Embeddings
 from langchain_core.language_models.chat_models import BaseChatModel
+from langchain_core.messages import BaseMessage
+from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
 from pydantic import AliasChoices, BaseModel, ConfigDict, Field
 
 from uipath_langchain_client.settings import (
     UiPathAPIConfig,
     UiPathBaseSettings,
     get_default_client_settings,
 )
-from uipath_llm_client.httpx_client import UiPathHttpxAsyncClient, UiPathHttpxClient
+from uipath_llm_client.httpx_client import (
+    UiPathHttpxAsyncClient,
+    UiPathHttpxClient,
+)
+from uipath_llm_client.utils.headers import (
+    get_captured_response_headers,
+    set_captured_response_headers,
+)
 from uipath_llm_client.utils.retry import RetryConfig
 
 
@@ -99,6 +108,13 @@ class UiPathBaseLLMClient(BaseModel, ABC):
         },
         description="Default request headers to include in requests",
     )
+    captured_headers: tuple[str, ...] = Field(
+        default=("x-uipath-",),
+        description="Case-insensitive response header prefixes to capture from LLM Gateway responses. "
+        "Captured headers appear in response_metadata under the 'uipath_llmgateway_headers' key. "
+        "Set to an empty tuple to disable.",
+    )
+
     request_timeout: float | None = Field(
         alias="timeout",
         validation_alias=AliasChoices("timeout", "request_timeout", "default_request_timeout"),
@@ -113,6 +129,7 @@ class UiPathBaseLLMClient(BaseModel, ABC):
         default=None,
         description="Retry configuration for failed requests",
     )
+
     logger: logging.Logger | None = Field(
         default=None,
         description="Logger for request/response logging",
@@ -135,6 +152,7 @@ def uipath_sync_client(self) -> UiPathHttpxClient:
                     model_name=self.model_name, api_config=self.api_config
                 ),
             },
+            captured_headers=self.captured_headers,
             timeout=self.request_timeout,
             max_retries=self.max_retries,
             retry_config=self.retry_config,
@@ -158,6 +176,7 @@ def uipath_async_client(self) -> UiPathHttpxAsyncClient:
                     model_name=self.model_name, api_config=self.api_config
                 ),
             },
+            captured_headers=self.captured_headers,
             timeout=self.request_timeout,
             max_retries=self.max_retries,
             retry_config=self.retry_config,
@@ -283,7 +302,87 @@ async def uipath_astream(
 
 
 class UiPathBaseChatModel(UiPathBaseLLMClient, BaseChatModel):
-    pass
+    """Base chat model that captures LLM Gateway response headers into response_metadata.
+
+    Wraps _generate/_agenerate/_stream/_astream to automatically read captured headers
+    from the ContextVar (populated by the httpx client's send()) and inject them into
+    the AIMessage's response_metadata under the 'uipath_llmgateway_headers' key.
+
+    Passthrough clients that delegate to vendor SDKs should inherit from this class
+    so that headers are captured transparently.
+    """
+
+    def _generate(
+        self,
+        messages: list[BaseMessage],
+        *args: Any,
+        **kwargs: Any,
+    ) -> ChatResult:
+        set_captured_response_headers({})
+        try:
+            result = super()._generate(messages, *args, **kwargs)
+            self._inject_gateway_headers(result.generations)
+            return result
+        finally:
+            set_captured_response_headers({})
+
+    async def _agenerate(
+        self,
+        messages: list[BaseMessage],
+        *args: Any,
+        **kwargs: Any,
+    ) -> ChatResult:
+        set_captured_response_headers({})
+        try:
+            result = await super()._agenerate(messages, *args, **kwargs)
+            self._inject_gateway_headers(result.generations)
+            return result
+        finally:
+            set_captured_response_headers({})
+
+    def _stream(
+        self,
+        messages: list[BaseMessage],
+        *args: Any,
+        **kwargs: Any,
+    ) -> Iterator[ChatGenerationChunk]:
+        set_captured_response_headers({})
+        try:
+            first = True
+            for chunk in super()._stream(messages, *args, **kwargs):
+                if first:
+                    self._inject_gateway_headers([chunk])
+                    first = False
+                yield chunk
+        finally:
+            set_captured_response_headers({})
+
+    async def _astream(
+        self,
+        messages: list[BaseMessage],
+        *args: Any,
+        **kwargs: Any,
+    ) -> AsyncIterator[ChatGenerationChunk]:
+        set_captured_response_headers({})
+        try:
+            first = True
+            async for chunk in super()._astream(messages, *args, **kwargs):
+                if first:
+                    self._inject_gateway_headers([chunk])
+                    first = False
+                yield chunk
+        finally:
+            set_captured_response_headers({})
+
+    def _inject_gateway_headers(self, generations: Sequence[ChatGeneration]) -> None:
+        """Inject captured gateway headers into each generation's response_metadata."""
+        if not self.captured_headers:
+            return
+        headers = get_captured_response_headers()
+        if not headers:
+            return
+        for generation in generations:
+            generation.message.response_metadata["uipath_llmgateway_headers"] = headers
 
 
 class UiPathBaseEmbeddings(UiPathBaseLLMClient, Embeddings):

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

@@ -4,7 +4,7 @@
 from pydantic import Field, model_validator
 from typing_extensions import override
 
-from uipath_langchain_client.base_client import UiPathBaseLLMClient
+from uipath_langchain_client.base_client import UiPathBaseChatModel
 from uipath_langchain_client.settings import UiPathAPIConfig
 
 try:
@@ -26,7 +26,7 @@
     ) from e
 
 
-class UiPathChatAnthropic(UiPathBaseLLMClient, ChatAnthropic):
+class UiPathChatAnthropic(UiPathBaseChatModel, ChatAnthropic):
     api_config: UiPathAPIConfig = UiPathAPIConfig(
         api_type="completions",
         client_type="passthrough",
@@ -65,7 +65,7 @@ def _anthropic_client(
                     api_key="PLACEHOLDER",
                     base_url=str(self.uipath_sync_client.base_url),
                     default_headers=dict(self.uipath_sync_client.headers),
-                    max_retries=0,  # handled by the UiPathBaseLLMClient
+                    max_retries=0,  # handled by the UiPathBaseChatModel
                     http_client=self.uipath_sync_client,
                 )
             case "vertexai":
@@ -75,7 +75,7 @@ def _anthropic_client(
                     access_token="PLACEHOLDER",
                     base_url=str(self.uipath_sync_client.base_url),
                     default_headers=dict(self.uipath_sync_client.headers),
-                    max_retries=0,  # handled by the UiPathBaseLLMClient
+                    max_retries=0,  # handled by the UiPathBaseChatModel
                     http_client=self.uipath_sync_client,
                 )
             case "awsbedrock":
@@ -85,15 +85,15 @@ def _anthropic_client(
                     aws_region="PLACEHOLDER",
                     base_url=str(self.uipath_sync_client.base_url),
                     default_headers=dict(self.uipath_sync_client.headers),
-                    max_retries=0,  # handled by the UiPathBaseLLMClient
+                    max_retries=0,  # handled by the UiPathBaseChatModel
                     http_client=self.uipath_sync_client,
                 )
             case "anthropic":
                 return Anthropic(
                     api_key="PLACEHOLDER",
                     base_url=str(self.uipath_sync_client.base_url),
                     default_headers=dict(self.uipath_sync_client.headers),
-                    max_retries=0,  # handled by the UiPathBaseLLMClient
+                    max_retries=0,  # handled by the UiPathBaseChatModel
                     http_client=self.uipath_sync_client,
                 )
 
@@ -107,7 +107,7 @@ def _async_anthropic_client(
                     api_key="PLACEHOLDER",
                     base_url=str(self.uipath_async_client.base_url),
                     default_headers=dict(self.uipath_async_client.headers),
-                    max_retries=0,  # handled by the UiPathBaseLLMClient
+                    max_retries=0,  # handled by the UiPathBaseChatModel
                     http_client=self.uipath_async_client,
                 )
             case "vertexai":
@@ -117,7 +117,7 @@ def _async_anthropic_client(
                     access_token="PLACEHOLDER",
                     base_url=str(self.uipath_async_client.base_url),
                     default_headers=dict(self.uipath_async_client.headers),
-                    max_retries=0,  # handled by the UiPathBaseLLMClient
+                    max_retries=0,  # handled by the UiPathBaseChatModel
                     http_client=self.uipath_async_client,
                 )
             case "awsbedrock":
@@ -127,15 +127,15 @@ def _async_anthropic_client(
                     aws_region="PLACEHOLDER",
                     base_url=str(self.uipath_async_client.base_url),
                     default_headers=dict(self.uipath_async_client.headers),
-                    max_retries=0,  # handled by the UiPathBaseLLMClient
+                    max_retries=0,  # handled by the UiPathBaseChatModel
                     http_client=self.uipath_async_client,
                 )
             case _:
                 return AsyncAnthropic(
                     api_key="PLACEHOLDER",
                     base_url=str(self.uipath_async_client.base_url),
                     default_headers=dict(self.uipath_async_client.headers),
-                    max_retries=0,  # handled by the UiPathBaseLLMClient
+                    max_retries=0,  # handled by the UiPathBaseChatModel
                     http_client=self.uipath_async_client,
                 )
 

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

@@ -2,7 +2,7 @@
 
 from pydantic import model_validator
 
-from uipath_langchain_client.base_client import UiPathBaseLLMClient
+from uipath_langchain_client.base_client import UiPathBaseChatModel
 from uipath_langchain_client.settings import UiPathAPIConfig
 
 try:
@@ -17,7 +17,7 @@
     ) from e
 
 
-class UiPathAzureAIChatCompletionsModel(UiPathBaseLLMClient, AzureAIChatCompletionsModel):  # type: ignore[override]
+class UiPathAzureAIChatCompletionsModel(UiPathBaseChatModel, AzureAIChatCompletionsModel):  # type: ignore[override]
     api_config: UiPathAPIConfig = UiPathAPIConfig(
         api_type="completions",
         client_type="passthrough",