feat(*): support ark batch chat

jizhongBao · jizhongBao · commit c0a21a724abe · 2025-01-15T11:20:04.000+08:00
diff --git a/volcenginesdkarkruntime/_base_client.py b/volcenginesdkarkruntime/_base_client.py
@@ -478,6 +478,57 @@ def _request(
             stream_cls=stream_cls,
         )
 
+    def _request_without_retry(self,
+            *,
+            cast_to: Type[ResponseT],
+            options: RequestOptions,
+            stream: bool,
+            stream_cls: type[_StreamT] | None,
+    ) -> ResponseT | _StreamT:
+        request = self._build_request(options)
+        req_id = request.headers.get(CLIENT_REQUEST_HEADER, "")
+        try:
+            response = self._client.send(
+                request,
+                stream=stream or self._should_stream_response_body(request=request),
+            )
+        except httpx.TimeoutException as err:
+            log.debug("Raising timeout error")
+            raise ArkAPITimeoutError(request=request, request_id=req_id) from err
+        except Exception as err:
+            log.debug("Encountered Exception", exc_info=True)
+            log.debug("Raising connection error")
+            raise ArkAPIConnectionError(request=request, request_id=req_id) from err
+
+        log.debug(
+            'HTTP Request: %s %s "%i %s"',
+            request.method,
+            request.url,
+            response.status_code,
+            response.reason_phrase,
+        )
+
+        try:
+            response.raise_for_status()
+        except httpx.HTTPStatusError as err:  # thrown on 4xx and 5xx status code
+            log.debug("Encountered httpx.HTTPStatusError", exc_info=True)
+            # If the response is streamed then we need to explicitly read the response
+            # to completion before attempting to access the response text.
+            if not err.response.is_closed:
+                err.response.read()
+
+            log.debug("Re-raising status error")
+            raise self._make_status_error_from_response(
+                err.response, request_id=req_id
+            ) from None
+
+        return self._process_response(
+            cast_to=cast_to,
+            response=response,
+            stream=stream,
+            stream_cls=stream_cls,
+        )
+
     def _retry_request(
             self,
             options: RequestOptions,
@@ -595,6 +646,28 @@ def delete(
 
         return cast(ResponseT, self.request(cast_to, opts))
 
+    def post_without_retry(
+            self,
+            path: str,
+            *,
+            cast_to: Type[ResponseT],
+            body: Dict | None = None,
+            options: ExtraRequestOptions = {},
+            files: RequestFiles | None = None,
+            stream: bool = False,
+            stream_cls: type[_StreamT] | None = None,
+    ) -> ResponseT | _StreamT:
+        opts = RequestOptions.construct(  # type: ignore
+            method="post",
+            url=path,
+            body=body,
+            **options,
+        )
+
+        return cast(
+            ResponseT, self.request_without_retry(cast_to, opts, stream=stream, stream_cls=stream_cls)
+        )
+
     def request(
             self,
             cast_to: Type[ResponseT],
@@ -612,6 +685,21 @@ def request(
             remaining_retries=remaining_retries,
         )
 
+    def request_without_retry(
+            self,
+            cast_to: Type[ResponseT],
+            options: RequestOptions,
+            *,
+            stream: bool = False,
+            stream_cls: type[_StreamT] | None = None,
+    ) -> ResponseT | _StreamT:
+        return self._request_without_retry(
+            cast_to=cast_to,
+            options=options,
+            stream=stream,
+            stream_cls=stream_cls,
+        )
+
     def is_closed(self) -> bool:
         return self._client.is_closed
 
@@ -755,6 +843,27 @@ async def delete(
 
         return await self.request(cast_to, opts)
 
+
+    async def post_without_retry(
+            self,
+            path: str,
+            *,
+            cast_to: Type[ResponseT],
+            body: Dict | None = None,
+            options: ExtraRequestOptions = {},
+            files: RequestFiles | None = None,
+            stream: bool = False,
+            stream_cls: type[_AsyncStreamT] | None = None,
+    ) -> ResponseT | _AsyncStreamT:
+        opts = RequestOptions.construct(
+            method="post",
+            url=path,
+            body=body,
+            **options,
+        )
+
+        return await self.request_without_retry(cast_to, opts, stream=stream, stream_cls=stream_cls)
+
     async def request(
             self,
             cast_to: Type[ResponseT],
@@ -772,6 +881,21 @@ async def request(
             remaining_retries=remaining_retries,
         )
 
+    async def request_without_retry(
+            self,
+            cast_to: Type[ResponseT],
+            options: RequestOptions,
+            *,
+            stream: bool = False,
+            stream_cls: type[_StreamT] | None = None,
+    ) -> ResponseT | _StreamT:
+        return await self._request_without_retry(
+            cast_to=cast_to,
+            options=options,
+            stream=stream,
+            stream_cls=stream_cls,
+        )
+
     async def _request(
             self,
             *,
@@ -859,6 +983,57 @@ async def _request(
             stream_cls=stream_cls,
         )
 
+    async def _request_without_retry(
+            self,
+            *,
+            cast_to: Type[ResponseT],
+            options: RequestOptions,
+            stream: bool,
+            stream_cls: type[_AsyncStreamT] | None,
+    ) -> ResponseT | _AsyncStreamT:
+        request = self._build_request(options)
+        req_id = request.headers.get(CLIENT_REQUEST_HEADER, "")
+        try:
+            response = await self._client.send(
+                request,
+                stream=stream or self._should_stream_response_body(request=request),
+            )
+        except httpx.TimeoutException as err:
+            log.debug("Raising timeout error")
+            raise ArkAPITimeoutError(request=request, request_id=req_id) from err
+        except Exception as err:
+            log.debug("Encountered Exception", exc_info=True)
+            log.debug("Raising connection error")
+            raise ArkAPIConnectionError(request=request, request_id=req_id) from err
+        log.debug(
+            'HTTP Request: %s %s "%i %s"',
+            request.method,
+            request.url,
+            response.status_code,
+            response.reason_phrase,
+        )
+        try:
+            response.raise_for_status()
+        except httpx.HTTPStatusError as err:  # thrown on 4xx and 5xx status code
+            log.debug("Encountered httpx.HTTPStatusError", exc_info=True)
+
+            # If the response is streamed then we need to explicitly read the response
+            # to completion before attempting to access the response text.
+            if not err.response.is_closed:
+                await err.response.aread()
+
+            log.debug("Re-raising status error")
+            raise self._make_status_error_from_response(
+                err.response, request_id=req_id
+            ) from None
+
+        return await self._process_response(
+            cast_to=cast_to,
+            response=response,
+            stream=stream,
+            stream_cls=stream_cls,
+        )
+
     async def _retry_request(
             self,
             options: RequestOptions,
diff --git a/volcenginesdkarkruntime/_client.py b/volcenginesdkarkruntime/_client.py
@@ -28,6 +28,7 @@
 from ._streaming import Stream
 
 from ._utils._key_agreement import key_agreement_client
+from ._utils._model_breaker import ModelBreaker
 
 __all__ = ["Ark", "AsyncArk"]
 
@@ -39,6 +40,9 @@ class Ark(SyncAPIClient):
     tokenization: resources.Tokenization
     context: resources.Context
     content_generation: resources.ContentGeneration
+    batch_chat: resources.BatchChat
+    model_breaker_map: dict[str, ModelBreaker]
+    model_breaker_lock: threading.Lock
 
     def __init__(
         self,
@@ -98,6 +102,9 @@ def __init__(
         self.tokenization = resources.Tokenization(self)
         self.context = resources.Context(self)
         self.content_generation = resources.ContentGeneration(self)
+        self.batch_chat = resources.BatchChat(self)
+        self.model_breaker_map = defaultdict(ModelBreaker)
+        self.model_breaker_lock = threading.Lock()
         # self.classification = resources.Classification(self)
 
     def _get_endpoint_sts_token(self, endpoint_id: str):
@@ -128,6 +135,9 @@ def auth_headers(self) -> dict[str, str]:
         api_key = self.api_key
         return {"Authorization": f"Bearer {api_key}"}
 
+    def get_model_breaker(self, model_name: str) -> ModelBreaker:
+        with self.model_breaker_lock:
+            return self.model_breaker_map[model_name]
 
 class AsyncArk(AsyncAPIClient):
     chat: resources.AsyncChat
@@ -136,6 +146,9 @@ class AsyncArk(AsyncAPIClient):
     tokenization: resources.AsyncTokenization
     context: resources.AsyncContext
     content_generation: resources.AsyncContentGeneration
+    batch_chat: resources.AsyncBatchChat
+    model_breaker_map: dict[str, ModelBreaker]
+    model_breaker_lock: threading.Lock
 
     def __init__(
         self,
@@ -194,6 +207,9 @@ def __init__(
         self.tokenization = resources.AsyncTokenization(self)
         self.context = resources.AsyncContext(self)
         self.content_generation = resources.AsyncContentGeneration(self)
+        self.batch_chat = resources.AsyncBatchChat(self)
+        self.model_breaker_map = defaultdict(ModelBreaker)
+        self.model_breaker_lock = threading.Lock()
         # self.classification = resources.AsyncClassification(self)
 
     def _get_endpoint_sts_token(self, endpoint_id: str):
@@ -217,6 +233,10 @@ def auth_headers(self) -> dict[str, str]:
         api_key = self.api_key
         return {"Authorization": f"Bearer {api_key}"}
 
+    def get_model_breaker(self, model_name: str) -> ModelBreaker:
+        with self.model_breaker_lock:
+            return self.model_breaker_map[model_name]
+
 
 class StsTokenManager(object):
     # The time at which we'll attempt to refresh, but not
diff --git a/volcenginesdkarkruntime/_constants.py b/volcenginesdkarkruntime/_constants.py
@@ -10,8 +10,11 @@
 SERVER_REQUEST_HEADER = "X-Request-Id"
 ARK_E2E_ENCRYPTION_HEADER = "x-is-encrypted"
 
+DEFAULT_TIMEOUT_SECONDS = 600.0
+DEFAULT_CONNECT_TIMEOUT_SECONDS = 60.0
 # default timeout is 1 minutes
-DEFAULT_TIMEOUT = httpx.Timeout(timeout=600.0, connect=60.0)
+DEFAULT_TIMEOUT = httpx.Timeout(timeout=DEFAULT_TIMEOUT_SECONDS, connect=DEFAULT_CONNECT_TIMEOUT_SECONDS)
+
 DEFAULT_MAX_RETRIES = 2
 DEFAULT_CONNECTION_LIMITS = httpx.Limits(
     max_connections=1000, max_keepalive_connections=100
diff --git a/volcenginesdkarkruntime/_resource.py b/volcenginesdkarkruntime/_resource.py
@@ -12,6 +12,7 @@ def __init__(self, client: "Ark") -> None:
         self._post = client.post
         self._get = client.get
         self._delete = client.delete
+        self._post_without_retry = client.post_without_retry
 
 
 class AsyncAPIResource:
@@ -22,3 +23,4 @@ def __init__(self, client: "AsyncArk") -> None:
         self._post = client.post
         self._get = client.get
         self._delete = client.delete
+        self._post_without_retry = client.post_without_retry
diff --git a/volcenginesdkarkruntime/_utils/_model_breaker.py b/volcenginesdkarkruntime/_utils/_model_breaker.py
@@ -0,0 +1,23 @@
+from datetime import datetime, timedelta
+
+
+class ModelBreaker:
+    def __init__(self):
+        # 初始化 allow_time 为当前时间
+        self.allow_time = datetime.now()
+
+    def allow(self):
+        # 检查当前时间是否在 allow_time 之后
+        return datetime.now() > self.allow_time
+
+    def reset(self, duration):
+        # 将 allow_time 重置为当前时间加上指定的持续时间
+        self.allow_time = datetime.now() + timedelta(seconds=duration.total_seconds())
+
+    def get_allowed_duration(self):
+        # 计算当前时间与 allow_time 之间的持续时间
+        allow_duration = self.allow_time - datetime.now()
+        # 如果持续时间为负，则返回一个零时长的 timedelta 对象
+        if allow_duration.total_seconds() < 0:
+            return timedelta(0)
+        return allow_duration
diff --git a/volcenginesdkarkruntime/resources/__init__.py b/volcenginesdkarkruntime/resources/__init__.py
@@ -5,6 +5,7 @@
 from .bot import BotChat, AsyncBotChat
 from .context import Context, AsyncContext
 from .content_generation import ContentGeneration, AsyncContentGeneration
+from .batch_chat import BatchChat, AsyncBatchChat
 
 __all__ = [
     "Chat",
@@ -19,4 +20,6 @@
     "AsyncContext",
     "ContentGeneration",
     "AsyncContentGeneration"
+    "BatchChat",
+    "AsyncBatchChat"
 ]
diff --git a/volcenginesdkarkruntime/resources/batch_chat/__init__.py b/volcenginesdkarkruntime/resources/batch_chat/__init__.py
@@ -0,0 +1,3 @@
+from .chat import BatchChat, AsyncBatchChat
+
+__all__ = ["BatchChat", "AsyncBatchChat"]
diff --git a/volcenginesdkarkruntime/resources/batch_chat/chat.py b/volcenginesdkarkruntime/resources/batch_chat/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__ = ["BatchChat", "AsyncBatchChat"]
+
+
+class BatchChat(SyncAPIResource):
+    @cached_property
+    def completions(self) -> Completions:
+        return Completions(self._client)
+
+
+class AsyncBatchChat(AsyncAPIResource):
+    @cached_property
+    def completions(self) -> AsyncCompletions:
+        return AsyncCompletions(self._client)
diff --git a/volcenginesdkarkruntime/resources/batch_chat/completions.py b/volcenginesdkarkruntime/resources/batch_chat/completions.py
diff --git a/volcenginesdkexamples/volcenginesdkarkruntime/async_batch_chat_completions.py b/volcenginesdkexamples/volcenginesdkarkruntime/async_batch_chat_completions.py
diff --git a/volcenginesdkexamples/volcenginesdkarkruntime/batch_chat_completions.py b/volcenginesdkexamples/volcenginesdkarkruntime/batch_chat_completions.py

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+from .chat import BatchChat, AsyncBatchChat`
	`2`	`+`
	`3`	`+__all__ = ["BatchChat", "AsyncBatchChat"]`