fix: wenxin model name invalid when llm call

langgenius · Sep 27, 2023 · 5c366b6 · 5c366b6
1 parent 9dbb8ac
commit 5c366b6
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/api/core/model_providers/models/llm/wenxin_model.py b/api/core/model_providers/models/llm/wenxin_model.py
@@ -18,6 +18,7 @@ def _init_client(self) -> Any:
         provider_model_kwargs = self._to_model_kwargs_input(self.model_rules, self.model_kwargs)
         # TODO load price_config from configs(db)
         return Wenxin(
+            model=self.name,
             streaming=self.streaming,
             callbacks=self.callbacks,
             **self.credentials,

diff --git a/api/core/model_providers/providers/wenxin_provider.py b/api/core/model_providers/providers/wenxin_provider.py
@@ -61,13 +61,18 @@ def get_model_parameter_rules(self, model_name: str, model_type: ModelType) -> M
         :param model_type:
         :return:
         """
+        model_max_tokens = {
+            'ernie-bot': 4800,
+            'ernie-bot-turbo': 11200,
+        }
+
         if model_name in ['ernie-bot', 'ernie-bot-turbo']:
             return ModelKwargsRules(
                 temperature=KwargRule[float](min=0.01, max=1, default=0.95, precision=2),
                 top_p=KwargRule[float](min=0.01, max=1, default=0.8, precision=2),
                 presence_penalty=KwargRule[float](enabled=False),
                 frequency_penalty=KwargRule[float](enabled=False),
-                max_tokens=KwargRule[int](enabled=False),
+                max_tokens=KwargRule[int](enabled=False, max=model_max_tokens.get(model_name)),
             )
         else:
             return ModelKwargsRules(