From 6a77523bf0d18096561dc809f31b9eb88179ea50 Mon Sep 17 00:00:00 2001 From: buua436 Date: Thu, 9 Jul 2026 14:02:08 +0800 Subject: [PATCH] refa: resolve tenant model refs consistently (#16744) --- agent/component/agent_with_tools.py | 6 +- agent/component/browser.py | 6 +- agent/component/categorize.py | 4 +- agent/component/llm.py | 10 +-- agent/tools/retrieval.py | 6 +- api/apps/restful_apis/bot_api.py | 10 +-- api/apps/restful_apis/chat_api.py | 10 +-- api/apps/restful_apis/chunk_api.py | 10 +-- api/apps/restful_apis/dify_retrieval_api.py | 4 +- api/apps/restful_apis/models_api.py | 23 ++++++- api/apps/restful_apis/openai_api.py | 6 +- api/apps/services/dataset_api_service.py | 32 ++++++---- api/apps/services/models_api_service.py | 8 ++- api/apps/services/provider_api_service.py | 4 +- .../joint_services/memory_message_service.py | 16 ++--- api/db/joint_services/tenant_model_service.py | 40 +++++++++--- api/db/services/dialog_service.py | 64 +++++++++---------- api/utils/api_utils.py | 4 +- api/utils/validation_utils.py | 9 ++- rag/app/naive.py | 12 ++-- rag/app/picture.py | 4 +- rag/benchmark.py | 4 +- rag/flow/parser/parser.py | 20 +++--- rag/flow/parser/utils.py | 4 +- rag/flow/tokenizer/tokenizer.py | 8 +-- rag/graphrag/general/smoke.py | 8 +-- rag/graphrag/light/smoke.py | 8 +-- rag/graphrag/search.py | 8 +-- rag/prompts/generator.py | 16 ++--- rag/svr/task_executor.py | 24 +++---- .../chunk_post_processor.py | 12 ++-- .../dataflow_service.py | 8 +-- .../dataset_wiki_generator.py | 4 +- .../task_executor_refactor/task_handler.py | 20 ++++-- test/testcases/restful_api/test_chats.py | 3 +- .../test_dify_retrieval_routes_unit.py | 1 + .../test_user_tenant_routes_unit.py | 1 + .../test_chat_sdk_routes_unit.py | 1 + .../test_dify_retrieval_routes_unit.py | 1 + .../test_doc_sdk_routes_unit.py | 1 + .../test_session_sdk_routes_unit.py | 2 + .../test_iteration_runtime_unit.py | 1 + .../test_chunk_app/test_chunk_routes_unit.py | 1 + .../test_agentbots_access_control.py | 8 ++- .../test_openai_stream_no_duplicate.py | 1 + .../api/apps/sdk/test_dify_retrieval.py | 1 + .../api/apps/services/test_delete_datasets.py | 2 + ...ls_api_service_list_tenant_added_models.py | 31 +++------ .../test_dialog_service_final_answer.py | 24 +++---- ...t_dialog_service_use_sql_source_columns.py | 4 +- .../svr/task_executor_refactor/conftest.py | 10 +-- 51 files changed, 300 insertions(+), 225 deletions(-) diff --git a/agent/component/agent_with_tools.py b/agent/component/agent_with_tools.py index 93b9f2eff2..dc4a0cf0de 100644 --- a/agent/component/agent_with_tools.py +++ b/agent/component/agent_with_tools.py @@ -27,7 +27,7 @@ import json_repair from agent.component.llm import LLM, LLMParam from agent.tools.base import LLMToolPluginCallSession, ToolBase, ToolMeta, ToolParamBase -from api.db.joint_services.tenant_model_service import get_model_config_from_provider_instance, get_model_type_by_name +from api.db.joint_services.tenant_model_service import resolve_model_config, resolve_model_type from api.db.services.llm_service import LLMBundle from api.db.services.mcp_server_service import MCPServerService from common.connection_utils import timeout @@ -82,9 +82,9 @@ class Agent(LLM, ToolBase): original_name = cpn.get_meta()["function"]["name"] indexed_name = f"{original_name}_{idx}" self.tools[indexed_name] = cpn - model_types = get_model_type_by_name(self._canvas.get_tenant_id(), self._param.llm_id) + model_types = resolve_model_type(self._canvas.get_tenant_id(), self._param.llm_id) model_type = "chat" if "chat" in model_types else model_types[0] - chat_model_config = get_model_config_from_provider_instance(self._canvas.get_tenant_id(), model_type, self._param.llm_id) + chat_model_config = resolve_model_config(self._canvas.get_tenant_id(), model_type, self._param.llm_id) self.chat_mdl = LLMBundle( self._canvas.get_tenant_id(), chat_model_config, diff --git a/agent/component/browser.py b/agent/component/browser.py index 19a8a91380..67116ed781 100644 --- a/agent/component/browser.py +++ b/agent/component/browser.py @@ -33,7 +33,7 @@ from urllib.request import Request, urlopen from agent.component.base import ComponentBase from agent.component.llm import LLMParam from api.db import FileType -from api.db.joint_services.tenant_model_service import get_model_config_from_provider_instance, get_model_type_by_name +from api.db.joint_services.tenant_model_service import resolve_model_config, resolve_model_type from api.db.services import duplicate_name from api.db.services.file_service import FileService from api.utils.file_utils import filename_type @@ -400,9 +400,9 @@ class Browser(ComponentBase, ABC): def _build_browser_llm(self): from browser_use.llm import ChatBrowserUse, ChatOpenAI - chat_model_config = get_model_config_from_provider_instance( + chat_model_config = resolve_model_config( self._canvas.get_tenant_id(), - get_model_type_by_name(self._canvas.get_tenant_id(), self._param.llm_id), + resolve_model_type(self._canvas.get_tenant_id(), self._param.llm_id), self._param.llm_id, ) cfg = self._as_model_config_dict(chat_model_config) diff --git a/agent/component/categorize.py b/agent/component/categorize.py index 7ad5315fdf..5361b2493f 100644 --- a/agent/component/categorize.py +++ b/agent/component/categorize.py @@ -21,7 +21,7 @@ from abc import ABC from common.constants import LLMType from api.db.services.llm_service import LLMBundle -from api.db.joint_services.tenant_model_service import get_model_config_from_provider_instance +from api.db.joint_services.tenant_model_service import resolve_model_config from agent.component.llm import LLMParam, LLM from common.connection_utils import timeout from rag.llm.chat_model import ERROR_PREFIX @@ -115,7 +115,7 @@ class Categorize(LLM, ABC): msg[-1]["content"] = query_value self.set_input_value(query_key, msg[-1]["content"]) self._param.update_prompt() - chat_model_config = get_model_config_from_provider_instance(self._canvas.get_tenant_id(), LLMType.CHAT, self._param.llm_id) + chat_model_config = resolve_model_config(self._canvas.get_tenant_id(), LLMType.CHAT, self._param.llm_id) chat_mdl = LLMBundle(self._canvas.get_tenant_id(), chat_model_config) user_prompt = """ diff --git a/agent/component/llm.py b/agent/component/llm.py index ae0e300ced..4b836f61fc 100644 --- a/agent/component/llm.py +++ b/agent/component/llm.py @@ -25,7 +25,7 @@ from functools import partial from common.constants import LLMType from api.db.services.dialog_service import _stream_with_think_delta from api.db.services.llm_service import LLMBundle -from api.db.joint_services.tenant_model_service import get_model_config_from_provider_instance, get_model_type_by_name +from api.db.joint_services.tenant_model_service import resolve_model_config, resolve_model_type from agent.component.base import ComponentBase, ComponentParamBase from common.connection_utils import timeout from rag.prompts.generator import tool_call_summary, message_fit_in, citation_prompt, structured_output_prompt @@ -88,9 +88,9 @@ class LLM(ComponentBase): def __init__(self, canvas, component_id, param: ComponentParamBase): super().__init__(canvas, component_id, param) - model_types = get_model_type_by_name(self._canvas.get_tenant_id(), self._param.llm_id) + model_types = resolve_model_type(self._canvas.get_tenant_id(), self._param.llm_id) model_type = "chat" if "chat" in model_types else model_types[0] - chat_model_config = get_model_config_from_provider_instance(self._canvas.get_tenant_id(), model_type, self._param.llm_id) + chat_model_config = resolve_model_config(self._canvas.get_tenant_id(), model_type, self._param.llm_id) self.chat_mdl = LLMBundle(self._canvas.get_tenant_id(), chat_model_config, max_retries=self._param.max_retries, retry_interval=self._param.delay_after_error) self.imgs = [] @@ -318,14 +318,14 @@ class LLM(ComponentBase): len(sys_file_imgs), max(0, prev_img_count + len(sys_file_imgs) - len(self.imgs)), ) - model_types = get_model_type_by_name(self._canvas.get_tenant_id(), self._param.llm_id) + model_types = resolve_model_type(self._canvas.get_tenant_id(), self._param.llm_id) if self.imgs and LLMType.IMAGE2TEXT.value in model_types: model_type = LLMType.IMAGE2TEXT.value elif LLMType.CHAT.value in model_types: model_type = LLMType.CHAT.value else: model_type = model_types[0] - model_config = get_model_config_from_provider_instance(self._canvas.get_tenant_id(), model_type, self._param.llm_id) + model_config = resolve_model_config(self._canvas.get_tenant_id(), model_type, self._param.llm_id) if self.imgs: self.chat_mdl = LLMBundle(self._canvas.get_tenant_id(), model_config, max_retries=self._param.max_retries, retry_interval=self._param.delay_after_error) diff --git a/agent/tools/retrieval.py b/agent/tools/retrieval.py index 2fb649e869..813ee22030 100644 --- a/agent/tools/retrieval.py +++ b/agent/tools/retrieval.py @@ -27,7 +27,7 @@ from api.db.services.knowledgebase_service import KnowledgebaseService from api.db.services.llm_service import LLMBundle from api.db.services.memory_service import MemoryService from api.db.joint_services import memory_message_service -from api.db.joint_services.tenant_model_service import get_tenant_default_model_by_type, get_model_config_from_provider_instance +from api.db.joint_services.tenant_model_service import get_tenant_default_model_by_type, resolve_model_config from common import settings from common.connection_utils import timeout from rag.app.tag import label_question @@ -116,12 +116,12 @@ class Retrieval(ToolBase, ABC): embd_mdl = None if embd_nms: tenant_id = self._canvas.get_tenant_id() - embd_model_config = get_model_config_from_provider_instance(tenant_id, LLMType.EMBEDDING, embd_nms[0]) + embd_model_config = resolve_model_config(tenant_id, LLMType.EMBEDDING, embd_nms[0]) embd_mdl = LLMBundle(tenant_id, embd_model_config) rerank_mdl = None if self._param.rerank_id: - rerank_model_config = get_model_config_from_provider_instance(kbs[0].tenant_id, LLMType.RERANK, self._param.rerank_id) + rerank_model_config = resolve_model_config(kbs[0].tenant_id, LLMType.RERANK, self._param.rerank_id) rerank_mdl = LLMBundle(kbs[0].tenant_id, rerank_model_config) vars = self.get_input_elements_from_text(query_text) diff --git a/api/apps/restful_apis/bot_api.py b/api/apps/restful_apis/bot_api.py index f2ee95ab21..dd2e94b3e8 100644 --- a/api/apps/restful_apis/bot_api.py +++ b/api/apps/restful_apis/bot_api.py @@ -37,7 +37,7 @@ from api.db.services.user_service import TenantService from common.metadata_utils import apply_meta_data_filter from api.db.services.search_service import SearchService from api.db.services.user_service import UserTenantService -from api.db.joint_services.tenant_model_service import get_tenant_default_model_by_type, get_model_config_from_provider_instance +from api.db.joint_services.tenant_model_service import get_tenant_default_model_by_type, resolve_model_config from common.misc_utils import thread_pool_exec from api.utils.api_utils import get_error_data_result, get_json_result, add_tenant_id_to_kwargs, get_result, get_request_json, server_error_response, validate_request from rag.app.tag import label_question @@ -405,7 +405,7 @@ async def retrieval_test_embedded(tenant_id=None): if meta_data_filter.get("method") in ["auto", "semi_auto"]: chat_id = search_config.get("chat_id", "") if chat_id: - chat_model_config = await thread_pool_exec(get_model_config_from_provider_instance, tenant_id, LLMType.CHAT, chat_id) + chat_model_config = await thread_pool_exec(resolve_model_config, tenant_id, LLMType.CHAT, chat_id) else: chat_model_config = await thread_pool_exec(get_tenant_default_model_by_type, tenant_id, LLMType.CHAT) chat_mdl = LLMBundle(tenant_id, chat_model_config) @@ -450,12 +450,12 @@ async def retrieval_test_embedded(tenant_id=None): if langs: _question = await cross_languages(kb.tenant_id, None, _question, langs) - embd_model_config = await thread_pool_exec(get_model_config_from_provider_instance, kb.tenant_id, LLMType.EMBEDDING, kb.embd_id) + embd_model_config = await thread_pool_exec(resolve_model_config, kb.tenant_id, LLMType.EMBEDDING, kb.embd_id) embd_mdl = LLMBundle(kb.tenant_id, embd_model_config) rerank_mdl = None if rerank_id: - rerank_model_config = await thread_pool_exec(get_model_config_from_provider_instance, tenant_id, LLMType.RERANK, rerank_id) + rerank_model_config = await thread_pool_exec(resolve_model_config, tenant_id, LLMType.RERANK, rerank_id) rerank_mdl = LLMBundle(kb.tenant_id, rerank_model_config) if req.get("keyword", False): @@ -523,7 +523,7 @@ async def related_questions_embedded(tenant_id=None): chat_id = search_config.get("chat_id", "") if chat_id: - chat_model_config = await thread_pool_exec(get_model_config_from_provider_instance, tenant_id, LLMType.CHAT, chat_id) + chat_model_config = await thread_pool_exec(resolve_model_config, tenant_id, LLMType.CHAT, chat_id) else: chat_model_config = await thread_pool_exec(get_tenant_default_model_by_type, tenant_id, LLMType.CHAT) chat_mdl = LLMBundle(tenant_id, chat_model_config) diff --git a/api/apps/restful_apis/chat_api.py b/api/apps/restful_apis/chat_api.py index 46e904e230..945baca77e 100644 --- a/api/apps/restful_apis/chat_api.py +++ b/api/apps/restful_apis/chat_api.py @@ -27,7 +27,7 @@ from quart import Response, request from api.apps import current_user, login_required from api.apps.restful_apis._generation_params import merge_generation_config, pop_generation_config -from api.db.joint_services.tenant_model_service import get_tenant_default_model_by_type, get_model_config_from_provider_instance, get_api_key +from api.db.joint_services.tenant_model_service import get_api_key, get_tenant_default_model_by_type, resolve_model_config from api.db.services.chunk_feedback_service import ChunkFeedbackService from api.db.services.conversation_service import ConversationService, structure_answer from api.db.services.dialog_service import DialogService, async_chat, gen_mindmap @@ -277,10 +277,10 @@ async def _validate_llm_id(llm_id, tenant_id, llm_setting=None): model_type = "chat" try: await thread_pool_exec( - get_model_config_from_provider_instance, + resolve_model_config, tenant_id=tenant_id, - model_name=llm_id, model_type=model_type, + model_ref=llm_id, ) except Exception as e: logging.error(f"Fail to get model config for {llm_id}: {e}") @@ -298,7 +298,7 @@ async def _validate_rerank_id(rerank_id, tenant_id): return None try: await thread_pool_exec( - get_model_config_from_provider_instance, + resolve_model_config, tenant_id=tenant_id, model_name=rerank_id, model_type="rerank", @@ -1127,7 +1127,7 @@ async def recommendation(): chat_id = search_config.get("chat_id", "") if chat_id: - chat_model_config = get_model_config_from_provider_instance(current_user.id, LLMType.CHAT, chat_id) + chat_model_config = resolve_model_config(current_user.id, LLMType.CHAT, chat_id) else: chat_model_config = get_tenant_default_model_by_type(current_user.id, LLMType.CHAT) chat_mdl = LLMBundle(current_user.id, chat_model_config) diff --git a/api/apps/restful_apis/chunk_api.py b/api/apps/restful_apis/chunk_api.py index b91519db0c..3de5463629 100644 --- a/api/apps/restful_apis/chunk_api.py +++ b/api/apps/restful_apis/chunk_api.py @@ -27,7 +27,7 @@ from quart import request from api.apps import login_required from api.db.joint_services.tenant_model_service import ( split_model_name, - get_model_config_from_provider_instance, + resolve_model_config, get_tenant_default_model_by_type, ) from api.db.db_models import Document, Task @@ -374,12 +374,12 @@ async def retrieval_test(tenant_id): e, kb = KnowledgebaseService.get_by_id(kb_ids[0]) if not e: return get_error_data_result(message="Dataset not found!") - embd_model_config = get_model_config_from_provider_instance(kb.tenant_id, LLMType.EMBEDDING, kb.embd_id) + embd_model_config = resolve_model_config(kb.tenant_id, LLMType.EMBEDDING, kb.embd_id) embd_mdl = LLMBundle(kb.tenant_id, embd_model_config) rerank_mdl = None if req.get("rerank_id"): - rerank_model_config = get_model_config_from_provider_instance(kb.tenant_id, LLMType.RERANK, req["rerank_id"]) + rerank_model_config = resolve_model_config(kb.tenant_id, LLMType.RERANK, req["rerank_id"]) rerank_mdl = LLMBundle(kb.tenant_id, rerank_model_config) if langs: @@ -902,7 +902,7 @@ async def add_chunk(tenant_id, dataset_id, document_id): d["doc_type_kwd"] = "image" embd_id = DocumentService.get_embd_id(document_id) - model_config = get_model_config_from_provider_instance(dataset_tenant_id, LLMType.EMBEDDING.value, embd_id) + model_config = resolve_model_config(dataset_tenant_id, LLMType.EMBEDDING.value, embd_id) embd_mdl = TenantLLMService.model_instance(model_config) v, c = embd_mdl.encode([doc.name, req["content"] if not d["question_kwd"] else "\n".join(d["question_kwd"])]) v = 0.1 * v[0] + 0.9 * v[1] @@ -1048,7 +1048,7 @@ async def update_chunk(tenant_id, dataset_id, document_id, chunk_id): d["doc_type_kwd"] = "image" embd_id = DocumentService.get_embd_id(document_id) - model_config = get_model_config_from_provider_instance(dataset_tenant_id, LLMType.EMBEDDING.value, embd_id) + model_config = resolve_model_config(dataset_tenant_id, LLMType.EMBEDDING.value, embd_id) embd_mdl = TenantLLMService.model_instance(model_config) if doc.parser_id == ParserType.QA: arr = [t for t in re.split(r"[\n\t]", d["content_with_weight"]) if len(t) > 1] diff --git a/api/apps/restful_apis/dify_retrieval_api.py b/api/apps/restful_apis/dify_retrieval_api.py index 6ede32cc40..d0245f52c3 100644 --- a/api/apps/restful_apis/dify_retrieval_api.py +++ b/api/apps/restful_apis/dify_retrieval_api.py @@ -27,7 +27,7 @@ from api.db.services.document_service import DocumentService from api.db.services.doc_metadata_service import DocMetadataService from api.db.services.knowledgebase_service import KnowledgebaseService from api.db.services.llm_service import LLMBundle -from api.db.joint_services.tenant_model_service import get_tenant_default_model_by_type, get_model_config_from_provider_instance +from api.db.joint_services.tenant_model_service import get_tenant_default_model_by_type, resolve_model_config from common.metadata_utils import meta_filter, convert_conditions from api.apps import login_required from api.utils.api_utils import add_tenant_id_to_kwargs, build_error_result, get_request_json, get_json_result @@ -261,7 +261,7 @@ async def retrieval(tenant_id): kb_id, ) return build_error_result(message="No authorization.", code=RetCode.AUTHENTICATION_ERROR) - model_config = get_model_config_from_provider_instance(kb.tenant_id, LLMType.EMBEDDING, kb.embd_id) + model_config = resolve_model_config(kb.tenant_id, LLMType.EMBEDDING, kb.embd_id) embd_mdl = LLMBundle(kb.tenant_id, model_config) if metadata_condition: doc_ids.extend(meta_filter(metas, convert_conditions(metadata_condition), metadata_condition.get("logic", "and"))) diff --git a/api/apps/restful_apis/models_api.py b/api/apps/restful_apis/models_api.py index c19d0f7d34..aa6e267179 100644 --- a/api/apps/restful_apis/models_api.py +++ b/api/apps/restful_apis/models_api.py @@ -19,6 +19,7 @@ from quart import request from api.apps import login_required from api.apps.services import models_api_service +from api.db.services.user_service import TenantService from api.utils.api_utils import ( add_tenant_id_to_kwargs, get_error_argument_result, @@ -39,6 +40,16 @@ def get_added_models(tenant_id: str): security: - ApiKeyAuth: [] parameters: + - in: query + name: owner_tenant_id + type: string + required: false + description: "If provided, list models from the owner tenant's scope after access validation." + - in: query + name: type + type: string + required: false + description: "Model type filter (chat, embedding, rerank, asr, vision, tts, ocr)." - in: header name: Authorization type: string @@ -70,8 +81,18 @@ def get_added_models(tenant_id: str): type: boolean """ model_type_filter = request.args.get("type") + owner_tenant_id = request.args.get("owner_tenant_id") try: - success, result = models_api_service.list_tenant_added_models(tenant_id, model_type_filter) + target_tenant_id = tenant_id + if owner_tenant_id: + if owner_tenant_id != tenant_id: + joined_tenants = TenantService.get_joined_tenants_by_user_id(tenant_id) + allowed_tenant_ids = {tenant_id, *(tenant["tenant_id"] for tenant in joined_tenants)} + if owner_tenant_id not in allowed_tenant_ids: + return get_error_data_result(message="Permission denied") + target_tenant_id = owner_tenant_id + + success, result = models_api_service.list_tenant_added_models(target_tenant_id, model_type_filter) if success: return get_result(data=result) else: diff --git a/api/apps/restful_apis/openai_api.py b/api/apps/restful_apis/openai_api.py index 4d913dccdd..8f345563b5 100644 --- a/api/apps/restful_apis/openai_api.py +++ b/api/apps/restful_apis/openai_api.py @@ -23,7 +23,7 @@ from api.apps import current_user, login_required from api.apps.restful_apis._generation_params import extract_generation_config, merge_generation_config from api.db.services.dialog_service import DialogService, async_chat from api.db.services.doc_metadata_service import DocMetadataService -from api.db.joint_services.tenant_model_service import get_model_config_from_provider_instance, get_api_key +from api.db.joint_services.tenant_model_service import resolve_model_config, get_api_key from api.utils.api_utils import get_error_data_result, get_request_json, validate_request from common.constants import RetCode, StatusEnum from common.metadata_utils import convert_conditions, meta_filter @@ -40,9 +40,9 @@ def _validate_llm_id(llm_id, tenant_id, llm_setting=None): model_type = "chat" try: - get_model_config_from_provider_instance( + resolve_model_config( tenant_id=tenant_id, - model_name=llm_id, + model_ref=llm_id, model_type=model_type, ) except Exception as e: diff --git a/api/apps/services/dataset_api_service.py b/api/apps/services/dataset_api_service.py index 4a0d5d9a08..a1a67f2ed5 100644 --- a/api/apps/services/dataset_api_service.py +++ b/api/apps/services/dataset_api_service.py @@ -18,8 +18,8 @@ import json import os import re -from api.db.joint_services.tenant_model_service import get_model_config_from_provider_instance -from common.constants import PAGERANK_FLD +from api.db.joint_services.tenant_model_service import resolve_model_config, resolve_model_id +from common.constants import PAGERANK_FLD, LLMType from common import settings from api.db.db_models import File from api.db.services.document_service import DocumentService, queue_raptor_o_graphrag_tasks @@ -28,6 +28,7 @@ from api.db.services.file_service import FileService from api.db.services.knowledgebase_service import KnowledgebaseService, validate_dataset_embedding_models from api.db.services.connector_service import Connector2KbService from api.db.services.task_service import GRAPH_RAPTOR_FAKE_DOC_ID, TaskService +from api.db.services.tenant_model_service import TenantModelService from api.db.services.user_service import TenantService, UserService, UserTenantService from common.constants import FileSource, StatusEnum from api.utils.api_utils import deep_merge, get_parser_config, remap_dictionary_keys, verify_embedding_availability @@ -328,6 +329,11 @@ async def update_dataset(tenant_id: str, dataset_id: str, req: dict): ok, err = verify_embedding_availability(req["embd_id"], tenant_id) if not ok: return False, err + ok, _ = TenantModelService.get_by_id(req["embd_id"]) + if ok: + req["tenant_embd_id"] = req["embd_id"] + else: + req["tenant_embd_id"] = resolve_model_id(tenant_id, LLMType.EMBEDDING, req["embd_id"]) if "pagerank" in req and req["pagerank"] != kb.pagerank: if os.environ.get("DOC_ENGINE", "elasticsearch") == "infinity": @@ -1011,7 +1017,7 @@ async def search(dataset_id: str, tenant_id: str, req: dict): if meta_data_filter.get("method") in ["auto", "semi_auto"]: chat_id = search_config.get("chat_id", "") if chat_id: - chat_model_config = get_model_config_from_provider_instance(tenant_id, LLMType.CHAT, search_config["chat_id"]) + chat_model_config = resolve_model_config(tenant_id, LLMType.CHAT, search_config["chat_id"]) else: chat_model_config = get_tenant_default_model_by_type(tenant_id, LLMType.CHAT) chat_mdl = LLMBundle(tenant_id, chat_model_config) @@ -1045,15 +1051,15 @@ async def search(dataset_id: str, tenant_id: str, req: dict): if langs: _question = await cross_languages(kb.tenant_id, None, _question, langs) if kb.embd_id: - embd_model_config = get_model_config_from_provider_instance(kb.tenant_id, LLMType.EMBEDDING, kb.embd_id) + embd_model_config = resolve_model_config(kb.tenant_id, LLMType.EMBEDDING, kb.embd_id) else: embd_model_config = get_tenant_default_model_by_type(kb.tenant_id, LLMType.EMBEDDING) embd_mdl = LLMBundle(kb.tenant_id, embd_model_config) rerank_mdl = None - rerank_id = search_config.get("rerank_id") or req.get("rerank_id") + rerank_id = req.get("rerank_id") or search_config.get("rerank_id") if rerank_id: - rerank_model_config = get_model_config_from_provider_instance(kb.tenant_id, LLMType.RERANK.value, rerank_id) + rerank_model_config = resolve_model_config(kb.tenant_id, LLMType.RERANK.value, rerank_id) rerank_mdl = LLMBundle(kb.tenant_id, rerank_model_config) if search_config.get("keyword", req.get("keyword", False)): @@ -1243,7 +1249,7 @@ def check_embedding(dataset_id: str, tenant_id: str, req: dict): if not ok: return False, err - embd_model_config = get_model_config_from_provider_instance(kb.tenant_id, LLMType.EMBEDDING, embd_id) + embd_model_config = resolve_model_config(kb.tenant_id, LLMType.EMBEDDING, embd_id) emb_mdl = LLMBundle(kb.tenant_id, embd_model_config) n = int(req.get("check_num", 5)) @@ -1400,7 +1406,7 @@ async def search_datasets(tenant_id: str, req: dict): if meta_data_filter.get("method") in ["auto", "semi_auto"]: chat_id = search_config.get("chat_id", "") if chat_id: - chat_model_config = get_model_config_from_provider_instance(tenant_id, LLMType.CHAT, search_config["chat_id"]) + chat_model_config = resolve_model_config(tenant_id, LLMType.CHAT, search_config["chat_id"]) else: chat_model_config = get_tenant_default_model_by_type(tenant_id, LLMType.CHAT) chat_mdl = LLMBundle(tenant_id, chat_model_config) @@ -1438,13 +1444,15 @@ async def search_datasets(tenant_id: str, req: dict): embd_mdl = None if kb.embd_id: - embd_model_config = get_model_config_from_provider_instance(kb.tenant_id, LLMType.EMBEDDING, kb.embd_id) - embd_mdl = LLMBundle(kb.tenant_id, embd_model_config) + embd_model_config = resolve_model_config(kb.tenant_id, LLMType.EMBEDDING, kb.embd_id) + else: + embd_model_config = get_tenant_default_model_by_type(kb.tenant_id, LLMType.EMBEDDING) + embd_mdl = LLMBundle(kb.tenant_id, embd_model_config) rerank_mdl = None - rerank_id = search_config.get("rerank_id") or req.get("rerank_id") + rerank_id = req.get("rerank_id") or search_config.get("rerank_id") if rerank_id: - rerank_model_config = get_model_config_from_provider_instance(kb.tenant_id, LLMType.RERANK.value, rerank_id) + rerank_model_config = resolve_model_config(kb.tenant_id, LLMType.RERANK.value, rerank_id) rerank_mdl = LLMBundle(kb.tenant_id, rerank_model_config) if search_config.get("keyword", req.get("keyword", False)): diff --git a/api/apps/services/models_api_service.py b/api/apps/services/models_api_service.py index 5b16e7dd99..c79028e904 100644 --- a/api/apps/services/models_api_service.py +++ b/api/apps/services/models_api_service.py @@ -285,14 +285,17 @@ def list_tenant_added_models(tenant_id: str, model_type_filter: str = None): model_records = TenantModelService.get_models_by_provider_ids_and_instance_ids(provider_ids, list({instance.id for instance in instances})) target_type_records = [record for record in model_records if record.model_type & model_type_filter_bin] if model_type_filter_bin else model_records - factory_rank_mapping = {factory["name"]: -_to_int(factory.get("rank", "500")) for factory in FACTORY_LLM_INFOS} model_rank_map: dict = {} for factory in FACTORY_LLM_INFOS: for llm in factory.get("llm", []): model_rank_map[(factory["name"], llm["llm_name"])] = _to_int(llm.get("rank", 500)) + factory_rank_mapping = {factory["name"]: -_to_int(factory.get("rank", "500")) for factory in FACTORY_LLM_INFOS} added_models = [ { + "model_id": model_record.id, + "tenant_id": provider_info_map[model_record.provider_id].tenant_id, + "tenant_name": tenant.name, "model_type": get_model_type_human(model_record.model_type), "name": model_record.model_name, "provider_id": model_record.provider_id, @@ -313,6 +316,9 @@ def list_tenant_added_models(tenant_id: str, model_type_filter: str = None): if not tei_already_added: added_models.append( { + "model_id": "", + "tenant_id": tenant.id, + "tenant_name": tenant.name, "model_type": ["embedding"], "name": tei_model, "provider_id": "", diff --git a/api/apps/services/provider_api_service.py b/api/apps/services/provider_api_service.py index 9f3f075895..31a25512e9 100644 --- a/api/apps/services/provider_api_service.py +++ b/api/apps/services/provider_api_service.py @@ -20,7 +20,7 @@ import asyncio from common.constants import LLMType, ActiveStatusEnum, ModelVerifyStatusEnum from common.settings import FACTORY_LLM_INFOS -from api.db.joint_services.tenant_model_service import get_model_config_from_provider_instance, delete_models_by_instance_ids, delete_instances_by_provider_ids +from api.db.joint_services.tenant_model_service import resolve_model_config, delete_models_by_instance_ids, delete_instances_by_provider_ids from api.db.services.tenant_model_provider_service import TenantModelProviderService from api.db.services.tenant_model_instance_service import TenantModelInstanceService from api.db.services.tenant_model_service import TenantModelService @@ -1141,7 +1141,7 @@ async def chat_to_model(tenant_id: str, provider_id_or_name: str, instance_id_or # Get model config composite_name = f"{model_name}@{instance_name}@{provider_name}" try: - model_config = get_model_config_from_provider_instance(tenant_id, LLMType.CHAT, composite_name) + model_config = resolve_model_config(tenant_id, LLMType.CHAT, composite_name) except LookupError: return False, f"Model '{composite_name}' not authorized" diff --git a/api/db/joint_services/memory_message_service.py b/api/db/joint_services/memory_message_service.py index d6014c8ce3..1e0d3bcf0e 100644 --- a/api/db/joint_services/memory_message_service.py +++ b/api/db/joint_services/memory_message_service.py @@ -27,7 +27,7 @@ from api.db.db_models import Task from api.db.services.task_service import TaskService from api.db.services.memory_service import MemoryService from api.db.services.llm_service import LLMBundle -from api.db.joint_services.tenant_model_service import get_model_config_from_provider_instance, get_model_config_by_id +from api.db.joint_services.tenant_model_service import resolve_model_config, get_model_config_by_id from api.utils.memory_utils import get_memory_type_human from memory.services.messages import MessageService from memory.services.query import MsgTextQuery, get_vector @@ -169,11 +169,11 @@ async def extract_by_llm(tenant_id: str, tenant_llm_id: str | None, extract_conf user_prompts.append({"role": "user", "content": PromptAssembler.assemble_user_prompt(conversation_content, conversation_time, conversation_time)}) if tenant_llm_id: try: - llm_config = get_model_config_by_id(tenant_id, tenant_llm_id) + llm_config = get_model_config_by_id(tenant_id, LLMType.CHAT, tenant_llm_id) except LookupError: - llm_config = get_model_config_from_provider_instance(tenant_id, LLMType.CHAT, llm_id) + llm_config = resolve_model_config(tenant_id, LLMType.CHAT, llm_id) else: - llm_config = get_model_config_from_provider_instance(tenant_id, LLMType.CHAT, llm_id) + llm_config = resolve_model_config(tenant_id, LLMType.CHAT, llm_id) with LLMBundle(tenant_id, llm_config) as llm: if task_id: TaskService.update_progress(task_id, {"progress": 0.15, "progress_msg": timestamp_to_date(current_timestamp()) + " " + "Prepared prompts and LLM."}) @@ -196,11 +196,11 @@ 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): if memory.tenant_embd_id: try: - embd_model_config = get_model_config_by_id(memory.tenant_id, memory.tenant_embd_id) + embd_model_config = get_model_config_by_id(memory.tenant_id, LLMType.EMBEDDING, memory.tenant_embd_id) except LookupError: - embd_model_config = get_model_config_from_provider_instance(memory.tenant_id, LLMType.EMBEDDING, memory.embd_id) + embd_model_config = resolve_model_config(memory.tenant_id, LLMType.EMBEDDING, memory.embd_id) else: - embd_model_config = get_model_config_from_provider_instance(memory.tenant_id, LLMType.EMBEDDING, memory.embd_id) + embd_model_config = resolve_model_config(memory.tenant_id, LLMType.EMBEDDING, memory.embd_id) with LLMBundle(memory.tenant_id, embd_model_config) as embedding_model: if task_id: TaskService.update_progress(task_id, {"progress": 0.65, "progress_msg": timestamp_to_date(current_timestamp()) + " " + "Prepared embedding model."}) @@ -270,7 +270,7 @@ def query_message(filter_dict: dict, params: dict): question = params["query"] question = question.strip() memory = memory_list[0] - embd_model_config = get_model_config_from_provider_instance(memory.tenant_id, LLMType.EMBEDDING, memory.embd_id) + embd_model_config = resolve_model_config(memory.tenant_id, LLMType.EMBEDDING, memory.embd_id) embd_model = LLMBundle(memory.tenant_id, embd_model_config) match_dense = get_vector(question, embd_model, similarity=params["similarity_threshold"]) match_text, _ = MsgTextQuery().question(question, min_match=params["similarity_threshold"]) diff --git a/api/db/joint_services/tenant_model_service.py b/api/db/joint_services/tenant_model_service.py index 8304dc4308..040c4636c4 100644 --- a/api/db/joint_services/tenant_model_service.py +++ b/api/db/joint_services/tenant_model_service.py @@ -195,10 +195,10 @@ def get_tenant_default_model_by_type(tenant_id: str, model_type: str | enum.Enum # Prefer resolving by tenant_model.id when available if model_id: try: - return get_model_config_by_id(tenant_id, model_id) + return get_model_config_by_id(tenant_id, model_type, model_id) except LookupError: logger.warning("tenant_model id=%s not found, falling back to model_name lookup for %s", model_id, model_name) - return get_model_config_from_provider_instance(tenant_id, model_type, model_name) + return resolve_model_config(tenant_id, model_type, model_name) def split_model_name(model_name: str): @@ -249,6 +249,11 @@ 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) @@ -312,16 +317,22 @@ def get_model_config_from_provider_instance(tenant_id, model_type: str | enum.En raise LookupError(f"Model {model_name} not found for model {model_type_val}") -def get_model_config_by_id(tenant_id: str, model_id: str): +def get_model_config_by_id(tenant_id: str, model_type: str | enum.Enum, model_id: str): """Get model config from tenant_model by its id (CharField PK).""" + model_type_val = model_type if isinstance(model_type, str) else model_type.value + model_type_bin = calculate_model_type(model_type_val) exist, model_obj = TenantModelService.get_by_id(model_id) if not exist: raise LookupError(f"TenantModel id={model_id} not found.") - if model_obj.status != ActiveStatusEnum.ACTIVE.value: + if model_obj.status == ActiveStatusEnum.INACTIVE.value: raise LookupError(f"TenantModel id={model_id} is disabled.") + if model_obj.status == ActiveStatusEnum.UNSUPPORTED.value: + raise LookupError(f"TenantModel id={model_id} cannot be used as {model_type_val} model.") + if not (model_obj.model_type & model_type_bin): + raise LookupError(f"TenantModel id={model_id} cannot be used as {model_type_val} model.") - provider_obj = TenantModelProviderService.get_by_id(model_obj.provider_id) - if not provider_obj: + ok, provider_obj = TenantModelProviderService.get_by_id(model_obj.provider_id) + if not ok: raise LookupError(f"Provider id={model_obj.provider_id} not found for model id={model_id}.") # Validate that tenant_id owns the provider or is a joined tenant of the provider's owner. @@ -331,8 +342,8 @@ def get_model_config_by_id(tenant_id: str, model_id: str): if provider_obj.tenant_id not in joined_tenant_ids: raise LookupError(f"Tenant {tenant_id} has no access to provider owned by tenant {provider_obj.tenant_id}.") - instance_obj = TenantModelInstanceService.get_by_id(model_obj.instance_id) - if not instance_obj: + ok, instance_obj = TenantModelInstanceService.get_by_id(model_obj.instance_id) + if not ok: raise LookupError(f"Instance id={model_obj.instance_id} not found for model id={model_id}.") api_key, is_tool, api_key_payload = _decode_api_key_config(instance_obj.api_key) @@ -344,7 +355,7 @@ def get_model_config_by_id(tenant_id: str, model_id: str): "api_key": api_key, "llm_name": model_obj.model_name, "api_base": extra_fields.get("base_url", ""), - "model_type": model_obj.model_type, + "model_type": model_type_val, "is_tools": model_extra.get("is_tools", is_tool), "max_tokens": model_extra.get("max_tokens") or 8192, } @@ -432,6 +443,17 @@ 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) diff --git a/api/db/services/dialog_service.py b/api/db/services/dialog_service.py index fb7ff21652..8987f25581 100644 --- a/api/db/services/dialog_service.py +++ b/api/db/services/dialog_service.py @@ -39,7 +39,7 @@ from api.utils.reference_metadata_utils import ( enrich_chunks_with_document_metadata, resolve_reference_metadata_preferences, ) -from api.db.joint_services.tenant_model_service import get_tenant_default_model_by_type, get_model_config_from_provider_instance, get_model_type_by_name, get_model_config_by_id +from api.db.joint_services.tenant_model_service import get_tenant_default_model_by_type, resolve_model_config, resolve_model_type, get_model_config_by_id from common.time_utils import current_timestamp, datetime_format from common.text_utils import normalize_arabic_digits from rag.advanced_rag.knowlege_compile.mind_map_extractor import MindMapExtractor @@ -295,23 +295,23 @@ async def async_chat_solo(dialog, messages, stream=True, session_id=None): if dialog.llm_id: if dialog.tenant_llm_id: try: - llm_types = get_model_type_by_name(dialog.tenant_id, dialog.llm_id) + llm_types = resolve_model_type(dialog.tenant_id, dialog.llm_id) if "chat" in llm_types: - model_config = get_model_config_by_id(dialog.tenant_id, dialog.tenant_llm_id) + model_config = get_model_config_by_id(dialog.tenant_id, LLMType.CHAT, dialog.tenant_llm_id) else: - model_config = get_model_config_from_provider_instance(dialog.tenant_id, LLMType.IMAGE2TEXT, dialog.llm_id) + model_config = resolve_model_config(dialog.tenant_id, LLMType.IMAGE2TEXT, dialog.llm_id) except LookupError: - llm_types = get_model_type_by_name(dialog.tenant_id, dialog.llm_id) + llm_types = resolve_model_type(dialog.tenant_id, dialog.llm_id) if "chat" in llm_types: - model_config = get_model_config_from_provider_instance(dialog.tenant_id, LLMType.CHAT, dialog.llm_id) + model_config = resolve_model_config(dialog.tenant_id, LLMType.CHAT, dialog.llm_id) else: - model_config = get_model_config_from_provider_instance(dialog.tenant_id, LLMType.IMAGE2TEXT, dialog.llm_id) + model_config = resolve_model_config(dialog.tenant_id, LLMType.IMAGE2TEXT, dialog.llm_id) else: - llm_types = get_model_type_by_name(dialog.tenant_id, dialog.llm_id) + llm_types = resolve_model_type(dialog.tenant_id, dialog.llm_id) if "chat" in llm_types: - model_config = get_model_config_from_provider_instance(dialog.tenant_id, LLMType.CHAT, dialog.llm_id) + model_config = resolve_model_config(dialog.tenant_id, LLMType.CHAT, dialog.llm_id) else: - model_config = get_model_config_from_provider_instance(dialog.tenant_id, LLMType.IMAGE2TEXT, dialog.llm_id) + model_config = resolve_model_config(dialog.tenant_id, LLMType.IMAGE2TEXT, dialog.llm_id) else: model_config = get_tenant_default_model_by_type(dialog.tenant_id, LLMType.CHAT) @@ -364,7 +364,7 @@ def get_models(dialog, trace_context=None, langfuse_session_id=None): if kbs and kbs[0].embd_id: embd_owner_tenant_id = kbs[0].tenant_id - embd_model_config = get_model_config_from_provider_instance(embd_owner_tenant_id, LLMType.EMBEDDING, kbs[0].embd_id) + embd_model_config = resolve_model_config(embd_owner_tenant_id, LLMType.EMBEDDING, kbs[0].embd_id) embd_mdl = LLMBundle(embd_owner_tenant_id, embd_model_config, trace_context=trace_context, langfuse_session_id=langfuse_session_id) if not embd_mdl: raise LookupError("Embedding model(%s) not found" % kbs[0].embd_id) @@ -372,11 +372,11 @@ def get_models(dialog, trace_context=None, langfuse_session_id=None): if dialog.llm_id: if dialog.tenant_llm_id: try: - chat_model_config = get_model_config_by_id(dialog.tenant_id, dialog.tenant_llm_id) + chat_model_config = get_model_config_by_id(dialog.tenant_id, LLMType.CHAT, dialog.tenant_llm_id) except LookupError: - chat_model_config = get_model_config_from_provider_instance(dialog.tenant_id, LLMType.CHAT, dialog.llm_id) + chat_model_config = resolve_model_config(dialog.tenant_id, LLMType.CHAT, dialog.llm_id) else: - chat_model_config = get_model_config_from_provider_instance(dialog.tenant_id, LLMType.CHAT, dialog.llm_id) + chat_model_config = resolve_model_config(dialog.tenant_id, LLMType.CHAT, dialog.llm_id) else: chat_model_config = get_tenant_default_model_by_type(dialog.tenant_id, LLMType.CHAT) @@ -385,11 +385,11 @@ def get_models(dialog, trace_context=None, langfuse_session_id=None): if dialog.rerank_id: if dialog.tenant_rerank_id: try: - rerank_model_config = get_model_config_by_id(dialog.tenant_id, dialog.tenant_rerank_id) + rerank_model_config = get_model_config_by_id(dialog.tenant_id, LLMType.RERANK, dialog.tenant_rerank_id) except LookupError: - rerank_model_config = get_model_config_from_provider_instance(dialog.tenant_id, LLMType.RERANK, dialog.rerank_id) + rerank_model_config = resolve_model_config(dialog.tenant_id, LLMType.RERANK, dialog.rerank_id) else: - rerank_model_config = get_model_config_from_provider_instance(dialog.tenant_id, LLMType.RERANK, dialog.rerank_id) + rerank_model_config = resolve_model_config(dialog.tenant_id, LLMType.RERANK, dialog.rerank_id) rerank_mdl = LLMBundle(dialog.tenant_id, rerank_model_config, trace_context=trace_context, langfuse_session_id=langfuse_session_id) if dialog.prompt_config.get("tts"): @@ -583,23 +583,23 @@ async def async_chat(dialog, messages, stream=True, **kwargs): if dialog.llm_id: if dialog.tenant_llm_id: try: - llm_types = get_model_type_by_name(dialog.tenant_id, dialog.llm_id) + llm_types = resolve_model_type(dialog.tenant_id, dialog.llm_id) if "chat" in llm_types: - llm_model_config = get_model_config_by_id(dialog.tenant_id, dialog.tenant_llm_id) + llm_model_config = get_model_config_by_id(dialog.tenant_id, LLMType.CHAT, dialog.tenant_llm_id) else: - llm_model_config = get_model_config_from_provider_instance(dialog.tenant_id, LLMType.IMAGE2TEXT, dialog.llm_id) + llm_model_config = resolve_model_config(dialog.tenant_id, LLMType.IMAGE2TEXT, dialog.llm_id) except LookupError: - llm_types = get_model_type_by_name(dialog.tenant_id, dialog.llm_id) + llm_types = resolve_model_type(dialog.tenant_id, dialog.llm_id) if "chat" in llm_types: - llm_model_config = get_model_config_from_provider_instance(dialog.tenant_id, LLMType.CHAT, dialog.llm_id) + llm_model_config = resolve_model_config(dialog.tenant_id, LLMType.CHAT, dialog.llm_id) else: - llm_model_config = get_model_config_from_provider_instance(dialog.tenant_id, LLMType.IMAGE2TEXT, dialog.llm_id) + llm_model_config = resolve_model_config(dialog.tenant_id, LLMType.IMAGE2TEXT, dialog.llm_id) else: - llm_types = get_model_type_by_name(dialog.tenant_id, dialog.llm_id) + llm_types = resolve_model_type(dialog.tenant_id, dialog.llm_id) if "chat" in llm_types: - llm_model_config = get_model_config_from_provider_instance(dialog.tenant_id, LLMType.CHAT, dialog.llm_id) + llm_model_config = resolve_model_config(dialog.tenant_id, LLMType.CHAT, dialog.llm_id) else: - llm_model_config = get_model_config_from_provider_instance(dialog.tenant_id, LLMType.IMAGE2TEXT, dialog.llm_id) + llm_model_config = resolve_model_config(dialog.tenant_id, LLMType.IMAGE2TEXT, dialog.llm_id) else: llm_model_config = get_tenant_default_model_by_type(dialog.tenant_id, LLMType.CHAT) @@ -1683,12 +1683,12 @@ async def async_ask(question, kb_ids, tenant_id, chat_llm_name=None, search_conf is_knowledge_graph = all([kb.parser_id == ParserType.KG for kb in kbs]) retriever = settings.retriever if not is_knowledge_graph else settings.kg_retriever embd_owner_tenant_id = kbs[0].tenant_id - embd_model_config = get_model_config_from_provider_instance(embd_owner_tenant_id, LLMType.EMBEDDING, embedding_list[0]) + embd_model_config = resolve_model_config(embd_owner_tenant_id, LLMType.EMBEDDING, embedding_list[0]) embd_mdl = LLMBundle(embd_owner_tenant_id, embd_model_config) - chat_model_config = get_model_config_from_provider_instance(tenant_id, LLMType.CHAT, chat_llm_name) + chat_model_config = resolve_model_config(tenant_id, LLMType.CHAT, chat_llm_name) chat_mdl = LLMBundle(tenant_id, chat_model_config) if rerank_id: - rerank_model_config = get_model_config_from_provider_instance(tenant_id, LLMType.RERANK, rerank_id) + rerank_model_config = resolve_model_config(tenant_id, LLMType.RERANK, rerank_id) rerank_mdl = LLMBundle(tenant_id, rerank_model_config) max_tokens = chat_mdl.max_length tenant_ids = list(set([kb.tenant_id for kb in kbs])) @@ -1794,16 +1794,16 @@ async def gen_mindmap(question, kb_ids, tenant_id, search_config={}): return {"error": "No KB selected"} tenant_ids = list(set([kb.tenant_id for kb in kbs])) embd_owner_tenant_id = kbs[0].tenant_id - embd_model_config = get_model_config_from_provider_instance(embd_owner_tenant_id, LLMType.EMBEDDING, kbs[0].embd_id) + embd_model_config = resolve_model_config(embd_owner_tenant_id, LLMType.EMBEDDING, kbs[0].embd_id) embd_mdl = LLMBundle(embd_owner_tenant_id, embd_model_config) chat_id = search_config.get("chat_id", "") if chat_id: - chat_model_config = get_model_config_from_provider_instance(tenant_id, LLMType.CHAT, chat_id) + chat_model_config = resolve_model_config(tenant_id, LLMType.CHAT, chat_id) else: chat_model_config = get_tenant_default_model_by_type(tenant_id, LLMType.CHAT) chat_mdl = LLMBundle(tenant_id, chat_model_config) if rerank_id: - rerank_model_config = get_model_config_from_provider_instance(tenant_id, LLMType.RERANK, rerank_id) + rerank_model_config = resolve_model_config(tenant_id, LLMType.RERANK, rerank_id) rerank_mdl = LLMBundle(tenant_id, rerank_model_config) if meta_data_filter: diff --git a/api/utils/api_utils.py b/api/utils/api_utils.py index 97398b3ed0..acb3849726 100644 --- a/api/utils/api_utils.py +++ b/api/utils/api_utils.py @@ -498,7 +498,7 @@ def check_duplicate_ids(ids, id_type="item"): def verify_embedding_availability(embd_id: str, tenant_id: str) -> tuple[bool, str | None]: - from api.db.joint_services.tenant_model_service import get_model_config_from_provider_instance + from api.db.joint_services.tenant_model_service import resolve_model_config """ Verifies availability of an embedding model for a specific tenant. @@ -534,7 +534,7 @@ def verify_embedding_availability(embd_id: str, tenant_id: str) -> tuple[bool, s (False, {'code': 101, 'message': "Unsupported model: "}) """ try: - get_model_config_from_provider_instance(tenant_id, LLMType.EMBEDDING, embd_id) + resolve_model_config(tenant_id, LLMType.EMBEDDING, embd_id) except LookupError as e: return False, str(e) except OperationalError as e: diff --git a/api/utils/validation_utils.py b/api/utils/validation_utils.py index fdbc31889a..31ba61b845 100644 --- a/api/utils/validation_utils.py +++ b/api/utils/validation_utils.py @@ -617,13 +617,14 @@ class CreateDatasetReq(Base): Validation pipeline: 1. Structural format verification 2. Component non-empty check - 3. Value normalization + 3. Tenant model id passthrough + 4. Value normalization Args: v (str): Raw model identifier Returns: - str: Validated @ format + str: Validated @ format or tenant_model id Raises: PydanticCustomError: For these violations: @@ -633,11 +634,15 @@ class CreateDatasetReq(Base): Examples: Valid: "text-embedding-3-large@openai" + Valid: "2f3c0f9c7b1d11f0a1b2c3d4e5f67890" # tenant_model.id Invalid: "invalid_model" (no @) Invalid: "@openai" (empty model_name) Invalid: "text-embedding-3-large@" (empty provider) """ if isinstance(v, str): + if re.fullmatch(r"[0-9a-fA-F]{32}", v): + return v + if "@" not in v: raise PydanticCustomError("format_invalid", "Embedding model identifier must follow @ format") diff --git a/rag/app/naive.py b/rag/app/naive.py index 024c46b154..fed97d1a63 100644 --- a/rag/app/naive.py +++ b/rag/app/naive.py @@ -36,7 +36,7 @@ from api.db.joint_services.tenant_model_service import ( ensure_opendataloader_from_env, ensure_paddleocr_from_env, get_first_provider_model_name, - get_model_config_from_provider_instance, + resolve_model_config, get_tenant_default_model_by_type, ) from rag.utils.file_utils import extract_embed_file, extract_links_from_pdf, extract_links_from_docx, extract_html @@ -150,7 +150,7 @@ def by_mineru( if mineru_llm_name: try: - ocr_model_config = get_model_config_from_provider_instance(tenant_id, LLMType.OCR, mineru_llm_name) + ocr_model_config = resolve_model_config(tenant_id, LLMType.OCR, mineru_llm_name) ocr_model = LLMBundle(tenant_id=tenant_id, model_config=ocr_model_config, lang=lang) pdf_parser = ocr_model.mdl @@ -226,7 +226,7 @@ def by_opendataloader( if opendataloader_llm_name: try: - ocr_model_config = get_model_config_from_provider_instance(tenant_id, LLMType.OCR, opendataloader_llm_name) + ocr_model_config = resolve_model_config(tenant_id, LLMType.OCR, opendataloader_llm_name) ocr_model = LLMBundle(tenant_id=tenant_id, model_config=ocr_model_config, lang=lang) pdf_parser = ocr_model.mdl parse_options = {k: kwargs[k] for k in ("hybrid", "image_output", "sanitize") if k in kwargs} @@ -280,7 +280,7 @@ def by_paddleocr( if paddleocr_llm_name: try: - ocr_model_config = get_model_config_from_provider_instance(tenant_id, LLMType.OCR, paddleocr_llm_name) + ocr_model_config = resolve_model_config(tenant_id, LLMType.OCR, paddleocr_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( @@ -327,7 +327,7 @@ def by_somark( if somark_llm_name: try: try: - ocr_model_config = get_model_config_from_provider_instance(tenant_id, LLMType.OCR, somark_llm_name) + ocr_model_config = resolve_model_config(tenant_id, LLMType.OCR, somark_llm_name) except Exception: if "@" in somark_llm_name: raise @@ -363,7 +363,7 @@ def by_plaintext(filename, binary=None, from_page=0, to_page=MAXIMUM_PAGE_NUMBER tenant_id = kwargs.get("tenant_id") if not tenant_id: raise ValueError("tenant_id is required when using vision layout recognizer") - vision_model_config = get_model_config_from_provider_instance(tenant_id, LLMType.IMAGE2TEXT, layout_recognizer) + vision_model_config = resolve_model_config(tenant_id, LLMType.IMAGE2TEXT, layout_recognizer) vision_model = LLMBundle( tenant_id, model_config=vision_model_config, diff --git a/rag/app/picture.py b/rag/app/picture.py index 348df9abc1..db6e76de49 100644 --- a/rag/app/picture.py +++ b/rag/app/picture.py @@ -25,7 +25,7 @@ import numpy as np from PIL import Image from api.db.services.llm_service import LLMBundle -from api.db.joint_services.tenant_model_service import get_tenant_default_model_by_type, get_first_provider_model_name, get_model_config_from_provider_instance, ensure_paddleocr_from_env +from api.db.joint_services.tenant_model_service import get_tenant_default_model_by_type, get_first_provider_model_name, resolve_model_config, ensure_paddleocr_from_env from common.constants import LLMType from common.parser_config_utils import normalize_layout_recognizer from common.string_utils import clean_markdown_block @@ -124,7 +124,7 @@ def _try_paddleocr_image(filename, binary, tenant_id, parser_config, callback): if not paddleocr_llm_name: return "" - ocr_model_config = get_model_config_from_provider_instance(tenant_id, LLMType.OCR, paddleocr_llm_name) + ocr_model_config = resolve_model_config(tenant_id, LLMType.OCR, paddleocr_llm_name) ocr_model = LLMBundle(tenant_id=tenant_id, model_config=ocr_model_config) pdf_parser = ocr_model.mdl diff --git a/rag/benchmark.py b/rag/benchmark.py index 7866ab6930..1a3d4e8f96 100644 --- a/rag/benchmark.py +++ b/rag/benchmark.py @@ -25,7 +25,7 @@ from common import settings from common.constants import LLMType from api.db.services.llm_service import LLMBundle from api.db.services.knowledgebase_service import KnowledgebaseService -from api.db.joint_services.tenant_model_service import get_model_config_from_provider_instance +from api.db.joint_services.tenant_model_service import resolve_model_config from common.misc_utils import get_uuid from rag.nlp import tokenize, search from ranx import evaluate @@ -43,7 +43,7 @@ class Benchmark: e, self.kb = KnowledgebaseService.get_by_id(kb_id) self.similarity_threshold = self.kb.similarity_threshold self.vector_similarity_weight = self.kb.vector_similarity_weight - embd_model_config = get_model_config_from_provider_instance(self.kb.tenant_id, LLMType.EMBEDDING, self.kb.embd_id) + embd_model_config = resolve_model_config(self.kb.tenant_id, LLMType.EMBEDDING, self.kb.embd_id) self.embd_mdl = LLMBundle(self.kb.tenant_id, embd_model_config, lang=self.kb.language) self.tenant_id = "" self.index_name = "" diff --git a/rag/flow/parser/parser.py b/rag/flow/parser/parser.py index 920df678a6..ab2bb6197f 100644 --- a/rag/flow/parser/parser.py +++ b/rag/flow/parser/parser.py @@ -32,7 +32,7 @@ from api.db.joint_services.tenant_model_service import ( ensure_opendataloader_from_env, ensure_paddleocr_from_env, get_first_provider_model_name, - get_model_config_from_provider_instance, + resolve_model_config, get_tenant_default_model_by_type, ) from common import settings @@ -351,7 +351,7 @@ class Parser(ProcessBase): elif lowered.endswith("@somark"): # Keep the full 3-segment ``@@`` # form produced by the new Tenant LLM Provider UI (#14595); - # ``get_model_config_from_provider_instance`` -> ``split_model_name`` + # ``resolve_model_config`` -> ``split_model_name`` # downstream requires all three segments. parser_model_name = raw_parse_method parse_method = "SoMark" @@ -389,7 +389,7 @@ class Parser(ProcessBase): raise RuntimeError("MinerU model not configured. Please add MinerU in Model Providers or set MINERU_* env.") tenant_id = self._canvas._tenant_id - ocr_model_config = get_model_config_from_provider_instance(tenant_id, LLMType.OCR, parser_model_name) + ocr_model_config = resolve_model_config(tenant_id, LLMType.OCR, parser_model_name) ocr_model = LLMBundle(tenant_id, ocr_model_config, lang=conf.get("lang", "Chinese")) pdf_parser = ocr_model.mdl @@ -462,7 +462,7 @@ class Parser(ProcessBase): raise RuntimeError("OpenDataLoader model not configured. Please add OpenDataLoader in Model Providers.") tenant_id = self._canvas._tenant_id - ocr_model_config = get_model_config_from_provider_instance(tenant_id, LLMType.OCR, 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 @@ -524,7 +524,7 @@ class Parser(ProcessBase): tenant_id = self._canvas._tenant_id try: - ocr_model_config = get_model_config_from_provider_instance(tenant_id, LLMType.OCR, parser_model_name) + ocr_model_config = resolve_model_config(tenant_id, LLMType.OCR, parser_model_name) except Exception: if "@" in parser_model_name: raise @@ -614,7 +614,7 @@ class Parser(ProcessBase): raise RuntimeError("PaddleOCR model not configured. Please add PaddleOCR in Model Providers or set PADDLEOCR_* env.") tenant_id = self._canvas._tenant_id - ocr_model_config = get_model_config_from_provider_instance(tenant_id, LLMType.OCR, 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 @@ -644,7 +644,7 @@ class Parser(ProcessBase): # Vision parser treats each page as a large image block. else: if conf.get("parse_method"): - vision_model_config = get_model_config_from_provider_instance(self._canvas._tenant_id, LLMType.IMAGE2TEXT, conf["parse_method"]) + vision_model_config = resolve_model_config(self._canvas._tenant_id, LLMType.IMAGE2TEXT, conf["parse_method"]) else: vision_model_config = get_tenant_default_model_by_type(self._canvas._tenant_id, LLMType.IMAGE2TEXT) vision_model = LLMBundle(self._canvas._tenant_id, vision_model_config, lang=self._param.setups["pdf"].get("lang")) @@ -1117,7 +1117,7 @@ class Parser(ProcessBase): else: lang = conf["lang"] # use VLM to describe the picture - cv_model_config = get_model_config_from_provider_instance(self._canvas.get_tenant_id(), LLMType.IMAGE2TEXT, conf["parse_method"]) + cv_model_config = resolve_model_config(self._canvas.get_tenant_id(), LLMType.IMAGE2TEXT, conf["parse_method"]) cv_model = LLMBundle(self._canvas.get_tenant_id(), cv_model_config, lang=lang) img_binary = io.BytesIO() img.save(img_binary, format="JPEG") @@ -1153,7 +1153,7 @@ class Parser(ProcessBase): tmpf.write(blob) tmpf.flush() tmp_path = os.path.abspath(tmpf.name) - seq2txt_model_config = get_model_config_from_provider_instance(self._canvas.get_tenant_id(), LLMType.SPEECH2TEXT, vlm["llm_id"]) + seq2txt_model_config = resolve_model_config(self._canvas.get_tenant_id(), LLMType.SPEECH2TEXT, vlm["llm_id"]) seq2txt_mdl = LLMBundle(self._canvas.get_tenant_id(), seq2txt_model_config) txt = seq2txt_mdl.transcription(tmp_path) @@ -1166,7 +1166,7 @@ class Parser(ProcessBase): conf = self._param.setups["video"] vlm = conf.get("vlm") self.set_output("output_format", conf["output_format"]) - cv_model_config = get_model_config_from_provider_instance(self._canvas.get_tenant_id(), LLMType.IMAGE2TEXT, vlm["llm_id"]) + cv_model_config = resolve_model_config(self._canvas.get_tenant_id(), LLMType.IMAGE2TEXT, vlm["llm_id"]) cv_mdl = LLMBundle(self._canvas.get_tenant_id(), cv_model_config) video_prompt = str(conf.get("prompt", "") or "") txt = asyncio.run(cv_mdl.async_chat(system="", history=[], gen_conf={}, video_bytes=blob, filename=name, video_prompt=video_prompt)) diff --git a/rag/flow/parser/utils.py b/rag/flow/parser/utils.py index 48df57ea3c..4185bb2bfb 100644 --- a/rag/flow/parser/utils.py +++ b/rag/flow/parser/utils.py @@ -21,7 +21,7 @@ from docx import Document from api.db.services.llm_service import LLMBundle from api.db.joint_services.tenant_model_service import ( get_tenant_default_model_by_type, - get_model_config_from_provider_instance, + resolve_model_config, ) from common.constants import LLMType from deepdoc.parser.figure_parser import VisionFigureParser @@ -170,7 +170,7 @@ def enhance_media_sections_with_vision( try: try: - vision_model_config = get_model_config_from_provider_instance(tenant_id, LLMType.IMAGE2TEXT, vlm_conf["llm_id"]) + vision_model_config = resolve_model_config(tenant_id, LLMType.IMAGE2TEXT, vlm_conf["llm_id"]) except Exception: vision_model_config = get_tenant_default_model_by_type(tenant_id, LLMType.IMAGE2TEXT) vision_model = LLMBundle(tenant_id, vision_model_config) diff --git a/rag/flow/tokenizer/tokenizer.py b/rag/flow/tokenizer/tokenizer.py index 63b9df0f96..dab4cf536e 100644 --- a/rag/flow/tokenizer/tokenizer.py +++ b/rag/flow/tokenizer/tokenizer.py @@ -21,7 +21,7 @@ import numpy as np from common.constants import LLMType from api.db.services.knowledgebase_service import KnowledgebaseService from api.db.services.llm_service import LLMBundle -from api.db.joint_services.tenant_model_service import get_tenant_default_model_by_type, get_model_config_from_provider_instance, get_model_config_by_id +from api.db.joint_services.tenant_model_service import get_tenant_default_model_by_type, resolve_model_config, get_model_config_by_id from common.connection_utils import timeout from rag.flow.base import ProcessBase, ProcessParamBase from rag.flow.parser.pdf_chunk_metadata import finalize_pdf_chunk @@ -63,11 +63,11 @@ class Tokenizer(ProcessBase): e, kb = KnowledgebaseService.get_by_id(self._canvas._kb_id) if kb.tenant_embd_id: try: - embd_model_config = get_model_config_by_id(self._canvas._tenant_id, kb.tenant_embd_id) + embd_model_config = get_model_config_by_id(self._canvas._tenant_id, LLMType.EMBEDDING, kb.tenant_embd_id) except LookupError: - embd_model_config = get_model_config_from_provider_instance(self._canvas._tenant_id, LLMType.EMBEDDING, kb.embd_id) + embd_model_config = resolve_model_config(self._canvas._tenant_id, LLMType.EMBEDDING, kb.embd_id) else: - embd_model_config = get_model_config_from_provider_instance(self._canvas._tenant_id, LLMType.EMBEDDING, kb.embd_id) + embd_model_config = resolve_model_config(self._canvas._tenant_id, LLMType.EMBEDDING, kb.embd_id) else: embd_model_config = get_tenant_default_model_by_type(self._canvas._tenant_id, LLMType.EMBEDDING) embedding_model = LLMBundle(self._canvas._tenant_id, embd_model_config) diff --git a/rag/graphrag/general/smoke.py b/rag/graphrag/general/smoke.py index 687ef2ba47..c0db81d553 100644 --- a/rag/graphrag/general/smoke.py +++ b/rag/graphrag/general/smoke.py @@ -24,7 +24,7 @@ from common.constants import LLMType from api.db.services.document_service import DocumentService from api.db.services.knowledgebase_service import KnowledgebaseService from api.db.services.llm_service import LLMBundle -from api.db.joint_services.tenant_model_service import get_tenant_default_model_by_type, get_model_config_from_provider_instance, get_model_config_by_id +from api.db.joint_services.tenant_model_service import get_tenant_default_model_by_type, resolve_model_config, get_model_config_by_id from rag.graphrag.general.graph_extractor import GraphExtractor from rag.graphrag.general.index import update_graph, with_resolution, with_community from common import settings @@ -76,11 +76,11 @@ async def main(): _, kb = KnowledgebaseService.get_by_id(kb_id) if kb.tenant_embd_id: try: - embd_model_config = get_model_config_by_id(args.tenant_id, kb.tenant_embd_id) + embd_model_config = get_model_config_by_id(args.tenant_id, LLMType.EMBEDDING, kb.tenant_embd_id) except LookupError: - embd_model_config = get_model_config_from_provider_instance(args.tenant_id, LLMType.EMBEDDING, kb.embd_id) + embd_model_config = resolve_model_config(args.tenant_id, LLMType.EMBEDDING, kb.embd_id) else: - embd_model_config = get_model_config_from_provider_instance(args.tenant_id, LLMType.EMBEDDING, kb.embd_id) + embd_model_config = resolve_model_config(args.tenant_id, LLMType.EMBEDDING, kb.embd_id) embed_bdl = LLMBundle(args.tenant_id, embd_model_config) graph, doc_ids = await update_graph( diff --git a/rag/graphrag/light/smoke.py b/rag/graphrag/light/smoke.py index 0084aca377..83a52a290e 100644 --- a/rag/graphrag/light/smoke.py +++ b/rag/graphrag/light/smoke.py @@ -24,7 +24,7 @@ from common.constants import LLMType from api.db.services.document_service import DocumentService from api.db.services.knowledgebase_service import KnowledgebaseService from api.db.services.llm_service import LLMBundle -from api.db.joint_services.tenant_model_service import get_tenant_default_model_by_type, get_model_config_from_provider_instance, get_model_config_by_id +from api.db.joint_services.tenant_model_service import get_tenant_default_model_by_type, resolve_model_config, get_model_config_by_id from rag.graphrag.general.index import update_graph from rag.graphrag.light.graph_extractor import GraphExtractor from common import settings @@ -77,11 +77,11 @@ async def main(): _, kb = KnowledgebaseService.get_by_id(kb_id) if kb.tenant_embd_id: try: - embd_model_config = get_model_config_by_id(args.tenant_id, kb.tenant_embd_id) + embd_model_config = get_model_config_by_id(args.tenant_id, LLMType.EMBEDDING, kb.tenant_embd_id) except LookupError: - embd_model_config = get_model_config_from_provider_instance(args.tenant_id, LLMType.EMBEDDING, kb.embd_id) + embd_model_config = resolve_model_config(args.tenant_id, LLMType.EMBEDDING, kb.embd_id) else: - embd_model_config = get_model_config_from_provider_instance(args.tenant_id, LLMType.EMBEDDING, kb.embd_id) + embd_model_config = resolve_model_config(args.tenant_id, LLMType.EMBEDDING, kb.embd_id) embed_bdl = LLMBundle(args.tenant_id, embd_model_config) graph, doc_ids = await update_graph( diff --git a/rag/graphrag/search.py b/rag/graphrag/search.py index 70d846a17e..d5a5efaf62 100644 --- a/rag/graphrag/search.py +++ b/rag/graphrag/search.py @@ -300,7 +300,7 @@ if __name__ == "__main__": from common.constants import LLMType from api.db.services.knowledgebase_service import KnowledgebaseService from api.db.services.llm_service import LLMBundle - from api.db.joint_services.tenant_model_service import get_tenant_default_model_by_type, get_model_config_from_provider_instance, get_model_config_by_id + from api.db.joint_services.tenant_model_service import get_tenant_default_model_by_type, resolve_model_config, get_model_config_by_id from rag.nlp import search settings.init_settings() @@ -316,11 +316,11 @@ if __name__ == "__main__": _, kb = KnowledgebaseService.get_by_id(kb_id) if kb.tenant_embd_id: try: - embd_model_config = get_model_config_by_id(args.tenant_id, kb.tenant_embd_id) + embd_model_config = get_model_config_by_id(args.tenant_id, LLMType.EMBEDDING, kb.tenant_embd_id) except LookupError: - embd_model_config = get_model_config_from_provider_instance(args.tenant_id, LLMType.EMBEDDING, kb.embd_id) + embd_model_config = resolve_model_config(args.tenant_id, LLMType.EMBEDDING, kb.embd_id) else: - embd_model_config = get_model_config_from_provider_instance(args.tenant_id, LLMType.EMBEDDING, kb.embd_id) + embd_model_config = resolve_model_config(args.tenant_id, LLMType.EMBEDDING, kb.embd_id) embed_bdl = LLMBundle(args.tenant_id, embd_model_config) kg = KGSearch(settings.docStoreConn) diff --git a/rag/prompts/generator.py b/rag/prompts/generator.py index 49103d56a0..ab2902c66d 100644 --- a/rag/prompts/generator.py +++ b/rag/prompts/generator.py @@ -254,14 +254,14 @@ async def question_proposal(chat_mdl, content, topn=3): async def full_question(tenant_id=None, llm_id=None, messages=[], language=None, chat_mdl=None): from common.constants import LLMType from api.db.services.llm_service import LLMBundle - from api.db.joint_services.tenant_model_service import get_model_config_from_provider_instance, get_model_type_by_name + from api.db.joint_services.tenant_model_service import resolve_model_config, resolve_model_type if not chat_mdl: - model_types = get_model_type_by_name(tenant_id, llm_id) + model_types = resolve_model_type(tenant_id, llm_id) if "image2text" in model_types: - chat_model_config = get_model_config_from_provider_instance(tenant_id, LLMType.IMAGE2TEXT, llm_id) + chat_model_config = resolve_model_config(tenant_id, LLMType.IMAGE2TEXT, llm_id) else: - chat_model_config = get_model_config_from_provider_instance(tenant_id, LLMType.CHAT, llm_id) + chat_model_config = resolve_model_config(tenant_id, LLMType.CHAT, llm_id) chat_mdl = LLMBundle(tenant_id, chat_model_config) conv = [] for m in messages: @@ -290,15 +290,15 @@ async def full_question(tenant_id=None, llm_id=None, messages=[], language=None, async def cross_languages(tenant_id, llm_id, query, languages=[]): from common.constants import LLMType from api.db.services.llm_service import LLMBundle - from api.db.joint_services.tenant_model_service import get_model_config_from_provider_instance, get_tenant_default_model_by_type, get_model_type_by_name + from api.db.joint_services.tenant_model_service import resolve_model_config, get_tenant_default_model_by_type, resolve_model_type - if llm_id and "image2text" in get_model_type_by_name(tenant_id, llm_id): - chat_model_config = get_model_config_from_provider_instance(tenant_id, LLMType.IMAGE2TEXT, llm_id) + if llm_id and "image2text" in resolve_model_type(tenant_id, llm_id): + chat_model_config = resolve_model_config(tenant_id, LLMType.IMAGE2TEXT, llm_id) else: if not llm_id: chat_model_config = get_tenant_default_model_by_type(tenant_id, LLMType.CHAT) else: - chat_model_config = get_model_config_from_provider_instance(tenant_id, LLMType.CHAT, llm_id) + chat_model_config = resolve_model_config(tenant_id, LLMType.CHAT, llm_id) chat_mdl = LLMBundle(tenant_id, chat_model_config) rendered_sys_prompt = PROMPT_JINJA_ENV.from_string(CROSS_LANGUAGES_SYS_PROMPT_TEMPLATE).render() rendered_user_prompt = PROMPT_JINJA_ENV.from_string(CROSS_LANGUAGES_USER_PROMPT_TEMPLATE).render(query=query, languages=languages) diff --git a/rag/svr/task_executor.py b/rag/svr/task_executor.py index 5dd2ab98a1..2f785cce96 100644 --- a/rag/svr/task_executor.py +++ b/rag/svr/task_executor.py @@ -79,7 +79,7 @@ from api.db.services.doc_metadata_service import DocMetadataService from api.db.services.llm_service import LLMBundle from api.db.services.task_service import TaskService, has_canceled, CANVAS_DEBUG_DOC_ID, GRAPH_RAPTOR_FAKE_DOC_ID from api.db.services.file2document_service import File2DocumentService -from api.db.joint_services.tenant_model_service import get_tenant_default_model_by_type, get_model_config_from_provider_instance, get_model_config_by_id +from api.db.joint_services.tenant_model_service import get_tenant_default_model_by_type, resolve_model_config, get_model_config_by_id from common.versions import get_ragflow_version from api.db.db_models import close_connection from rag.app import laws, paper, presentation, manual, qa, table, book, resume, picture, naive, one, audio, email, tag @@ -424,7 +424,7 @@ async def build_chunks(task, progress_callback): if task["parser_config"].get("auto_keywords", 0): st = timer() progress_callback(msg="Start to generate keywords for every chunk ...") - chat_model_config = get_model_config_from_provider_instance(task["tenant_id"], LLMType.CHAT, task["llm_id"]) + chat_model_config = resolve_model_config(task["tenant_id"], LLMType.CHAT, task["llm_id"]) chat_mdl = LLMBundle(task["tenant_id"], chat_model_config, lang=task["language"]) async def doc_keyword_extraction(chat_mdl, d, topn): @@ -461,7 +461,7 @@ async def build_chunks(task, progress_callback): if task["parser_config"].get("auto_questions", 0): st = timer() progress_callback(msg="Start to generate questions for every chunk ...") - chat_model_config = get_model_config_from_provider_instance(task["tenant_id"], LLMType.CHAT, task["llm_id"]) + chat_model_config = resolve_model_config(task["tenant_id"], LLMType.CHAT, task["llm_id"]) chat_mdl = LLMBundle(task["tenant_id"], chat_model_config, lang=task["language"]) async def doc_question_proposal(chat_mdl, d, topn): @@ -497,7 +497,7 @@ async def build_chunks(task, progress_callback): if task["parser_config"].get("enable_metadata", False) and (task["parser_config"].get("metadata") or task["parser_config"].get("built_in_metadata")): st = timer() progress_callback(msg="Start to generate meta-data for every chunk ...") - chat_model_config = get_model_config_from_provider_instance(task["tenant_id"], LLMType.CHAT, task["llm_id"]) + chat_model_config = resolve_model_config(task["tenant_id"], LLMType.CHAT, task["llm_id"]) chat_mdl = LLMBundle(task["tenant_id"], chat_model_config, lang=task["language"]) async def gen_metadata_task(chat_mdl, d): @@ -570,7 +570,7 @@ async def build_chunks(task, progress_callback): set_tags_to_cache(kb_ids, all_tags) else: all_tags = json.loads(all_tags) - chat_model_config = get_model_config_from_provider_instance(tenant_id, LLMType.CHAT, task["llm_id"]) + chat_model_config = resolve_model_config(tenant_id, LLMType.CHAT, task["llm_id"]) chat_mdl = LLMBundle(task["tenant_id"], chat_model_config, lang=task["language"]) docs_to_tag = [] @@ -635,7 +635,7 @@ async def build_chunks(task, progress_callback): @timed_with_recording def build_TOC(task, docs, progress_callback): progress_callback(msg="Start to generate table of content ...") - chat_model_config = get_model_config_from_provider_instance(task["tenant_id"], LLMType.CHAT, task["llm_id"]) + chat_model_config = resolve_model_config(task["tenant_id"], LLMType.CHAT, task["llm_id"]) chat_mdl = LLMBundle(task["tenant_id"], chat_model_config, lang=task["language"]) docs = sorted( docs, @@ -820,11 +820,11 @@ async def run_dataflow(task: dict): embedding_id = kb.embd_id if kb.tenant_embd_id: try: - embd_model_config = get_model_config_by_id(task["tenant_id"], kb.tenant_embd_id) + embd_model_config = get_model_config_by_id(task["tenant_id"], LLMType.EMBEDDING, kb.tenant_embd_id) except LookupError: - embd_model_config = get_model_config_from_provider_instance(task["tenant_id"], LLMType.EMBEDDING, embedding_id) + embd_model_config = resolve_model_config(task["tenant_id"], LLMType.EMBEDDING, embedding_id) else: - embd_model_config = get_model_config_from_provider_instance(task["tenant_id"], LLMType.EMBEDDING, embedding_id) + embd_model_config = resolve_model_config(task["tenant_id"], LLMType.EMBEDDING, embedding_id) embedding_model = LLMBundle(task["tenant_id"], embd_model_config) @timeout(60) @@ -1426,7 +1426,7 @@ async def do_handle_task(task): try: # bind embedding model if task_embedding_id: - embd_model_config = get_model_config_from_provider_instance(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) @@ -1474,7 +1474,7 @@ async def do_handle_task(task): return # bind LLM for raptor - chat_model_config = get_model_config_from_provider_instance(task_tenant_id, LLMType.CHAT, kb_task_llm_id) + chat_model_config = resolve_model_config(task_tenant_id, LLMType.CHAT, kb_task_llm_id) chat_model = LLMBundle(task_tenant_id, chat_model_config, lang=task_language) # run RAPTOR async with kg_limiter: @@ -1533,7 +1533,7 @@ async def do_handle_task(task): graphrag_conf = kb_parser_config.get("graphrag", {}) start_ts = timer() - chat_model_config = get_model_config_from_provider_instance(task_tenant_id, LLMType.CHAT, kb_task_llm_id) + chat_model_config = resolve_model_config(task_tenant_id, LLMType.CHAT, kb_task_llm_id) chat_model = LLMBundle(task_tenant_id, chat_model_config, lang=task_language) with_resolution = graphrag_conf.get("resolution", False) with_community = graphrag_conf.get("community", False) diff --git a/rag/svr/task_executor_refactor/chunk_post_processor.py b/rag/svr/task_executor_refactor/chunk_post_processor.py index c52e246c17..555656f4ba 100644 --- a/rag/svr/task_executor_refactor/chunk_post_processor.py +++ b/rag/svr/task_executor_refactor/chunk_post_processor.py @@ -40,7 +40,7 @@ from rag.svr.task_executor_refactor.task_context import TaskContext from api.db.services.doc_metadata_service import DocMetadataService from api.db.services.llm_service import LLMBundle -from api.db.joint_services.tenant_model_service import get_model_config_from_provider_instance +from api.db.joint_services.tenant_model_service import resolve_model_config from rag.prompts.generator import gen_metadata, keyword_extraction, question_proposal, content_tagging from rag.graphrag.utils import get_llm_cache, set_llm_cache, get_tags_from_cache, set_tags_to_cache @@ -56,7 +56,7 @@ async def extract_keywords(docs: List[Dict], ctx: TaskContext) -> None: st = timer() ctx.progress_cb(msg="Start to generate keywords for every chunk ...") - chat_model_config = get_model_config_from_provider_instance(ctx.tenant_id, LLMType.CHAT, ctx.llm_id) + chat_model_config = resolve_model_config(ctx.tenant_id, LLMType.CHAT, ctx.llm_id) with LLMBundle(ctx.tenant_id, chat_model_config, lang=ctx.language) as chat_model: async def doc_keyword_extraction(chat_mdl, d, topn): @@ -98,7 +98,7 @@ async def generate_questions(docs: List[Dict], ctx: TaskContext) -> None: st = timer() ctx.progress_cb(msg="Start to generate questions for every chunk ...") - chat_model_config = get_model_config_from_provider_instance(ctx.tenant_id, LLMType.CHAT, ctx.llm_id) + chat_model_config = resolve_model_config(ctx.tenant_id, LLMType.CHAT, ctx.llm_id) with LLMBundle(ctx.tenant_id, chat_model_config, lang=ctx.language) as chat_model: async def doc_question_proposal(chat_mdl, d, topn): @@ -178,7 +178,7 @@ async def generate_metadata(docs: List[Dict], ctx: TaskContext) -> None: st = timer() ctx.progress_cb(msg="Start to generate meta-data for every chunk ...") - chat_model_config = get_model_config_from_provider_instance(ctx.tenant_id, LLMType.CHAT, ctx.llm_id) + chat_model_config = resolve_model_config(ctx.tenant_id, LLMType.CHAT, ctx.llm_id) with LLMBundle(ctx.tenant_id, chat_model_config, lang=ctx.language) as chat_model: metadata_conf = build_metadata_config(ctx.parser_config) @@ -266,7 +266,7 @@ async def apply_tags(docs: List[Dict], ctx: TaskContext) -> None: set_tags_to_cache(kb_ids, all_tags) else: all_tags = json.loads(all_tags) - chat_model_config = get_model_config_from_provider_instance(tenant_id, LLMType.CHAT, ctx.llm_id) + chat_model_config = resolve_model_config(tenant_id, LLMType.CHAT, ctx.llm_id) with LLMBundle(ctx.tenant_id, chat_model_config, lang=ctx.language) as chat_model: docs_to_tag = [] for doc in docs: @@ -940,7 +940,7 @@ async def run_document_structure_compile(handler, embedding_model: LLMBundle) -> chat_llm_id = _resolve_template_chat_llm_id(parser_cfg, ctx) if chat_llm_id not in llm_bundle_cache: try: - cfg = get_model_config_from_provider_instance( + cfg = resolve_model_config( ctx.tenant_id, LLMType.CHAT, chat_llm_id, diff --git a/rag/svr/task_executor_refactor/dataflow_service.py b/rag/svr/task_executor_refactor/dataflow_service.py index c4912655f0..280916a500 100644 --- a/rag/svr/task_executor_refactor/dataflow_service.py +++ b/rag/svr/task_executor_refactor/dataflow_service.py @@ -38,7 +38,7 @@ from api.db.services.canvas_service import UserCanvasService from api.db.services.document_service import DocumentService from api.db.services.doc_metadata_service import DocMetadataService from api.db.services.pipeline_operation_log_service import PipelineOperationLogService -from api.db.joint_services.tenant_model_service import get_model_config_from_provider_instance, get_model_config_by_id +from api.db.joint_services.tenant_model_service import resolve_model_config, get_model_config_by_id from common.connection_utils import timeout from common.constants import LLMType, PipelineTaskType from common.metadata_utils import update_metadata_to @@ -244,13 +244,13 @@ class DataflowService: embedding_id = kb.embd_id if kb.tenant_embd_id: try: - embd_model_config = get_model_config_by_id(ctx.tenant_id, kb.tenant_embd_id) + embd_model_config = get_model_config_by_id(ctx.tenant_id, LLMType.EMBEDDING, kb.tenant_embd_id) except LookupError: - embd_model_config = get_model_config_from_provider_instance( + embd_model_config = resolve_model_config( ctx.tenant_id, LLMType.EMBEDDING, embedding_id ) else: - embd_model_config = get_model_config_from_provider_instance( + embd_model_config = resolve_model_config( ctx.tenant_id, LLMType.EMBEDDING, embedding_id ) from api.db.services.llm_service import LLMBundle diff --git a/rag/svr/task_executor_refactor/dataset_wiki_generator.py b/rag/svr/task_executor_refactor/dataset_wiki_generator.py index a747266abe..5823fd7254 100644 --- a/rag/svr/task_executor_refactor/dataset_wiki_generator.py +++ b/rag/svr/task_executor_refactor/dataset_wiki_generator.py @@ -645,7 +645,7 @@ async def run_wiki( from api.db.services.llm_service import LLMBundle from api.db.joint_services.tenant_model_service import ( get_tenant_default_model_by_type, - get_model_config_from_provider_instance, + resolve_model_config, ) from api.apps.restful_apis.chunk_api import _compilation_template_kind @@ -705,7 +705,7 @@ async def run_wiki( if key == "__tenant_default__": cfg = get_tenant_default_model_by_type(ctx.tenant_id, LLMType.CHAT) else: - cfg = get_model_config_from_provider_instance( + cfg = resolve_model_config( ctx.tenant_id, LLMType.CHAT, key, diff --git a/rag/svr/task_executor_refactor/task_handler.py b/rag/svr/task_executor_refactor/task_handler.py index cbb37e1399..0e6e3e2fe5 100644 --- a/rag/svr/task_executor_refactor/task_handler.py +++ b/rag/svr/task_executor_refactor/task_handler.py @@ -42,7 +42,7 @@ from api.db.services.compilation_template_group_service import CompilationTempla from api.db.joint_services.memory_message_service import handle_save_to_memory_task from api.db.joint_services.tenant_model_service import ( get_tenant_default_model_by_type, - get_model_config_from_provider_instance, + resolve_model_config, get_model_config_by_id, ) from api.db.services.llm_service import LLMBundle @@ -321,11 +321,17 @@ class TaskHandler: try: if ctx.tenant_embd_id: try: - embd_model_config = get_model_config_by_id(task_tenant_id, 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 = get_model_config_from_provider_instance(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 = get_model_config_from_provider_instance(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) @@ -381,7 +387,7 @@ class TaskHandler: return # Bind LLM for raptor - chat_model_config = get_model_config_from_provider_instance(task_tenant_id, LLMType.CHAT, kb_task_llm_id) + chat_model_config = resolve_model_config(task_tenant_id, LLMType.CHAT, kb_task_llm_id) with LLMBundle(task_tenant_id, chat_model_config, lang=ctx.language) as chat_model: # Run RAPTOR raptor_service = RaptorService(ctx=ctx) @@ -497,7 +503,7 @@ class TaskHandler: graphrag_conf = kb_parser_config.get("graphrag", {}) start_ts = timer() - chat_model_config = get_model_config_from_provider_instance(task_tenant_id, LLMType.CHAT, kb_task_llm_id) + chat_model_config = resolve_model_config(task_tenant_id, LLMType.CHAT, kb_task_llm_id) with LLMBundle(task_tenant_id, chat_model_config, lang=task_language) as chat_model: with_resolution = graphrag_conf.get("resolution", False) with_community = graphrag_conf.get("community", False) @@ -780,7 +786,7 @@ class TaskHandler: def _build_toc(cls, ctx: TaskContext, docs: List[Dict], progress_cb: Callable) -> Optional[Dict]: """Build table of contents.""" progress_cb(msg="Start to generate table of content ...") - chat_model_config = get_model_config_from_provider_instance(ctx.tenant_id, LLMType.CHAT, ctx.llm_id) + chat_model_config = resolve_model_config(ctx.tenant_id, LLMType.CHAT, ctx.llm_id) with LLMBundle(ctx.tenant_id, chat_model_config, lang=ctx.language) as chat_mdl: docs = sorted( docs, diff --git a/test/testcases/restful_api/test_chats.py b/test/testcases/restful_api/test_chats.py index d39df0c06e..84463ffdeb 100644 --- a/test/testcases/restful_api/test_chats.py +++ b/test/testcases/restful_api/test_chats.py @@ -764,6 +764,7 @@ def _load_chat_routes_unit_module(monkeypatch): tenant_model_service_mod = ModuleType("api.db.joint_services.tenant_model_service") tenant_model_service_mod.get_model_config_from_provider_instance = lambda *_args, **_kwargs: {} + tenant_model_service_mod.resolve_model_config = lambda *_args, **_kwargs: {} tenant_model_service_mod.get_tenant_default_model_by_type = lambda *_args, **_kwargs: {} tenant_model_service_mod.get_api_key = lambda *_args, **_kwargs: SimpleNamespace(id=1) tenant_model_service_mod.split_model_name = lambda model: (model.split("@")[0], "default", "factory") @@ -1151,7 +1152,7 @@ def test_chat_create_accepts_provider_scoped_rerank_id_unit(monkeypatch): query_calls.append(kwargs) return {} - monkeypatch.setattr(module, "get_model_config_from_provider_instance", _get_model_config_from_provider_instance) + monkeypatch.setattr(module, "resolve_model_config", _get_model_config_from_provider_instance) def _save(**kwargs): saved.update(kwargs) diff --git a/test/testcases/restful_api/test_dify_retrieval_routes_unit.py b/test/testcases/restful_api/test_dify_retrieval_routes_unit.py index 95650fb2f3..0f912a7844 100644 --- a/test/testcases/restful_api/test_dify_retrieval_routes_unit.py +++ b/test/testcases/restful_api/test_dify_retrieval_routes_unit.py @@ -239,6 +239,7 @@ def _load_dify_retrieval_module(monkeypatch): tenant_model_service_mod.get_model_config_by_id = _get_model_config_by_id tenant_model_service_mod.get_tenant_default_model_by_type = _get_tenant_default_model_by_type tenant_model_service_mod.get_model_config_from_provider_instance = _get_model_config_from_provider_instance + tenant_model_service_mod.resolve_model_config = _get_model_config_from_provider_instance monkeypatch.setitem(sys.modules, "api.db.joint_services.tenant_model_service", tenant_model_service_mod) module_name = "test_dify_retrieval_routes_unit_module" diff --git a/test/testcases/restful_api/test_user_tenant_routes_unit.py b/test/testcases/restful_api/test_user_tenant_routes_unit.py index 13be318323..9bd9f45f19 100644 --- a/test/testcases/restful_api/test_user_tenant_routes_unit.py +++ b/test/testcases/restful_api/test_user_tenant_routes_unit.py @@ -1502,6 +1502,7 @@ def _load_chat_routes_unit_module(monkeypatch): tenant_model_provider_mod = ModuleType("api.db.joint_services.tenant_model_service") tenant_model_provider_mod.get_model_config_from_provider_instance = lambda *_args, **_kwargs: {} + tenant_model_provider_mod.resolve_model_config = lambda *_args, **_kwargs: {} tenant_model_provider_mod.get_tenant_default_model_by_type = lambda *_args, **_kwargs: {} def _split_model_name(model_name): diff --git a/test/testcases/test_http_api/test_chat_assistant_management/test_chat_sdk_routes_unit.py b/test/testcases/test_http_api/test_chat_assistant_management/test_chat_sdk_routes_unit.py index c7eba82293..16af98fa4b 100644 --- a/test/testcases/test_http_api/test_chat_assistant_management/test_chat_sdk_routes_unit.py +++ b/test/testcases/test_http_api/test_chat_assistant_management/test_chat_sdk_routes_unit.py @@ -368,6 +368,7 @@ def _load_chat_module(monkeypatch): tenant_model_service_mod = ModuleType("api.db.joint_services.tenant_model_service") tenant_model_service_mod.get_model_config_from_provider_instance = lambda *_args, **_kwargs: {} + tenant_model_service_mod.resolve_model_config = lambda *_args, **_kwargs: {} tenant_model_service_mod.get_tenant_default_model_by_type = lambda *_args, **_kwargs: {} monkeypatch.setitem(sys.modules, "api.db.joint_services.tenant_model_service", tenant_model_service_mod) diff --git a/test/testcases/test_http_api/test_dataset_management/test_dify_retrieval_routes_unit.py b/test/testcases/test_http_api/test_dataset_management/test_dify_retrieval_routes_unit.py index 9573eab268..54bcfd7208 100644 --- a/test/testcases/test_http_api/test_dataset_management/test_dify_retrieval_routes_unit.py +++ b/test/testcases/test_http_api/test_dataset_management/test_dify_retrieval_routes_unit.py @@ -246,6 +246,7 @@ def _load_dify_retrieval_module(monkeypatch): tenant_model_service_mod.get_model_config_by_id = _get_model_config_by_id tenant_model_service_mod.get_model_config_from_provider_instance = _get_model_config_from_provider_instance + tenant_model_service_mod.resolve_model_config = _get_model_config_from_provider_instance tenant_model_service_mod.get_tenant_default_model_by_type = _get_tenant_default_model_by_type monkeypatch.setitem(sys.modules, "api.db.joint_services.tenant_model_service", tenant_model_service_mod) diff --git a/test/testcases/test_http_api/test_file_management_within_dataset/test_doc_sdk_routes_unit.py b/test/testcases/test_http_api/test_file_management_within_dataset/test_doc_sdk_routes_unit.py index b21f3db6b7..552288692c 100644 --- a/test/testcases/test_http_api/test_file_management_within_dataset/test_doc_sdk_routes_unit.py +++ b/test/testcases/test_http_api/test_file_management_within_dataset/test_doc_sdk_routes_unit.py @@ -472,6 +472,7 @@ def _load_doc_module(monkeypatch, module_basename="chunk_api"): tenant_model_service_mod.get_model_config_by_id = _get_model_config_by_id tenant_model_service_mod.get_model_config_from_provider_instance = _get_model_config_from_provider_instance + tenant_model_service_mod.resolve_model_config = _get_model_config_from_provider_instance tenant_model_service_mod.get_tenant_default_model_by_type = _get_tenant_default_model_by_type monkeypatch.setitem(sys.modules, "api.db.joint_services.tenant_model_service", tenant_model_service_mod) diff --git a/test/testcases/test_http_api/test_session_management/test_session_sdk_routes_unit.py b/test/testcases/test_http_api/test_session_management/test_session_sdk_routes_unit.py index 7530b893f5..219d77c341 100644 --- a/test/testcases/test_http_api/test_session_management/test_session_sdk_routes_unit.py +++ b/test/testcases/test_http_api/test_session_management/test_session_sdk_routes_unit.py @@ -520,6 +520,7 @@ def _load_session_module(monkeypatch): tenant_model_service_mod.get_model_config_by_id = _get_model_config_by_id tenant_model_service_mod.get_model_config_from_provider_instance = _get_model_config_from_provider_instance + tenant_model_service_mod.resolve_model_config = _get_model_config_from_provider_instance tenant_model_service_mod.get_tenant_default_model_by_type = _get_tenant_default_model_by_type tenant_model_service_mod.get_api_key = _get_api_key tenant_model_service_mod.split_model_name = _split_model_name @@ -2243,6 +2244,7 @@ def _load_chat_api_module(monkeypatch): tenant_model_svc = ModuleType("api.db.joint_services.tenant_model_service") tenant_model_svc.get_tenant_default_model_by_type = lambda *_a, **_k: {} tenant_model_svc.get_model_config_from_provider_instance = lambda **_k: {} + tenant_model_svc.resolve_model_config = lambda **_k: {} tenant_model_svc.get_api_key = lambda **_k: "fake-api-key" tenant_model_svc.split_model_name = lambda model_name: (model_name, "", "") monkeypatch.setitem(sys.modules, "api.db.joint_services.tenant_model_service", tenant_model_svc) diff --git a/test/testcases/test_web_api/test_canvas_app/test_iteration_runtime_unit.py b/test/testcases/test_web_api/test_canvas_app/test_iteration_runtime_unit.py index 74796312ae..8f3c90585e 100644 --- a/test/testcases/test_web_api/test_canvas_app/test_iteration_runtime_unit.py +++ b/test/testcases/test_web_api/test_canvas_app/test_iteration_runtime_unit.py @@ -96,6 +96,7 @@ def _load_canvas_runtime(monkeypatch): tenant_model_service = ModuleType("api.db.joint_services.tenant_model_service") tenant_model_service.get_tenant_default_model_by_type = lambda *_a, **_kw: None + tenant_model_service.resolve_model_config = lambda *_a, **_kw: None monkeypatch.setitem( sys.modules, "api.db.joint_services.tenant_model_service", diff --git a/test/testcases/test_web_api/test_chunk_app/test_chunk_routes_unit.py b/test/testcases/test_web_api/test_chunk_app/test_chunk_routes_unit.py index d483c61860..7b40e7f184 100644 --- a/test/testcases/test_web_api/test_chunk_app/test_chunk_routes_unit.py +++ b/test/testcases/test_web_api/test_chunk_app/test_chunk_routes_unit.py @@ -310,6 +310,7 @@ def _load_chunk_module(monkeypatch): tenant_model_service_mod = ModuleType("api.db.joint_services.tenant_model_service") tenant_model_service_mod.get_model_config_by_id = lambda *_args, **_kwargs: {"llm_name": "embed", "model_type": "embedding"} tenant_model_service_mod.get_model_config_from_provider_instance = lambda *_args, **_kwargs: {"llm_name": "embed", "model_type": "embedding"} + tenant_model_service_mod.resolve_model_config = lambda *_args, **_kwargs: {"llm_name": "embed", "model_type": "embedding"} tenant_model_service_mod.get_tenant_default_model_by_type = lambda *_args, **_kwargs: {"llm_name": "chat", "model_type": "chat"} monkeypatch.setitem(sys.modules, "api.db.joint_services.tenant_model_service", tenant_model_service_mod) diff --git a/test/unit_test/api/apps/restful_apis/test_agentbots_access_control.py b/test/unit_test/api/apps/restful_apis/test_agentbots_access_control.py index 938f3c5ed7..9174d8c9f3 100644 --- a/test/unit_test/api/apps/restful_apis/test_agentbots_access_control.py +++ b/test/unit_test/api/apps/restful_apis/test_agentbots_access_control.py @@ -90,7 +90,13 @@ def _load_bot_api(monkeypatch, *, accessible, calls): TenantService=SimpleNamespace(), UserTenantService=SimpleNamespace(), ) - _stub(monkeypatch, "api.db.joint_services.tenant_model_service", get_tenant_default_model_by_type=lambda *_a, **_k: None, get_model_config_from_provider_instance=lambda *_a, **_k: None) + _stub( + monkeypatch, + "api.db.joint_services.tenant_model_service", + get_tenant_default_model_by_type=lambda *_a, **_k: None, + get_model_config_from_provider_instance=lambda *_a, **_k: None, + resolve_model_config=lambda *_a, **_k: None, + ) _stub(monkeypatch, "common.misc_utils", get_uuid=lambda: "uuid", thread_pool_exec=_passthrough_thread_pool_exec) _stub( monkeypatch, diff --git a/test/unit_test/api/apps/restful_apis/test_openai_stream_no_duplicate.py b/test/unit_test/api/apps/restful_apis/test_openai_stream_no_duplicate.py index 77748cc8b2..82be1dfd73 100644 --- a/test/unit_test/api/apps/restful_apis/test_openai_stream_no_duplicate.py +++ b/test/unit_test/api/apps/restful_apis/test_openai_stream_no_duplicate.py @@ -73,6 +73,7 @@ def _load_openai_api(monkeypatch): monkeypatch, "api.db.joint_services.tenant_model_service", get_model_config_from_provider_instance=lambda *_a, **_k: {}, + resolve_model_config=lambda *_a, **_k: {}, get_api_key=lambda *_a, **_k: "key", ) _stub( diff --git a/test/unit_test/api/apps/sdk/test_dify_retrieval.py b/test/unit_test/api/apps/sdk/test_dify_retrieval.py index 0b286544d4..271a2624e9 100644 --- a/test/unit_test/api/apps/sdk/test_dify_retrieval.py +++ b/test/unit_test/api/apps/sdk/test_dify_retrieval.py @@ -135,6 +135,7 @@ def _load_dify_retrieval(monkeypatch, *, kb, accessible, request_body, tenant_id "api.db.joint_services.tenant_model_service", get_tenant_default_model_by_type=lambda *_a, **_k: {}, get_model_config_from_provider_instance=lambda *_a, **_k: {}, + resolve_model_config=lambda *_a, **_k: {}, ) _stub( diff --git a/test/unit_test/api/apps/services/test_delete_datasets.py b/test/unit_test/api/apps/services/test_delete_datasets.py index bc1b04d4e3..d8ac3263a7 100644 --- a/test/unit_test/api/apps/services/test_delete_datasets.py +++ b/test/unit_test/api/apps/services/test_delete_datasets.py @@ -103,6 +103,8 @@ def _load_delete_datasets_module(monkeypatch, *, f2d_rows, file_filter_delete): monkeypatch, "api.db.joint_services.tenant_model_service", get_model_config_from_provider_instance=MagicMock(), + resolve_model_config=MagicMock(), + resolve_model_id=MagicMock(), ) _stub( monkeypatch, diff --git a/test/unit_test/api/apps/services/test_models_api_service_list_tenant_added_models.py b/test/unit_test/api/apps/services/test_models_api_service_list_tenant_added_models.py index 48c04baaac..e72e3947f7 100644 --- a/test/unit_test/api/apps/services/test_models_api_service_list_tenant_added_models.py +++ b/test/unit_test/api/apps/services/test_models_api_service_list_tenant_added_models.py @@ -21,12 +21,10 @@ the quantization tag, such as `text-embedding-nomic-embed-text-v1.5@q8_0`). """ import importlib.util -import logging import sys from pathlib import Path from enum import IntEnum from types import ModuleType, SimpleNamespace -from unittest.mock import MagicMock import pytest @@ -35,6 +33,7 @@ pytestmark = pytest.mark.p2 class _StubModelTypeBinary(IntEnum): """Mimics common.constants.ModelTypeBinary for the stubbed environment.""" + CHAT = 1 EMBEDDING = 2 SPEECH2TEXT = 4 @@ -86,10 +85,11 @@ def _load_module(monkeypatch, *, tenant_model_records, factory_llm_infos=None): additional behaviour at runtime. """ - tenant = SimpleNamespace(id="tenant-1") + tenant = SimpleNamespace(id="tenant-1", name="tenant-1") provider = SimpleNamespace( id="provider-1", provider_name="LM-Studio", + tenant_id="tenant-1", ) instance = SimpleNamespace( id="instance-1", @@ -110,9 +110,7 @@ def _load_module(monkeypatch, *, tenant_model_records, factory_llm_infos=None): # Default no-op; tests that need to observe the resolved provider # name passed by `_get_model_info` override this with # monkeypatch.setattr on the loaded stub. - get_by_tenant_id_and_provider_name=lambda tenant_id, provider_name: SimpleNamespace( - id="provider-1", provider_name=provider_name - ), + get_by_tenant_id_and_provider_name=lambda tenant_id, provider_name: SimpleNamespace(id="provider-1", provider_name=provider_name), ), ) _stub( @@ -120,9 +118,7 @@ def _load_module(monkeypatch, *, tenant_model_records, factory_llm_infos=None): "api.db.services.tenant_model_instance_service", TenantModelInstanceService=SimpleNamespace( get_by_provider_ids=lambda provider_ids: [instance], - get_by_provider_id_and_instance_name=lambda provider_id, instance_name: SimpleNamespace( - id="instance-1", provider_id=provider_id, instance_name=instance_name - ), + get_by_provider_id_and_instance_name=lambda provider_id, instance_name: SimpleNamespace(id="instance-1", provider_id=provider_id, instance_name=instance_name), ), ) _stub( @@ -132,12 +128,8 @@ def _load_module(monkeypatch, *, tenant_model_records, factory_llm_infos=None): get_models_by_provider_ids_and_instance_ids=lambda p_ids, i_ids: list(tenant_model_records), # Default returns an "active" model so `_get_model_info` treats # the row as enabled when exercising the bare-model branch. - get_by_provider_id_and_instance_id_and_model_type_and_model_name=lambda *args: SimpleNamespace( - status=1 - ), - get_by_provider_id_and_instance_id_and_model_name=lambda *args: SimpleNamespace( - status=1 - ), + get_by_provider_id_and_instance_id_and_model_type_and_model_name=lambda *args: SimpleNamespace(status=1), + get_by_provider_id_and_instance_id_and_model_name=lambda *args: SimpleNamespace(status=1), ), ) @@ -169,13 +161,7 @@ def _load_module(monkeypatch, *, tenant_model_records, factory_llm_infos=None): FACTORY_LLM_INFOS=factory_llm_infos if factory_llm_infos is not None else [], ) - module_path = ( - Path(__file__).resolve().parents[5] - / "api" - / "apps" - / "services" - / "models_api_service.py" - ) + module_path = Path(__file__).resolve().parents[5] / "api" / "apps" / "services" / "models_api_service.py" spec = importlib.util.spec_from_file_location( "test_models_api_service_list_tenant_added_models", module_path, @@ -213,6 +199,7 @@ def _make_model_record(model_name, model_type="embedding", status=1): """ model_type_bin = _MODEL_TYPE_TO_BIN.get(model_type, model_type) if isinstance(model_type, str) else model_type return SimpleNamespace( + id=f"{model_name}-id", provider_id="provider-1", instance_id="instance-1", model_name=model_name, diff --git a/test/unit_test/api/db/services/test_dialog_service_final_answer.py b/test/unit_test/api/db/services/test_dialog_service_final_answer.py index 6234fc10c6..33218622ec 100644 --- a/test/unit_test/api/db/services/test_dialog_service_final_answer.py +++ b/test/unit_test/api/db/services/test_dialog_service_final_answer.py @@ -224,7 +224,7 @@ def test_async_ask_final_event_carries_decorated_answer(monkeypatch): monkeypatch.setattr(dialog_service.KnowledgebaseService, "get_by_ids", lambda _ids: [_KB]) monkeypatch.setattr( dialog_service, - "get_model_config_from_provider_instance", + "resolve_model_config", lambda _tid, _type, _name: _LLM_CONFIG, ) monkeypatch.setattr(dialog_service, "LLMBundle", lambda _tid, _cfg: chat_mdl) @@ -269,7 +269,7 @@ def test_async_ask_delta_events_carry_incremental_text_only(monkeypatch): monkeypatch.setattr(dialog_service.KnowledgebaseService, "get_by_ids", lambda _ids: [_KB]) monkeypatch.setattr( dialog_service, - "get_model_config_from_provider_instance", + "resolve_model_config", lambda _tid, _type, _name: _LLM_CONFIG, ) monkeypatch.setattr(dialog_service, "LLMBundle", lambda _tid, _cfg: chat_mdl) @@ -381,7 +381,7 @@ def _make_dialog(chat_mdl_stub): top_n=6, top_k=1024, rerank_id="", - tenant_rerank_id=None + tenant_rerank_id=None, ) @@ -400,10 +400,10 @@ def test_async_chat_final_event_carries_decorated_answer(monkeypatch): retriever = _StubRetriever() # Stub out the heavy service/model calls - monkeypatch.setattr(dialog_service, "get_model_type_by_name", lambda _tid, _llm_id: ["chat"]) + monkeypatch.setattr(dialog_service, "resolve_model_type", lambda _tid, _llm_id: ["chat"]) monkeypatch.setattr( dialog_service, - "get_model_config_from_provider_instance", + "resolve_model_config", lambda _tid, _type, _llm_id: _LLM_CONFIG, ) monkeypatch.setattr( @@ -452,10 +452,10 @@ def test_async_chat_langfuse_uses_start_observation(monkeypatch): chat_mdl = _StreamingChatModel(llm_answer) retriever = _StubRetriever() - monkeypatch.setattr(dialog_service, "get_model_type_by_name", lambda _tid, _llm_id: ["chat"]) + monkeypatch.setattr(dialog_service, "resolve_model_type", lambda _tid, _llm_id: ["chat"]) monkeypatch.setattr( dialog_service, - "get_model_config_from_provider_instance", + "resolve_model_config", lambda _tid, _type, _llm_id: _LLM_CONFIG, ) monkeypatch.setattr( @@ -515,10 +515,10 @@ def test_async_chat_langfuse_observation_includes_session_id(monkeypatch): chat_mdl = _StreamingChatModel("Session traces should be grouped.") retriever = _StubRetriever() - monkeypatch.setattr(dialog_service, "get_model_type_by_name", lambda _tid, _llm_id: ["chat"]) + monkeypatch.setattr(dialog_service, "resolve_model_type", lambda _tid, _llm_id: ["chat"]) monkeypatch.setattr( dialog_service, - "get_model_config_from_provider_instance", + "resolve_model_config", lambda _tid, _type, _llm_id: _LLM_CONFIG, ) monkeypatch.setattr( @@ -573,7 +573,7 @@ def test_get_models_passes_langfuse_trace_context_to_llm_bundles(monkeypatch): monkeypatch.setattr(dialog_service.KnowledgebaseService, "get_by_ids", lambda _ids: [_KB]) monkeypatch.setattr( dialog_service, - "get_model_config_from_provider_instance", + "resolve_model_config", lambda _tenant_id, model_type, _model_id: {**_LLM_CONFIG, "model_type": model_type}, ) monkeypatch.setattr( @@ -614,10 +614,10 @@ def test_async_chat_continues_when_langfuse_observation_start_fails(monkeypatch) chat_mdl = _StreamingChatModel(llm_answer) retriever = _StubRetriever() - monkeypatch.setattr(dialog_service, "get_model_type_by_name", lambda _tid, _llm_id: ["chat"]) + monkeypatch.setattr(dialog_service, "resolve_model_type", lambda _tid, _llm_id: ["chat"]) monkeypatch.setattr( dialog_service, - "get_model_config_from_provider_instance", + "resolve_model_config", lambda _tid, _type, _llm_id: _LLM_CONFIG, ) monkeypatch.setattr( diff --git a/test/unit_test/api/db/services/test_dialog_service_use_sql_source_columns.py b/test/unit_test/api/db/services/test_dialog_service_use_sql_source_columns.py index 5da77c6615..7c4c602ffa 100644 --- a/test/unit_test/api/db/services/test_dialog_service_use_sql_source_columns.py +++ b/test/unit_test/api/db/services/test_dialog_service_use_sql_source_columns.py @@ -317,10 +317,10 @@ def test_async_chat_uses_all_docs_when_no_doc_ids_selected(monkeypatch): ) monkeypatch.setattr(dialog_service.settings, "retriever", retriever, raising=False) - monkeypatch.setattr(dialog_service, "get_model_type_by_name", lambda _tid, _llm_id: ["chat"]) + monkeypatch.setattr(dialog_service, "resolve_model_type", lambda _tid, _llm_id: ["chat"]) monkeypatch.setattr( dialog_service, - "get_model_config_from_provider_instance", + "resolve_model_config", lambda *_args, **_kwargs: {"llm_factory": "unit", "max_tokens": 4096, "model_type": "chat"}, ) monkeypatch.setattr(dialog_service.TenantLangfuseService, "filter_by_tenant", lambda **_kwargs: None) diff --git a/test/unit_test/rag/svr/task_executor_refactor/conftest.py b/test/unit_test/rag/svr/task_executor_refactor/conftest.py index f36092dec7..363f227c29 100644 --- a/test/unit_test/rag/svr/task_executor_refactor/conftest.py +++ b/test/unit_test/rag/svr/task_executor_refactor/conftest.py @@ -324,7 +324,7 @@ def create_patch_embedding_model(vectors=None, vector_size=128): return ( patch( - "rag.svr.task_executor_refactor.task_handler.get_model_config_from_provider_instance", + "rag.svr.task_executor_refactor.task_handler.resolve_model_config", return_value=MagicMock(), ), patch( @@ -579,7 +579,7 @@ def mock_raptor_context(): class patch_embedding_binding: """Context manager that patches embedding model binding at the external boundary. - Patches ``LLMBundle``, ``get_model_config_from_provider_instance``, and + Patches ``LLMBundle``, ``resolve_model_config``, and ``get_tenant_default_model_by_type`` so that ``TaskHandler._bind_embedding_model`` executes its real logic without making actual API calls. @@ -609,7 +609,7 @@ class patch_embedding_binding: self._patches = [ patch( - "rag.svr.task_executor_refactor.task_handler.get_model_config_from_provider_instance", + "rag.svr.task_executor_refactor.task_handler.resolve_model_config", return_value=MagicMock(), ), patch( @@ -711,7 +711,7 @@ class patch_pipeline_mocks: Usage:: with patch_pipeline_mocks() as m: - m.get_model_config_from_provider_instance.return_value = MagicMock() + m.resolve_model_config.return_value = MagicMock() handler = TaskHandler(ctx) await handler.handle() """ @@ -723,7 +723,7 @@ class patch_pipeline_mocks: # (module_key, attr_name, use_AsyncMock) _COMMON = [ - ("task_handler", "get_model_config_from_provider_instance", False), + ("task_handler", "resolve_model_config", False), ("task_handler", "LLMBundle", False), ("task_handler", "get_tenant_default_model_by_type", False), ("task_handler", "search.index_name", False),