From 381219aa41d9304573d7d6cfde1685b4ec895d3f Mon Sep 17 00:00:00 2001 From: Zhichang Yu Date: Fri, 29 Nov 2024 17:02:49 +0800 Subject: [PATCH] Fixed increase_usage for builtin models (#3748) ### What problem does this PR solve? Fixed increase_usage for builtin models. Close #1803 ### Type of change - [x] Bug Fix (non-breaking change which fixes an issue) --- api/db/services/llm_service.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/api/db/services/llm_service.py b/api/db/services/llm_service.py index 1a91c78b6..98d8fd21d 100644 --- a/api/db/services/llm_service.py +++ b/api/db/services/llm_service.py @@ -172,12 +172,17 @@ class TenantLLMService(CommonService): num = 0 try: - for u in cls.query(tenant_id=tenant_id, llm_name=llm_name): - num += cls.model.update(used_tokens=u.used_tokens + used_tokens)\ - .where(cls.model.tenant_id == tenant_id, cls.model.llm_name == llm_name)\ + tenant_llms = cls.query(tenant_id=tenant_id, llm_name=llm_name) + if tenant_llms: + tenant_llm = tenant_llms[0] + num = cls.model.update(used_tokens=tenant_llm.used_tokens + used_tokens)\ + .where(cls.model.tenant_id == tenant_id, cls.model.llm_factory == tenant_llm.llm_factory, cls.model.llm_name == llm_name)\ .execute() - except Exception as e: - pass + else: + llm_factory = llm_name.split("/")[0] if "/" in llm_name else llm_name + num = cls.model.create(tenant_id=tenant_id, llm_factory=llm_factory, llm_name=llm_name, used_tokens=used_tokens) + except Exception: + logging.exception("TenantLLMService.increase_usage got exception") return num @classmethod