fix: Added azure api version in the url

AI21Labs · Jun 21, 2024 · bd49dbc · bd49dbc
1 parent c509b7e
commit bd49dbc
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/ai21/clients/azure/ai21_azure_client.py b/ai21/clients/azure/ai21_azure_client.py
@@ -8,6 +8,7 @@
 from ai21.clients.studio.resources.studio_chat import StudioChat, AsyncStudioChat
 
 AzureADTokenProvider = Callable[[], str]
+_DEFAULT_AZURE_VERSION = "v1"
 
 
 class BaseAzureClient(ABC):
@@ -47,6 +48,7 @@ class AsyncAI21AzureClient(BaseAzureClient, AsyncAI21HTTPClient):
     def __init__(
         self,
         base_url: str,
+        api_version: str = _DEFAULT_AZURE_VERSION,
         api_key: Optional[str] = None,
         azure_ad_token: str | None = None,
         azure_ad_token_provider: AzureADTokenProvider | None = None,
@@ -63,6 +65,9 @@ def __init__(
 
         headers = self._prepare_headers(headers=default_headers or {})
 
+        if api_version:
+            base_url += f"/{api_version}"
+
         super().__init__(
             api_key=api_key,
             base_url=base_url,
@@ -81,6 +86,7 @@ class AI21AzureClient(BaseAzureClient, AI21HTTPClient):
     def __init__(
         self,
         base_url: str,
+        api_version: str = _DEFAULT_AZURE_VERSION,
         api_key: Optional[str] = None,
         azure_ad_token: str | None = None,
         azure_ad_token_provider: AzureADTokenProvider | None = None,
@@ -97,6 +103,9 @@ def __init__(
 
         headers = self._prepare_headers(headers=default_headers or {})
 
+        if api_version:
+            base_url += f"/{api_version}"
+
         super().__init__(
             api_key=api_key,
             base_url=base_url,