feat: via param

AI21Labs · Dec 14, 2023 · d51f101 · d51f101
1 parent 1d820a8
commit d51f101
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 6 deletions.
diff --git a/ai21/ai21_studio_client.py b/ai21/ai21_studio_client.py
@@ -18,6 +18,7 @@ def __init__(
         timeout_sec: Optional[int] = None,
         num_retries: Optional[int] = None,
         organization: Optional[str] = None,
+        via: Optional[str] = None,
         env_config: _AI21EnvConfig = AI21EnvConfig,
     ):
         self._env_config = env_config
@@ -33,6 +34,7 @@ def __init__(
         self._num_retries = num_retries or self._env_config.num_retries
         self._organization = organization or self._env_config.organization
         self._application = self._env_config.application
+        self._via = via
 
         headers = self._build_headers(passed_headers=headers)
 
@@ -54,15 +56,15 @@ def _build_headers(self, passed_headers: Optional[Dict[str, Any]]) -> Dict[str,
 
     def _build_user_agent(self) -> str:
         user_agent = f"ai21 studio SDK {VERSION}"
-        organization = self._organization
 
-        if organization is not None:
-            user_agent = f"{user_agent} organization: {organization}"
+        if self._organization is not None:
+            user_agent = f"{user_agent} organization: {self._organization}"
 
-        application = self._application
+        if self._application is not None:
+            user_agent = f"{user_agent} application: {self._application}"
 
-        if application is not None:
-            user_agent = f"{user_agent} application: {application}"
+        if self._via is not None:
+            user_agent = f"{user_agent} via: {self._via}"
 
         return user_agent
 

diff --git a/ai21/clients/studio/ai21_client.py b/ai21/clients/studio/ai21_client.py
@@ -33,6 +33,7 @@ def __init__(
         headers: Optional[Dict[str, Any]] = None,
         timeout_sec: Optional[float] = None,
         num_retries: Optional[int] = None,
+        via: Optional[str] = None,
         **kwargs,
     ):
         super().__init__(
@@ -42,6 +43,7 @@ def __init__(
             headers=headers,
             timeout_sec=timeout_sec,
             num_retries=num_retries,
+            via=via,
         )
         self.completion = StudioCompletion(self)
         self.chat = StudioChat(self)