fix: add embeddings; misc

Author: hexiaochun

volcenginesdkarkruntime/_base_client.py

@@ -49,7 +49,6 @@ def __init__(self, **kwargs: Any) -> None:
         kwargs.setdefault("timeout", DEFAULT_TIMEOUT)
         kwargs.setdefault("limits", DEFAULT_CONNECTION_LIMITS)
         kwargs.setdefault("follow_redirects", True)
-        kwargs.setdefault("base_url", BASE_URL)
         super().__init__(**kwargs)
 
 

volcenginesdkarkruntime/_client.py

@@ -39,6 +39,7 @@ def __init__(
         ak: str | None = None,
         sk: str | None = None,
         api_key: str | None = None,
+        region: str = "cn-beijing",
         timeout: float | Timeout | None = DEFAULT_TIMEOUT,
         max_retries: int = DEFAULT_MAX_RETRIES,
         http_client: Client | None = None,
@@ -66,6 +67,7 @@ def __init__(
         self.ak = ak
         self.sk = sk
         self.api_key = api_key
+        self.region = region
 
         assert (api_key is not None) or (ak is not None and sk is not None), "you need to support api_key or ak&sk"
 
@@ -81,12 +83,15 @@ def __init__(
         self._sts_token_manager: StsTokenManager | None = None
 
         self.chat = resources.Chat(self)
+        self.embeddings = resources.Embeddings(self)
+        # self.tokenization = resources.Tokenization(self)
+        # self.classification = resources.Classification(self)
 
     def _get_endpoint_sts_token(self, endpoint_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._sts_token_manager = StsTokenManager(self.ak, self.sk, self.region)
         return self._sts_token_manager.get(endpoint_id)
 
     @property
@@ -105,6 +110,7 @@ def __init__(
         sk: str | None = None,
         api_key: str | None = None,
         base_url: str | URL = BASE_URL,
+        region: str = "cn-beijing",
         timeout: float | Timeout | None = DEFAULT_TIMEOUT,
         max_retries: int = DEFAULT_MAX_RETRIES,
         http_client: AsyncClient | None = None,
@@ -131,6 +137,7 @@ def __init__(
         self.ak = ak
         self.sk = sk
         self.api_key = api_key
+        self.region = region
 
         assert (api_key is not None) or (ak is not None and sk is not None), "you need to support api_key or ak&sk"
 
@@ -146,12 +153,15 @@ def __init__(
         self._sts_token_manager: StsTokenManager | None = None
 
         self.chat = resources.AsyncChat(self)
+        self.embeddings = resources.AsyncEmbeddings(self)
+        # self.tokenization = resources.AsyncTokenization(self)
+        # self.classification = resources.AsyncClassification(self)
 
     def _get_endpoint_sts_token(self, endpoint_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._sts_token_manager = StsTokenManager(self.ak, self.sk, self.region)
         return self._sts_token_manager.get(endpoint_id)
 
     @property
@@ -169,7 +179,7 @@ class StsTokenManager(object):
     # refreshed credentials.
     _mandatory_refresh_timeout: int = _DEFAULT_MANDATORY_REFRESH_TIMEOUT
 
-    def __init__(self, ak: str, sk: str):
+    def __init__(self, ak: str, sk: str, region: str):
         self._endpoint_sts_tokens: Dict[str, Tuple[str, int]] = defaultdict(lambda: ("", 0))
         self._refresh_lock = threading.Lock()
 
@@ -178,7 +188,8 @@ def __init__(self, ak: str, sk: str):
         configuration = volcenginesdkcore.Configuration()
         configuration.ak = ak
         configuration.sk = sk
-        configuration.region = "cn-beijing"
+        configuration.region = region
+        configuration.schema = "https"
 
         volcenginesdkcore.Configuration.set_default(configuration)
         self.api_instance = volcenginesdkark.ARKApi()
@@ -190,8 +201,8 @@ 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):
-        if ttl < _DEFAULT_ADVISORY_REFRESH_TIMEOUT * 2:
-            raise ArkAPIError("ttl should not be under {} seconds.".format(_DEFAULT_ADVISORY_REFRESH_TIMEOUT * 2))
+        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(

volcenginesdkarkruntime/_constants.py

@@ -11,7 +11,7 @@
 
 # default timeout is 1 minutes
 DEFAULT_TIMEOUT = httpx.Timeout(timeout=60.0, connect=60.0)
-DEFAULT_MAX_RETRIES = 1
+DEFAULT_MAX_RETRIES = 2
 DEFAULT_CONNECTION_LIMITS = httpx.Limits(
     max_connections=1000, max_keepalive_connections=100
 )

volcenginesdkarkruntime/_exceptions.py

@@ -66,7 +66,7 @@ def __init__(
             self.type = None
 
     def __str__(self):
-        return f"{self.message} (request_id: {self.request_id})"
+        return f"{self.message}, request_id: {self.request_id}"
 
 
 class ArkAPIResponseValidationError(ArkAPIError):

volcenginesdkarkruntime/_types.py

@@ -62,3 +62,7 @@ def __repr__(self) -> str:
 
 
 NOT_GIVEN = NotGiven()
+
+Headers = Dict[str, str]
+Query = Dict[str, object]
+Body = Dict[str, object]

volcenginesdkarkruntime/_utils/_utils.py

@@ -82,4 +82,5 @@ def _insert_sts_token(args, kwargs):
     model = kwargs.get("model", "")
     if ark_client.api_key is None and model and model.startswith("ep-") and ark_client.ak and ark_client.sk:
         default_auth_header = {"Authorization": "Bearer " + ark_client._get_endpoint_sts_token(model)}
-        kwargs["extra_headers"] = {**default_auth_header, **kwargs.get("extra_headers", {})}
+        extra_headers = kwargs.get("extra_headers") if kwargs.get("extra_headers") else {}
+        kwargs["extra_headers"] = {**default_auth_header, **extra_headers}

volcenginesdkarkruntime/resources/__init__.py

@@ -1,3 +1,13 @@
 from .chat import Chat, AsyncChat
+from .embeddings import Embeddings, AsyncEmbeddings
+from .tokenization import Tokenization, AsyncTokenization
+from .classification import Classification, AsyncClassification
 
-__all__ = ["Chat", "AsyncChat"]
+__all__ = [
+    "Chat",
+    "AsyncChat",
+    "Embeddings",
+    "AsyncEmbeddings",
+    "Tokenization",
+    "AsyncTokenization"
+]

volcenginesdkarkruntime/resources/chat/completions.py

@@ -5,6 +5,7 @@
 import httpx
 from typing_extensions import Literal
 
+from ..._types import Body, Query, Headers
 from ..._utils._utils import with_sts_token, async_with_sts_token
 from ..._base_client import make_request_options
 from ..._resource import SyncAPIResource, AsyncAPIResource
@@ -59,9 +60,9 @@ def create(
         top_logprobs: Optional[int] | None = None,
         top_p: Optional[float] | None = None,
         user: str | None = None,
-        extra_headers: Dict[str, str] | None = None,
-        extra_query: Dict[str, object] | None = None,
-        extra_body: Dict[str, object] | None = None,
+        extra_headers: Headers | None = None,
+        extra_query: Query | None = None,
+        extra_body: Body | None = None,
         timeout: float | httpx.Timeout | None = None,
     ) -> ChatCompletion | Stream[ChatCompletionChunk]:
         return self._post(
@@ -127,9 +128,9 @@ async def create(
         top_logprobs: Optional[int] | None = None,
         top_p: Optional[float] | None = None,
         user: str | None = None,
-        extra_headers: Dict[str, str] | None = None,
-        extra_query: Dict[str, object] | None = None,
-        extra_body: Dict[str, object] | None = None,
+        extra_headers: Headers | None = None,
+        extra_query: Query | None = None,
+        extra_body: Body | None = None,
         timeout: float | httpx.Timeout | None = None,
     ) -> ChatCompletion | AsyncStream[ChatCompletionChunk]:
         return await self._post(

volcenginesdkarkruntime/resources/classification.py

@@ -0,0 +1,109 @@
+from __future__ import annotations
+
+from typing import List
+
+import httpx
+
+from .._base_client import (
+    make_request_options,
+)
+from .._compat import cached_property
+from .._resource import SyncAPIResource, AsyncAPIResource
+from .._response import to_raw_response_wrapper, async_to_raw_response_wrapper
+from .._types import Body, Query, Headers
+from .._utils._utils import with_sts_token, async_with_sts_token
+from ..types.create_classification_response import CreateClassificationResponse
+
+__all__ = ["Classification", "AsyncClassification"]
+
+
+class Classification(SyncAPIResource):
+    @cached_property
+    def with_raw_response(self) -> ClassificationWithRawResponse:
+        return ClassificationWithRawResponse(self)
+
+    @with_sts_token
+    def create(
+        self,
+        *,
+        query: str,
+        model: str,
+        labels: List[str],
+        user: str | None = None,
+        # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
+        # The extra values given here take precedence over values defined on the client or passed to this method.
+        extra_headers: Headers | None = None,
+        extra_query: Query | None = None,
+        extra_body: Body | None = None,
+        timeout: float | httpx.Timeout | None = None,
+    ) -> CreateClassificationResponse:
+        return self._post(
+            "/classification",
+            body={
+                "query": query,
+                "model": model,
+                "labels": labels
+            },
+            options=make_request_options(
+                extra_headers=extra_headers,
+                extra_query=extra_query,
+                extra_body=extra_body,
+                timeout=timeout,
+            ),
+            cast_to=CreateClassificationResponse,
+        )
+
+
+class AsyncClassification(AsyncAPIResource):
+    @cached_property
+    def with_raw_response(self) -> AsyncClassificationWithRawResponse:
+        return AsyncClassificationWithRawResponse(self)
+
+    @async_with_sts_token
+    async def create(
+        self,
+        *,
+        query: str,
+        model: str,
+        labels: List[str],
+        user: str | None = None,
+        # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs.
+        # The extra values given here take precedence over values defined on the client or passed to this method.
+        extra_headers: Headers | None = None,
+        extra_query: Query | None = None,
+        extra_body: Body | None = None,
+        timeout: float | httpx.Timeout | None = None,
+    ) -> CreateClassificationResponse:
+        return await self._post(
+            "/classification",
+            body={
+                "query": query,
+                "model": model,
+                "labels": labels
+            },
+            options=make_request_options(
+                extra_headers=extra_headers,
+                extra_query=extra_query,
+                extra_body=extra_body,
+                timeout=timeout,
+            ),
+            cast_to=CreateClassificationResponse,
+        )
+
+
+class ClassificationWithRawResponse:
+    def __init__(self, classification: Classification) -> None:
+        self._classification = classification
+
+        self.create = to_raw_response_wrapper(
+            classification.create,
+        )
+
+
+class AsyncClassificationWithRawResponse:
+    def __init__(self, classification: AsyncClassification) -> None:
+        self._classification = classification
+
+        self.create = async_to_raw_response_wrapper(
+            classification.create,
+        )