youngkingdom/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-07-15 15:47:15 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix: align model default handling (#16782)

Browse Source
This commit is contained in:
buua436
2026-07-10 10:34:19 +08:00
committed by GitHub
parent d48a5622df
commit 74bbbba3e0
15 changed files with 169 additions and 71 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
api/apps/restful_apis/chat_api.py
View File

@ -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)

10
api/apps/restful_apis/models_api.py
View File

@ -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)

70
api/apps/services/models_api_service.py
View File

@ -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"

28
api/db/services/compilation_template_service.py
View File

@ -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

7
api/db/services/user_service.py
View File

@ -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,
]

25
common/misc_utils.py
View File

@ -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)

37
test/testcases/restful_api/test_chats.py
View File

@ -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)

7
test/testcases/restful_api/test_datasets.py
View File

@ -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

4
test/testcases/restful_api/test_user_tenant_routes_unit.py
View File

@ -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

4
test/testcases/test_sdk_api/test_dataset_mangement/test_create_dataset.py
View File

@ -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(

4
test/testcases/test_sdk_api/test_dataset_mangement/test_update_dataset.py
View File

@ -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(

23
test/unit_test/api/apps/services/test_delete_datasets.py
View File

@ -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,

2
test/unit_test/api/apps/services/test_models_api_service_list_tenant_added_models.py
View File

@ -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,

11
tools/scripts/mysql_migration.py
View File

@ -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)

2
tools/scripts/run_migrations.sh
View File

@ -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."