Implement get_async_http_client function for async Foundry HTTP clients (#341)

Author: daankolthof

assets/language_model_utils.py

@@ -19,7 +19,7 @@
 from foundry_sdk._core.config import Config
 from foundry_sdk._core.context_and_environment_vars import HOSTNAME_VAR
 from foundry_sdk._core.context_and_environment_vars import TOKEN_VAR
-from foundry_sdk._core.http_client import HttpClient
+from foundry_sdk._core.http_client import HttpClient, AsyncHttpClient
 
 
 def _get_api_gateway_base_url(*, preview: bool = False) -> str:
@@ -125,20 +125,58 @@ def get_http_client(*, preview: bool = False, config: Optional[Config] = None) -
 
     Raises:
         ValueError: If preview is not set to True.
-        RuntimeError: If the Foundry API gateway base URL or token is not available in the current context.
+        RuntimeError: If the Foundry token is not available in the current context.
     """
     if not preview:
         raise ValueError(
             "get_http_client() is in beta. " "Please set the preview parameter to True to use it."
         )
-    token = get_foundry_token(preview=True)
 
-    # Merge auth header with any user-provided headers
+    return HttpClient(config=_create_http_client_config(config))
+
+
+def get_async_http_client(
+    *, preview: bool = False, config: Optional[Config] = None
+) -> AsyncHttpClient:
+    """Get an async HTTP client configured for the current Foundry environment.
+
+    Args:
+        preview: Must be set to True to use this beta feature.
+        config: Optional configuration for the async HTTP client.
+
+    Returns:
+        An AsyncHttpClient instance configured with the Foundry hostname and authentication.
+
+    Raises:
+        ValueError: If preview is not set to True.
+        RuntimeError: If the Foundry token is not available in the current context.
+    """
+    if not preview:
+        raise ValueError(
+            "get_async_http_client() is in beta. "
+            "Please set the preview parameter to True to use it."
+        )
+
+    return AsyncHttpClient(config=_create_http_client_config(config))
+
+
+def _create_http_client_config(config: Optional[Config]) -> Config:
+    """Create a Config object for the HTTP client, ensuring that the Foundry token is included in the headers.
+
+    Args:
+        config: An optional Config object provided by the user.
+
+    Returns:
+        A Config object with the Foundry token included in the default headers.
+
+    Raises:
+        RuntimeError: If the Foundry token is not available in the current context.
+    """
+    token = get_foundry_token(preview=True)
     auth_header = {"Authorization": f"Bearer {token}"}
+
     if config is None:
-        config = Config(default_headers=auth_header)
+        return Config(default_headers=auth_header)
     else:
         merged_headers = {**auth_header, **(config.default_headers or {})}
-        config = replace(config, default_headers=merged_headers)
-
-    return HttpClient(config=config)
+        return replace(config, default_headers=merged_headers)

assets/test_language_model_utils.py

@@ -17,11 +17,12 @@
 
 from foundry_sdk._core.context_and_environment_vars import HOSTNAME_VAR
 from foundry_sdk._core.context_and_environment_vars import TOKEN_VAR
-from foundry_sdk._core.http_client import HttpClient
+from foundry_sdk._core.http_client import HttpClient, AsyncHttpClient
 from foundry_sdk.v2.language_models import (
     get_anthropic_base_url,
     get_foundry_token,
     get_http_client,
+    get_async_http_client,
     get_openai_base_url,
 )
 from foundry_sdk.v2.language_models.utils import _get_api_gateway_base_url
@@ -177,3 +178,17 @@ def test_raises_runtime_error_when_token_not_in_context(self):
                 get_http_client(preview=True)
         finally:
             HOSTNAME_VAR.reset(hostname_token)
+
+
+class TestGetAsyncHttpClient:
+    """Test get_async_http_client function."""
+
+    def test_returns_async_http_client(self):
+        hostname_token = HOSTNAME_VAR.set("test.palantirfoundry.com")
+        auth_token = TOKEN_VAR.set("test-token-12345")
+        try:
+            result = get_async_http_client(preview=True)
+            assert isinstance(result, AsyncHttpClient)
+        finally:
+            HOSTNAME_VAR.reset(hostname_token)
+            TOKEN_VAR.reset(auth_token)

config.json

@@ -186,7 +186,8 @@
             "from foundry_sdk.v2.language_models.utils import get_foundry_token",
             "from foundry_sdk.v2.language_models.utils import get_openai_base_url",
             "from foundry_sdk.v2.language_models.utils import get_anthropic_base_url",
-            "from foundry_sdk.v2.language_models.utils import get_http_client"
+            "from foundry_sdk.v2.language_models.utils import get_http_client",
+            "from foundry_sdk.v2.language_models.utils import get_async_http_client"
           ]
         },
         "MapRendering": false,