Perf: set timeout of some steps in KG. (#8873)

### What problem does this PR solve?

### Type of change


- [x] Performance Improvement
This commit is contained in:
Kevin Hu
2025-07-16 18:06:03 +08:00
committed by GitHub
parent b3018a455f
commit fbd115773b
9 changed files with 28 additions and 29 deletions

View File

@ -22,7 +22,6 @@ from collections import defaultdict
import jinja2
import json_repair
from api import settings
from rag.prompt_template import load_prompt
from rag.settings import TAG_FLD
from rag.utils import encoder, num_tokens_from_string
@ -51,18 +50,6 @@ def chunks_format(reference):
]
def llm_id2llm_type(llm_id):
from api.db.services.llm_service import TenantLLMService
llm_id, *_ = TenantLLMService.split_model_name_and_factory(llm_id)
llm_factories = settings.FACTORY_LLM_INFOS
for llm_factory in llm_factories:
for llm in llm_factory["llm"]:
if llm_id == llm["llm_name"]:
return llm["model_type"].strip(",")[-1]
def message_fit_in(msg, max_length=4000):
def count():
nonlocal msg
@ -188,8 +175,9 @@ def question_proposal(chat_mdl, content, topn=3):
def full_question(tenant_id, llm_id, messages, language=None):
from api.db import LLMType
from api.db.services.llm_service import LLMBundle
from api.db.services.llm_service import TenantLLMService
if llm_id2llm_type(llm_id) == "image2text":
if TenantLLMService.llm_id2llm_type(llm_id) == "image2text":
chat_mdl = LLMBundle(tenant_id, LLMType.IMAGE2TEXT, llm_id)
else:
chat_mdl = LLMBundle(tenant_id, LLMType.CHAT, llm_id)
@ -220,8 +208,9 @@ def full_question(tenant_id, llm_id, messages, language=None):
def cross_languages(tenant_id, llm_id, query, languages=[]):
from api.db import LLMType
from api.db.services.llm_service import LLMBundle
from api.db.services.llm_service import TenantLLMService
if llm_id and llm_id2llm_type(llm_id) == "image2text":
if llm_id and TenantLLMService.llm_id2llm_type(llm_id) == "image2text":
chat_mdl = LLMBundle(tenant_id, LLMType.IMAGE2TEXT, llm_id)
else:
chat_mdl = LLMBundle(tenant_id, LLMType.CHAT, llm_id)

View File

@ -506,7 +506,7 @@ async def run_raptor(row, chat_mdl, embd_mdl, vector_size, callback=None):
return res, tk_count
@timeout(60*60*1.5)
@timeout(60*60, 1)
async def do_handle_task(task):
task_id = task["id"]
task_from_page = task["from_page"]