From 3ac0a147e50a5829d727c0c9ae8c104c10f55f79 Mon Sep 17 00:00:00 2001 From: ionmincu Date: Tue, 16 Dec 2025 18:24:20 +0200 Subject: [PATCH] fix(llm): add headers for agenthub config --- .gitignore | 2 + pyproject.toml | 2 +- src/uipath_langchain/_utils/_request_mixin.py | 8 ++ src/uipath_langchain/chat/bedrock.py | 16 +++ src/uipath_langchain/chat/openai.py | 82 ++++++++----- src/uipath_langchain/chat/vertex.py | 108 ++++++++++-------- tests/chat/__init__.py | 0 tests/chat/test_openai_url_rewrite.py | 82 +++++++++++++ uv.lock | 2 +- 9 files changed, 228 insertions(+), 74 deletions(-) create mode 100644 tests/chat/__init__.py create mode 100644 tests/chat/test_openai_url_rewrite.py diff --git a/.gitignore b/.gitignore index 0b66de238..ada74fd4c 100644 --- a/.gitignore +++ b/.gitignore @@ -180,3 +180,5 @@ cython_debug/ **/__uipath/ **/.langgraph_api **/testcases/**/uipath.json + +/playground.py diff --git a/pyproject.toml b/pyproject.toml index 13f634c39..851e49d09 100644 --- a/pyproject.toml +++ b/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" diff --git a/src/uipath_langchain/_utils/_request_mixin.py b/src/uipath_langchain/_utils/_request_mixin.py index b84f509a1..317877de6 100644 --- a/src/uipath_langchain/_utils/_request_mixin.py +++ b/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 diff --git a/src/uipath_langchain/chat/bedrock.py b/src/uipath_langchain/chat/bedrock.py index e62817550..ccd4ad84d 100644 --- a/src/uipath_langchain/chat/bedrock.py +++ b/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() diff --git a/src/uipath_langchain/chat/openai.py b/src/uipath_langchain/chat/openai.py index b51c97399..048bd07b1 100644 --- a/src/uipath_langchain/chat/openai.py +++ b/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: diff --git a/src/uipath_langchain/chat/vertex.py b/src/uipath_langchain/chat/vertex.py index 84836f936..6fc21a5c2 100644 --- a/src/uipath_langchain/chat/vertex.py +++ b/src/uipath_langchain/chat/vertex.py @@ -40,60 +40,59 @@ def _check_genai_dependencies() -> None: from pydantic import PrivateAttr -def _rewrite_request_for_gateway( - request: httpx.Request, gateway_url: str -) -> httpx.Request: - """Rewrite a request to redirect to the UiPath gateway.""" - url_str = str(request.url) - if "generateContent" in url_str or "streamGenerateContent" in url_str: - is_streaming = "alt=sse" in url_str - - headers = dict(request.headers) - - headers["X-UiPath-Streaming-Enabled"] = "true" if is_streaming else "false" - - gateway_url_parsed = httpx.URL(gateway_url) - if gateway_url_parsed.host: - headers["host"] = gateway_url_parsed.host - - return httpx.Request( - method=request.method, - url=gateway_url, - headers=headers, - content=request.content, - extensions=request.extensions, - ) - return request +def _rewrite_vertex_url(original_url: str, gateway_url: str) -> httpx.URL | None: + """Rewrite Google GenAI URLs to UiPath gateway endpoint. + + Handles URL patterns containing generateContent or streamGenerateContent. + Returns the gateway URL, or None if no rewrite needed. + """ + if "generateContent" in original_url or "streamGenerateContent" in original_url: + return httpx.URL(gateway_url + "?api-version=v1") + return None -class _UrlRewriteTransport(httpx.BaseTransport): +class _UrlRewriteTransport(httpx.HTTPTransport): """Transport that rewrites URLs to redirect to UiPath gateway.""" - def __init__(self, gateway_url: str): + def __init__(self, gateway_url: str, verify: bool = True): + super().__init__(verify=verify) self.gateway_url = gateway_url - self._transport = httpx.HTTPTransport() def handle_request(self, request: httpx.Request) -> httpx.Response: - request = _rewrite_request_for_gateway(request, self.gateway_url) - return self._transport.handle_request(request) - - def close(self) -> None: - self._transport.close() + original_url = str(request.url) + new_url = _rewrite_vertex_url(original_url, self.gateway_url) + if new_url: + # Set streaming header based on original URL before modifying + is_streaming = "alt=sse" in original_url + request.headers["X-UiPath-Streaming-Enabled"] = ( + "true" if is_streaming else "false" + ) + # Update host header to match the new URL + request.headers["host"] = new_url.host + request.url = new_url + return super().handle_request(request) -class _AsyncUrlRewriteTransport(httpx.AsyncBaseTransport): +class _AsyncUrlRewriteTransport(httpx.AsyncHTTPTransport): """Async transport that rewrites URLs to redirect to UiPath gateway.""" - def __init__(self, gateway_url: str): + def __init__(self, gateway_url: str, verify: bool = True): + super().__init__(verify=verify) self.gateway_url = gateway_url - self._transport = httpx.AsyncHTTPTransport() async def handle_async_request(self, request: httpx.Request) -> httpx.Response: - request = _rewrite_request_for_gateway(request, self.gateway_url) - return await self._transport.handle_async_request(request) - - async def aclose(self) -> None: - await self._transport.aclose() + original_url = str(request.url) + new_url = _rewrite_vertex_url(original_url, self.gateway_url) + if new_url: + # Set streaming header based on original URL before modifying + is_streaming = "alt=sse" in original_url + request.headers["X-UiPath-Streaming-Enabled"] = ( + "true" if is_streaming else "false" + ) + # Update host header to match the new URL + request.headers["host"] = new_url.host + request.url = new_url + return await super().handle_async_request(request) class UiPathChatVertex(ChatGoogleGenerativeAI): @@ -103,6 +102,8 @@ class UiPathChatVertex(ChatGoogleGenerativeAI): _model_name: str = PrivateAttr() _uipath_token: str = PrivateAttr() _uipath_llmgw_url: Optional[str] = PrivateAttr(default=None) + _agenthub_config: Optional[str] = PrivateAttr(default=None) + _byo_connection_id: Optional[str] = PrivateAttr(default=None) def __init__( self, @@ -111,6 +112,8 @@ def __init__( token: Optional[str] = None, model_name: str = GeminiModels.gemini_2_5_flash, temperature: Optional[float] = None, + agenthub_config: Optional[str] = None, + byo_connection_id: Optional[str] = None, **kwargs: Any, ): org_id = org_id or os.getenv("UIPATH_ORGANIZATION_ID") @@ -131,18 +134,21 @@ def __init__( ) uipath_url = self._build_base_url(model_name) - headers = self._build_headers(token) + headers = self._build_headers(token, agenthub_config, byo_connection_id) + + client_kwargs = get_httpx_client_kwargs() + verify = client_kwargs.get("verify", True) http_options = genai_types.HttpOptions( httpx_client=httpx.Client( - transport=_UrlRewriteTransport(uipath_url), + transport=_UrlRewriteTransport(uipath_url, verify=verify), headers=headers, - **get_httpx_client_kwargs(), + **client_kwargs, ), httpx_async_client=httpx.AsyncClient( - transport=_AsyncUrlRewriteTransport(uipath_url), + transport=_AsyncUrlRewriteTransport(uipath_url, verify=verify), headers=headers, - **get_httpx_client_kwargs(), + **client_kwargs, ), ) @@ -168,6 +174,8 @@ def __init__( self._model_name = model_name self._uipath_token = token self._uipath_llmgw_url = uipath_url + self._agenthub_config = agenthub_config + self._byo_connection_id = byo_connection_id if self.temperature is not None and not 0 <= self.temperature <= 2.0: raise ValueError("temperature must be in the range [0.0, 2.0]") @@ -182,11 +190,19 @@ def __init__( self.default_metadata = tuple(additional_headers.items()) @staticmethod - def _build_headers(token: str) -> dict[str, str]: + def _build_headers( + token: str, + agenthub_config: Optional[str] = None, + byo_connection_id: Optional[str] = None, + ) -> dict[str, str]: """Build HTTP headers for UiPath Gateway requests.""" headers = { "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 job_key := os.getenv("UIPATH_JOB_KEY"): headers["X-UiPath-JobKey"] = job_key if process_key := os.getenv("UIPATH_PROCESS_KEY"): diff --git a/tests/chat/__init__.py b/tests/chat/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/chat/test_openai_url_rewrite.py b/tests/chat/test_openai_url_rewrite.py new file mode 100644 index 000000000..89fd0a53f --- /dev/null +++ b/tests/chat/test_openai_url_rewrite.py @@ -0,0 +1,82 @@ +import httpx + +from uipath_langchain.chat.openai import _rewrite_openai_url + + +class TestRewriteOpenAIUrl: + """Tests for the _rewrite_openai_url function.""" + + def test_rewrite_deployments_url(self): + """Test rewriting URLs with /openai/deployments/ pattern (responses: false).""" + original_url = "https://cloud.uipath.com/account/tenant/agenthub_/llm/raw/vendor/openai/model/gpt-5-mini-2025-08-07/openai/deployments/gpt-5-mini-2025-08-07/chat/completions?api-version=2024-12-01-preview" + params = httpx.QueryParams({"api-version": "2024-12-01-preview"}) + + result = _rewrite_openai_url(original_url, params) + + assert result is not None + assert ( + str(result) + == "https://cloud.uipath.com/account/tenant/agenthub_/llm/raw/vendor/openai/model/gpt-5-mini-2025-08-07/completions?api-version=2024-12-01-preview" + ) + + def test_rewrite_responses_url(self): + """Test rewriting URLs with /openai/responses pattern (responses: true).""" + original_url = "https://cloud.uipath.com/account/tenant/agenthub_/llm/raw/vendor/openai/model/gpt-5-mini-2025-08-07/openai/responses?api-version=2024-12-01-preview" + params = httpx.QueryParams({"api-version": "2024-12-01-preview"}) + + result = _rewrite_openai_url(original_url, params) + + assert result is not None + assert ( + str(result) + == "https://cloud.uipath.com/account/tenant/agenthub_/llm/raw/vendor/openai/model/gpt-5-mini-2025-08-07/completions?api-version=2024-12-01-preview" + ) + + def test_rewrite_base_url_with_query_params(self): + """Test rewriting base URL with query params (responses API base case).""" + original_url = "https://cloud.uipath.com/account/tenant/agenthub_/llm/raw/vendor/openai/model/gpt-5-mini-2025-08-07?api-version=2024-12-01-preview" + params = httpx.QueryParams({"api-version": "2024-12-01-preview"}) + + result = _rewrite_openai_url(original_url, params) + + assert result is not None + assert ( + str(result) + == "https://cloud.uipath.com/account/tenant/agenthub_/llm/raw/vendor/openai/model/gpt-5-mini-2025-08-07/completions?api-version=2024-12-01-preview" + ) + + def test_rewrite_without_query_params(self): + """Test rewriting URL without query parameters.""" + original_url = "https://cloud.uipath.com/account/tenant/agenthub_/llm/raw/vendor/openai/model/gpt-5-mini-2025-08-07/openai/responses" + params = httpx.QueryParams() + + result = _rewrite_openai_url(original_url, params) + + assert result is not None + assert ( + str(result) + == "https://cloud.uipath.com/account/tenant/agenthub_/llm/raw/vendor/openai/model/gpt-5-mini-2025-08-07/completions" + ) + + def test_rewrite_localhost_url(self): + """Test rewriting localhost URL.""" + original_url = "https://localhost:7024/account/tenant/llm/raw/vendor/openai/model/gpt-5-mini-2025-08-07/openai/deployments/gpt-5-mini-2025-08-07/chat/completions" + params = httpx.QueryParams() + + result = _rewrite_openai_url(original_url, params) + + assert result is not None + assert ( + str(result) + == "https://localhost:7024/account/tenant/llm/raw/vendor/openai/model/gpt-5-mini-2025-08-07/completions" + ) + + def test_rewrite_preserves_different_api_versions(self): + """Test that different api-version values are preserved.""" + original_url = "https://cloud.uipath.com/account/tenant/agenthub_/llm/raw/vendor/openai/model/gpt-5-mini-2025-08-07/openai/responses?api-version=2025-04-01-preview" + params = httpx.QueryParams({"api-version": "2025-04-01-preview"}) + + result = _rewrite_openai_url(original_url, params) + + assert result is not None + assert "api-version=2025-04-01-preview" in str(result) diff --git a/uv.lock b/uv.lock index 8e10dda29..2ab5af58e 100644 --- a/uv.lock +++ b/uv.lock @@ -3260,7 +3260,7 @@ wheels = [ [[package]] name = "uipath-langchain" -version = "0.1.31" +version = "0.1.32" source = { editable = "." } dependencies = [ { name = "aiosqlite" },