Fix: allow datasets sharing a base embedding model to be searched together (#18166)

api/apps/restful_apis/chunk_api.py

@@ -26,17 +26,16 @@ from quart import request
 
 from api.apps import login_required
 from api.apps.services import structure_graph_common as sgc
-from api.db.joint_services.tenant_model_service import (
-    split_model_name,
-    resolve_model_config,
-    get_tenant_default_model_by_type,
-)
 from api.db.db_models import Document, Task
+from api.db.joint_services.tenant_model_service import (
+    get_tenant_default_model_by_type,
+    resolve_model_config,
+)
 from api.db.services.doc_metadata_service import DocMetadataService
 from api.db.services.document_counter_service import release_reparse_counters
 from api.db.services.document_service import DocumentService
 from api.db.services.file2document_service import File2DocumentService
-from api.db.services.knowledgebase_service import KnowledgebaseService
+from api.db.services.knowledgebase_service import KnowledgebaseService, validate_dataset_embedding_models
 from api.db.services.llm_service import LLMBundle
 from api.db.services.task_service import TaskService, cancel_all_task_of, queue_tasks
 from api.db.services.tenant_llm_service import TenantLLMService
@@ -49,8 +48,8 @@ from api.utils.api_utils import (
     get_result,
     server_error_response,
 )
-from api.utils.pagination_utils import DEFAULT_PAGE, DEFAULT_PAGE_SIZE, validate_rest_api_ids, validate_rest_api_page, validate_rest_api_page_size
 from api.utils.image_utils import store_chunk_image
+from api.utils.pagination_utils import DEFAULT_PAGE, DEFAULT_PAGE_SIZE, validate_rest_api_ids, validate_rest_api_page, validate_rest_api_page_size
 from api.utils.reference_metadata_utils import (
     enrich_chunks_with_document_metadata,
     resolve_reference_metadata_preferences,
@@ -66,7 +65,6 @@ from rag.app.tag import label_question
 from rag.nlp import search
 from rag.prompts.generator import cross_languages, keyword_extraction
 
-
 DOC_STOP_PARSING_INVALID_STATE_MESSAGE = "Can't stop parsing document that has not started or already completed"
 DOC_STOP_PARSING_INVALID_STATE_ERROR_CODE = "DOC_STOP_PARSING_INVALID_STATE"
 
@@ -342,9 +340,9 @@ async def retrieval_test(tenant_id):
         if not KnowledgebaseService.accessible(kb_id=id, user_id=tenant_id):
             return get_error_data_result(f"You don't own the dataset {id}.")
     kbs = KnowledgebaseService.get_by_ids(kb_ids)
-    embd_nms = list(set([split_model_name(kb.embd_id)[0] for kb in kbs]))
-    if len(embd_nms) != 1:
-        return get_result(message="Datasets use different embedding models.", code=RetCode.DATA_ERROR)
+    embd_err = validate_dataset_embedding_models(kbs)
+    if embd_err:
+        return get_result(message=embd_err, code=RetCode.DATA_ERROR)
     if "question" not in req:
         return get_error_data_result("`question` is required.")
     page = validate_rest_api_page(req.get("page", DEFAULT_PAGE))
@@ -605,9 +603,9 @@ async def get_document_structure_graph(tenant_id, dataset_id, document_id):
     migration doesn't drop their data on the floor. Empty templates
     (zero entities AND zero relations) are filtered out.
     """
-    from rag.nlp import search
     from api.db.services.compilation_template_group_service import CompilationTemplateGroupService
     from api.db.services.compilation_template_service import CompilationTemplateService
+    from rag.nlp import search
 
     if not KnowledgebaseService.accessible(kb_id=dataset_id, user_id=tenant_id):
         return get_error_data_result(message=f"You don't own the dataset {dataset_id}.")
@@ -855,7 +853,7 @@ async def get_document_structure_graph(tenant_id, dataset_id, document_id):
     for tid in configured_ids:
         if tid in grouped and tid not in ordered_ids:
             ordered_ids.append(tid)
-    for bucket_id in grouped.keys():
+    for bucket_id in grouped:
         if bucket_id not in ordered_ids:
             ordered_ids.append(bucket_id)
 

api/db/services/knowledgebase_service.py

@@ -15,18 +15,19 @@
 #
 from datetime import datetime
 
-from peewee import fn, JOIN
+from peewee import JOIN, fn
 
+from api.constants import DATASET_NAME_LIMIT
 from api.db import TenantPermission
 from api.db.db_models import DB, Document, Knowledgebase, User, UserCanvas
-from api.db.services.common_service import CommonService
-from common.time_utils import current_timestamp, datetime_format
+from api.db.joint_services.tenant_model_service import get_composite_model_name_by_ids
 from api.db.services import duplicate_name
+from api.db.services.common_service import CommonService
 from api.db.services.user_service import TenantService
-from common.misc_utils import get_uuid
+from api.utils.api_utils import get_data_error_result, get_parser_config
 from common.constants import StatusEnum
-from api.constants import DATASET_NAME_LIMIT
-from api.utils.api_utils import get_parser_config, get_data_error_result
+from common.misc_utils import get_uuid
+from common.time_utils import current_timestamp, datetime_format
 
 
 def _base_model_name(embd_id: str) -> str:
@@ -35,9 +36,37 @@ def _base_model_name(embd_id: str) -> str:
     return parts[0]
 
 
+def _kb_embedding_base_name(kb, resolved_names) -> str:
+    """Resolve a dataset's embedding reference to its base model name.
+
+    ``tenant_embd_id`` — or ``embd_id`` itself when it stores a raw
+    tenant_model id — is resolved through ``resolved_names`` (id to
+    ``model@instance@provider``). An id that no longer resolves falls back to
+    the composite base name when ``embd_id`` holds one, otherwise to the id
+    itself so only exact matches group together.
+    """
+    embd_id = (kb.embd_id or "").strip()
+    ref = (getattr(kb, "tenant_embd_id", None) or "").strip()
+    if not ref and "@" not in embd_id:
+        ref = embd_id
+    if not ref:
+        return _base_model_name(embd_id)
+    composite = resolved_names.get(ref)
+    if composite:
+        return _base_model_name(composite)
+    if embd_id and embd_id != ref:
+        return _base_model_name(embd_id)
+    return ref
+
+
 def validate_dataset_embedding_models(kbs):
     """Validate that all given datasets use the same embedding model (or all use none).
 
+    Embedding references are resolved through tenant_model first, so datasets
+    storing a raw tenant_model id and datasets storing a legacy
+    ``model@instance@provider`` composite compare equal when they point at the
+    same model.
+
     Returns an error message string on failure, or ``None`` on success.
     """
     # Either all datasets have an embedding model, or none do. Mixing is not allowed.
@@ -46,7 +75,20 @@ def validate_dataset_embedding_models(kbs):
     if has_embd and len(embd_ids) != len(kbs):
         return "Cannot search across datasets where some have embedding models and others do not."
     if has_embd:
-        embd_nms = list({_base_model_name(eid) for eid in embd_ids})
+        candidates = []
+        for kb in kbs:
+            if not kb.embd_id:
+                continue
+            ref = (getattr(kb, "tenant_embd_id", None) or "").strip()
+            if not ref and "@" not in kb.embd_id:
+                ref = kb.embd_id.strip()
+            if ref:
+                candidates.append(ref)
+        try:
+            resolved_names = get_composite_model_name_by_ids(candidates)
+        except Exception:  # noqa: BLE001 - resolution is best-effort; unresolvable ids keep their raw value
+            resolved_names = {}
+        embd_nms = {_kb_embedding_base_name(kb, resolved_names) for kb in kbs if kb.embd_id}
         if len(embd_nms) > 1:
             return f"Datasets use different embedding models: {[kb.embd_id for kb in kbs]}"
     return None
@@ -129,8 +171,8 @@ class KnowledgebaseService(CommonService):
         # Returns:
         #     If all documents are parsed successfully, returns (True, None)
         #     If any document is not fully parsed, returns (False, error_message)
-        from common.constants import TaskStatus
         from api.db.services.document_service import DocumentService
+        from common.constants import TaskStatus
 
         # Get dataset information
         kbs = cls.query(id=kb_id)