fix(llm): add headers for agenthub config

Author: ionmincu

.gitignore

@@ -180,3 +180,5 @@ cython_debug/
 **/__uipath/
 **/.langgraph_api
 **/testcases/**/uipath.json
+
+/playground.py

pyproject.toml

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

@@ -137,6 +137,8 @@ class UiPathRequestMixin(BaseModel):
     max_tokens: int | None = 1000
     frequency_penalty: float | None = None
     presence_penalty: float | None = None
+    agenthub_config: str | None = None
+    byo_connection_id: str | None = None
 
     logger: logging.Logger | None = None
     max_retries: int | None = 5
@@ -748,6 +750,12 @@ def auth_headers(self) -> dict[str, str]:
                 "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
+                )
             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

@@ -48,10 +48,14 @@ def __init__(
         model: str,
         token: str,
         api_flavor: str,
+        agenthub_config: Optional[str] = None,
+        byo_connection_id: Optional[str] = None,
     ):
         self.model = model
         self.token = token
         self.api_flavor = api_flavor
+        self.agenthub_config = agenthub_config
+        self.byo_connection_id = byo_connection_id
         self._vendor = "awsbedrock"
         self._url: Optional[str] = None
 
@@ -101,6 +105,10 @@ def _modify_request(self, request, **kwargs):
             "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:
@@ -118,6 +126,8 @@ def __init__(
         tenant_id: Optional[str] = None,
         token: Optional[str] = None,
         model_name: str = BedrockModels.anthropic_claude_haiku_4_5,
+        agenthub_config: Optional[str] = None,
+        byo_connection_id: Optional[str] = None,
         **kwargs,
     ):
         org_id = org_id or os.getenv("UIPATH_ORGANIZATION_ID")
@@ -141,6 +151,8 @@ def __init__(
             model=model_name,
             token=token,
             api_flavor="converse",
+            agenthub_config=agenthub_config,
+            byo_connection_id=byo_connection_id,
         )
 
         client = passthrough_client.get_client()
@@ -156,6 +168,8 @@ def __init__(
         tenant_id: Optional[str] = None,
         token: Optional[str] = None,
         model_name: str = BedrockModels.anthropic_claude_haiku_4_5,
+        agenthub_config: Optional[str] = None,
+        byo_connection_id: Optional[str] = None,
         **kwargs,
     ):
         org_id = org_id or os.getenv("UIPATH_ORGANIZATION_ID")
@@ -179,6 +193,8 @@ def __init__(
             model=model_name,
             token=token,
             api_flavor="invoke",
+            agenthub_config=agenthub_config,
+            byo_connection_id=byo_connection_id,
         )
 
         client = passthrough_client.get_client()

src/uipath_langchain/chat/openai.py

@@ -12,21 +12,41 @@
 logger = logging.getLogger(__name__)
 
 
+def _rewrite_openai_url(
+    original_url: str, params: httpx.QueryParams
+) -> httpx.URL | None:
+    """Rewrite OpenAI URLs to UiPath gateway completions endpoint.
+
+    Handles three URL patterns:
+    - responses: false -> .../openai/deployments/.../chat/completions?api-version=...
+    - responses: true  -> .../openai/responses?api-version=...
+    - responses API base -> .../{model}?api-version=... (no /openai/ path)
+
+    All are rewritten to .../completions
+    """
+    if "/openai/deployments/" in original_url:
+        base_url = original_url.split("/openai/deployments/")[0]
+    elif "/openai/responses" in original_url:
+        base_url = original_url.split("/openai/responses")[0]
+    else:
+        # Handle base URL case (no /openai/ path appended yet)
+        # Strip query string to get base URL
+        base_url = original_url.split("?")[0]
+
+    new_url_str = f"{base_url}/completions"
+    if params:
+        return httpx.URL(new_url_str, params=params)
+    return httpx.URL(new_url_str)
+
+
 class UiPathURLRewriteTransport(httpx.AsyncHTTPTransport):
     def __init__(self, verify: bool = True, **kwargs):
         super().__init__(verify=verify, **kwargs)
 
     async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
