Merge branch 'feat/add_bot_sdk' into 'master'

feat(ark): add bot chat sdk

See merge request iaasng/volcengine-python-sdk!328

Author: 修杰

meta.json

@@ -1,4 +1,4 @@
 {
-	"lasted": "1.0.88",
+	"lasted": "1.0.89",
 	"meta_commit": "daee4996c2a9d1fb908eeaabcace8611ebae30f3"
 }

setup.py

@@ -3,7 +3,7 @@
 from setuptools import setup, find_packages  # noqa: H301
 
 NAME = "volcengine-python-sdk"
-VERSION = "1.0.88"
+VERSION = "1.0.89"
 # To install the library, run the following
 #
 # python setup.py install

volcenginesdkarkruntime/_client.py

@@ -22,6 +22,7 @@
     _DEFAULT_ADVISORY_REFRESH_TIMEOUT,
     _DEFAULT_MANDATORY_REFRESH_TIMEOUT,
     _DEFAULT_STS_TIMEOUT,
+    _DEFAULT_RESOURCE_TYPE,
     DEFAULT_TIMEOUT
 )
 from ._streaming import Stream
@@ -31,6 +32,8 @@
 
 class Ark(SyncAPIClient):
     chat: resources.Chat
+    bot_chat: resources.BotChat
+    embeddings: resources.Embeddings
 
     def __init__(
         self,
@@ -83,6 +86,7 @@ def __init__(
         self._sts_token_manager: StsTokenManager | None = None
 
         self.chat = resources.Chat(self)
+        self.bot_chat = resources.BotChat(self)
         self.embeddings = resources.Embeddings(self)
         # self.tokenization = resources.Tokenization(self)
         # self.classification = resources.Classification(self)
@@ -94,6 +98,13 @@ def _get_endpoint_sts_token(self, endpoint_id: str):
             self._sts_token_manager = StsTokenManager(self.ak, self.sk, self.region)
         return self._sts_token_manager.get(endpoint_id)
 
+    def _get_bot_sts_token(self, bot_id: str):
+        if self._sts_token_manager is None:
+            if self.ak is None or self.sk is None:
+                raise ArkAPIError("must set ak and sk before get endpoint token.")
+            self._sts_token_manager = StsTokenManager(self.ak, self.sk, self.region)
+        return self._sts_token_manager.get(bot_id, resource_type="bot")
+
     @property
     def auth_headers(self) -> dict[str, str]:
         api_key = self.api_key
@@ -102,6 +113,8 @@ def auth_headers(self) -> dict[str, str]:
 
 class AsyncArk(AsyncAPIClient):
     chat: resources.AsyncChat
+    bot_chat: resources.AsyncBotChat
+    embeddings: resources.AsyncEmbeddings
 
     def __init__(
         self,
@@ -153,6 +166,7 @@ def __init__(
         self._sts_token_manager: StsTokenManager | None = None
 
         self.chat = resources.AsyncChat(self)
+        self.bot_chat = resources.AsyncBotChat(self)
         self.embeddings = resources.AsyncEmbeddings(self)
         # self.tokenization = resources.AsyncTokenization(self)
         # self.classification = resources.AsyncClassification(self)
@@ -171,7 +185,6 @@ def auth_headers(self) -> dict[str, str]:
 
 
 class StsTokenManager(object):
-
     # The time at which we'll attempt to refresh, but not
     # block if someone else is refreshing.
     _advisory_refresh_timeout: int = _DEFAULT_ADVISORY_REFRESH_TIMEOUT
@@ -200,13 +213,14 @@ def _need_refresh(self, ep: str, refresh_in: int | None = None) -> bool:
 
         return self._endpoint_sts_tokens[ep][1] - time.time() < refresh_in
 
-    def _protected_refresh(self, ep: str, ttl: int = _DEFAULT_STS_TIMEOUT, is_mandatory: bool = False):
+    def _protected_refresh(self, ep: str, ttl: int = _DEFAULT_STS_TIMEOUT, is_mandatory: bool = False,
+                           resource_type: str = _DEFAULT_RESOURCE_TYPE):
         if ttl < self._advisory_refresh_timeout * 2:
             raise ArkAPIError("ttl should not be under {} seconds.".format(self._advisory_refresh_timeout * 2))
 
         try:
             api_key, expired_time = self._load_api_key(
-                ep, ttl
+                ep, ttl, resource_type=resource_type
             )
             self._endpoint_sts_tokens[ep] = (api_key, expired_time)
         except ApiException as e:
@@ -215,7 +229,7 @@ def _protected_refresh(self, ep: str, ttl: int = _DEFAULT_STS_TIMEOUT, is_mandat
             else:
                 logging.error("load api key cause error: e={}".format(e))
 
-    def _refresh(self, ep: str):
+    def _refresh(self, ep: str, resource_type: str = _DEFAULT_RESOURCE_TYPE):
         if not self._need_refresh(ep, self._advisory_refresh_timeout):
             return
 
@@ -228,7 +242,7 @@ def _refresh(self, ep: str):
                     ep, self._mandatory_refresh_timeout
                 )
 
-                self._protected_refresh(ep, is_mandatory=is_mandatory_refresh)
+                self._protected_refresh(ep, is_mandatory=is_mandatory_refresh, resource_type=resource_type)
                 return
             finally:
                 self._refresh_lock.release()
@@ -237,16 +251,17 @@ def _refresh(self, ep: str):
                 if not self._need_refresh(ep, self._mandatory_refresh_timeout):
                     return
 
-                self._protected_refresh(ep, is_mandatory=True)
+                self._protected_refresh(ep, is_mandatory=True, resource_type=resource_type)
 
-    def get(self, ep: str) -> str:
-        self._refresh(ep)
+    def get(self, ep: str, resource_type: str = _DEFAULT_RESOURCE_TYPE) -> str:
+        self._refresh(ep, resource_type=resource_type)
         return self._endpoint_sts_tokens[ep][0]
 
-    def _load_api_key(self, ep: str, duration_seconds: int) -> Tuple[str, int]:
+    def _load_api_key(self, ep: str, duration_seconds: int,
+                      resource_type: str = _DEFAULT_RESOURCE_TYPE) -> Tuple[str, int]:
         get_api_key_request = volcenginesdkark.GetApiKeyRequest(
             duration_seconds=duration_seconds,
-            resource_type="endpoint",
+            resource_type=resource_type,
             resource_ids=[ep],
         )
         resp: volcenginesdkark.GetApiKeyResponse = self.api_instance.get_api_key(

volcenginesdkarkruntime/_constants.py

@@ -22,3 +22,5 @@
 _DEFAULT_MANDATORY_REFRESH_TIMEOUT = 10 * 60  # 10 min
 _DEFAULT_ADVISORY_REFRESH_TIMEOUT = 30 * 60  # 30 min
 _DEFAULT_STS_TIMEOUT = 7 * 24 * 60 * 60  # 7 days
+
+_DEFAULT_RESOURCE_TYPE = "endpoint"

volcenginesdkarkruntime/_utils/_utils.py

@@ -84,3 +84,7 @@ def _insert_sts_token(args, kwargs):
         default_auth_header = {"Authorization": "Bearer " + ark_client._get_endpoint_sts_token(model)}
         extra_headers = kwargs.get("extra_headers") if kwargs.get("extra_headers") else {}
         kwargs["extra_headers"] = {**default_auth_header, **extra_headers}
+    elif ark_client.api_key is None and model and model.startswith("bot-") and ark_client.ak and ark_client.sk:
+        default_auth_header = {"Authorization": "Bearer " + ark_client._get_bot_sts_token(model)}
+        extra_headers = kwargs.get("extra_headers") if kwargs.get("extra_headers") else {}
+        kwargs["extra_headers"] = {**default_auth_header, **extra_headers}

volcenginesdkarkruntime/resources/__init__.py

@@ -2,10 +2,13 @@
 from .embeddings import Embeddings, AsyncEmbeddings
 from .tokenization import Tokenization, AsyncTokenization
 from .classification import Classification, AsyncClassification
+from .bot import BotChat, AsyncBotChat
 
 __all__ = [
     "Chat",
+    "BotChat",
     "AsyncChat",
+    "AsyncBotChat",
     "Embeddings",
     "AsyncEmbeddings",
     "Tokenization",

volcenginesdkarkruntime/resources/bot/__init__.py

@@ -0,0 +1,3 @@
+from .chat import BotChat, AsyncBotChat
+
+__all__ = ["BotChat", "AsyncBotChat"]

volcenginesdkarkruntime/resources/bot/chat.py

@@ -0,0 +1,21 @@
+# File generated from our OpenAPI spec by Stainless. See CONTRIBUTING.md for details.
+
+from __future__ import annotations
+
+from .completions import Completions, AsyncCompletions
+from ..._compat import cached_property
+from ..._resource import SyncAPIResource, AsyncAPIResource
+
+__all__ = ["BotChat", "AsyncBotChat"]
+
+
+class BotChat(SyncAPIResource):
+    @cached_property
+    def completions(self) -> Completions:
+        return Completions(self._client)
+
+
+class AsyncBotChat(AsyncAPIResource):
+    @cached_property
+    def completions(self) -> AsyncCompletions:
+        return AsyncCompletions(self._client)