From cc9463955520f08c2097a9eb83cf5d9973f37bbe Mon Sep 17 00:00:00 2001 From: Lynn Date: Thu, 9 Jul 2026 14:52:41 +0800 Subject: [PATCH] Fix: get_by_id (#16765) --- .../joint_services/memory_message_service.py | 16 +++++++++++--- api/db/joint_services/tenant_model_service.py | 21 +++++++++++-------- rag/app/naive.py | 9 +------- rag/flow/parser/parser.py | 9 +------- rag/nlp/__init__.py | 3 ++- .../dataflow_service.py | 8 ++----- .../task_executor_refactor/task_handler.py | 12 +++-------- 7 files changed, 34 insertions(+), 44 deletions(-) diff --git a/api/db/joint_services/memory_message_service.py b/api/db/joint_services/memory_message_service.py index 1e0d3bcf0e..e1abfc7cfa 100644 --- a/api/db/joint_services/memory_message_service.py +++ b/api/db/joint_services/memory_message_service.py @@ -155,8 +155,18 @@ async def save_extracted_to_memory_only(memory_id: str, message_dict, source_mes return await embed_and_save(memory, message_list, task_id) -async def extract_by_llm(tenant_id: str, tenant_llm_id: str | None, extract_conf: dict, memory_type: List[str], user_input: str, - agent_response: str, system_prompt: str = "", user_prompt: str="", task_id: str=None, llm_id: str = "") -> List[dict]: +async def extract_by_llm( + tenant_id: str, + tenant_llm_id: str | None, + extract_conf: dict, + memory_type: List[str], + user_input: str, + agent_response: str, + system_prompt: str = "", + user_prompt: str = "", + task_id: str = None, + llm_id: str = "", +) -> List[dict]: if not system_prompt: system_prompt = PromptAssembler.assemble_system_prompt({"memory_type": memory_type}) conversation_content = f"User Input: {user_input}\nAgent Response: {agent_response}" @@ -193,7 +203,7 @@ async def extract_by_llm(tenant_id: str, tenant_llm_id: str | None, extract_conf ] -async def embed_and_save(memory, message_list: list[dict], task_id: str=None): +async def embed_and_save(memory, message_list: list[dict], task_id: str = None): if memory.tenant_embd_id: try: embd_model_config = get_model_config_by_id(memory.tenant_id, LLMType.EMBEDDING, memory.tenant_embd_id) diff --git a/api/db/joint_services/tenant_model_service.py b/api/db/joint_services/tenant_model_service.py index 040c4636c4..eead80a10d 100644 --- a/api/db/joint_services/tenant_model_service.py +++ b/api/db/joint_services/tenant_model_service.py @@ -249,12 +249,14 @@ def _resolve_instance_for_model(provider_obj, instance_name: str, model_name: st raise LookupError(f"Instance {instance_name} not found for model {model_name}.") + def resolve_model_config(tenant_id, model_type: str | enum.Enum, model_ref: str): try: return get_model_config_by_id(tenant_id, model_type, model_ref) except LookupError: return get_model_config_from_provider_instance(tenant_id, model_type, model_ref) + def get_model_config_from_provider_instance(tenant_id, model_type: str | enum.Enum, model_name: str): pure_model_name, instance_name, provider_name = split_model_name(model_name) model_type_val = model_type if isinstance(model_type, str) else model_type.value @@ -390,9 +392,7 @@ def resolve_model_id(tenant_id: str, model_type: str | enum.Enum, model_name: st raise LookupError(f"Provider {provider_name} not found for model {model_name}.") instance_obj = _resolve_instance_for_model(provider_obj, instance_name, model_name) - model_obj = TenantModelService.get_by_provider_id_and_instance_id_and_model_type_and_model_name( - provider_obj.id, instance_obj.id, model_type_val, pure_model_name - ) + model_obj = TenantModelService.get_by_provider_id_and_instance_id_and_model_type_and_model_name(provider_obj.id, instance_obj.id, model_type_val, pure_model_name) if not model_obj: raise LookupError(f"Model {model_name} not found for type {model_type_val}.") return model_obj.id @@ -400,12 +400,12 @@ def resolve_model_id(tenant_id: str, model_type: str | enum.Enum, model_name: st # Mapping from model-name field → (LLMType, tenant_model id field) _MODEL_NAME_TO_ID_FIELD_MAP: dict[str, tuple[str, str]] = { - "llm_id": (LLMType.CHAT, "tenant_llm_id"), - "embd_id": (LLMType.EMBEDDING, "tenant_embd_id"), - "rerank_id": (LLMType.RERANK, "tenant_rerank_id"), - "asr_id": (LLMType.SPEECH2TEXT, "tenant_asr_id"), - "img2txt_id": (LLMType.IMAGE2TEXT, "tenant_img2txt_id"), - "tts_id": (LLMType.TTS, "tenant_tts_id"), + "llm_id": (LLMType.CHAT, "tenant_llm_id"), + "embd_id": (LLMType.EMBEDDING, "tenant_embd_id"), + "rerank_id": (LLMType.RERANK, "tenant_rerank_id"), + "asr_id": (LLMType.SPEECH2TEXT, "tenant_asr_id"), + "img2txt_id": (LLMType.IMAGE2TEXT, "tenant_img2txt_id"), + "tts_id": (LLMType.TTS, "tenant_tts_id"), } @@ -443,18 +443,21 @@ def get_api_key(tenant_id: str, model_name: str): instance_obj = _resolve_instance_for_model(provider_obj, instance_name, model_name) return instance_obj.api_key + def get_model_type_by_id(model_id: str): exist, model_obj = TenantModelService.get_by_id(model_id) if not exist: raise LookupError(f"TenantModel id={model_id} not found.") return get_model_type_human(model_obj.model_type) + def resolve_model_type(tenant_id: str, model_ref: str): try: return get_model_type_by_id(model_ref) except LookupError: return get_model_type_by_name(tenant_id, model_ref) + def get_model_type_by_name(tenant_id: str, model_name: str): pure_model_name, instance_name, provider_name = split_model_name(model_name) provider_obj = TenantModelProviderService.get_by_tenant_id_and_provider_name(tenant_id, provider_name) diff --git a/rag/app/naive.py b/rag/app/naive.py index fed97d1a63..302b5899b6 100644 --- a/rag/app/naive.py +++ b/rag/app/naive.py @@ -326,14 +326,7 @@ def by_somark( if somark_llm_name: try: - try: - ocr_model_config = resolve_model_config(tenant_id, LLMType.OCR, somark_llm_name) - except Exception: - if "@" in somark_llm_name: - raise - from api.db.services.tenant_llm_service import TenantLLMService - - ocr_model_config = TenantLLMService.get_model_config(tenant_id, LLMType.OCR.value, somark_llm_name) + ocr_model_config = resolve_model_config(tenant_id, LLMType.OCR, somark_llm_name) ocr_model = LLMBundle(tenant_id=tenant_id, model_config=ocr_model_config, lang=lang) pdf_parser = ocr_model.mdl sections, tables = pdf_parser.parse_pdf( diff --git a/rag/flow/parser/parser.py b/rag/flow/parser/parser.py index ab2bb6197f..a35827fc25 100644 --- a/rag/flow/parser/parser.py +++ b/rag/flow/parser/parser.py @@ -523,14 +523,7 @@ class Parser(ProcessBase): raise RuntimeError("SoMark model not configured. Please add SoMark in Model Providers or set SOMARK_* env.") tenant_id = self._canvas._tenant_id - try: - ocr_model_config = resolve_model_config(tenant_id, LLMType.OCR, parser_model_name) - except Exception: - if "@" in parser_model_name: - raise - from api.db.services.tenant_llm_service import TenantLLMService - - ocr_model_config = TenantLLMService.get_model_config(tenant_id, LLMType.OCR.value, parser_model_name) + ocr_model_config = resolve_model_config(tenant_id, LLMType.OCR, parser_model_name) ocr_model = LLMBundle(tenant_id, ocr_model_config) pdf_parser = ocr_model.mdl diff --git a/rag/nlp/__init__.py b/rag/nlp/__init__.py index fa4f5a5faf..597b6b462b 100644 --- a/rag/nlp/__init__.py +++ b/rag/nlp/__init__.py @@ -354,6 +354,7 @@ def is_chinese(text): def tokenize(d, txt, eng, language="English"): from . import rag_tokenizer + rag_tokenizer.tokenizer.set_language(language) d["content_with_weight"] = txt t = re.sub(r"]{0,12})?>", " ", txt) @@ -483,7 +484,7 @@ def tokenize_table(tbls, doc, eng, batch_size=10, language="English"): de = "； " if lang_key in {"chinese", "japanese"} else "; " for i in range(0, len(rows), batch_size): d = copy.deepcopy(doc) - r = de.join(rows[i:i + batch_size]) + r = de.join(rows[i : i + batch_size]) tokenize(d, r, eng, language=language) d["doc_type_kwd"] = "table" if img: diff --git a/rag/svr/task_executor_refactor/dataflow_service.py b/rag/svr/task_executor_refactor/dataflow_service.py index 280916a500..5491c9cb73 100644 --- a/rag/svr/task_executor_refactor/dataflow_service.py +++ b/rag/svr/task_executor_refactor/dataflow_service.py @@ -246,13 +246,9 @@ class DataflowService: try: embd_model_config = get_model_config_by_id(ctx.tenant_id, LLMType.EMBEDDING, kb.tenant_embd_id) except LookupError: - embd_model_config = resolve_model_config( - ctx.tenant_id, LLMType.EMBEDDING, embedding_id - ) + embd_model_config = resolve_model_config(ctx.tenant_id, LLMType.EMBEDDING, embedding_id) else: - embd_model_config = resolve_model_config( - ctx.tenant_id, LLMType.EMBEDDING, embedding_id - ) + embd_model_config = resolve_model_config(ctx.tenant_id, LLMType.EMBEDDING, embedding_id) from api.db.services.llm_service import LLMBundle with LLMBundle(ctx.tenant_id, embd_model_config) as embedding_model: diff --git a/rag/svr/task_executor_refactor/task_handler.py b/rag/svr/task_executor_refactor/task_handler.py index 0e6e3e2fe5..983cd01cb1 100644 --- a/rag/svr/task_executor_refactor/task_handler.py +++ b/rag/svr/task_executor_refactor/task_handler.py @@ -321,17 +321,11 @@ class TaskHandler: try: if ctx.tenant_embd_id: try: - embd_model_config = get_model_config_by_id( - task_tenant_id, LLMType.EMBEDDING, ctx.tenant_embd_id - ) + embd_model_config = get_model_config_by_id(task_tenant_id, LLMType.EMBEDDING, ctx.tenant_embd_id) except LookupError: - embd_model_config = resolve_model_config( - task_tenant_id, LLMType.EMBEDDING, task_embedding_id - ) + embd_model_config = resolve_model_config(task_tenant_id, LLMType.EMBEDDING, task_embedding_id) elif task_embedding_id: - embd_model_config = resolve_model_config( - task_tenant_id, LLMType.EMBEDDING, task_embedding_id - ) + embd_model_config = resolve_model_config(task_tenant_id, LLMType.EMBEDDING, task_embedding_id) else: embd_model_config = get_tenant_default_model_by_type(task_tenant_id, LLMType.EMBEDDING) embedding_model = LLMBundle(task_tenant_id, embd_model_config, lang=task_language)