-        original_url = str(request.url)
-
-        if "/openai/deployments/" in original_url:
-            base_url = original_url.split("/openai/deployments/")[0]
-            query_string = request.url.params
-            new_url_str = f"{base_url}/completions"
-            if query_string:
-                request.url = httpx.URL(new_url_str, params=query_string)
-            else:
-                request.url = httpx.URL(new_url_str)
+        new_url = _rewrite_openai_url(str(request.url), request.url.params)
+        if new_url:
+            request.url = new_url
 
         return await super().handle_async_request(request)
 
@@ -36,16 +56,9 @@ def __init__(self, verify: bool = True, **kwargs):
         super().__init__(verify=verify, **kwargs)
 
     def handle_request(self, request: httpx.Request) -> httpx.Response:
-        original_url = str(request.url)
-
-        if "/openai/deployments/" in original_url:
-            base_url = original_url.split("/openai/deployments/")[0]
-            query_string = request.url.params
-            new_url_str = f"{base_url}/completions"
-            if query_string:
-                request.url = httpx.URL(new_url_str, params=query_string)
-            else:
-                request.url = httpx.URL(new_url_str)
+        new_url = _rewrite_openai_url(str(request.url), request.url.params)
+        if new_url:
+            request.url = new_url
 
         return super().handle_request(request)
 
@@ -58,6 +71,9 @@ def __init__(
         api_version: str = "2024-12-01-preview",
         org_id: Optional[str] = None,
         tenant_id: Optional[str] = None,
+        agenthub_config: Optional[str] = None,
+        extra_headers: Optional[dict[str, str]] = None,
+        byo_connection_id: Optional[str] = None,
         **kwargs,
     ):
         org_id = org_id or os.getenv("UIPATH_ORGANIZATION_ID")
@@ -81,18 +97,24 @@ def __init__(
         self._vendor = "openai"
         self._model_name = model_name
         self._url: Optional[str] = None
+        self._agenthub_config = agenthub_config
+        self._byo_connection_id = byo_connection_id
+        self._extra_headers = extra_headers or {}
+
+        client_kwargs = get_httpx_client_kwargs()
+        verify = client_kwargs.get("verify", True)
 
         super().__init__(
             azure_endpoint=self._build_base_url(),
             model_name=model_name,
             default_headers=self._build_headers(token),
             http_async_client=httpx.AsyncClient(
-                transport=UiPathURLRewriteTransport(verify=True),
-                **get_httpx_client_kwargs(),
+                transport=UiPathURLRewriteTransport(verify=verify),
+                **client_kwargs,
             ),
             http_client=httpx.Client(
-                transport=UiPathSyncURLRewriteTransport(verify=True),
-                **get_httpx_client_kwargs(),
+                transport=UiPathSyncURLRewriteTransport(verify=verify),
+                **client_kwargs,
             ),
             api_key=token,
             api_version=api_version,
@@ -105,10 +127,18 @@ def _build_headers(self, token: str) -> dict[str, str]:
             "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 job_key := os.getenv("UIPATH_JOB_KEY"):
             headers["X-UiPath-JobKey"] = job_key
         if process_key := os.getenv("UIPATH_PROCESS_KEY"):
             headers["X-UiPath-ProcessKey"] = process_key
+
+        # Allow extra_headers to override defaults
+        headers.update(self._extra_headers)
         return headers
 
     @property
@@ -117,9 +147,9 @@ def endpoint(self) -> str:
         formatted_endpoint = vendor_endpoint.format(
             vendor=self._vendor,
             model=self._model_name,
-            api_version=self._openai_api_version,
         )
-        return formatted_endpoint.replace("/completions", "")
+        base_endpoint = formatted_endpoint.replace("/completions", "")
+        return f"{base_endpoint}?api-version={self._openai_api_version}"
 
     def _build_base_url(self) -> str:
         if not self._url: