Merge branch 'feat/ark-runtime-batch' into 'integration_2025-01-16_675380848130'

feat: [development task] ark-runtime-manual-Python (978354)

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

Author: BitsAdmin

volcenginesdkarkruntime/_base_client.py

@@ -595,6 +595,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(cast_to, opts, remaining_retries=0, stream=stream, stream_cls=stream_cls)
+        )
+
     def request(
             self,
             cast_to: Type[ResponseT],
@@ -755,6 +777,26 @@ 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(cast_to, opts, remaining_retries=0, stream=stream, stream_cls=stream_cls)
+
     async def request(
             self,
             cast_to: Type[ResponseT],

volcenginesdkarkruntime/_client.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 import logging
 import os
 import threading
@@ -28,6 +29,7 @@
 from ._streaming import Stream
 
 from ._utils._key_agreement import key_agreement_client
+from ._utils._model_breaker import ModelBreaker
 
 __all__ = ["Ark", "AsyncArk"]
 
@@ -39,6 +41,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 +103,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 +136,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 +147,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: asyncio.Lock
 
     def __init__(
         self,
@@ -194,6 +208,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 = asyncio.Lock()
         # self.classification = resources.AsyncClassification(self)
 
     def _get_endpoint_sts_token(self, endpoint_id: str):
@@ -217,6 +234,10 @@ def auth_headers(self) -> dict[str, str]:
         api_key = self.api_key
         return {"Authorization": f"Bearer {api_key}"}
 
+    async def get_model_breaker(self, model_name: str) -> ModelBreaker:
+        async 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

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

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

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

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"
 ]

volcenginesdkarkruntime/resources/batch_chat/__init__.py

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

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)