Handle searching dataset without embedding model (#16742)

### Summary

Handle searching dataset without embedding model

In this PR, Searching datasets with different embedding models or
searching dataset with/without embedding models are not allowed. We will
improve the behavior later.

api/apps/restful_apis/chat_api.py

@@ -27,11 +27,11 @@ 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, split_model_name
+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.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
-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.search_service import SearchService
 from api.db.services.user_service import TenantService, UserTenantService
@@ -337,9 +337,9 @@ async def _validate_dataset_ids(dataset_ids, tenant_id):
             return f"The dataset {dataset_id} doesn't own parsed file"
         kbs.append(kb)
 
-    embd_ids = [split_model_name(kb.embd_id)[0] for kb in kbs]
-    if len(set(embd_ids)) > 1:
-        return f"Datasets use different embedding models: {[kb.embd_id for kb in kbs]}"
+    err = validate_dataset_embedding_models(kbs)
+    if err:
+        return err
 
     return normalized_ids
 

api/apps/services/dataset_api_service.py

@@ -25,7 +25,7 @@ from api.db.db_models import File
 from api.db.services.document_service import DocumentService, queue_raptor_o_graphrag_tasks
 from api.db.services.file2document_service import File2DocumentService
 from api.db.services.file_service import FileService
-from api.db.services.knowledgebase_service import KnowledgebaseService
+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.user_service import TenantService, UserService, UserTenantService
@@ -1167,8 +1167,7 @@ def check_embedding(dataset_id: str, tenant_id: str, req: dict):
         except Exception as e:
             if "not_found_exception" in repr(e) or "index_not_found_exception" in repr(e):
                 logging.info(
-                    "sample_random_chunks_with_vectors: index %s not yet created for tenant %s; "
-                    "returning empty sample set",
+                    "sample_random_chunks_with_vectors: index %s not yet created for tenant %s; returning empty sample set",
                     index_nm,
                     tenant_id,
                 )
@@ -1327,7 +1326,7 @@ async def search_datasets(tenant_id: str, req: dict):
     :param req: search request containing dataset_ids and other params
     :return: (success, result) or (success, error_message)
     """
-    from api.db.joint_services.tenant_model_service import get_tenant_default_model_by_type, split_model_name
+    from api.db.joint_services.tenant_model_service import get_tenant_default_model_by_type
     from api.db.services.doc_metadata_service import DocMetadataService
     from api.db.services.llm_service import LLMBundle
     from api.db.services.search_service import SearchService
@@ -1365,10 +1364,9 @@ async def search_datasets(tenant_id: str, req: dict):
     if not kbs:
         return False, "Datasets not found!"
 
-    # All datasets must use the same embedding model
-    embd_nms = list(set([split_model_name(kb.embd_id)[0] for kb in kbs]))
-    if len(embd_nms) != 1:
-        return False, "Datasets use different embedding models."
+    err = validate_dataset_embedding_models(kbs)
+    if err:
+        return False, err
 
     if doc_ids is not None and not isinstance(doc_ids, list):
         return False, "`doc_ids` should be a list"
@@ -1437,11 +1435,11 @@ async def search_datasets(tenant_id: str, req: dict):
     _question = question
     if langs:
         _question = await cross_languages(kb.tenant_id, None, _question, langs)
+
+    embd_mdl = None
     if kb.embd_id:
         embd_model_config = get_model_config_from_provider_instance(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)
+        embd_mdl = LLMBundle(kb.tenant_id, embd_model_config)
 
     rerank_mdl = None
     rerank_id = search_config.get("rerank_id") or req.get("rerank_id")

api/db/services/dialog_service.py

@@ -31,7 +31,7 @@ from common.constants import LLMType, ParserType, StatusEnum
 from api.db.db_models import DB, Dialog
 from api.db.services.common_service import CommonService
 from api.db.services.doc_metadata_service import DocMetadataService
-from api.db.services.knowledgebase_service import KnowledgebaseService
+from api.db.services.knowledgebase_service import KnowledgebaseService, validate_dataset_embedding_models
 from api.db.services.langfuse_service import TenantLangfuseService
 from api.db.services.llm_service import LLMBundle
 from common.metadata_utils import apply_meta_data_filter
@@ -358,16 +358,16 @@ async def async_chat_solo(dialog, messages, stream=True, session_id=None):
 def get_models(dialog, trace_context=None, langfuse_session_id=None):
     embd_mdl, chat_mdl, rerank_mdl, tts_mdl = None, None, None, None
     kbs = KnowledgebaseService.get_by_ids(dialog.kb_ids)
-    embedding_list = list(set([kb.embd_id for kb in kbs]))
-    if len(embedding_list) > 1:
-        raise Exception("**ERROR**: Knowledge bases use different embedding models.")
+    err = validate_dataset_embedding_models(kbs)
+    if err:
+        raise Exception(err)
 
-    if embedding_list:
+    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, embedding_list[0])
+        embd_model_config = get_model_config_from_provider_instance(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" % embedding_list[0])
+            raise LookupError("Embedding model(%s) not found" % kbs[0].embd_id)
 
     if dialog.llm_id:
         if dialog.tenant_llm_id:
@@ -721,8 +721,8 @@ async def async_chat(dialog, messages, stream=True, **kwargs):
                 prompt_config,
                 partial(
                     retriever.retrieval,
-                    embd_mdl = embd_mdl,
-                    tenant_ids = tenant_ids,
+                    embd_mdl=embd_mdl,
+                    tenant_ids=tenant_ids,
                     kb_ids=dialog.kb_ids,
                     page=1,
                     page_size=dialog.top_n,

api/db/services/knowledgebase_service.py

@@ -29,6 +29,29 @@ from api.constants import DATASET_NAME_LIMIT
 from api.utils.api_utils import get_parser_config, get_data_error_result
 
 
+def _base_model_name(embd_id: str) -> str:
+    """Return the base model name by stripping provider/instance suffix from an embd_id."""
+    parts = embd_id.rsplit("@", 2)
+    return parts[0]
+
+
+def validate_dataset_embedding_models(kbs):
+    """Validate that all given datasets use the same embedding model (or all use none).
+
+    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.
+    embd_ids = [kb.embd_id for kb in kbs if kb.embd_id]
+    has_embd = len(embd_ids) > 0
+    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})
+        if len(embd_nms) > 1:
+            return f"Datasets use different embedding models: {[kb.embd_id for kb in kbs]}"
+    return None
+
+
 class KnowledgebaseService(CommonService):
     """Service class for managing dataset operations.