feat: support json schema

Author: 潘婉宁

setup.py

@@ -29,7 +29,8 @@
             "httpx>=0.23.0, <1",
             "anyio>=3.5.0, <5",
             "cached-property; python_version < '3.8'",
-            "cryptography>=43.0.3, <43.0.4"
+            "cryptography>=43.0.3, <43.0.4",
+            "jiter>=0.4.0, <1"
         ]
     },
 )

volcenginesdkarkruntime/_base_client.py

@@ -42,11 +42,11 @@
     ArkAPIResponseValidationError,
 )
 from ._models import construct_type
+from ._request_options import RequestOptions, ExtraRequestOptions
 from ._response import ArkAPIResponse, ArkAsyncAPIResponse
 from ._streaming import SSEDecoder, SSEBytesDecoder, Stream, AsyncStream
-from ._types import ResponseT, NotGiven, NOT_GIVEN
-from ._request_options import RequestOptions, ExtraRequestOptions
-from ._utils._utils import _gen_request_id
+from ._types import ResponseT, NotGiven, NOT_GIVEN, PostParser
+from ._utils._utils import _gen_request_id, is_given
 
 _T = TypeVar("_T")
 _StreamT = TypeVar("_StreamT", bound=Stream[Any])
@@ -90,6 +90,7 @@ def make_request_options(
     extra_query: Dict[str, Any] | None = None,
     extra_body: Dict[str, Any] | None = None,
     timeout: float | httpx.Timeout | None = None,
+    post_parser: PostParser | NotGiven = NOT_GIVEN,
 ) -> ExtraRequestOptions:
     options: ExtraRequestOptions = {}
     if extra_headers is not None:
