KMS/ragflow
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

optimize TenantLLMService.increase_usage for "can't update token usag… (#4755)

Browse Source 
…e error " message

### What problem does this PR solve?

optimize TenantLLMService.increase_usage Performance

### Type of change

- [x] Performance Improvement

Co-authored-by: che_shuai <che_shuai@massclouds.com>
This commit is contained in:
davidche 2025-02-07 12:16:17 +08:00 committed by GitHub
parent 2aa0cdde8f
commit 588207d7c1
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 24 additions and 27 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

51
api/db/services/llm_service.py
View File

@ -175,38 +175,35 @@ class TenantLLMService(CommonService):
if not e:
raise LookupError("Tenant not found")
if llm_type == LLMType.EMBEDDING.value:
mdlnm = tenant.embd_id
elif llm_type == LLMType.SPEECH2TEXT.value:
mdlnm = tenant.asr_id
elif llm_type == LLMType.IMAGE2TEXT.value:
mdlnm = tenant.img2txt_id
elif llm_type == LLMType.CHAT.value:
mdlnm = tenant.llm_id if not llm_name else llm_name
elif llm_type == LLMType.RERANK:
mdlnm = tenant.rerank_id if not llm_name else llm_name
elif llm_type == LLMType.TTS:
mdlnm = tenant.tts_id if not llm_name else llm_name
else:
assert False, "LLM type error"
llm_map = {
LLMType.EMBEDDING.value: tenant.embd_id,
LLMType.SPEECH2TEXT.value: tenant.asr_id,
LLMType.IMAGE2TEXT.value: tenant.img2txt_id,
LLMType.CHAT.value: tenant.llm_id if not llm_name else llm_name,
LLMType.RERANK.value: tenant.rerank_id if not llm_name else llm_name,
LLMType.TTS.value: tenant.tts_id if not llm_name else llm_name
}
mdlnm = llm_map.get(llm_type)
if mdlnm is None:
raise ValueError("LLM type error")
llm_name, llm_factory = TenantLLMService.split_model_name_and_factory(mdlnm)
num = 0
try:
if llm_factory:
tenant_llms = cls.query(tenant_id=tenant_id, llm_name=llm_name, llm_factory=llm_factory)
else:
tenant_llms = cls.query(tenant_id=tenant_id, llm_name=llm_name)
if not tenant_llms:
return num
else:
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()
num = cls.model.update(
used_tokens=cls.model.used_tokens + used_tokens
).where(
cls.model.tenant_id == tenant_id,
cls.model.llm_name == llm_name,
cls.model.llm_factory == llm_factory if llm_factory else True
).execute()
except Exception:
logging.exception("TenantLLMService.increase_usage got exception")
logging.exception(
"TenantLLMService.increase_usage got exception,Failed to update used_tokens for tenant_id=%s, llm_name=%s",
tenant_id, llm_name)
return 0
return num
@classmethod