Cherry-pick 4aca819: Add ernie-speed (#315)

PaddlePaddle · Feb 8, 2024 · b78704b · b78704b
1 parent 7a51518
commit b78704b
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 5 deletions.
diff --git a/erniebot/src/erniebot/constants.py b/erniebot/src/erniebot/constants.py
@@ -12,11 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Final
+
 LOGGER_NAME: str = "erniebot"
 
 STREAM_RESPONSE_PREFIX: bytes = b"data: "
 
-DEFAULT_REQUEST_TIMEOUT_SECS: float = 600
+DEFAULT_REQUEST_TIMEOUT_SECS: Final[float] = 600
 
-POLLING_INTERVAL_SECS: float = 5
-POLLING_TIMEOUT_SECS: float = 20
+POLLING_INTERVAL_SECS: Final[float] = 5
+POLLING_TIMEOUT_SECS: Final[float] = 600
diff --git 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