refactor: uniform headers into chat/http_client module

Move all shared HTTP concerns (headers, header capture, retryers) into
a new chat/http_client subpackage. Introduce build_uipath_headers() as
the single entry point for building UiPath LLM Gateway headers across
OpenAI, Vertex, and Bedrock providers. Context headers (job_key,
folder_key, trace_id) now read directly from env vars like process_key.

Co-Authored-By: Claude Opus 4.6 <[email protected]>

Author: cotovanu-cristian

src/uipath_langchain/_utils/_request_mixin.py

@@ -4,7 +4,6 @@
 import os
 import time
 from typing import Any, AsyncIterator, Dict, Iterator, Mapping
-from urllib.parse import quote
 
 import httpx
 import openai
@@ -34,7 +33,7 @@
     get_uipath_token_header,
 )
 from uipath_langchain._utils._sleep_policy import before_sleep_log
-from uipath_langchain.chat._headers import build_uipath_context_headers
+from uipath_langchain.chat.http_client import build_uipath_headers
 from uipath_langchain.runtime.errors import (
     LangGraphErrorCode,
     LangGraphRuntimeError,
@@ -80,7 +79,6 @@ class UiPathRequestMixin(BaseModel):
 
     default_headers: Mapping[str, str] | None = {
         "X-UiPath-Streaming-Enabled": "false",
-        "X-UiPath-ProcessKey": quote(os.getenv("UIPATH_PROCESS_KEY", ""), safe=""),
     }
     model_name: str | None = Field(
         default_factory=lambda: os.getenv(
@@ -756,23 +754,24 @@ def auth_headers(self) -> dict[str, str]:
         if not self._auth_headers:
             self._auth_headers = {
                 **self.default_headers,  # type: ignore
-                "Authorization": f"Bearer {self.access_token}",
-                "X-UiPath-LlmGateway-TimeoutSeconds": str(self.default_request_timeout),
             }
-            if self.agenthub_config:
-                self._auth_headers["X-UiPath-AgentHub-Config"] = self.agenthub_config
-            if self.byo_connection_id:
-                self._auth_headers["X-UiPath-LlmGateway-ByoIsConnectionId"] = (
-                    self.byo_connection_id
+            self._auth_headers.update(
+                build_uipath_headers(
+                    self.access_token,
+                    agenthub_config=self.agenthub_config,
+                    byo_connection_id=self.byo_connection_id,
                 )
+            )
+            self._auth_headers["X-UiPath-LlmGateway-TimeoutSeconds"] = str(
+                self.default_request_timeout
+            )
             if self.is_normalized and self.model_name:
                 self._auth_headers["X-UiPath-LlmGateway-NormalizedApi-ModelName"] = (
                     self.model_name
                 )
             if self.include_account_id:
                 self._auth_headers["x-uipath-internal-accountid"] = self.org_id
                 self._auth_headers["x-uipath-internal-tenantid"] = self.tenant_id
-            self._auth_headers.update(build_uipath_context_headers())
         return self._auth_headers
 
     def _get_llm_string(self, stop: list[str] | None = None, **kwargs: Any) -> str:

src/uipath_langchain/chat/_headers.py

@@ -2,17 +2,16 @@
 import os
 from collections.abc import Iterator
 from typing import Any, Optional
-from urllib.parse import quote
 
 from langchain_core.callbacks import CallbackManagerForLLMRun
 from langchain_core.messages import BaseMessage
 from langchain_core.outputs import ChatGenerationChunk, ChatResult
 from tenacity import AsyncRetrying, Retrying
 from uipath.platform.common import EndpointManager, resource_override
 
-from ._headers import build_uipath_context_headers
-from .header_capture import HeaderCapture
-from .retryers.bedrock import AsyncBedrockRetryer, BedrockRetryer
+from .http_client import build_uipath_headers
+from .http_client.header_capture import HeaderCapture
+from .http_client.retryers.bedrock import AsyncBedrockRetryer, BedrockRetryer
 from .supported_models import BedrockModels
 from .types import APIFlavor, LLMProvider
 
@@ -130,20 +129,13 @@ def _modify_request(self, request, **kwargs):
         streaming = "true" if request.url.endswith("-stream") else "false"
         request.url = self._build_base_url()
 
-        headers = {
-            "Authorization": f"Bearer {self.token}",
-            "X-UiPath-LlmGateway-ApiFlavor": self.api_flavor,
-            "X-UiPath-Streaming-Enabled": streaming,
-        }
-
-        if self.agenthub_config:
-            headers["X-UiPath-AgentHub-Config"] = self.agenthub_config
-        if self.byo_connection_id:
-            headers["X-UiPath-LlmGateway-ByoIsConnectionId"] = self.byo_connection_id
-        if process_key := os.getenv("UIPATH_PROCESS_KEY"):
-            headers["X-UiPath-ProcessKey"] = quote(process_key, safe="")
-
-        headers.update(build_uipath_context_headers())
+        headers = build_uipath_headers(
+            self.token,
+            agenthub_config=self.agenthub_config,
+            byo_connection_id=self.byo_connection_id,
+        )
+        headers["X-UiPath-LlmGateway-ApiFlavor"] = self.api_flavor
+        headers["X-UiPath-Streaming-Enabled"] = streaming
 
         request.headers.update(headers)
 

src/uipath_langchain/chat/bedrock.py

@@ -0,0 +1,3 @@
+from .headers import build_uipath_headers
+
+__all__ = ["build_uipath_headers"]

src/uipath_langchain/chat/http_client/__init__.py

@@ -0,0 +1,33 @@
+"""Shared UiPath HTTP headers for LLM Gateway requests."""
+
+import os
+from urllib.parse import quote
+
+
+def build_uipath_headers(
+    token: str,
+    *,
+    agenthub_config: str | None = None,
+    byo_connection_id: str | None = None,
+) -> dict[str, str]:
+    """Build common UiPath headers for LLM Gateway requests.
+
+    Reads process_key, job_key, folder_key, and trace_id directly from
+    environment variables when set.
+    """
+    headers: dict[str, str] = {
+        "Authorization": f"Bearer {token}",
+    }
+    if agenthub_config:
+        headers["X-UiPath-AgentHub-Config"] = agenthub_config
+    if byo_connection_id:
+        headers["X-UiPath-LlmGateway-ByoIsConnectionId"] = byo_connection_id
+    if process_key := os.getenv("UIPATH_PROCESS_KEY"):
+        headers["X-UiPath-ProcessKey"] = quote(process_key, safe="")
+    if job_key := os.getenv("UIPATH_JOB_KEY"):
+        headers["x-uipath-jobkey"] = job_key
+    if folder_key := os.getenv("UIPATH_FOLDER_KEY"):
+        headers["x-uipath-folderkey"] = folder_key
+    if trace_id := os.getenv("UIPATH_TRACE_ID"):
+        headers["x-uipath-traceid"] = trace_id
+    return headers

…/uipath_langchain/chat/header_capture.py …chain/chat/http_client/header_capture.pysrc/uipath_langchain/chat/header_capture.py renamed to src/uipath_langchain/chat/http_client/header_capture.py src/uipath_langchain/chat/header_capture.py renamed to src/uipath_langchain/chat/http_client/header_capture.py

@@ -1,7 +1,6 @@
 import logging
 import os
 from typing import Optional
-from urllib.parse import quote
 
 import httpx
 from langchain_openai import AzureChatOpenAI
@@ -12,7 +11,7 @@
     resource_override,
 )
 
-from ._headers import build_uipath_context_headers
+from .http_client import build_uipath_headers
 from .supported_models import OpenAIModels
 from .types import APIFlavor, LLMProvider
 
@@ -152,21 +151,12 @@ def __init__(
         self._api_flavor = api_flavor
 
     def _build_headers(self, token: str) -> dict[str, str]:
-        headers = {
-            "X-UiPath-LlmGateway-ApiFlavor": "auto",
-            "Authorization": f"Bearer {token}",
-        }
-
-        if self._agenthub_config:
-            headers["X-UiPath-AgentHub-Config"] = self._agenthub_config
-        if self._byo_connection_id:
-            headers["X-UiPath-LlmGateway-ByoIsConnectionId"] = self._byo_connection_id
-        if process_key := os.getenv("UIPATH_PROCESS_KEY"):
-            headers["X-UiPath-ProcessKey"] = quote(process_key, safe="")
-
-        headers.update(build_uipath_context_headers())
-
-        # Allow extra_headers to override defaults
+        headers = build_uipath_headers(
+            token,
+            agenthub_config=self._agenthub_config,
+            byo_connection_id=self._byo_connection_id,
+        )
+        headers["X-UiPath-LlmGateway-ApiFlavor"] = "auto"
         headers.update(self._extra_headers)
         return headers