From 7d3861fa4347f0558833bd5e08e150f7ab830e91 Mon Sep 17 00:00:00 2001 From: Bobholamovic Date: Thu, 8 Feb 2024 10:47:39 +0800 Subject: [PATCH] Add ernie-speed --- erniebot/src/erniebot/constants.py | 2 +- erniebot/src/erniebot/resources/chat_completion.py | 13 +++++++++++-- 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/erniebot/src/erniebot/constants.py b/erniebot/src/erniebot/constants.py index 0267411a0..e305f4f69 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 ee0982836..ff3bfe2de 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