feat: add local URL override support

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)
 
 

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

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()