update 8

Author: ionmincu

src/uipath_llamaindex/llms/vertex.py

@@ -41,6 +41,31 @@ def _check_vertex_dependencies() -> None:
 from llama_index.llms.google_genai import GoogleGenAI
 
 
+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 = "streamGenerateContent" in url_str
+        # Build headers with streaming flag and correct host
+        headers = dict(request.headers)
+        if is_streaming:
+            headers["X-UiPath-Streaming-Enabled"] = "true"
+        # Update host header to match the gateway URL
+        gateway_url_parsed = httpx.URL(gateway_url)
+        headers["host"] = gateway_url_parsed.host
+        # Create new request with rewritten URL
+        return httpx.Request(
+            method=request.method,
+            url=gateway_url,
+            headers=headers,
+            content=request.content,
+            extensions=request.extensions,
+        )
+    return request
+
+
 class UiPathHttpxClient(httpx.Client):
     """Custom httpx client that redirects generateContent requests to UiPath gateway."""
 
@@ -57,27 +82,30 @@ def request(self, method: str, url: Any, **kwargs) -> httpx.Response:
 
     def send(self, request: httpx.Request, **kwargs) -> httpx.Response:
         """Override send to redirect generateContent/streamGenerateContent to UiPath gateway."""
-        url_str = str(request.url)
-        if "generateContent" in url_str or "streamGenerateContent" in url_str:
-            is_streaming = "streamGenerateContent" in url_str
-            # Build headers with streaming flag and correct host
-            headers = dict(request.headers)
-            if is_streaming:
-                headers["X-UiPath-Streaming-Enabled"] = "true"
-            # Update host header to match the gateway URL
-            gateway_url_parsed = httpx.URL(self.gateway_url)
-            headers["host"] = gateway_url_parsed.host
-            # Create new request with rewritten URL
-            request = httpx.Request(
-                method=request.method,
-                url=self.gateway_url,
-                headers=headers,
-                content=request.content,
-                extensions=request.extensions,
-            )
+        request = _rewrite_request_for_gateway(request, self.gateway_url)
         return super().send(request, **kwargs)
 
 
+class UiPathHttpxAsyncClient(httpx.AsyncClient):
+    """Custom async httpx client that redirects generateContent requests to UiPath gateway."""
+
+    def __init__(self, gateway_url: str, **kwargs):
+        self.gateway_url = gateway_url
+        super().__init__(**kwargs)
+
+    async def request(self, method: str, url: Any, **kwargs) -> httpx.Response:
+        """Override request to redirect generateContent/streamGenerateContent to UiPath gateway."""
+        url_str = str(url)
+        if "generateContent" in url_str or "streamGenerateContent" in url_str:
+            url = self.gateway_url
+        return await super().request(method, url, **kwargs)
+
+    async def send(self, request: httpx.Request, **kwargs) -> httpx.Response:
+        """Override send to redirect generateContent/streamGenerateContent to UiPath gateway."""
+        request = _rewrite_request_for_gateway(request, self.gateway_url)
+        return await super().send(request, **kwargs)
+
+
 class UiPathVertex(GoogleGenAI):
     """
     UiPath Vertex AI LLM that routes requests through UiPath's LLM Gateway.
@@ -144,16 +172,22 @@ def __init__(
         uipath_url = self._build_base_url_static(model)
         headers = self._build_headers_static(token)
 
-        # Create custom httpx client that redirects requests to UiPath gateway
+        # Create custom httpx clients that redirect requests to UiPath gateway
         custom_httpx_client = UiPathHttpxClient(
             gateway_url=uipath_url,
             headers=headers,
             follow_redirects=True,
         )
+        custom_httpx_async_client = UiPathHttpxAsyncClient(
+            gateway_url=uipath_url,
+            headers=headers,
+            follow_redirects=True,
+        )
 
-        # Configure HTTP options with our custom client
+        # Configure HTTP options with our custom clients (sync and async)
         http_options = genai_types.HttpOptions(
             httpxClient=custom_httpx_client,
+            httpxAsyncClient=custom_httpx_async_client,
         )
 
         # Create google.genai client with custom httpx client