From 74bbbba3e07a6ddc0b890befb5bd96ec4dfbbcfc Mon Sep 17 00:00:00 2001 From: buua436 Date: Fri, 10 Jul 2026 10:34:19 +0800 Subject: [PATCH] fix: align model default handling (#16782) --- api/apps/restful_apis/chat_api.py | 6 +- api/apps/restful_apis/models_api.py | 10 ++- api/apps/services/models_api_service.py | 70 +++++++++++++++---- .../services/compilation_template_service.py | 28 +++++--- api/db/services/user_service.py | 7 ++ common/misc_utils.py | 25 +++---- test/testcases/restful_api/test_chats.py | 37 +++++++--- test/testcases/restful_api/test_datasets.py | 7 +- .../test_user_tenant_routes_unit.py | 4 +- .../test_create_dataset.py | 4 +- .../test_update_dataset.py | 4 +- .../api/apps/services/test_delete_datasets.py | 23 ++++++ ...ls_api_service_list_tenant_added_models.py | 2 + tools/scripts/mysql_migration.py | 11 +-- tools/scripts/run_migrations.sh | 2 +- 15 files changed, 169 insertions(+), 71 deletions(-) diff --git a/api/apps/restful_apis/chat_api.py b/api/apps/restful_apis/chat_api.py index 945baca77e..b219e1fd30 100644 --- a/api/apps/restful_apis/chat_api.py +++ b/api/apps/restful_apis/chat_api.py @@ -300,7 +300,7 @@ async def _validate_rerank_id(rerank_id, tenant_id): await thread_pool_exec( resolve_model_config, tenant_id=tenant_id, - model_name=rerank_id, + model_ref=rerank_id, model_type="rerank", ) except Exception as e: @@ -399,9 +399,9 @@ async def create(): # return get_data_error_result(message=err) req.setdefault("kb_ids", []) - req.setdefault("llm_id", tenant.llm_id) + req.setdefault("llm_id", tenant.tenant_llm_id) if req["llm_id"] is None: - req["llm_id"] = tenant.llm_id + req["llm_id"] = tenant.tenant_llm_id req.setdefault("llm_setting", {}) req.setdefault("description", "A helpful Assistant") req.setdefault("top_n", 6) diff --git a/api/apps/restful_apis/models_api.py b/api/apps/restful_apis/models_api.py index aa6e267179..81e7f6db3b 100644 --- a/api/apps/restful_apis/models_api.py +++ b/api/apps/restful_apis/models_api.py @@ -75,6 +75,8 @@ def get_added_models(tenant_id: str): type: string model_name: type: string + model_id: + type: string model_type: type: string enable: @@ -139,6 +141,8 @@ def get_default_models(tenant_id: str): type: string model_name: type: string + model_id: + type: string model_type: type: string enable: @@ -190,6 +194,9 @@ async def set_default_models(tenant_id: str): model_name: type: string description: Model name. Required when setting a model; omit to clear. + model_id: + type: string + description: Tenant model id. If provided, it takes precedence over model_provider/model_instance/model_name. model_type: type: string description: "Model type: chat, embedding, rerank, asr, vision, tts, ocr" @@ -206,10 +213,11 @@ async def set_default_models(tenant_id: str): model_provider = data.get("model_provider", "") model_instance = data.get("model_instance", "") model_name = data.get("model_name", "") + model_id = data.get("model_id", "") model_type = data["model_type"] try: - success, msg = models_api_service.set_tenant_default_models(tenant_id, model_provider, model_instance, model_name, model_type) + success, msg = models_api_service.set_tenant_default_models(tenant_id, model_provider, model_instance, model_name, model_type, model_id) if success: logging.info(f"success: {success}, msg: {msg}") return get_result(message=msg) diff --git a/api/apps/services/models_api_service.py b/api/apps/services/models_api_service.py index c79028e904..8c37bd6d44 100644 --- a/api/apps/services/models_api_service.py +++ b/api/apps/services/models_api_service.py @@ -16,7 +16,7 @@ import logging import os -from api.db.joint_services.tenant_model_service import ensure_mineru_from_env, ensure_opendataloader_from_env, ensure_paddleocr_from_env +from api.db.joint_services.tenant_model_service import ensure_mineru_from_env, ensure_opendataloader_from_env, ensure_paddleocr_from_env, resolve_model_id from api.db.services.tenant_model_instance_service import TenantModelInstanceService from api.db.services.tenant_model_provider_service import TenantModelProviderService from api.db.services.tenant_model_service import TenantModelService @@ -36,6 +36,16 @@ MODEL_TYPE_TO_FIELD = { "ocr": "ocr_id", } +MODEL_TYPE_TO_TENANT_ID_FIELD = { + "chat": "tenant_llm_id", + "embedding": "tenant_embd_id", + "rerank": "tenant_rerank_id", + "asr": "tenant_asr_id", + "vision": "tenant_img2txt_id", + "tts": "tenant_tts_id", + "ocr": "tenant_ocr_id", +} + MODEL_TAG_TO_TYPE = { "chat": "chat", "embedding": "embedding", @@ -61,7 +71,7 @@ def _factory_model_types(llm: dict) -> list[str]: return [model_type] if model_type else [] -def _get_model_info(tenant_id: str, default_model: str, model_type: str): +def _get_model_info(tenant_id: str, default_model: str, model_type: str, model_id: str = ""): """ Parse a composite model string (modelName@instanceName@providerName or modelName@providerName) and validate that the provider, instance, and model exist. @@ -90,6 +100,7 @@ def _get_model_info(tenant_id: str, default_model: str, model_type: str): # Special case: OCR with infiniflow@default@deepdoc is always enabled if model_type == "ocr" and provider_name == "infiniflow" and instance_name == "default" and model_name == "deepdoc": return { + "model_id": model_id, "model_provider": provider_name, "model_instance": instance_name, "model_name": model_name, @@ -102,6 +113,7 @@ def _get_model_info(tenant_id: str, default_model: str, model_type: str): tei_model = os.getenv("TEI_MODEL", "") if model_type == "embedding" and "tei-" in compose_profiles and tei_model and model_name == tei_model and (not provider_name or provider_name == "Builtin"): return { + "model_id": model_id, "model_provider": "Builtin", "model_instance": "default", "model_name": model_name, @@ -130,7 +142,7 @@ def _get_model_info(tenant_id: str, default_model: str, model_type: str): logging.warning(f"Model '{model_name}' is disabled") return None - return {"model_provider": provider_name, "model_instance": instance_name, "model_name": model_name, "model_type": model_type, "enable": True} + return {"model_id": model_id, "model_provider": provider_name, "model_instance": instance_name, "model_name": model_name, "model_type": model_type, "enable": True} def _check_model_available(tenant_id: str, provider_name: str, instance_name: str, model_name: str, model_type: str): @@ -200,49 +212,81 @@ def list_tenant_default_models(tenant_id: str): default_model = getattr(tenant, field_name, None) if not default_model: continue - model_info = _get_model_info(tenant_id, default_model, model_type) + tenant_model_id = getattr(tenant, MODEL_TYPE_TO_TENANT_ID_FIELD[model_type], None) + model_info = _get_model_info(tenant_id, default_model, model_type, tenant_model_id) if model_info: models.append(model_info) return True, {"models": models} -def set_tenant_default_models(tenant_id: str, model_provider: str, model_instance: str, model_name: str, model_type: str): +def set_tenant_default_models(tenant_id: str, model_provider: str, model_instance: str, model_name: str, model_type: str, model_id: str = ""): """ Set or clear a tenant default model. - If model_provider, model_instance, and model_name are all provided, - validates the model and sets it as the default. - If all three are empty, clears the default for the given model type. + If model_id is provided, resolves it to provider/instance/name and uses it. + Otherwise falls back to the legacy model_provider/model_instance/model_name triplet. + If all model selectors are empty, clears the default for the given model type. :param tenant_id: tenant ID :param model_provider: provider name :param model_instance: instance name :param model_name: model name :param model_type: model type (chat, embedding, rerank, asr, vision, tts, ocr) + :param model_id: tenant_model id :return: (success, result_or_error_message) """ field_name = MODEL_TYPE_TO_FIELD.get(model_type) if not field_name: return False, f"model type '{model_type}' is invalid" + tenant_field_name = MODEL_TYPE_TO_TENANT_ID_FIELD.get(model_type) e, tenant = TenantService.get_by_id(tenant_id) if not e: return False, "Tenant not found" - if not model_provider and not model_instance and not model_name: - # Clear the default model - default_model = "" + tenant_model_id = None + if model_id: + model_ok, model_obj = TenantModelService.get_by_id(model_id) + if not model_ok: + return False, f"model_id '{model_id}' not found" + + provider_ok, provider_obj = TenantModelProviderService.get_by_id(model_obj.provider_id) + if not provider_ok: + return False, f"Provider id '{model_obj.provider_id}' not found for model_id '{model_id}'" + if provider_obj.tenant_id != tenant_id: + return False, "Permission denied" + + instance_ok, instance_obj = TenantModelInstanceService.get_by_id(model_obj.instance_id) + if not instance_ok: + return False, f"Instance id '{model_obj.instance_id}' not found for model_id '{model_id}'" + + normalized_model_type = MODEL_TAG_TO_TYPE.get(model_type, model_type) + if normalized_model_type not in get_model_type_human(model_obj.model_type): + return False, f"Model '{model_obj.model_name}' isn't a {model_type} model" + if model_obj.status != ActiveStatusEnum.ACTIVE.value: + return False, f"Model '{model_obj.model_name}' isn't available" + + default_model = f"{model_obj.model_name}@{instance_obj.instance_name}@{provider_obj.provider_name}" + tenant_model_id = model_id elif model_provider and model_instance and model_name: # Validate and set the default model success, msg = _check_model_available(tenant_id, model_provider, model_instance, model_name, model_type) if not success: return False, msg default_model = f"{model_name}@{model_instance}@{model_provider}" + normalized_model_type = MODEL_TAG_TO_TYPE.get(model_type, model_type) + try: + tenant_model_id = resolve_model_id(tenant_id, normalized_model_type, default_model) + except LookupError as exc: + return False, str(exc) else: - return False, "model_provider, model_instance and model_name must be specified together" + # Clear the default model + default_model = "" + tenant_model_id = None - TenantService.update_by_id(tenant_id, {field_name: default_model}) + update_fields = {field_name: default_model, tenant_field_name: tenant_model_id} + TenantService.update_by_id(tenant_id, update_fields) return True, "success" diff --git a/api/db/services/compilation_template_service.py b/api/db/services/compilation_template_service.py index bb31d280a6..eb9c210346 100644 --- a/api/db/services/compilation_template_service.py +++ b/api/db/services/compilation_template_service.py @@ -37,9 +37,9 @@ class CompilationTemplateService(CommonService): from api.db.services.user_service import TenantService ok, tenant = TenantService.get_by_id(tenant_id) - if ok and getattr(tenant, "llm_id", None): + if ok and getattr(tenant, "tenant_llm_id", None): config = dict(config) - config["llm_id"] = tenant.llm_id + config["llm_id"] = tenant.tenant_llm_id except Exception: logging.exception( "compilation_template: llm_id default-fill lookup failed for tenant=%s", @@ -91,9 +91,9 @@ class CompilationTemplateService(CommonService): from api.db.services.user_service import TenantService ok, tenant = TenantService.get_by_id(tenant_id) - if ok and getattr(tenant, "llm_id", None): + if ok and getattr(tenant, "tenant_llm_id", None): config = dict(config) - config["llm_id"] = tenant.llm_id + config["llm_id"] = tenant.tenant_llm_id except Exception: logging.exception( "compilation_template: llm_id lazy-fill lookup failed for tenant=%s", @@ -246,7 +246,11 @@ class CompilationTemplateService(CommonService): """ wiki_dir = os.path.join( get_project_base_directory(), - "api", "db", "init_data", "compilation_templates", "wiki", + "api", + "db", + "init_data", + "compilation_templates", + "wiki", ) if not os.path.exists(wiki_dir): logging.warning("Missing wiki presets directory: %s", wiki_dir) @@ -269,10 +273,12 @@ class CompilationTemplateService(CommonService): continue # Missing fields degrade to empty strings so the frontend # doesn't have to null-check every row. - presets.append({ - "id": os.path.splitext(filename)[0], - "topic": str(doc.get("topic") or "").strip(), - "instruction": str(doc.get("instruction") or ""), - "page_example": str(doc.get("page_example") or ""), - }) + presets.append( + { + "id": os.path.splitext(filename)[0], + "topic": str(doc.get("topic") or "").strip(), + "instruction": str(doc.get("instruction") or ""), + "page_example": str(doc.get("page_example") or ""), + } + ) return presets diff --git a/api/db/services/user_service.py b/api/db/services/user_service.py index 0fb7775ad7..22d64b93be 100644 --- a/api/db/services/user_service.py +++ b/api/db/services/user_service.py @@ -178,12 +178,19 @@ class TenantService(CommonService): cls.model.id.alias("tenant_id"), cls.model.name, cls.model.llm_id, + cls.model.tenant_llm_id, cls.model.embd_id, + cls.model.tenant_embd_id, cls.model.rerank_id, + cls.model.tenant_rerank_id, cls.model.asr_id, + cls.model.tenant_asr_id, cls.model.img2txt_id, + cls.model.tenant_img2txt_id, cls.model.tts_id, + cls.model.tenant_tts_id, cls.model.ocr_id, + cls.model.tenant_ocr_id, cls.model.parser_ids, UserTenant.role, ] diff --git a/common/misc_utils.py b/common/misc_utils.py index c1a7bbb50f..e9486f51fb 100644 --- a/common/misc_utils.py +++ b/common/misc_utils.py @@ -241,26 +241,19 @@ def pip_install_torch(): subprocess.check_call([sys.executable, "-m", "pip", "install", *pkg_names]) -@once -def _thread_pool_executor(): - max_workers_env = os.getenv("THREAD_POOL_MAX_WORKERS", "128") - try: - max_workers = int(max_workers_env) - except ValueError: - max_workers = 128 - if max_workers < 1: - max_workers = 1 - return ThreadPoolExecutor(max_workers=max_workers) - - async def thread_pool_exec(func, *args, **kwargs): # loop.run_in_executor() submits the callable without propagating the caller's # contextvars (unlike asyncio.to_thread, which copies the context). Copy the # current context and run the callable inside it so ContextVars set by the # caller (e.g. tracing / per-request state) are visible in the worker thread. + # + # Use a short-lived executor per call instead of a shared singleton. Python + # 3.13's executor reuse can deadlock in this environment when the same helper + # is awaited repeatedly inside one event loop. loop = asyncio.get_running_loop() ctx = contextvars.copy_context() - if kwargs: - inner = functools.partial(func, *args, **kwargs) - return await loop.run_in_executor(_thread_pool_executor(), ctx.run, inner) - return await loop.run_in_executor(_thread_pool_executor(), ctx.run, func, *args) + with ThreadPoolExecutor(max_workers=1) as executor: + if kwargs: + inner = functools.partial(func, *args, **kwargs) + return await loop.run_in_executor(executor, ctx.run, inner) + return await loop.run_in_executor(executor, ctx.run, func, *args) diff --git a/test/testcases/restful_api/test_chats.py b/test/testcases/restful_api/test_chats.py index 84463ffdeb..95d726588c 100644 --- a/test/testcases/restful_api/test_chats.py +++ b/test/testcases/restful_api/test_chats.py @@ -16,6 +16,7 @@ import asyncio import importlib.util +import re import sys from copy import deepcopy from concurrent.futures import ThreadPoolExecutor @@ -775,7 +776,7 @@ def _load_chat_routes_unit_module(monkeypatch): class _StubTenantService: @staticmethod def get_by_id(_tenant_id): - return True, SimpleNamespace(llm_id="glm-4") + return True, SimpleNamespace(llm_id="glm-4", tenant_llm_id="tenant-llm-id") @staticmethod def get_joined_tenants_by_user_id(_user_id): @@ -1142,7 +1143,7 @@ def test_chat_create_accepts_provider_scoped_rerank_id_unit(monkeypatch): "vector_similarity_weight": 0.25, }, ) - monkeypatch.setattr(module.TenantService, "get_by_id", lambda _tid: (True, SimpleNamespace(llm_id="glm-4@CI@ZHIPU-AI"))) + monkeypatch.setattr(module.TenantService, "get_by_id", lambda _tid: (True, SimpleNamespace(llm_id="glm-4@CI@ZHIPU-AI", tenant_llm_id="tenant-llm-id"))) monkeypatch.setattr(module.DialogService, "query", lambda **_kwargs: []) monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: [SimpleNamespace(id="kb-1")]) monkeypatch.setattr(module.KnowledgebaseService, "query", lambda **_kwargs: [_DummyKB()]) @@ -1166,7 +1167,7 @@ def test_chat_create_accepts_provider_scoped_rerank_id_unit(monkeypatch): assert saved["rerank_id"] == "custom-reranker@OpenAI" assert { "tenant_id": "tenant-1", - "model_name": "custom-reranker@OpenAI", + "model_ref": "custom-reranker@OpenAI", "model_type": "rerank", } in query_calls @@ -1176,7 +1177,7 @@ def test_chat_create_allows_default_knowledge_placeholder_without_sources_unit(m module = _load_chat_routes_unit_module(monkeypatch) saved = {} _set_route_unit_request_json(monkeypatch, module, {"name": "chat-a"}) - monkeypatch.setattr(module.TenantService, "get_by_id", lambda _tid: (True, SimpleNamespace(llm_id="glm-4"))) + monkeypatch.setattr(module.TenantService, "get_by_id", lambda _tid: (True, SimpleNamespace(llm_id="glm-4", tenant_llm_id="tenant-llm-id"))) monkeypatch.setattr(module.DialogService, "query", lambda **_kwargs: []) monkeypatch.setattr(module, "get_api_key", lambda *_args, **_kwargs: SimpleNamespace(id=1)) @@ -1215,7 +1216,7 @@ def test_chat_create_uses_direct_chat_fields_unit(monkeypatch): "vector_similarity_weight": 0.25, }, ) - monkeypatch.setattr(module.TenantService, "get_by_id", lambda _tid: (True, SimpleNamespace(llm_id="glm-4"))) + monkeypatch.setattr(module.TenantService, "get_by_id", lambda _tid: (True, SimpleNamespace(llm_id="glm-4", tenant_llm_id="tenant-llm-id"))) monkeypatch.setattr(module.DialogService, "query", lambda **_kwargs: []) monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: [SimpleNamespace(id="kb-1")]) monkeypatch.setattr(module.KnowledgebaseService, "query", lambda **_kwargs: [_DummyKB()]) @@ -1371,7 +1372,7 @@ def test_patch_chat_drops_response_only_fields_before_update_unit(monkeypatch): _set_route_unit_request_json(monkeypatch, module, payload) monkeypatch.setattr(module.DialogService, "query", lambda **kwargs: [] if "name" in kwargs else [SimpleNamespace(id="chat-1")]) monkeypatch.setattr(module.DialogService, "get_by_id", lambda _id: (True, _DummyDialogRecord(existing))) - monkeypatch.setattr(module.TenantService, "get_by_id", lambda _tid: (True, SimpleNamespace(llm_id="glm-4"))) + monkeypatch.setattr(module.TenantService, "get_by_id", lambda _tid: (True, SimpleNamespace(llm_id="glm-4", tenant_llm_id="tenant-llm-id"))) monkeypatch.setattr(module.KnowledgebaseService, "accessible", lambda **_kwargs: [SimpleNamespace(id="kb-1")]) monkeypatch.setattr(module.KnowledgebaseService, "query", lambda **_kwargs: [_DummyKB()]) monkeypatch.setattr(module, "get_api_key", lambda *args, **kwargs: SimpleNamespace(id=1)) @@ -1399,7 +1400,7 @@ def test_patch_chat_merges_prompt_and_llm_settings_unit(monkeypatch): ) monkeypatch.setattr(module.DialogService, "query", lambda **_kwargs: [SimpleNamespace(id="chat-1")]) monkeypatch.setattr(module.DialogService, "get_by_id", lambda _id: (True, _DummyDialogRecord(existing))) - monkeypatch.setattr(module.TenantService, "get_by_id", lambda _tid: (True, SimpleNamespace(llm_id="glm-4"))) + monkeypatch.setattr(module.TenantService, "get_by_id", lambda _tid: (True, SimpleNamespace(llm_id="glm-4", tenant_llm_id="tenant-llm-id"))) def _update(_chat_id, payload): updated.update(payload) @@ -1442,7 +1443,7 @@ def test_update_chat_allows_knowledge_placeholder_without_sources_unit(monkeypat ) monkeypatch.setattr(module.DialogService, "query", lambda **_kwargs: [SimpleNamespace(id="chat-1")]) monkeypatch.setattr(module.DialogService, "get_by_id", lambda _id: (True, _DummyDialogRecord(existing))) - monkeypatch.setattr(module.TenantService, "get_by_id", lambda _tid: (True, SimpleNamespace(llm_id="glm-4"))) + monkeypatch.setattr(module.TenantService, "get_by_id", lambda _tid: (True, SimpleNamespace(llm_id="glm-4", tenant_llm_id="tenant-llm-id"))) updated = {} def _update(_chat_id, payload): @@ -1540,7 +1541,15 @@ def test_chat_create_llm_contract(rest_client, clear_chats, ensure_parsed_docume body = res.json() assert body["code"] == expected_code, (scenario_name, body) if expected_code == 0: - assert body["data"]["llm_id"] == expected_llm_id, (scenario_name, body) + actual_llm_id = body["data"]["llm_id"] + tenant_llm_id = body["data"].get("tenant_llm_id") + if expected_llm_id == "glm-4-flash@CI@ZHIPU-AI": + if tenant_llm_id: + assert actual_llm_id == tenant_llm_id, (scenario_name, body) + else: + assert re.fullmatch(r"[0-9a-f]{32}", actual_llm_id), (scenario_name, body) + else: + assert actual_llm_id == expected_llm_id, (scenario_name, body) assert body["data"]["llm_setting"] == expected_llm_setting, (scenario_name, body) else: assert body["message"] == expected_message, (scenario_name, body) @@ -1819,7 +1828,15 @@ def test_chat_update_llm_contract(rest_client, clear_chats, ensure_parsed_docume get_payload = get_res.json() assert get_payload["code"] == 0, (scenario_name, get_payload) assert get_payload["data"]["name"] == updated_name, (scenario_name, get_payload) - assert get_payload["data"]["llm_id"] == expected_llm_id, (scenario_name, get_payload) + actual_llm_id = get_payload["data"]["llm_id"] + tenant_llm_id = get_payload["data"].get("tenant_llm_id") + if expected_llm_id == "glm-4-flash@CI@ZHIPU-AI": + if tenant_llm_id: + assert actual_llm_id == tenant_llm_id, (scenario_name, get_payload) + else: + assert re.fullmatch(r"[0-9a-f]{32}", actual_llm_id), (scenario_name, get_payload) + else: + assert actual_llm_id == expected_llm_id, (scenario_name, get_payload) assert get_payload["data"]["llm_setting"] == expected_llm_setting, (scenario_name, get_payload) else: assert body["message"] == expected_message, (scenario_name, body) diff --git a/test/testcases/restful_api/test_datasets.py b/test/testcases/restful_api/test_datasets.py index 3f951478c4..133573e5bb 100644 --- a/test/testcases/restful_api/test_datasets.py +++ b/test/testcases/restful_api/test_datasets.py @@ -832,7 +832,7 @@ def test_dataset_update_embedding_model_invalid_and_none_contract(rest_client, c assert list_res.status_code == 200 list_payload = list_res.json() assert list_payload["code"] == 0, list_payload - assert list_payload["data"][0]["embedding_model"] == "BAAI/bge-small-en-v1.5@Local@Builtin", list_payload + assert list_payload["data"][0]["embedding_model"].startswith("BAAI/bge-small-en-v1.5"), list_payload @pytest.mark.p2 @@ -1197,7 +1197,10 @@ def test_dataset_create_embedding_model_contract(rest_client, clear_datasets, na pytest.xfail(f"Environment has no authorized tenant model for {embedding_model}: {payload}") assert payload["code"] == expected_code, payload if expected_embedding_model is not None: - assert payload["data"]["embedding_model"] == expected_embedding_model, payload + if name in {"embedding_model_unset", "embedding_model_none"}: + assert payload["data"]["embedding_model"].startswith("BAAI/bge-small-en-v1.5"), payload + else: + assert payload["data"]["embedding_model"] == expected_embedding_model, payload if expected_message is not None: assert payload["message"] == expected_message, payload 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 9bd9f45f19..81a286d986 100644 --- a/test/testcases/restful_api/test_user_tenant_routes_unit.py +++ b/test/testcases/restful_api/test_user_tenant_routes_unit.py @@ -1532,7 +1532,7 @@ def _load_chat_routes_unit_module(monkeypatch): "TenantService", (), { - "get_by_id": staticmethod(lambda _tenant_id: (True, SimpleNamespace(llm_id="glm-4"))), + "get_by_id": staticmethod(lambda _tenant_id: (True, SimpleNamespace(llm_id="glm-4", tenant_llm_id="tenant-llm-id"))), "get_joined_tenants_by_user_id": staticmethod(lambda _user_id: [{"tenant_id": "tenant-1"}, {"tenant_id": "team-tenant-2"}]), }, ) @@ -1599,7 +1599,7 @@ def test_create_chat_uses_tenant_default_llm_when_llm_id_is_null_unit(monkeypatc res = _run(module.create.__wrapped__()) assert res["code"] == 0 - assert saved["llm_id"] == "glm-4" + assert saved["llm_id"] == "tenant-llm-id" assert saved["llm_setting"]["temperature"] == 0.8 diff --git a/test/testcases/test_sdk_api/test_dataset_mangement/test_create_dataset.py b/test/testcases/test_sdk_api/test_dataset_mangement/test_create_dataset.py index 91673746e9..3686b3d970 100644 --- a/test/testcases/test_sdk_api/test_dataset_mangement/test_create_dataset.py +++ b/test/testcases/test_sdk_api/test_dataset_mangement/test_create_dataset.py @@ -237,13 +237,13 @@ class TestDatasetCreate: def test_embedding_model_unset(self, client): payload = {"name": "embedding_model_unset"} dataset = client.create_dataset(**payload) - assert dataset.embedding_model == "BAAI/bge-small-en-v1.5@Local@Builtin", str(dataset) + assert dataset.embedding_model.split("@", 1)[0] == "BAAI/bge-small-en-v1.5", str(dataset) @pytest.mark.p2 def test_embedding_model_none(self, client): payload = {"name": "embedding_model_none", "embedding_model": None} dataset = client.create_dataset(**payload) - assert dataset.embedding_model == "BAAI/bge-small-en-v1.5@Local@Builtin", str(dataset) + assert dataset.embedding_model.split("@", 1)[0] == "BAAI/bge-small-en-v1.5", str(dataset) @pytest.mark.p2 @pytest.mark.parametrize( diff --git a/test/testcases/test_sdk_api/test_dataset_mangement/test_update_dataset.py b/test/testcases/test_sdk_api/test_dataset_mangement/test_update_dataset.py index a0dfd5cdf5..a26749e697 100644 --- a/test/testcases/test_sdk_api/test_dataset_mangement/test_update_dataset.py +++ b/test/testcases/test_sdk_api/test_dataset_mangement/test_update_dataset.py @@ -229,10 +229,10 @@ class TestDatasetUpdate: def test_embedding_model_none(self, client, add_dataset_func): dataset = add_dataset_func dataset.update({"embedding_model": None}) - assert dataset.embedding_model == "BAAI/bge-small-en-v1.5@Local@Builtin", str(dataset) + assert dataset.embedding_model.split("@", 1)[0] == "BAAI/bge-small-en-v1.5", str(dataset) retrieved_dataset = client.get_dataset(name=dataset.name) - assert retrieved_dataset.embedding_model == "BAAI/bge-small-en-v1.5@Local@Builtin", str(retrieved_dataset) + assert retrieved_dataset.embedding_model.split("@", 1)[0] == "BAAI/bge-small-en-v1.5", str(retrieved_dataset) @pytest.mark.p2 @pytest.mark.parametrize( 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 d8ac3263a7..7e5c021c71 100644 --- a/test/unit_test/api/apps/services/test_delete_datasets.py +++ b/test/unit_test/api/apps/services/test_delete_datasets.py @@ -18,6 +18,7 @@ import importlib.util import sys from pathlib import Path +from enum import IntEnum from types import ModuleType, SimpleNamespace from unittest.mock import MagicMock @@ -26,6 +27,16 @@ import pytest pytestmark = pytest.mark.p2 +class _StubModelTypeBinary(IntEnum): + CHAT = 1 + EMBEDDING = 2 + SPEECH2TEXT = 4 + IMAGE2TEXT = 8 + RERANK = 16 + TTS = 32 + OCR = 64 + + def _stub(monkeypatch, name, **attrs): mod = ModuleType(name) for key, value in attrs.items(): @@ -122,6 +133,8 @@ def _load_delete_datasets_module(monkeypatch, *, f2d_rows, file_filter_delete): _stub( monkeypatch, "api.db.db_models", + DB=SimpleNamespace(connection_context=lambda: lambda func: func), + TenantModel=SimpleNamespace(), File=SimpleNamespace(source_type="source_type", id="id", type="type", name="name"), ) _stub( @@ -130,8 +143,18 @@ def _load_delete_datasets_module(monkeypatch, *, f2d_rows, file_filter_delete): PAGERANK_FLD="pagerank", TAG_FLD="tag", FileSource=SimpleNamespace(KNOWLEDGEBASE="knowledgebase"), + PipelineTaskType=SimpleNamespace( + PARSE="parse", + DOWNLOAD="download", + RAPTOR="raptor", + GRAPH_RAG="graph_rag", + MINDMAP="mindmap", + ARTIFACT="artifact", + SKILL="skill", + ), StatusEnum=SimpleNamespace(), LLMType=SimpleNamespace(), + ModelTypeBinary=_StubModelTypeBinary, ) _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 e72e3947f7..c0d426f000 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 @@ -25,6 +25,7 @@ import sys from pathlib import Path from enum import IntEnum from types import ModuleType, SimpleNamespace +from unittest.mock import MagicMock import pytest @@ -141,6 +142,7 @@ def _load_module(monkeypatch, *, tenant_model_records, factory_llm_infos=None): ensure_mineru_from_env=lambda *a, **kw: None, ensure_paddleocr_from_env=lambda *a, **kw: None, ensure_opendataloader_from_env=lambda *a, **kw: None, + resolve_model_id=MagicMock(), ) _stub( monkeypatch, diff --git a/tools/scripts/mysql_migration.py b/tools/scripts/mysql_migration.py index 4d871548c6..55a61537d3 100644 --- a/tools/scripts/mysql_migration.py +++ b/tools/scripts/mysql_migration.py @@ -1777,7 +1777,7 @@ class TenantModelIdMigrationStage(MigrationStage): has_work = True break # Check if there are NULL values that should be populated - cursor = self.db.execute_sql(f"SELECT COUNT(*) FROM `{table_name}` WHERE `{tenant_id_col}` IS NULL OR `{tenant_id_col}` = ''") + cursor = self.db.execute_sql(f"SELECT COUNT(*) FROM `{table_name}` WHERE `{tenant_id_col}` IS NULL OR `{tenant_id_col}` = '' OR LENGTH(`{tenant_id_col}`) <> 32") null_count = cursor.fetchone()[0] if null_count > 0: has_work = True @@ -1837,17 +1837,12 @@ class TenantModelIdMigrationStage(MigrationStage): logger.info(f"Column {table_name}.{legacy_col} does not exist, skipping") continue - # Get rows where tenant_*_id is NULL or empty - cursor = self.db.execute_sql( - f"SELECT id, `{legacy_col}` FROM `{table_name}` WHERE (`{tenant_id_col}` IS NULL OR `{tenant_id_col}` = '') AND `{legacy_col}` IS NOT NULL AND `{legacy_col}` != ''" - ) - # Also need tenant_id for model lookup # For tenant table, the PK is the tenant_id # For knowledgebase, dialog, memory we also need tenant_id if table_name == "tenant": cursor = self.db.execute_sql( - f"SELECT id, `{legacy_col}` FROM `{table_name}` WHERE (`{tenant_id_col}` IS NULL OR `{tenant_id_col}` = '') AND `{legacy_col}` IS NOT NULL AND `{legacy_col}` != ''" + f"SELECT id, `{legacy_col}` FROM `{table_name}` WHERE (`{tenant_id_col}` IS NULL OR `{tenant_id_col}` = '' OR LENGTH(`{tenant_id_col}`) <> 32) AND `{legacy_col}` IS NOT NULL AND `{legacy_col}` != ''" ) while True: rows = cursor.fetchmany(self.scan_batch_size) @@ -1872,7 +1867,7 @@ class TenantModelIdMigrationStage(MigrationStage): tables_operated.add(table_name) else: cursor = self.db.execute_sql( - f"SELECT id, tenant_id, `{legacy_col}` FROM `{table_name}` WHERE (`{tenant_id_col}` IS NULL OR `{tenant_id_col}` = '') AND `{legacy_col}` IS NOT NULL AND `{legacy_col}` != ''" + f"SELECT id, tenant_id, `{legacy_col}` FROM `{table_name}` WHERE (`{tenant_id_col}` IS NULL OR `{tenant_id_col}` = '' OR LENGTH(`{tenant_id_col}`) <> 32) AND `{legacy_col}` IS NOT NULL AND `{legacy_col}` != ''" ) while True: rows = cursor.fetchmany(self.scan_batch_size) diff --git a/tools/scripts/run_migrations.sh b/tools/scripts/run_migrations.sh index e2492e077b..991d215bc2 100755 --- a/tools/scripts/run_migrations.sh +++ b/tools/scripts/run_migrations.sh @@ -32,7 +32,7 @@ echo "Running model provider table migrations..." --stages tenant_model_seeding,model_type_merge,tenant_model_id_migration \ --config "$CONFIG" \ --execute \ - --database-version "v0.27.0.dev0" \ + --database-version "v0.27.0.dev1" \ --mark-database-version-on-success echo "Model provider table migrations completed."