Add ernie-speed

PaddlePaddle · Feb 8, 2024 · 7d3861f · 7d3861f
1 parent be8c1af
commit 7d3861f
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 3 deletions.
diff --git 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
@@ -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