feat: add uipath headers llm calls (#663)

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: cotovanu-cristian

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath-langchain"
-version = "0.8.5"
+version = "0.8.6"
 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"

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,6 +33,7 @@
     get_uipath_token_header,
 )
 from uipath_langchain._utils._sleep_policy import before_sleep_log
+from uipath_langchain.chat.http_client import build_uipath_headers, resolve_gateway_url
 from uipath_langchain.runtime.errors import (
     LangGraphErrorCode,
     LangGraphRuntimeError,
@@ -79,8 +79,6 @@ class UiPathRequestMixin(BaseModel):
 
     default_headers: Mapping[str, str] | None = {
         "X-UiPath-Streaming-Enabled": "false",
-        "X-UiPath-JobKey": os.getenv("UIPATH_JOB_KEY", ""),
-        "X-UiPath-ProcessKey": quote(os.getenv("UIPATH_PROCESS_KEY", ""), safe=""),
     }
     model_name: str | None = Field(
         default_factory=lambda: os.getenv(
@@ -155,6 +153,7 @@ class UiPathRequestMixin(BaseModel):
     max_delay: float = 60.0
 
     _url: str | None = None
+    _is_override: bool = False
     _auth_headers: dict[str, str] | None = None
 
     # required to instantiate AzureChatOpenAI subclasses
@@ -732,17 +731,26 @@ def _prepare_url(self, url: str) -> httpx.URL:
     def _build_headers(self, options, retries_taken: int = 0) -> httpx.Headers:
         return httpx.Headers(self.auth_headers)
 
+    def _resolve_url_and_override(self) -> None:
+        """Resolve ``_url`` and ``_is_override`` idempotently."""
+        if self._url:
+            return
+        try:
+            self._url, self._is_override = resolve_gateway_url(self.endpoint)
+        except ValueError:
+            self._url = (
+                f"{self.base_url}/{self.org_id}/{self.tenant_id}/{self.endpoint}"
+            )
+        except NotImplementedError:
+            pass
+
     @property
     def url(self) -> str:
+        self._resolve_url_and_override()
         if not self._url:
-            env_uipath_url = os.getenv("UIPATH_URL")
-
-            if env_uipath_url:
-                self._url = f"{env_uipath_url.rstrip('/')}/{self.endpoint}"
-            else:
-                self._url = (
-                    f"{self.base_url}/{self.org_id}/{self.tenant_id}/{self.endpoint}"
-                )
+            raise NotImplementedError(
+                "The endpoint property is not implemented for this class."
+            )
         return self._url
 
     @property
@@ -754,17 +762,21 @@ def endpoint(self) -> str:
     @property
     def auth_headers(self) -> dict[str, str]:
         if not self._auth_headers:
+            self._resolve_url_and_override()
             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["Authorization"] = f"Bearer {self.access_token}"
+            self._auth_headers.update(
+                build_uipath_headers(
+                    agenthub_config=self.agenthub_config,
+                    byo_connection_id=self.byo_connection_id,
+                    inject_routing=self._is_override,
                 )
+            )
+            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

src/uipath_langchain/chat/bedrock.py

@@ -2,16 +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 .header_capture import HeaderCapture
-from .retryers.bedrock import AsyncBedrockRetryer, BedrockRetryer
+from .http_client import build_uipath_headers, resolve_gateway_url
+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
 
@@ -72,6 +72,7 @@ def __init__(
         self.byo_connection_id = byo_connection_id
         self._vendor = "awsbedrock"
         self._url: Optional[str] = None
+        self._is_override: bool = False
         self.header_capture = header_capture
 
     @property
@@ -83,16 +84,10 @@ def endpoint(self) -> str:
         )
         return formatted_endpoint
 
-    def _build_base_url(self) -> str:
+    def _resolve_url(self) -> tuple[str, bool]:
         if not self._url:
-            env_uipath_url = os.getenv("UIPATH_URL")
-
-            if env_uipath_url:
-                self._url = f"{env_uipath_url.rstrip('/')}/{self.endpoint}"
-            else:
-                raise ValueError("UIPATH_URL environment variable is required")
-
-        return self._url
+            self._url, self._is_override = resolve_gateway_url(self.endpoint)
+        return self._url, self._is_override
 
     def _capture_response_headers(self, parsed, model, **kwargs):
         if "ResponseMetadata" in parsed:
@@ -122,29 +117,24 @@ def get_client(self):
         return client
 
     def _modify_request(self, request, **kwargs):
-        """Intercept boto3 request and redirect to LLM Gateway"""
+        """Intercept boto3 request and redirect to LLM Gateway."""
         # Detect streaming based on URL suffix:
         # - converse-stream / invoke-with-response-stream -> streaming
         # - converse / invoke -> non-streaming
         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
-        job_key = os.getenv("UIPATH_JOB_KEY")
-        process_key = os.getenv("UIPATH_PROCESS_KEY")
-        if job_key:
-            headers["X-UiPath-JobKey"] = job_key
-        if process_key:
-            headers["X-UiPath-ProcessKey"] = quote(process_key, safe="")
+        url, is_override = self._resolve_url()
+        request.url = url
+
+        headers: dict[str, str] = {"Authorization": f"Bearer {self.token}"}
+        headers.update(
+            build_uipath_headers(
+                agenthub_config=self.agenthub_config,
+                byo_connection_id=self.byo_connection_id,
+                inject_routing=is_override,
+            )
+        )
+        headers["X-UiPath-LlmGateway-ApiFlavor"] = self.api_flavor
+        headers["X-UiPath-Streaming-Enabled"] = streaming
 
         request.headers.update(headers)
 

src/uipath_langchain/chat/http_client/__init__.py

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

…/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

@@ -0,0 +1,61 @@
+"""Shared UiPath HTTP headers for LLM Gateway requests."""
+
+import os
+from urllib.parse import quote
+
+from uipath.platform.common.constants import (
+    ENV_FOLDER_KEY,
+    ENV_JOB_KEY,
+    ENV_ORGANIZATION_ID,
+    ENV_PROCESS_KEY,
+    ENV_TENANT_ID,
+    ENV_UIPATH_TRACE_ID,
+    HEADER_AGENTHUB_CONFIG,
+    HEADER_FOLDER_KEY,
+    HEADER_INTERNAL_ACCOUNT_ID,
+    HEADER_INTERNAL_TENANT_ID,
+    HEADER_JOB_KEY,
+    HEADER_LLMGATEWAY_BYO_CONNECTION_ID,
+    HEADER_PROCESS_KEY,
+    HEADER_TRACE_ID,
+)
+
+
+def build_uipath_headers(
+    *,
+    agenthub_config: str | None = None,
+    byo_connection_id: str | None = None,
+    inject_routing: bool = False,
+) -> 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.
+
+    Args:
+        agenthub_config: Optional AgentHub configuration identifier.
+        byo_connection_id: Optional BYO connection identifier.
+        inject_routing: When True, adds tenant and account routing
+            headers that are normally injected by the platform routing
+            layer.  Set this when using a service URL override that
+            bypasses the platform.
+    """
+    headers: dict[str, str] = {}
+    if agenthub_config:
+        headers[HEADER_AGENTHUB_CONFIG] = agenthub_config
+    if byo_connection_id:
+        headers[HEADER_LLMGATEWAY_BYO_CONNECTION_ID] = byo_connection_id
+    if process_key := os.getenv(ENV_PROCESS_KEY):
+        headers[HEADER_PROCESS_KEY] = quote(process_key, safe="")
+    if job_key := os.getenv(ENV_JOB_KEY):
+        headers[HEADER_JOB_KEY] = job_key
+    if folder_key := os.getenv(ENV_FOLDER_KEY):
+        headers[HEADER_FOLDER_KEY] = folder_key
+    if trace_id := os.getenv(ENV_UIPATH_TRACE_ID):
+        headers[HEADER_TRACE_ID] = trace_id
+    if inject_routing:
+        if tenant_id := os.getenv(ENV_TENANT_ID):
+            headers[HEADER_INTERNAL_TENANT_ID] = tenant_id
+        if organization_id := os.getenv(ENV_ORGANIZATION_ID):
+            headers[HEADER_INTERNAL_ACCOUNT_ID] = organization_id
+    return headers

src/uipath_langchain/chat/http_client/headers.py

@@ -0,0 +1,46 @@
+"""Gateway URL resolution with service override support.
+
+Integrates chat models into the platform's per-service URL override system
+(UIPATH_SERVICE_URL_<SERVICE> env vars) so that local development redirects
+work consistently across all UiPath HTTP clients.
+"""
+
+import os
+from collections.abc import Callable
+
+import uipath.platform.common as _platform_common
+
+_resolve_service_url: Callable[[str], str | None] | None = getattr(
+    _platform_common, "resolve_service_url", None
+)
+
+
+def resolve_gateway_url(endpoint_path: str) -> tuple[str, bool]:
+    """Resolve the full gateway URL for a given endpoint path.
+
+    Checks for a per-service URL override first (e.g.
+    ``UIPATH_SERVICE_URL_AGENTHUB``). Falls back to building the URL
+    from ``UIPATH_URL``.
+
+    Args:
+        endpoint_path: Endpoint path with service prefix, e.g.
+            ``"agenthub_/llm/raw/vendor/openai/model/gpt-4/completions"``.
+
+    Returns:
+        Tuple of (resolved_url, is_override). When *is_override* is True
+        the caller should inject routing headers since the platform
+        routing layer is bypassed.
+
+    Raises:
+        ValueError: If neither a service override nor UIPATH_URL is set.
+    """
+    if _resolve_service_url is not None:
+        override_url = _resolve_service_url(endpoint_path)
+        if override_url:
+            return override_url, True
+
+    env_uipath_url = os.getenv("UIPATH_URL")
+    if not env_uipath_url:
+        raise ValueError("UIPATH_URL environment variable is required")
+
+    return f"{env_uipath_url.rstrip('/')}/{endpoint_path}", False