From 799eefe25ae48d3f767443fa610ecbcce108dc7a Mon Sep 17 00:00:00 2001 From: Andrei Ancuta Date: Tue, 3 Feb 2026 16:52:02 +0200 Subject: [PATCH 1/2] feat: add response headers in trace attributes --- src/uipath_langchain/chat/bedrock.py | 36 ++++++++- src/uipath_langchain/chat/constants.py | 9 +++ src/uipath_langchain/chat/header_capture.py | 41 ++++++++++ src/uipath_langchain/chat/openai.py | 1 + src/uipath_langchain/chat/vertex.py | 84 ++++++++++++++++++--- 5 files changed, 158 insertions(+), 13 deletions(-) create mode 100644 src/uipath_langchain/chat/constants.py create mode 100644 src/uipath_langchain/chat/header_capture.py diff --git a/src/uipath_langchain/chat/bedrock.py b/src/uipath_langchain/chat/bedrock.py index 108eebb0..a29631de 100644 --- a/src/uipath_langchain/chat/bedrock.py +++ b/src/uipath_langchain/chat/bedrock.py @@ -10,6 +10,7 @@ from uipath._utils import resource_override from uipath.utils import EndpointManager +from .header_capture import HeaderCapture from .retryers.bedrock import AsyncBedrockRetryer, BedrockRetryer from .supported_models import BedrockModels from .types import APIFlavor, LLMProvider @@ -62,6 +63,7 @@ def __init__( api_flavor: str, agenthub_config: Optional[str] = None, byo_connection_id: Optional[str] = None, + header_capture: HeaderCapture | None = None, ): self.model = model self.token = token @@ -70,6 +72,7 @@ def __init__( self.byo_connection_id = byo_connection_id self._vendor = "awsbedrock" self._url: Optional[str] = None + self.header_capture = header_capture @property def endpoint(self) -> str: @@ -91,6 +94,12 @@ def _build_base_url(self) -> str: return self._url + def _capture_response_headers(self, parsed, model, **kwargs): + if "ResponseMetadata" in parsed: + headers = parsed["ResponseMetadata"].get("HTTPHeaders", {}) + if self.header_capture: + self.header_capture.set(dict(headers)) + def get_client(self): client = boto3.client( "bedrock-runtime", @@ -106,6 +115,9 @@ def get_client(self): client.meta.events.register( "before-send.bedrock-runtime.*", self._modify_request ) + client.meta.events.register( + "after-call.bedrock-runtime.*", self._capture_response_headers + ) return client def _modify_request(self, request, **kwargs): @@ -203,6 +215,7 @@ class UiPathChatBedrock(ChatBedrock): model: str = "" # For tracing serialization retryer: Optional[Retrying] = None aretryer: Optional[AsyncRetrying] = None + header_capture: HeaderCapture def __init__( self, @@ -233,17 +246,21 @@ def __init__( "UIPATH_ACCESS_TOKEN environment variable or token parameter is required" ) + header_capture = HeaderCapture(name=f"bedrock_headers_{id(self)}") + passthrough_client = AwsBedrockCompletionsPassthroughClient( model=model_name, token=token, api_flavor="invoke", agenthub_config=agenthub_config, byo_connection_id=byo_connection_id, + header_capture=header_capture, ) client = passthrough_client.get_client() kwargs["client"] = client kwargs["model"] = model_name + kwargs["header_capture"] = header_capture super().__init__(**kwargs) self.model = model_name self.retryer = retryer @@ -297,7 +314,15 @@ def _generate( **kwargs: Any, ) -> ChatResult: messages = self._convert_file_blocks_to_anthropic_documents(messages) - return super()._generate(messages, stop=stop, run_manager=run_manager, **kwargs) + result = super()._generate( + messages, + stop=stop, + run_manager=run_manager, + **kwargs, + ) + self.header_capture.attach_to_chat_result(result) + self.header_capture.clear() + return result def _stream( self, @@ -307,9 +332,12 @@ def _stream( **kwargs: Any, ) -> Iterator[ChatGenerationChunk]: messages = self._convert_file_blocks_to_anthropic_documents(messages) - yield from super()._stream( - messages, stop=stop, run_manager=run_manager, **kwargs - ) + chunks = super()._stream(messages, stop=stop, run_manager=run_manager, **kwargs) + + for chunk in chunks: + self.header_capture.attach_to_chat_generation(chunk) + yield chunk + self.header_capture.clear() def _get_default_retryer() -> BedrockRetryer: diff --git a/src/uipath_langchain/chat/constants.py b/src/uipath_langchain/chat/constants.py new file mode 100644 index 00000000..e3d104c1 --- /dev/null +++ b/src/uipath_langchain/chat/constants.py @@ -0,0 +1,9 @@ +from enum import StrEnum + + +class LlmGatewayHeaders(StrEnum): + """LLM Gateway headers.""" + + IS_BYO_EXECUTION = "x-uipath-llmgateway-isbyoexecution" + EXECUTION_DEPLOYMENT_TYPE = "x-uipath-llmgateway-executiondeploymenttype" + IS_PII_MASKED = "x-uipath-llmgateway-ispiimasked" diff --git a/src/uipath_langchain/chat/header_capture.py b/src/uipath_langchain/chat/header_capture.py new file mode 100644 index 00000000..4d7db881 --- /dev/null +++ b/src/uipath_langchain/chat/header_capture.py @@ -0,0 +1,41 @@ +from contextvars import ContextVar +from typing import Optional + +from langchain_core.outputs import ChatGeneration, ChatResult + + +class HeaderCapture: + """Captures HTTP response headers and applies them to LangChain generations.""" + + def __init__(self, name: str = "response_headers"): + """Initialize with a new context var. + + Args: + name: Name for the context var.""" + self._headers: ContextVar[Optional[dict[str, str]]] = ContextVar( + name, default=None + ) + + def set(self, headers: dict[str, str]) -> None: + """Store headers in this instance's context var.""" + self._headers.set(headers) + + def clear(self) -> None: + """Clear stored headers from this instance's context var.""" + self._headers.set(None) + + def attach_to_chat_generation( + self, generation: ChatGeneration, metadata_key: str = "headers" + ) -> None: + """Attach captured headers to the generation message's response_metadata.""" + headers = self._headers.get() + if headers: + generation.message.response_metadata[metadata_key] = headers + + def attach_to_chat_result( + self, result: ChatResult, metadata_key: str = "headers" + ) -> ChatResult: + """Attach captured headers to the message response_metadata of each generation.""" + for generation in result.generations: + self.attach_to_chat_generation(generation, metadata_key) + return result diff --git a/src/uipath_langchain/chat/openai.py b/src/uipath_langchain/chat/openai.py index 39564b1e..9a4eb365 100644 --- a/src/uipath_langchain/chat/openai.py +++ b/src/uipath_langchain/chat/openai.py @@ -140,6 +140,7 @@ def __init__( api_version=api_version, validate_base_url=False, use_responses_api=use_responses_api, + include_response_headers=True, **kwargs, ) diff --git a/src/uipath_langchain/chat/vertex.py b/src/uipath_langchain/chat/vertex.py index 31f94d77..ce568ca3 100644 --- a/src/uipath_langchain/chat/vertex.py +++ b/src/uipath_langchain/chat/vertex.py @@ -1,5 +1,6 @@ import logging import os +from collections.abc import AsyncIterator, Iterator from typing import Any, Optional import httpx @@ -8,12 +9,13 @@ CallbackManagerForLLMRun, ) from langchain_core.messages import BaseMessage -from langchain_core.outputs import ChatResult +from langchain_core.outputs import ChatGenerationChunk, ChatResult from tenacity import AsyncRetrying, Retrying from uipath._utils import resource_override from uipath._utils._ssl_context import get_httpx_client_kwargs from uipath.utils import EndpointManager +from .header_capture import HeaderCapture from .retryers.vertex import AsyncVertexRetryer, VertexRetryer from .supported_models import GeminiModels from .types import APIFlavor, LLMProvider @@ -70,9 +72,15 @@ def _rewrite_vertex_url(original_url: str, gateway_url: str) -> httpx.URL | None class _UrlRewriteTransport(httpx.HTTPTransport): """Transport that rewrites URLs to redirect to UiPath gateway.""" - def __init__(self, gateway_url: str, verify: bool = True): + def __init__( + self, + gateway_url: str, + verify: bool = True, + header_capture: HeaderCapture | None = None, + ): super().__init__(verify=verify) self.gateway_url = gateway_url + self.header_capture = header_capture def handle_request(self, request: httpx.Request) -> httpx.Response: original_url = str(request.url) @@ -86,15 +94,26 @@ def handle_request(self, request: httpx.Request) -> httpx.Response: # Update host header to match the new URL request.headers["host"] = new_url.host request.url = new_url - return super().handle_request(request) + + response = super().handle_request(request) + if self.header_capture: + self.header_capture.set(dict(response.headers)) + + return response class _AsyncUrlRewriteTransport(httpx.AsyncHTTPTransport): """Async transport that rewrites URLs to redirect to UiPath gateway.""" - def __init__(self, gateway_url: str, verify: bool = True): + def __init__( + self, + gateway_url: str, + verify: bool = True, + header_capture: HeaderCapture | None = None, + ): super().__init__(verify=verify) self.gateway_url = gateway_url + self.header_capture = header_capture async def handle_async_request(self, request: httpx.Request) -> httpx.Response: original_url = str(request.url) @@ -108,7 +127,12 @@ async def handle_async_request(self, request: httpx.Request) -> httpx.Response: # Update host header to match the new URL request.headers["host"] = new_url.host request.url = new_url - return await super().handle_async_request(request) + + response = await super().handle_async_request(request) + if self.header_capture: + self.header_capture.set(dict(response.headers)) + + return response class UiPathChatVertex(ChatGoogleGenerativeAI): @@ -162,17 +186,22 @@ def __init__( uipath_url = self._build_base_url(model_name) headers = self._build_headers(token, agenthub_config, byo_connection_id) + header_capture = HeaderCapture(name=f"vertex_headers_{id(self)}") client_kwargs = get_httpx_client_kwargs() verify = client_kwargs.get("verify", True) http_options = genai_types.HttpOptions( httpx_client=httpx.Client( - transport=_UrlRewriteTransport(uipath_url, verify=verify), + transport=_UrlRewriteTransport( + uipath_url, verify=verify, header_capture=header_capture + ), headers=headers, **client_kwargs, ), httpx_async_client=httpx.AsyncClient( - transport=_AsyncUrlRewriteTransport(uipath_url, verify=verify), + transport=_AsyncUrlRewriteTransport( + uipath_url, verify=verify, header_capture=header_capture + ), headers=headers, **client_kwargs, ), @@ -205,6 +234,7 @@ def __init__( self._byo_connection_id = byo_connection_id self._retryer = retryer self._aretryer = aretryer + self._header_capture = header_capture if self.temperature is not None and not 0 <= self.temperature <= 2.0: raise ValueError("temperature must be in the range [0.0, 2.0]") @@ -295,7 +325,10 @@ def _generate( result = super()._generate( messages, stop=stop, run_manager=run_manager, **kwargs ) - return self._merge_finish_reason_to_response_metadata(result) + result = self._merge_finish_reason_to_response_metadata(result) + self._header_capture.attach_to_chat_result(result) + self._header_capture.clear() + return result async def _agenerate( self, @@ -308,7 +341,40 @@ async def _agenerate( result = await super()._agenerate( messages, stop=stop, run_manager=run_manager, **kwargs ) - return self._merge_finish_reason_to_response_metadata(result) + result = self._merge_finish_reason_to_response_metadata(result) + self._header_capture.attach_to_chat_result(result) + self._header_capture.clear() + return result + + def _stream( + self, + messages: list[BaseMessage], + stop: list[str] | None = None, + run_manager: CallbackManagerForLLMRun | None = None, + **kwargs: Any, + ) -> Iterator[ChatGenerationChunk]: + for chunk in super()._stream( + messages, stop=stop, run_manager=run_manager, **kwargs + ): + self._header_capture.attach_to_chat_generation(chunk) + yield chunk + + self._header_capture.clear() + + async def _astream( + self, + messages: list[BaseMessage], + stop: list[str] | None = None, + run_manager: AsyncCallbackManagerForLLMRun | None = None, + **kwargs: Any, + ) -> AsyncIterator[ChatGenerationChunk]: + async for chunk in super()._astream( + messages, stop=stop, run_manager=run_manager, **kwargs + ): + self._header_capture.attach_to_chat_generation(chunk) + yield chunk + + self._header_capture.clear() def _get_default_retryer() -> VertexRetryer: From 1416d6ed373806abf12004b7905f1d3d341d4741 Mon Sep 17 00:00:00 2001 From: cotovanu-cristian Date: Wed, 4 Feb 2026 22:02:51 +0200 Subject: [PATCH 2/2] chore: bump pyproject version --- pyproject.toml | 2 +- uv.lock | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index b5316423..5e56d63e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "uipath-langchain" -version = "0.5.26" +version = "0.5.27" 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/uv.lock b/uv.lock index 6a33ce26..216fb64b 100644 --- a/uv.lock +++ b/uv.lock @@ -3297,7 +3297,7 @@ wheels = [ [[package]] name = "uipath-langchain" -version = "0.5.26" +version = "0.5.27" source = { editable = "." } dependencies = [ { name = "httpx" },