From b733f8f55bb810ca48c63b2984ba17065d65b490 Mon Sep 17 00:00:00 2001 From: Rock Chin <1010553892@qq.com> Date: Sun, 5 Feb 2023 16:01:13 +0800 Subject: [PATCH] =?UTF-8?q?perf:=20=E6=B7=BB=E5=8A=A0=E5=AF=B9completion?= =?UTF-8?q?=5Fapi=5Fparams=E4=B8=ADengine=E5=8F=82=E6=95=B0=E5=80=BC?= =?UTF-8?q?=E7=9A=84=E7=BB=9F=E8=AE=A1=E6=94=AF=E6=8C=81?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- pkg/openai/manager.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/pkg/openai/manager.py b/pkg/openai/manager.py index a934685f..b21a545e 100644 --- a/pkg/openai/manager.py +++ b/pkg/openai/manager.py @@ -40,8 +40,12 @@ class OpenAIInteract: logging.debug("OpenAI response: %s", response) - self.audit_mgr.report_text_model_usage(config.completion_api_params['model'], - response['usage']['total_tokens']) + if 'model' in config.completion_api_params: + self.audit_mgr.report_text_model_usage(config.completion_api_params['model'], + response['usage']['total_tokens']) + elif 'engine' in config.completion_api_params: + self.audit_mgr.report_text_model_usage(config.completion_api_params['engine'], + response['usage']['total_tokens']) return response