@@ -107,6 +108,10 @@ def make_request_options(
     if timeout:
         options["timeout"] = timeout
 
+    if is_given(post_parser):
+        # internal
+        options["post_parser"] = post_parser  # type: ignore
+
     return options
 
 
@@ -524,6 +529,7 @@ def _process_response(
         self,
         *,
         cast_to: Type[ResponseT],
+        options: RequestOptions,
         response: httpx.Response,
         stream: bool,
         stream_cls: type[Stream[Any]] | type[AsyncStream[Any]] | None,
@@ -537,6 +543,7 @@ def _process_response(
             cast_to=cast("type[ResponseT]", cast_to),  # pyright: ignore[reportUnnecessaryCast]
             stream=stream,
             stream_cls=stream_cls,
+            options=options,
         )
         if bool(response.request.headers.get(RAW_RESPONSE_HEADER)):
             return cast(ResponseT, api_response)
@@ -557,6 +564,7 @@ def post(
         opts = RequestOptions.construct(  # type: ignore
             method="post",
             url=path,
+            files=files,
             body=body,
             **options,
         )
@@ -618,6 +626,7 @@ def post_without_retry(
             method="post",
             url=path,
             body=body,
+            files=files,
             **options,
         )
 
@@ -745,6 +754,7 @@ async def post(
             method="post",
             url=path,
             body=body,
+            files=files,
             **options,
         )
 
@@ -801,6 +811,7 @@ async def post_without_retry(
             method="post",
             url=path,
             body=body,
+            files=files,
             **options,
         )
 
@@ -910,6 +921,7 @@ async def _request(
             response=response,
             stream=stream,
             stream_cls=stream_cls,
+            options=options,
         )
 
     async def _retry_request(
@@ -945,6 +957,7 @@ async def _process_response(
         self,
         *,
         cast_to: Type[ResponseT],
+        options: RequestOptions,
         response: httpx.Response,
         stream: bool,
         stream_cls: type[Stream[Any]] | type[AsyncStream[Any]] | None,
@@ -958,6 +971,7 @@ async def _process_response(
             cast_to=cast("type[ResponseT]", cast_to),  # pyright: ignore[reportUnnecessaryCast]
             stream=stream,
             stream_cls=stream_cls,
+            options=options,
         )
         if bool(response.request.headers.get(RAW_RESPONSE_HEADER)):
             return cast(ResponseT, api_response)

volcenginesdkarkruntime/_client.py

@@ -16,6 +16,7 @@
 import volcenginesdkark
 
 from . import resources
+from .resources.beta import beta
 from ._base_client import SyncAPIClient, AsyncAPIClient
 from ._constants import (
     DEFAULT_MAX_RETRIES,
@@ -35,6 +36,7 @@
 
 
 class Ark(SyncAPIClient):
+    beta: beta.Beta
     chat: resources.Chat
     bot_chat: resources.BotChat
     embeddings: resources.Embeddings
@@ -100,6 +102,7 @@ def __init__(
         self._sts_token_manager: StsTokenManager | None = None
         self._certificate_manager: E2ECertificateManager | None = None
 
+        self.beta = beta.Beta(self)
         self.chat = resources.Chat(self)
         self.bot_chat = resources.BotChat(self)
         self.embeddings = resources.Embeddings(self)
@@ -155,6 +158,7 @@ def get_model_breaker(self, model_name: str) -> ModelBreaker:
 
 
 class AsyncArk(AsyncAPIClient):
+    beta: beta.AsyncBeta
     chat: resources.AsyncChat
     bot_chat: resources.AsyncBotChat
     embeddings: resources.AsyncEmbeddings
@@ -220,6 +224,7 @@ def __init__(
         self._sts_token_manager: StsTokenManager | None = None
         self._certificate_manager: E2ECertificateManager | None = None
 
+        self.beta = beta.AsyncBeta(self)
         self.chat = resources.AsyncChat(self)
         self.bot_chat = resources.AsyncBotChat(self)
         self.embeddings = resources.AsyncEmbeddings(self)

volcenginesdkarkruntime/_compat.py

@@ -2,10 +2,10 @@
 
 from typing import TYPE_CHECKING, Any, Union, Generic, TypeVar, Callable, cast, overload
 from datetime import date, datetime
-from typing_extensions import Self
-
+from typing_extensions import Self, Literal
 import pydantic
 from pydantic.fields import FieldInfo
+from ._types import IncEx
 
 _T = TypeVar("_T")
 _ModelT = TypeVar("_ModelT", bound=pydantic.BaseModel)
@@ -68,7 +68,6 @@ def is_typeddict(type_: type[Any]) -> bool:  # noqa: ARG001
             parse_datetime as parse_datetime,
         )
 
-
 # refactored config
 if TYPE_CHECKING:
     from pydantic import ConfigDict as ConfigDict
@@ -138,17 +137,25 @@ def model_json(model: pydantic.BaseModel, *, indent: int | None = None) -> str:
 def model_dump(
     model: pydantic.BaseModel,
     *,
+    exclude: IncEx | None = None,
     exclude_unset: bool = False,
     exclude_defaults: bool = False,
+    warnings: bool = True,
+    mode: Literal["json", "python"] = "python",
 ) -> dict[str, Any]:
-    if PYDANTIC_V2:
+    if PYDANTIC_V2 or hasattr(model, "model_dump"):
         return model.model_dump(
+            mode=mode,
+            exclude=exclude,
             exclude_unset=exclude_unset,
             exclude_defaults=exclude_defaults,
+            # warnings are not supported in Pydantic v1
+            warnings=warnings if PYDANTIC_V2 else True,
         )
     return cast(
         "dict[str, Any]",
         model.dict(  # pyright: ignore[reportDeprecated, reportUnnecessaryCast]
+            exclude=exclude,
             exclude_unset=exclude_unset,
             exclude_defaults=exclude_defaults,
         ),
@@ -161,6 +168,18 @@ def model_parse(model: type[_ModelT], data: Any) -> _ModelT:
     return model.parse_obj(data)  # pyright: ignore[reportDeprecated]
 
 
+def model_parse_json(model: type[_ModelT], data: str | bytes) -> _ModelT:
+    if PYDANTIC_V2:
+        return model.model_validate_json(data)
+    return model.parse_raw(data)  # pyright: ignore[reportDeprecated]
+
+
+def model_json_schema(model: type[_ModelT]) -> dict[str, Any]:
+    if PYDANTIC_V2:
+        return model.model_json_schema()
+    return model.schema()  # pyrigh
+
+
 # generic models
 if TYPE_CHECKING:
 

volcenginesdkarkruntime/_exceptions.py

@@ -2,11 +2,14 @@
 
 from __future__ import annotations
 
-from typing import Optional
+from typing import TYPE_CHECKING, Optional
 from typing_extensions import Literal
 
 import httpx
 
+if TYPE_CHECKING:
+    from .types.chat import ChatCompletion
+
 __all__ = [
     "ArkBadRequestError",
     "ArkAuthenticationError",
@@ -159,3 +162,27 @@ class ArkRateLimitError(ArkAPIStatusError):
 
 class ArkInternalServerError(ArkAPIStatusError):
     pass
+
+
+class ArkLengthFinishReasonError(ArkAPIStatusError):
+    completion: ChatCompletion
+    """The completion that caused this error.
+
+    Note: this will *not* be a complete `ChatCompletion` object when streaming as `usage`
+          will not be included.
+    """
+
+    def __init__(self, *, completion: ChatCompletion) -> None:
+        msg = "Could not parse response content as the length limit was reached"
+        if completion.usage:
+            msg += f" - {completion.usage}"
+
+        super().__init__(msg)
+        self.completion = completion
+
+
+class ArkContentFilterFinishReasonError(ArkAPIStatusError):
+    def __init__(self) -> None:
+        super().__init__(
+            "Could not parse response content as the request was rejected by the content filter",
+        )