From 8725e1ae322e04cf8ab0c71cf9f0327fcd6b6197 Mon Sep 17 00:00:00 2001 From: cotovanu-cristian Date: Sat, 14 Feb 2026 13:39:08 +0200 Subject: [PATCH] feat: add local URL override support --- src/uipath_langchain/chat/bedrock.py | 19 +++++++++++++------ src/uipath_langchain/chat/openai.py | 16 ++++++++++------ src/uipath_langchain/chat/vertex.py | 28 +++++++++++++++++++--------- 3 files changed, 42 insertions(+), 21 deletions(-) diff --git a/src/uipath_langchain/chat/bedrock.py b/src/uipath_langchain/chat/bedrock.py index a29631deb..c49b66633 100644 --- a/src/uipath_langchain/chat/bedrock.py +++ b/src/uipath_langchain/chat/bedrock.py @@ -7,7 +7,7 @@ from langchain_core.messages import BaseMessage from langchain_core.outputs import ChatGenerationChunk, ChatResult from tenacity import AsyncRetrying, Retrying -from uipath._utils import resource_override +from uipath._utils import resolve_endpoint_override, resource_override from uipath.utils import EndpointManager from .header_capture import HeaderCapture @@ -72,6 +72,7 @@ def __init__( self.byo_connection_id = byo_connection_id self._vendor = "awsbedrock" self._url: Optional[str] = None + self._override_headers: dict[str, str] = {} self.header_capture = header_capture @property @@ -85,12 +86,17 @@ def endpoint(self) -> str: def _build_base_url(self) -> str: if not self._url: - env_uipath_url = os.getenv("UIPATH_URL") - - if env_uipath_url: - self._url = f"{env_uipath_url.rstrip('/')}/{self.endpoint}" + override_url, self._override_headers = resolve_endpoint_override( + self.endpoint + ) + if override_url: + self._url = override_url else: - raise ValueError("UIPATH_URL environment variable is required") + 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 @@ -145,6 +151,7 @@ def _modify_request(self, request, **kwargs): if process_key: headers["X-UiPath-ProcessKey"] = process_key + headers.update(self._override_headers) request.headers.update(headers) diff --git a/src/uipath_langchain/chat/openai.py b/src/uipath_langchain/chat/openai.py index 9a4eb365a..52961e7fb 100644 --- a/src/uipath_langchain/chat/openai.py +++ b/src/uipath_langchain/chat/openai.py @@ -5,7 +5,7 @@ import httpx from langchain_openai import AzureChatOpenAI from pydantic import PrivateAttr -from uipath._utils import resource_override +from uipath._utils import resolve_endpoint_override, resource_override from uipath._utils._ssl_context import get_httpx_client_kwargs from uipath.utils import EndpointManager @@ -177,11 +177,15 @@ def endpoint(self) -> str: def _build_base_url(self) -> str: if not self._url: - env_uipath_url = os.getenv("UIPATH_URL") - - if env_uipath_url: - self._url = f"{env_uipath_url.rstrip('/')}/{self.endpoint}" + override_url, override_headers = resolve_endpoint_override(self.endpoint) + if override_url: + self._url = override_url + self._extra_headers.update(override_headers) else: - raise ValueError("UIPATH_URL environment variable is required") + 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 diff --git a/src/uipath_langchain/chat/vertex.py b/src/uipath_langchain/chat/vertex.py index ce568ca3d..423af8026 100644 --- a/src/uipath_langchain/chat/vertex.py +++ b/src/uipath_langchain/chat/vertex.py @@ -11,7 +11,7 @@ from langchain_core.messages import BaseMessage from langchain_core.outputs import ChatGenerationChunk, ChatResult from tenacity import AsyncRetrying, Retrying -from uipath._utils import resource_override +from uipath._utils import resolve_endpoint_override, resource_override from uipath._utils._ssl_context import get_httpx_client_kwargs from uipath.utils import EndpointManager @@ -183,8 +183,9 @@ def __init__( "UIPATH_ACCESS_TOKEN environment variable or token parameter is required" ) - uipath_url = self._build_base_url(model_name) + uipath_url, override_headers = self._build_base_url(model_name) headers = self._build_headers(token, agenthub_config, byo_connection_id) + headers.update(override_headers) header_capture = HeaderCapture(name=f"vertex_headers_{id(self)}") client_kwargs = get_httpx_client_kwargs() @@ -269,19 +270,28 @@ def _build_headers( return headers @staticmethod - def _build_base_url(model_name: str) -> str: - """Build the full URL for the UiPath LLM Gateway.""" - env_uipath_url = os.getenv("UIPATH_URL") - - if not env_uipath_url: - raise ValueError("UIPATH_URL environment variable is required") + def _build_base_url(model_name: str) -> tuple[str, dict[str, str]]: + """Build the full URL for the UiPath LLM Gateway. + Returns: + A tuple of (url, extra_headers). extra_headers contains routing + headers when a local service URL override is active. + """ vendor_endpoint = EndpointManager.get_vendor_endpoint() formatted_endpoint = vendor_endpoint.format( vendor="vertexai", model=model_name, ) - return f"{env_uipath_url.rstrip('/')}/{formatted_endpoint}" + + override_url, override_headers = resolve_endpoint_override(formatted_endpoint) + if override_url: + return override_url, override_headers + + 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('/')}/{formatted_endpoint}", {} def invoke(self, *args, **kwargs): retryer = self._retryer or _get_default_retryer()