[API ADD]Add max_out_token (#343)

PaddlePaddle · Apr 29, 2024 · 621ce45 · 621ce45
1 parent 8de8fdb
commit 621ce45
Showing 1 changed file with 14 additions and 6 deletions.
diff --git a/erniebot/src/erniebot/resources/chat_completion.py b/erniebot/src/erniebot/resources/chat_completion.py
@@ -156,6 +156,7 @@ def create(
         extra_params: Optional[dict] = ...,
         headers: Optional[HeadersType] = ...,
         request_timeout: Optional[float] = ...,
+        max_output_tokens: Optional[int] = ...,
         _config_: Optional[ConfigDictType] = ...,
     ) -> "ChatCompletionResponse":
         ...
@@ -182,6 +183,7 @@ def create(
         extra_params: Optional[dict] = ...,
         headers: Optional[HeadersType] = ...,
         request_timeout: Optional[float] = ...,
+        max_output_tokens: Optional[int] = ...,
         _config_: Optional[ConfigDictType] = ...,
     ) -> Iterator["ChatCompletionResponse"]:
         ...
@@ -208,6 +210,7 @@ def create(
         extra_params: Optional[dict] = ...,
         headers: Optional[HeadersType] = ...,
         request_timeout: Optional[float] = ...,
+        max_output_tokens: Optional[int] = ...,
         _config_: Optional[ConfigDictType] = ...,
     ) -> Union["ChatCompletionResponse", Iterator["ChatCompletionResponse"]]:
         ...
@@ -233,6 +236,7 @@ def create(
         extra_params: Optional[dict] = None,
         headers: Optional[HeadersType] = None,
         request_timeout: Optional[float] = None,
+        max_output_tokens: Optional[int] = None,
         _config_: Optional[ConfigDictType] = None,
     ) -> Union["ChatCompletionResponse", Iterator["ChatCompletionResponse"]]:
         """Creates a model response for the given conversation.
@@ -279,6 +283,7 @@ def create(
             user_id=user_id,
             tool_choice=tool_choice,
             stream=stream,
+            max_output_tokens=max_output_tokens,
         )
         kwargs["validate_functions"] = validate_functions
         if extra_params is not None:
@@ -313,6 +318,7 @@ async def acreate(
         extra_params: Optional[dict] = ...,
         headers: Optional[HeadersType] = ...,
         request_timeout: Optional[float] = ...,
+        max_output_tokens: Optional[int] = ...,
         _config_: Optional[ConfigDictType] = ...,
     ) -> EBResponse:
         ...
@@ -339,6 +345,7 @@ async def acreate(
         extra_params: Optional[dict] = ...,
         headers: Optional[HeadersType] = ...,
         request_timeout: Optional[float] = ...,
+        max_output_tokens: Optional[int] = ...,
         _config_: Optional[ConfigDictType] = ...,
     ) -> AsyncIterator["ChatCompletionResponse"]:
         ...
@@ -365,6 +372,7 @@ async def acreate(
         extra_params: Optional[dict] = ...,
         headers: Optional[HeadersType] = ...,
         request_timeout: Optional[float] = ...,
+        max_output_tokens: Optional[int] = ...,
         _config_: Optional[ConfigDictType] = ...,
     ) -> Union["ChatCompletionResponse", AsyncIterator["ChatCompletionResponse"]]:
         ...
@@ -390,6 +398,7 @@ async def acreate(
         extra_params: Optional[dict] = None,
         headers: Optional[HeadersType] = None,
         request_timeout: Optional[float] = None,
+        max_output_tokens: Optional[int] = None,
         _config_: Optional[ConfigDictType] = None,
     ) -> Union["ChatCompletionResponse", AsyncIterator["ChatCompletionResponse"]]:
         """Creates a model response for the given conversation.
@@ -436,6 +445,7 @@ async def acreate(
             user_id=user_id,
             tool_choice=tool_choice,
             stream=stream,
+            max_output_tokens=max_output_tokens,
         )
         kwargs["validate_functions"] = validate_functions
         if extra_params is not None:
@@ -450,12 +460,7 @@ async def acreate(
 
     def _check_model_kwargs(self, model_name: str, kwargs: Dict[str, Any]) -> None:
         if model_name in ("ernie-speed", "ernie-speed-128k", "ernie-char-8k", "ernie-tiny-8k", "ernie-lite"):
-            for arg in (
-                "functions",
-                "disable_search",
-                "enable_citation",
-                "tool_choice",
-            ):
+            for arg in ("functions", "disable_search", "enable_citation", "tool_choice"):
                 if arg in kwargs:
                     raise errors.InvalidArgumentError(f"`{arg}` is not supported by the {model_name} model.")
 
@@ -492,6 +497,7 @@ def _set_val_if_key_exists(src: dict, dst: dict, key: str) -> None:
             "extra_params",
             "headers",
             "request_timeout",
+            "max_output_tokens",
         }
 
         invalid_keys = kwargs.keys() - valid_keys
@@ -554,6 +560,8 @@ def _set_val_if_key_exists(src: dict, dst: dict, key: str) -> None:
             _set_val_if_key_exists(kwargs, params, "user_id")
         _set_val_if_key_exists(kwargs, params, "tool_choice")
         _set_val_if_key_exists(kwargs, params, "stream")
+        _set_val_if_key_exists(kwargs, params, "max_output_tokens")
+
         if "extra_params" in kwargs:
             params.update(kwargs["extra_params"])