diff --git a/api/apps/restful_apis/chunk_api.py b/api/apps/restful_apis/chunk_api.py index 751df6fd9e..8a7485b790 100644 --- a/api/apps/restful_apis/chunk_api.py +++ b/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) diff --git a/api/db/services/knowledgebase_service.py b/api/db/services/knowledgebase_service.py index dfa668f18d..2d906ad1e4 100644 --- a/api/db/services/knowledgebase_service.py +++ b/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) diff --git a/internal/agent/tool/retrieval_nlp.go b/internal/agent/tool/retrieval_nlp.go index 7434c5fd6e..dc4976cc7d 100644 --- a/internal/agent/tool/retrieval_nlp.go +++ b/internal/agent/tool/retrieval_nlp.go @@ -215,7 +215,7 @@ func (a *NLPRetrievalAdapter) Search(ctx context.Context, db *gorm.DB, req Retri if len(datasets.tenantIDs) != 1 { return nil, fmt.Errorf("retrieval: datasets span multiple tenants") } - if err := validateEmbeddingModels(datasets.kbs); err != nil { + if err := validateEmbeddingModels(ctx, db, datasets.kbs); err != nil { return nil, err } embeddingModel, err := a.resolveEmbeddingModel(ctx, datasets.kbs[0]) @@ -412,7 +412,7 @@ func (a *NLPRetrievalAdapter) resolveDatasets( return &resolvedDatasets{kbs: kbs, kbIDs: resolvedKBIDs, tenantIDs: tenantIDs}, nil } -func validateEmbeddingModels(kbs []*entity.Knowledgebase) error { +func validateEmbeddingModels(ctx context.Context, db *gorm.DB, kbs []*entity.Knowledgebase) error { if len(kbs) == 0 { return fmt.Errorf("retrieval: no datasets selected") } @@ -421,23 +421,26 @@ func validateEmbeddingModels(kbs []*entity.Knowledgebase) error { return fmt.Errorf("retrieval: dataset record is nil") } } - firstKey := knowledgebaseEmbeddingKey(kbs[0]) + embdNameCache := make(map[string]string) + firstKey := knowledgebaseEmbeddingKey(ctx, db, kbs[0], embdNameCache) for _, kb := range kbs[1:] { - if knowledgebaseEmbeddingKey(kb) != firstKey { + if knowledgebaseEmbeddingKey(ctx, db, kb, embdNameCache) != firstKey { return fmt.Errorf("retrieval: datasets use different embedding models") } } return nil } -func knowledgebaseEmbeddingKey(kb *entity.Knowledgebase) string { - if kb.TenantEmbdID != nil && strings.TrimSpace(*kb.TenantEmbdID) != "" { - return "tenant:" + strings.TrimSpace(*kb.TenantEmbdID) +// knowledgebaseEmbeddingKey groups datasets by their resolved base embedding +// model name (e.g. "BAAI/bge-m3"), matching the chat/dataset-search +// validation, so datasets pointing at the same model through different +// provider instances or storage forms (tenant_model id vs legacy composite +// name) retrieve together. +func knowledgebaseEmbeddingKey(ctx context.Context, db *gorm.DB, kb *entity.Knowledgebase, cache map[string]string) string { + if strings.TrimSpace(kb.EmbdID) == "" && (kb.TenantEmbdID == nil || strings.TrimSpace(*kb.TenantEmbdID) == "") { + return "default:" + strings.TrimSpace(kb.TenantID) } - if strings.TrimSpace(kb.EmbdID) != "" { - return "embedding:" + strings.TrimSpace(kb.EmbdID) - } - return "default:" + strings.TrimSpace(kb.TenantID) + return "embedding:" + dao.NewKnowledgebaseDAO().EmbeddingBaseName(ctx, db, kb, cache) } func (a *NLPRetrievalAdapter) resolveEmbeddingModel( diff --git a/internal/agent/tool/retrieval_nlp_test.go b/internal/agent/tool/retrieval_nlp_test.go index ed15ee6cd8..2de87c92fe 100644 --- a/internal/agent/tool/retrieval_nlp_test.go +++ b/internal/agent/tool/retrieval_nlp_test.go @@ -437,7 +437,7 @@ func TestNLPRetrievalAdapter_ResolveEmbeddingModelPriority(t *testing.T) { func TestValidateEmbeddingModelsRejectsDifferentTenantModels(t *testing.T) { firstID := "tenant-embedding-1" secondID := "tenant-embedding-2" - err := validateEmbeddingModels([]*entity.Knowledgebase{ + err := validateEmbeddingModels(t.Context(), nil, []*entity.Knowledgebase{ {ID: "kb-1", TenantID: "tenant-1", TenantEmbdID: &firstID}, {ID: "kb-2", TenantID: "tenant-1", TenantEmbdID: &secondID}, }) @@ -446,6 +446,28 @@ func TestValidateEmbeddingModelsRejectsDifferentTenantModels(t *testing.T) { } } +func TestValidateEmbeddingModelsAllowsSameBaseAcrossInstances(t *testing.T) { + // Datasets using the same base embedding model through different provider + // instances must validate together, matching the chat/dataset-search rule. + err := validateEmbeddingModels(t.Context(), nil, []*entity.Knowledgebase{ + {ID: "kb-1", TenantID: "tenant-1", EmbdID: "BAAI/bge-m3@renew@SILICONFLOW"}, + {ID: "kb-2", TenantID: "tenant-1", EmbdID: "BAAI/bge-m3@COPY@SILICONFLOW"}, + }) + if err != nil { + t.Fatalf("expected same-base composites to be accepted, got %v", err) + } +} + +func TestValidateEmbeddingModelsRejectsDifferentBases(t *testing.T) { + err := validateEmbeddingModels(t.Context(), nil, []*entity.Knowledgebase{ + {ID: "kb-1", TenantID: "tenant-1", EmbdID: "BAAI/bge-m3@renew@SILICONFLOW"}, + {ID: "kb-2", TenantID: "tenant-1", EmbdID: "Qwen/Qwen3-Embedding-0.6B@renew@SILICONFLOW"}, + }) + if err == nil { + t.Fatal("expected different base embedding models to be rejected") + } +} + type fakeKnowledgebaseLookup struct { kbs []*entity.Knowledgebase err error diff --git a/internal/common/format.go b/internal/common/format.go index 7d1ac423c4..5cd807a0e9 100644 --- a/internal/common/format.go +++ b/internal/common/format.go @@ -58,6 +58,22 @@ func IsUUID(uuid string) bool { return false } +// BaseModelName returns the bare model name of a (possibly composite) model +// reference by stripping the trailing "@instance@provider" (or "@provider") +// segments. The split is right-anchored so model names that legitimately +// contain '@' (e.g. LM Studio quant suffixes) are preserved. Mirrors Python's +// api/db/services/knowledgebase_service.py _base_model_name (rsplit("@", 2)[0]). +func BaseModelName(modelName string) string { + if idx := strings.LastIndex(modelName, "@"); idx > 0 { + base := modelName[:idx] + if idx2 := strings.LastIndex(base, "@"); idx2 > 0 { + return base[:idx2] + } + return base + } + return modelName +} + // ExtractCompositeName splits a composite model name into three parts. // Returns (modelName, instanceName, providerName, true) on success, // or ("", "", "", false) if the name is not a valid composite name. diff --git a/internal/common/format_test.go b/internal/common/format_test.go index 32c657f8d8..24bc1f0bcf 100644 --- a/internal/common/format_test.go +++ b/internal/common/format_test.go @@ -53,3 +53,25 @@ func TestChunkID_PreservesLeadingZero(t *testing.T) { t.Fatalf("ChunkID length = %d, want 16", len(got)) } } + +func TestBaseModelName(t *testing.T) { + tests := []struct { + name string + input string + want string + }{ + {"bare model name stays", "BAAI/bge-m3", "BAAI/bge-m3"}, + {"2-part composite strips provider", "BAAI/bge-m3@SILICONFLOW", "BAAI/bge-m3"}, + {"3-part composite strips instance and provider", "BAAI/bge-m3@renew@SILICONFLOW", "BAAI/bge-m3"}, + {"embedded @ in model name is preserved", "text-embedding-nomic-embed-text-v1.5@q8_0@lmstudio@LM-Studio", "text-embedding-nomic-embed-text-v1.5@q8_0"}, + {"opaque id without @ stays intact", "2d8ff0a97d75431c8c91526549939328", "2d8ff0a97d75431c8c91526549939328"}, + {"empty string", "", ""}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := BaseModelName(tt.input); got != tt.want { + t.Fatalf("BaseModelName(%q) = %q, want %q", tt.input, got, tt.want) + } + }) + } +} diff --git a/internal/dao/kb.go b/internal/dao/kb.go index bd265e2c96..c723437a53 100644 --- a/internal/dao/kb.go +++ b/internal/dao/kb.go @@ -21,6 +21,7 @@ import ( "errors" "fmt" "path" + "ragflow/internal/common" "ragflow/internal/entity" "strconv" @@ -59,6 +60,50 @@ func NewKnowledgebaseDAO() *KnowledgebaseDAO { return &KnowledgebaseDAO{} } +// EmbeddingBaseName resolves a knowledge base's embedding model reference to +// the bare model name (e.g. "BAAI/bge-m3") used to decide whether datasets +// can be selected and searched together. tenant_embd_id — or embd_id itself +// when it stores a raw tenant_model id — is resolved through tenant_model; +// legacy composite "model@instance@provider" values are reduced with +// common.BaseModelName. 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. +func (dao *KnowledgebaseDAO) EmbeddingBaseName(ctx context.Context, db *gorm.DB, kb *entity.Knowledgebase, cache map[string]string) string { + raw := strings.TrimSpace(kb.EmbdID) + id := "" + if kb.TenantEmbdID != nil { + id = strings.TrimSpace(*kb.TenantEmbdID) + } + if id == "" && raw != "" && !strings.Contains(raw, "@") { + id = raw + } + if id == "" { + return common.BaseModelName(raw) + } + if cache != nil { + if cached, ok := cache[id]; ok { + return cached + } + } + base := "" + if db != nil { + if model, err := NewTenantModelDAO().GetByID(ctx, db, id); err == nil && model != nil { + base = strings.TrimSpace(model.ModelName) + } + } + if base == "" { + if raw != "" && raw != id { + base = common.BaseModelName(raw) + } else { + base = id + } + } + if cache != nil { + cache[id] = base + } + return base +} + // Create creates a new knowledge base record func (dao *KnowledgebaseDAO) Create(ctx context.Context, db *gorm.DB, kb *entity.Knowledgebase) error { return db.WithContext(ctx).Create(kb).Error diff --git a/internal/service/chat.go b/internal/service/chat.go index a924a6c25c..13e2970bf4 100644 --- a/internal/service/chat.go +++ b/internal/service/chat.go @@ -401,7 +401,7 @@ func (s *ChatService) validateCreateDatasetIDs(ctx context.Context, value interf kbs = append(kbs, kb) } - if err := ValidateDatasetEmbeddingModels(kbs); err != nil { + if err := ValidateDatasetEmbeddingModels(ctx, dao.DB, kbs); err != nil { return nil, err } return normalizedIDs, nil @@ -774,21 +774,6 @@ const ( pyDefaultEmptyResponse = "Sorry! No relevant content was found in the knowledge base!" ) -// splitModelNameAndFactory extracts the base model name by stripping -// provider and instance suffixes, matching Python's rsplit("@", 2)[0]. -func (s *ChatService) splitModelNameAndFactory(embeddingModelID string) string { - if idx := strings.LastIndex(embeddingModelID, "@"); idx > 0 { - // Strip the provider segment. - base := embeddingModelID[:idx] - // Strip the instance segment (second-to-last @). - if idx2 := strings.LastIndex(base, "@"); idx2 > 0 { - return base[:idx2] - } - return base - } - return embeddingModelID -} - func (s *ChatService) getOwnedValidChat(ctx context.Context, userID, chatID string) (*entity.Chat, error) { chat, err := s.chatDAO.GetByIDAndStatus(ctx, dao.DB, chatID, string(entity.StatusValid)) if err != nil { @@ -1067,9 +1052,10 @@ func (s *ChatService) validateRESTDatasetIDs(ctx context.Context, value interfac embeddingModelIDs := make([]string, 0, len(kbs)) seenEmbedIDs := make(map[string]struct{}) + embdNameCache := make(map[string]string) for _, kb := range kbs { embeddingModelIDs = append(embeddingModelIDs, kb.EmbdID) - seenEmbedIDs[s.splitModelNameAndFactory(kb.EmbdID)] = struct{}{} + seenEmbedIDs[s.kbDAO.EmbeddingBaseName(ctx, dao.DB, kb, embdNameCache)] = struct{}{} } if len(seenEmbedIDs) > 1 { return nil, fmt.Errorf("datasets use different embedding models: %v", embeddingModelIDs) @@ -1112,7 +1098,7 @@ func (s *ChatService) resolveRESTRerankID(ctx context.Context, rerankID, tenantI if rerankID == "" { return "", nil } - baseName := s.splitModelNameAndFactory(rerankID) + baseName := common.BaseModelName(rerankID) if _, ok := defaultRerankModels[baseName]; ok { return "", nil } diff --git a/internal/service/chat_pipeline.go b/internal/service/chat_pipeline.go index f029701a8c..05f689eb32 100644 --- a/internal/service/chat_pipeline.go +++ b/internal/service/chat_pipeline.go @@ -1971,7 +1971,7 @@ func (s *ChatPipelineService) getModels(ctx context.Context, chat *entity.Chat) // Embedding model. var embModel *modelModule.EmbeddingModel if len(kbs) > 0 { - if err := ValidateDatasetEmbeddingModels(kbs); err != nil { + if err := ValidateDatasetEmbeddingModels(ctx, dao.DB, kbs); err != nil { return nil, nil, nil, nil, nil, err } if kbs[0].EmbdID != "" { diff --git a/internal/service/chunk/chunk.go b/internal/service/chunk/chunk.go index 8c14cebf88..0eebfa677e 100644 --- a/internal/service/chunk/chunk.go +++ b/internal/service/chunk/chunk.go @@ -192,11 +192,12 @@ func (s *ChunkService) RetrievalTest(ctx context.Context, req *service.Retrieval } } - // Check if all kbs have the same embedding model + // Check if all kbs resolve to the same base embedding model if len(kbRecords) > 1 { - firstEmbeddingKey := knowledgebaseEmbeddingKey(kbRecords[0], tenantIDs[0]) + embdNameCache := make(map[string]string) + firstEmbeddingKey := knowledgebaseEmbeddingKey(ctx, dao.DB, kbRecords[0], tenantIDs[0], embdNameCache) for i := 1; i < len(kbRecords); i++ { - if knowledgebaseEmbeddingKey(kbRecords[i], tenantIDs[i]) != firstEmbeddingKey { + if knowledgebaseEmbeddingKey(ctx, dao.DB, kbRecords[i], tenantIDs[i], embdNameCache) != firstEmbeddingKey { return nil, fmt.Errorf("cannot retrieve across datasets with different embedding models") } } @@ -444,14 +445,15 @@ func (s *ChunkService) RetrievalTest(ctx context.Context, req *service.Retrieval }, nil } -func knowledgebaseEmbeddingKey(kb *entity.Knowledgebase, tenantID string) string { - if kb.TenantEmbdID != nil && *kb.TenantEmbdID != "" { - return fmt.Sprintf("tenant:%s", *kb.TenantEmbdID) - } - if kb.EmbdID == "" { +// knowledgebaseEmbeddingKey groups datasets by their resolved base embedding +// model name (e.g. "BAAI/bge-m3") so datasets pointing at the same model +// through different provider instances or storage forms (tenant_model id vs +// legacy composite name) retrieve together. +func knowledgebaseEmbeddingKey(ctx context.Context, db *gorm.DB, kb *entity.Knowledgebase, tenantID string, cache map[string]string) string { + if strings.TrimSpace(kb.EmbdID) == "" && (kb.TenantEmbdID == nil || strings.TrimSpace(*kb.TenantEmbdID) == "") { return fmt.Sprintf("default:%s", tenantID) } - return fmt.Sprintf("embd:%s", kb.EmbdID) + return "embd:" + dao.NewKnowledgebaseDAO().EmbeddingBaseName(ctx, db, kb, cache) } // hydrateChunkVectors replaces zero (placeholder) vectors in chunks with real diff --git a/internal/service/chunk/chunk_test.go b/internal/service/chunk/chunk_test.go index 2a54adac26..fe529a0904 100644 --- a/internal/service/chunk/chunk_test.go +++ b/internal/service/chunk/chunk_test.go @@ -81,6 +81,7 @@ func TestHydrateChunkVectors_NoDim(t *testing.T) { func TestKnowledgebaseEmbeddingKey(t *testing.T) { tenantEmbdID := "42" + staleTenantEmbdID := "stale-id" tests := []struct { name string @@ -89,12 +90,20 @@ func TestKnowledgebaseEmbeddingKey(t *testing.T) { want string }{ { - name: "uses tenant embedding id before embd id", + name: "resolves tenant embedding id to base model name", kb: &entity.Knowledgebase{ EmbdID: "shared-model", TenantEmbdID: &tenantEmbdID, }, - want: "tenant:42", + want: "embd:BAAI/bge-m3", + }, + { + name: "stale tenant embedding id falls back to composite base name", + kb: &entity.Knowledgebase{ + EmbdID: "BAAI/bge-m3@1@SILICONFLOW", + TenantEmbdID: &staleTenantEmbdID, + }, + want: "embd:BAAI/bge-m3", }, { name: "uses embd id without tenant embedding id", @@ -117,11 +126,29 @@ func TestKnowledgebaseEmbeddingKey(t *testing.T) { }, want: "embd:shared-model", }, + { + name: "legacy composite reduces to base model name", + kb: &entity.Knowledgebase{ + EmbdID: "BAAI/bge-m3@2@SILICONFLOW", + }, + want: "embd:BAAI/bge-m3", + }, + } + + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{TranslateError: true}) + if err != nil { + t.Fatalf("failed to open sqlite: %v", err) + } + if err := db.AutoMigrate(&entity.TenantModel{}); err != nil { + t.Fatalf("failed to migrate test schema: %v", err) + } + if err := db.Create(&entity.TenantModel{ID: "42", ModelName: "BAAI/bge-m3"}).Error; err != nil { + t.Fatalf("failed to seed tenant_model: %v", err) } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - if got := knowledgebaseEmbeddingKey(tt.kb, tt.tenantID); got != tt.want { + if got := knowledgebaseEmbeddingKey(t.Context(), db, tt.kb, tt.tenantID, map[string]string{}); got != tt.want { t.Fatalf("knowledgebaseEmbeddingKey() = %q, want %q", got, tt.want) } }) diff --git a/internal/service/dataset/search.go b/internal/service/dataset/search.go index 272a10a894..4394af649e 100644 --- a/internal/service/dataset/search.go +++ b/internal/service/dataset/search.go @@ -102,7 +102,7 @@ func (d *DatasetService) SearchDatasets(ctx context.Context, req *service.Search } // Check if all kbs have the same embedding model - if err := service.ValidateDatasetEmbeddingModels(kbRecords); err != nil { + if err := service.ValidateDatasetEmbeddingModels(ctx, dao.DB, kbRecords); err != nil { return nil, err } diff --git a/internal/service/pipeline_params.go b/internal/service/pipeline_params.go index 88fbf0d23a..4adddf06ac 100644 --- a/internal/service/pipeline_params.go +++ b/internal/service/pipeline_params.go @@ -10,6 +10,8 @@ import ( "ragflow/internal/dao" "ragflow/internal/entity" pipelinepkg "ragflow/internal/ingestion/pipeline" + + "gorm.io/gorm" ) // loadCanvasDSLJSON returns the DSL JSON for a custom canvas pipeline. The @@ -97,24 +99,21 @@ func ResolveComponentParamsDefaults(ctx context.Context, parserID string, pipeli } // ValidateDatasetEmbeddingModels checks that all knowledge bases in the list -// either have an embedding model or none do, and that they all use the same model. -func ValidateDatasetEmbeddingModels(kbs []*entity.Knowledgebase) error { - embdIDs := make(map[string]struct{}) +// either have an embedding model or none do, and that they all resolve to the +// same base embedding model name (e.g. "BAAI/bge-m3"). 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. +func ValidateDatasetEmbeddingModels(ctx context.Context, db *gorm.DB, kbs []*entity.Knowledgebase) error { + embdNames := make(map[string]struct{}) hasEmbd := false noEmbd := false + cache := make(map[string]string) + kbDAO := dao.NewKnowledgebaseDAO() for _, kb := range kbs { if kb.EmbdID != "" { hasEmbd = true - baseName := kb.EmbdID - if idx := strings.LastIndex(kb.EmbdID, "@"); idx > 0 { - baseName = kb.EmbdID[:idx] - // Strip the second-to-last @-segment too (instance name), - // matching Python's _base_model_name which uses rsplit("@", 2). - if idx2 := strings.LastIndex(baseName, "@"); idx2 > 0 { - baseName = baseName[:idx2] - } - } - embdIDs[baseName] = struct{}{} + embdNames[kbDAO.EmbeddingBaseName(ctx, db, kb, cache)] = struct{}{} } else { noEmbd = true } @@ -122,7 +121,7 @@ func ValidateDatasetEmbeddingModels(kbs []*entity.Knowledgebase) error { if hasEmbd && noEmbd { return fmt.Errorf("cannot search across datasets where some have embedding models and others do not") } - if len(embdIDs) > 1 { + if len(embdNames) > 1 { return fmt.Errorf("datasets use different embedding models: %v", getEmbdIDs(kbs)) } return nil diff --git a/internal/service/pipeline_params_test.go b/internal/service/pipeline_params_test.go index a9da3ccbce..fcb56e9bbd 100644 --- a/internal/service/pipeline_params_test.go +++ b/internal/service/pipeline_params_test.go @@ -4,14 +4,32 @@ import ( "testing" "ragflow/internal/entity" + + "github.com/glebarez/sqlite" + "gorm.io/gorm" ) +func setupValidateEmbeddingTestDB(t *testing.T) *gorm.DB { + t.Helper() + + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{ + TranslateError: true, + }) + if err != nil { + t.Fatalf("failed to open sqlite: %v", err) + } + if err := db.AutoMigrate(&entity.TenantModel{}); err != nil { + t.Fatalf("failed to migrate test schema: %v", err) + } + return db +} + func TestValidateDatasetEmbeddingModels_AllHaveEmbeddingModel(t *testing.T) { kbs := []*entity.Knowledgebase{ {EmbdID: "BAAI/bge-large-zh-v1.5@Builtin"}, {EmbdID: "BAAI/bge-large-zh-v1.5@Builtin"}, } - if err := ValidateDatasetEmbeddingModels(kbs); err != nil { + if err := ValidateDatasetEmbeddingModels(t.Context(), nil, kbs); err != nil { t.Fatalf("expected nil, got %v", err) } } @@ -21,7 +39,7 @@ func TestValidateDatasetEmbeddingModels_NoneHasEmbeddingModel(t *testing.T) { {EmbdID: ""}, {EmbdID: ""}, } - if err := ValidateDatasetEmbeddingModels(kbs); err != nil { + if err := ValidateDatasetEmbeddingModels(t.Context(), nil, kbs); err != nil { t.Fatalf("expected nil, got %v", err) } } @@ -31,7 +49,7 @@ func TestValidateDatasetEmbeddingModels_MixedErrors(t *testing.T) { {EmbdID: "BAAI/bge-large-zh-v1.5@Builtin"}, {EmbdID: ""}, } - err := ValidateDatasetEmbeddingModels(kbs) + err := ValidateDatasetEmbeddingModels(t.Context(), nil, kbs) if err == nil { t.Fatal("expected error for mixed embedding") } @@ -45,7 +63,7 @@ func TestValidateDatasetEmbeddingModels_DifferentEmbeddingsErrors(t *testing.T) {EmbdID: "model-a@provider-1"}, {EmbdID: "model-b@provider-2"}, } - err := ValidateDatasetEmbeddingModels(kbs) + err := ValidateDatasetEmbeddingModels(t.Context(), nil, kbs) if err == nil { t.Fatal("expected error for different embeddings") } @@ -58,7 +76,7 @@ func TestValidateDatasetEmbeddingModels_SameBaseDifferentInstanceOK(t *testing.T {EmbdID: "BAAI/bge-large-zh-v1.5@instance1@provider1"}, {EmbdID: "BAAI/bge-large-zh-v1.5@instance2@provider2"}, } - if err := ValidateDatasetEmbeddingModels(kbs); err != nil { + if err := ValidateDatasetEmbeddingModels(t.Context(), nil, kbs); err != nil { t.Fatalf("expected nil, got %v", err) } } @@ -68,14 +86,84 @@ func TestValidateDatasetEmbeddingModels_DifferentBasesErrors(t *testing.T) { {EmbdID: "model-a@instance1@provider1"}, {EmbdID: "model-b@instance2@provider2"}, } - err := ValidateDatasetEmbeddingModels(kbs) + err := ValidateDatasetEmbeddingModels(t.Context(), nil, kbs) if err == nil { t.Fatal("expected error for different base models") } } func TestValidateDatasetEmbeddingModels_EmptyList(t *testing.T) { - if err := ValidateDatasetEmbeddingModels(nil); err != nil { + if err := ValidateDatasetEmbeddingModels(t.Context(), nil, nil); err != nil { t.Fatalf("expected nil for empty list, got %v", err) } } + +func TestValidateDatasetEmbeddingModels_TenantModelIDsResolveToSameBase(t *testing.T) { + db := setupValidateEmbeddingTestDB(t) + for _, m := range []entity.TenantModel{ + {ID: "11111111111111111111111111111111", ModelName: "BAAI/bge-m3"}, + {ID: "22222222222222222222222222222222", ModelName: "BAAI/bge-m3"}, + } { + if err := db.Create(&m).Error; err != nil { + t.Fatalf("failed to seed tenant_model: %v", err) + } + } + + // Two KBs referencing different tenant_model rows (e.g. different provider + // instances) that serve the same base model must validate together. + kbs := []*entity.Knowledgebase{ + {EmbdID: "11111111111111111111111111111111"}, + {EmbdID: "22222222222222222222222222222222"}, + } + if err := ValidateDatasetEmbeddingModels(t.Context(), db, kbs); err != nil { + t.Fatalf("expected nil, got %v", err) + } +} + +func TestValidateDatasetEmbeddingModels_TenantModelIDAndLegacyCompositeMatch(t *testing.T) { + db := setupValidateEmbeddingTestDB(t) + if err := db.Create(&entity.TenantModel{ID: "11111111111111111111111111111111", ModelName: "BAAI/bge-m3"}).Error; err != nil { + t.Fatalf("failed to seed tenant_model: %v", err) + } + + // One KB stores the tenant_model id, the other a legacy composite name + // pointing at the same base model through another instance. + kbs := []*entity.Knowledgebase{ + {EmbdID: "11111111111111111111111111111111"}, + {EmbdID: "BAAI/bge-m3@2@SILICONFLOW"}, + } + if err := ValidateDatasetEmbeddingModels(t.Context(), db, kbs); err != nil { + t.Fatalf("expected nil, got %v", err) + } +} + +func TestValidateDatasetEmbeddingModels_ResolvedDifferentBasesError(t *testing.T) { + db := setupValidateEmbeddingTestDB(t) + for _, m := range []entity.TenantModel{ + {ID: "11111111111111111111111111111111", ModelName: "BAAI/bge-m3"}, + {ID: "22222222222222222222222222222222", ModelName: "Qwen/Qwen3-Embedding-0.6B"}, + } { + if err := db.Create(&m).Error; err != nil { + t.Fatalf("failed to seed tenant_model: %v", err) + } + } + + kbs := []*entity.Knowledgebase{ + {EmbdID: "11111111111111111111111111111111"}, + {EmbdID: "22222222222222222222222222222222"}, + } + if err := ValidateDatasetEmbeddingModels(t.Context(), db, kbs); err == nil { + t.Fatal("expected error for different resolved base models") + } +} + +func TestValidateDatasetEmbeddingModels_UnresolvableIDsStayIsolated(t *testing.T) { + // Stale ids that no longer resolve must not suddenly compare equal. + kbs := []*entity.Knowledgebase{ + {EmbdID: "11111111111111111111111111111111"}, + {EmbdID: "22222222222222222222222222222222"}, + } + if err := ValidateDatasetEmbeddingModels(t.Context(), nil, kbs); err == nil { + t.Fatal("expected error for distinct unresolvable embedding ids") + } +} diff --git a/test/unit_test/api/db/services/test_knowledgebase_embedding_validation.py b/test/unit_test/api/db/services/test_knowledgebase_embedding_validation.py new file mode 100644 index 0000000000..4ad63ada47 --- /dev/null +++ b/test/unit_test/api/db/services/test_knowledgebase_embedding_validation.py @@ -0,0 +1,174 @@ +# +# Copyright 2026 The InfiniFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Regression tests for validate_dataset_embedding_models() in +api.db.services.knowledgebase_service. + +Datasets may store their embedding reference either as a tenant_model id +(hex) or as a legacy composite ``model@instance@provider`` string. Datasets +that resolve to the same base embedding model must validate together even +when the storage forms or provider instances differ (e.g. +``BAAI/bge-m3@renew@SILICONFLOW`` vs ``BAAI/bge-m3@COPY@SILICONFLOW``). + +This test loads only the relevant function definitions from the source file +via AST, so it doesn't pull in the full ``api.db.db_models`` import chain and +can run in any minimal pytest environment. +""" + +import ast +from pathlib import Path +from types import SimpleNamespace + +import pytest + +pytestmark = pytest.mark.p2 + + +def _load_validator(resolved_names): + """Exec the validator + helpers from the production source file with a + stubbed ``get_composite_model_name_by_ids``.""" + src_path = Path(__file__).resolve().parents[5] / "api" / "db" / "services" / "knowledgebase_service.py" + tree = ast.parse(src_path.read_text(encoding="utf-8")) + wanted = { + "_base_model_name", + "_kb_embedding_base_name", + "validate_dataset_embedding_models", + } + nodes = [node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name in wanted] + assert len(nodes) == len(wanted), f"missing functions: {wanted - {n.name for n in nodes}}" + module = ast.Module(body=nodes, type_ignores=[]) + ns = {"get_composite_model_name_by_ids": lambda _ids: resolved_names} + exec(compile(module, str(src_path), "exec"), ns) # noqa: S102 - exec'ing AST-extracted functions from our own source file + return ns["validate_dataset_embedding_models"] + + +def test_same_composite_validates(): + validate = _load_validator({}) + kbs = [ + SimpleNamespace(embd_id="BAAI/bge-m3@1@SILICONFLOW", tenant_embd_id=None), + SimpleNamespace(embd_id="BAAI/bge-m3@1@SILICONFLOW", tenant_embd_id=None), + ] + assert validate(kbs) is None + + +def test_same_base_different_instances_validate(): + """The search-settings bug: same model/provider, different instance name.""" + validate = _load_validator({}) + kbs = [ + SimpleNamespace(embd_id="BAAI/bge-m3@renew@SILICONFLOW", tenant_embd_id=None), + SimpleNamespace(embd_id="BAAI/bge-m3@COPY@SILICONFLOW", tenant_embd_id=None), + ] + assert validate(kbs) is None + + +def test_different_base_models_rejected(): + validate = _load_validator({}) + kbs = [ + SimpleNamespace(embd_id="BAAI/bge-m3@1@SILICONFLOW", tenant_embd_id=None), + SimpleNamespace(embd_id="Qwen/Qwen3-Embedding-0.6B@1@SILICONFLOW", tenant_embd_id=None), + ] + assert "different embedding models" in validate(kbs) + + +def test_tenant_model_id_and_composite_resolve_together(): + """A raw tenant_model id and a legacy composite pointing at the same base + model must validate together.""" + validate = _load_validator({"5fb939ba64d04dd2b6dd9c1c08775d52": "BAAI/bge-m3@1@SILICONFLOW"}) + kbs = [ + SimpleNamespace( + embd_id="5fb939ba64d04dd2b6dd9c1c08775d52", + tenant_embd_id="5fb939ba64d04dd2b6dd9c1c08775d52", + ), + SimpleNamespace(embd_id="BAAI/bge-m3@2@SILICONFLOW", tenant_embd_id=None), + ] + assert validate(kbs) is None + + +def test_two_tenant_model_ids_resolving_to_same_base_validate(): + validate = _load_validator( + { + "5fb939ba64d04dd2b6dd9c1c08775d52": "BAAI/bge-m3@1@SILICONFLOW", + "fdbbd8f9a0564a6f92d068bd596de861": "BAAI/bge-m3@2@SILICONFLOW", + } + ) + kbs = [ + SimpleNamespace( + embd_id="5fb939ba64d04dd2b6dd9c1c08775d52", + tenant_embd_id="5fb939ba64d04dd2b6dd9c1c08775d52", + ), + SimpleNamespace( + embd_id="fdbbd8f9a0564a6f92d068bd596de861", + tenant_embd_id="fdbbd8f9a0564a6f92d068bd596de861", + ), + ] + assert validate(kbs) is None + + +def test_resolved_different_bases_rejected(): + validate = _load_validator( + { + "5fb939ba64d04dd2b6dd9c1c08775d52": "BAAI/bge-m3@1@SILICONFLOW", + "b346fecb098e4305b58926f1bae6bbc9": "Qwen/Qwen3-Embedding-0.6B@1@SILICONFLOW", + } + ) + kbs = [ + SimpleNamespace( + embd_id="5fb939ba64d04dd2b6dd9c1c08775d52", + tenant_embd_id="5fb939ba64d04dd2b6dd9c1c08775d52", + ), + SimpleNamespace( + embd_id="b346fecb098e4305b58926f1bae6bbc9", + tenant_embd_id="b346fecb098e4305b58926f1bae6bbc9", + ), + ] + assert "different embedding models" in validate(kbs) + + +def test_unresolvable_ids_stay_isolated(): + validate = _load_validator({}) + kbs = [ + SimpleNamespace(embd_id="2d8ff0a97d75431c8c91526549939328", tenant_embd_id="2d8ff0a97d75431c8c91526549939328"), + SimpleNamespace(embd_id="BAAI/bge-m3@1@SILICONFLOW", tenant_embd_id=None), + ] + assert "different embedding models" in validate(kbs) + + +def test_stale_tenant_id_falls_back_to_composite_base(): + """A KB whose tenant_embd_id no longer resolves still groups by its + legacy composite embd_id.""" + validate = _load_validator({}) + kbs = [ + SimpleNamespace(embd_id="BAAI/bge-m3@1@SILICONFLOW", tenant_embd_id="2540bad87b7511f197c0d16229aaaa64"), + SimpleNamespace(embd_id="BAAI/bge-m3@2@SILICONFLOW", tenant_embd_id=None), + ] + assert validate(kbs) is None + + +def test_mixed_embedding_and_none_rejected(): + validate = _load_validator({}) + kbs = [ + SimpleNamespace(embd_id="BAAI/bge-m3@1@SILICONFLOW", tenant_embd_id=None), + SimpleNamespace(embd_id="", tenant_embd_id=None), + ] + assert "some have embedding models and others do not" in validate(kbs) + + +def test_all_without_embedding_validates(): + validate = _load_validator({}) + kbs = [ + SimpleNamespace(embd_id="", tenant_embd_id=None), + SimpleNamespace(embd_id="", tenant_embd_id=None), + ] + assert validate(kbs) is None diff --git a/web/src/components/knowledge-base-item.tsx b/web/src/components/knowledge-base-item.tsx index 3e96698575..b2e58b05f8 100644 --- a/web/src/components/knowledge-base-item.tsx +++ b/web/src/components/knowledge-base-item.tsx @@ -21,6 +21,7 @@ import { } from '@/hooks/use-knowledge-request'; import { IDataset } from '@/interfaces/database/dataset'; import { useBuildQueryVariableOptions } from '@/pages/agent/hooks/use-get-begin-query'; +import { getEmbeddingBaseName } from '@/utils/llm-util'; import { useDebounce } from 'ahooks'; import { toLower } from 'lodash'; import { type ReactNode, useCallback, useMemo, useRef, useState } from 'react'; @@ -94,9 +95,12 @@ export function useDisableDifferenceEmbeddingDataset(name: string) { ); }, [datasetListOrigin, selectedDatasetIds, missingDatasets]); - const selectedEmbedId = useMemo(() => { + // Datasets are mutually selectable when their embedding models resolve to + // the same base model name, even if they use different provider instances + // (e.g. "BAAI/bge-m3@renew@SILICONFLOW" vs "BAAI/bge-m3@COPY@SILICONFLOW"). + const selectedEmbedBaseName = useMemo(() => { const data = datasetList?.find((item) => item.id === datasetId?.[0]); - return data?.embedding_model ?? ''; + return getEmbeddingBaseName(data?.embedding_model); }, [datasetId, datasetList]); const nextOptions = useMemo(() => { @@ -126,12 +130,13 @@ export function useDisableDifferenceEmbeddingDataset(name: string) { disabled: item.chunk_count <= 0 || item.chunk_method === DocumentParserType.Tag || - (item.embedding_model !== selectedEmbedId && selectedEmbedId !== ''), + (selectedEmbedBaseName !== '' && + getEmbeddingBaseName(item.embedding_model) !== selectedEmbedBaseName), }; }); return datasetListMap; - }, [datasetList, selectedEmbedId]); + }, [datasetList, selectedEmbedBaseName]); const handleSearchChange = useCallback((value: string) => { setSearchString(value); diff --git a/web/src/utils/llm-util.ts b/web/src/utils/llm-util.ts index a1a12f4113..7f18b9d810 100644 --- a/web/src/utils/llm-util.ts +++ b/web/src/utils/llm-util.ts @@ -123,6 +123,21 @@ export function parseModelValue(val: string) { }; } +/** + * Base embedding model name used to decide whether two datasets can be + * selected and searched together. Composite references + * ("modelName@instanceName@providerName" or "modelName@providerName") reduce + * to the bare model name, so datasets using the same embedding model through + * different provider instances still group together. Opaque values without + * '@' (e.g. an unresolved tenant_model id) are returned unchanged, so only + * exact matches group together. Mirrors the backend's base-name comparison + * (Python `_base_model_name`, Go `common.BaseModelName`). + */ +export function getEmbeddingBaseName(embeddingModel?: string | null): string { + if (!embeddingModel) return ''; + return parseModelValue(embeddingModel)?.model_name ?? embeddingModel; +} + // Extract model name and factory ID from a model UUID // Supports both "model_name@factory_id" and "model_name@factory_id#instance_name". // Uses right-anchored split for the same reason as parseModelValue: diff --git a/web/src/utils/tests/llm-util.test.ts b/web/src/utils/tests/llm-util.test.ts index 776c421593..4b9392c26f 100644 --- a/web/src/utils/tests/llm-util.test.ts +++ b/web/src/utils/tests/llm-util.test.ts @@ -1,4 +1,9 @@ -import { buildModelValue, parseModelUuid, parseModelValue } from '../llm-util'; +import { + buildModelValue, + getEmbeddingBaseName, + parseModelUuid, + parseModelValue, +} from '../llm-util'; // Composite model keys are right-anchored: // "model_name@instance_name@provider_name" or "model_name@provider_name". @@ -117,3 +122,41 @@ describe('parseModelUuid — right-anchored', () => { }); }); }); + +describe('getEmbeddingBaseName — dataset co-selection grouping', () => { + test('3-part composite reduces to bare model name', () => { + expect(getEmbeddingBaseName('BAAI/bge-m3@renew@SILICONFLOW')).toBe( + 'BAAI/bge-m3', + ); + }); + + test('same model through different instances shares one base name', () => { + expect( + getEmbeddingBaseName('BAAI/bge-m3@renew@SILICONFLOW') === + getEmbeddingBaseName('BAAI/bge-m3@COPY@SILICONFLOW'), + ).toBe(true); + }); + + test('2-part composite strips the provider', () => { + expect(getEmbeddingBaseName('BAAI/bge-m3@SILICONFLOW')).toBe('BAAI/bge-m3'); + }); + + test('model names containing "@" keep their suffix', () => { + expect( + getEmbeddingBaseName( + 'text-embedding-nomic-embed-text-v1.5@q8_0@lmstudio@LM-Studio', + ), + ).toBe('text-embedding-nomic-embed-text-v1.5@q8_0'); + }); + + test('opaque tenant_model id is returned unchanged', () => { + expect(getEmbeddingBaseName('2d8ff0a97d75431c8c91526549939328')).toBe( + '2d8ff0a97d75431c8c91526549939328', + ); + }); + + test('empty and undefined values produce an empty base name', () => { + expect(getEmbeddingBaseName('')).toBe(''); + expect(getEmbeddingBaseName(undefined)).toBe(''); + }); +});