diff --git a/erniebot/src/erniebot/constants.py b/erniebot/src/erniebot/constants.py index 0267411a..e305f4f6 100644 --- a/erniebot/src/erniebot/constants.py +++ b/erniebot/src/erniebot/constants.py @@ -21,4 +21,4 @@ DEFAULT_REQUEST_TIMEOUT_SECS: Final[float] = 600 POLLING_INTERVAL_SECS: Final[float] = 5 -POLLING_TIMEOUT_SECS: Final[float] = 20 +POLLING_TIMEOUT_SECS: Final[float] = 600 diff --git a/erniebot/src/erniebot/resources/chat_completion.py b/erniebot/src/erniebot/resources/chat_completion.py index ee098283..ff3bfe2d 100644 --- a/erniebot/src/erniebot/resources/chat_completion.py +++ b/erniebot/src/erniebot/resources/chat_completion.py @@ -64,6 +64,9 @@ class ChatCompletion(EBResource, CreatableWithStreaming): "ernie-longtext": { "model_id": "ernie_bot_8k", }, + "ernie-speed": { + "model_id": "ernie_speed", + }, }, }, APIType.AISTUDIO: { @@ -478,8 +481,14 @@ def _set_val_if_key_exists(src: dict, dst: dict, key: str) -> None: # params params = {} - if model == "ernie-turbo": - for arg in ("functions", "stop", "disable_search", "enable_citation"): + if model in ("ernie-turbo", "ernie-speed"): + for arg in ( + "functions", + "stop", + "disable_search", + "enable_citation", + "tool_choice", + ): if arg in kwargs: raise errors.InvalidArgumentError(f"`{arg}` is not supported by the {model} model.") params["messages"] = messages