diff --git a/api/apps/restful_apis/dataset_api.py b/api/apps/restful_apis/dataset_api.py index 701c7340b7..55ded90e02 100644 --- a/api/apps/restful_apis/dataset_api.py +++ b/api/apps/restful_apis/dataset_api.py @@ -19,7 +19,7 @@ from peewee import OperationalError from quart import request from common.constants import RetCode from api.apps import login_required, current_user -from api.utils.api_utils import get_error_argument_result, get_error_data_result, get_result, add_tenant_id_to_kwargs +from api.utils.api_utils import get_error_argument_result, get_error_data_result, get_json_result, get_result, add_tenant_id_to_kwargs from api.utils.validation_utils import ( CreateDatasetReq, DeleteDatasetReq, @@ -653,6 +653,26 @@ async def run_embedding(tenant_id, dataset_id): return get_error_data_result(message="Internal server error") +@manager.route("/datasets//embedding/check", methods=["POST"]) # noqa: F821 +@login_required +@add_tenant_id_to_kwargs +async def check_embedding(tenant_id, dataset_id): + try: + req = await request.get_json() + if not req or not req.get("embd_id"): + return get_error_data_result(message="`embd_id` is required.") + status, result = dataset_api_service.check_embedding(dataset_id, tenant_id, req) + if status is True: + return get_result(data=result) + elif status == "not_effective": + return get_json_result(code=result["code"], message=result["message"], data=result["data"]) + else: + return get_error_data_result(message=result) + except Exception as e: + logging.exception(e) + return get_error_data_result(message="Internal server error") + + @manager.route("/datasets//ingestions", methods=["GET"]) # noqa: F821 @login_required @add_tenant_id_to_kwargs diff --git a/api/apps/services/dataset_api_service.py b/api/apps/services/dataset_api_service.py index 795e42b7b8..9e49596539 100644 --- a/api/apps/services/dataset_api_service.py +++ b/api/apps/services/dataset_api_service.py @@ -16,6 +16,7 @@ import logging import json import os +import re from common.constants import PAGERANK_FLD from common import settings from api.db.db_models import File @@ -306,8 +307,6 @@ async def update_dataset(tenant_id: str, dataset_id: str, req: dict): if "embd_id" in req: if not req["embd_id"]: req["embd_id"] = kb.embd_id - if kb.chunk_num != 0 and req["embd_id"] != kb.embd_id: - return False, f"When chunk_num ({kb.chunk_num}) > 0, embedding_model must remain {kb.embd_id}" ok, err = verify_embedding_availability(req["embd_id"], tenant_id) if not ok: return False, err @@ -1053,6 +1052,209 @@ async def search(dataset_id: str, tenant_id: str, req: dict): return True, ranks +def check_embedding(dataset_id: str, tenant_id: str, req: dict): + """ + Check embedding model compatibility by sampling random chunks, + re-embedding them with the new model, and computing cosine similarity. + + :param dataset_id: dataset ID + :param tenant_id: tenant ID + :param req: request body with embd_id + :return: (success, result) or (success, error_message) + """ + import random + + import numpy as np + from common.constants import RetCode + from common.doc_store.doc_store_base import OrderByExpr + from rag.nlp import search + + from api.db.joint_services.tenant_model_service import ( + get_model_config_by_type_and_name, + ) + from api.db.services.llm_service import LLMBundle + from common.constants import LLMType + + def _guess_vec_field(src: dict): + for k in src or {}: + if k.endswith("_vec"): + return k + return None + + def _as_float_vec(v): + if v is None: + return [] + if isinstance(v, str): + return [float(x) for x in v.split("\t") if x != ""] + if isinstance(v, (list, tuple, np.ndarray)): + return [float(x) for x in v] + return [] + + def _to_1d(x): + a = np.asarray(x, dtype=np.float32) + return a.reshape(-1) + + def _cos_sim(a, b, eps=1e-12): + a = _to_1d(a) + b = _to_1d(b) + na = np.linalg.norm(a) + nb = np.linalg.norm(b) + if na < eps or nb < eps: + return 0.0 + return float(np.dot(a, b) / (na * nb)) + + def sample_random_chunks_with_vectors( + docStoreConn, + tenant_id: str, + kb_id: str, + n: int = 5, + base_fields=("docnm_kwd", "doc_id", "content_with_weight", "page_num_int", "position_int", "top_int"), + ): + index_nm = search.index_name(tenant_id) + + res0 = docStoreConn.search( + select_fields=[], highlight_fields=[], + condition={"kb_id": kb_id, "available_int": 1}, + match_expressions=[], order_by=OrderByExpr(), + offset=0, limit=1, + index_names=index_nm, knowledgebase_ids=[kb_id], + ) + total = docStoreConn.get_total(res0) + if total <= 0: + return [] + + n = min(n, total) + offsets = sorted(random.sample(range(min(total, 1000)), n)) + out = [] + + for off in offsets: + res1 = docStoreConn.search( + select_fields=list(base_fields), + highlight_fields=[], + condition={"kb_id": kb_id, "available_int": 1}, + match_expressions=[], order_by=OrderByExpr(), + offset=off, limit=1, + index_names=index_nm, knowledgebase_ids=[kb_id], + ) + ids = docStoreConn.get_doc_ids(res1) + if not ids: + continue + + cid = ids[0] + full_doc = docStoreConn.get(cid, index_nm, [kb_id]) or {} + vec_field = _guess_vec_field(full_doc) + vec = _as_float_vec(full_doc.get(vec_field)) + + out.append({ + "chunk_id": cid, + "kb_id": kb_id, + "doc_id": full_doc.get("doc_id"), + "doc_name": full_doc.get("docnm_kwd"), + "vector_field": vec_field, + "vector_dim": len(vec), + "vector": vec, + "page_num_int": full_doc.get("page_num_int"), + "position_int": full_doc.get("position_int"), + "top_int": full_doc.get("top_int"), + "content_with_weight": full_doc.get("content_with_weight") or "", + "question_kwd": full_doc.get("question_kwd") or [], + }) + return out + + def _clean(s: str): + return re.sub(r"]{0,12})?>", " ", s or "").strip() + + if not dataset_id: + return False, 'Lack of "Dataset ID"' + + if not KnowledgebaseService.accessible(dataset_id, tenant_id): + return False, "No authorization." + + ok, kb = KnowledgebaseService.get_by_id(dataset_id) + if not ok: + return False, "Invalid Dataset ID" + + embd_id = req.get("embd_id", "") + if not embd_id: + return False, "`embd_id` is required." + + logging.info("check_embedding: dataset=%s tenant=%s embd_id=%s", dataset_id, tenant_id, embd_id) + + ok, err = verify_embedding_availability(embd_id, tenant_id) + if not ok: + return False, err + + embd_model_config = get_model_config_by_type_and_name(kb.tenant_id, LLMType.EMBEDDING, embd_id) + emb_mdl = LLMBundle(kb.tenant_id, embd_model_config) + + n = int(req.get("check_num", 5)) + samples = sample_random_chunks_with_vectors(settings.docStoreConn, tenant_id=kb.tenant_id, kb_id=dataset_id, n=n) + logging.info("check_embedding: dataset=%s sampled=%d chunks", dataset_id, len(samples)) + + results, eff_sims = [], [] + mode = "content_only" + for ck in samples: + title = ck.get("doc_name") or "Title" + + txt_in = "\n".join(ck.get("question_kwd") or []) or ck.get("content_with_weight") or "" + txt_in = _clean(txt_in) + if not txt_in: + results.append({"chunk_id": ck["chunk_id"], "reason": "no_text"}) + continue + + if not ck.get("vector"): + results.append({"chunk_id": ck["chunk_id"], "reason": "no_stored_vector"}) + continue + + try: + v, _ = emb_mdl.encode([title, txt_in]) + assert len(v[1]) == len(ck["vector"]), ( + f"The dimension ({len(v[1])}) of given embedding model is different from the original ({len(ck['vector'])})" + ) + sim_content = _cos_sim(v[1], ck["vector"]) + title_w = 0.1 + qv_mix = title_w * v[0] + (1 - title_w) * v[1] + sim_mix = _cos_sim(qv_mix, ck["vector"]) + sim = sim_content + mode = "content_only" + if sim_mix > sim: + sim = sim_mix + mode = "title+content" + except Exception as e: + return False, f"Embedding failure. {e}" + + eff_sims.append(sim) + results.append({ + "chunk_id": ck["chunk_id"], + "doc_id": ck["doc_id"], + "doc_name": ck["doc_name"], + "vector_field": ck["vector_field"], + "vector_dim": ck["vector_dim"], + "cos_sim": round(sim, 6), + }) + + summary = { + "kb_id": dataset_id, + "model": embd_id, + "sampled": len(samples), + "valid": len(eff_sims), + "avg_cos_sim": round(float(np.mean(eff_sims)) if eff_sims else 0.0, 6), + "min_cos_sim": round(float(np.min(eff_sims)) if eff_sims else 0.0, 6), + "max_cos_sim": round(float(np.max(eff_sims)) if eff_sims else 0.0, 6), + "match_mode": mode, + } + + data = {"summary": summary, "results": results} + if not eff_sims: + logging.warning("check_embedding: dataset=%s no comparable chunks", dataset_id) + return False, "No embedded chunks are available to compare." + if summary["avg_cos_sim"] >= 0.9: + logging.info("check_embedding: dataset=%s compatible avg_cos_sim=%s valid=%d", dataset_id, summary["avg_cos_sim"], len(eff_sims)) + return True, data + logging.warning("check_embedding: dataset=%s not_effective avg_cos_sim=%s valid=%d", dataset_id, summary["avg_cos_sim"], len(eff_sims)) + return "not_effective", {"code": RetCode.NOT_EFFECTIVE, "message": "Embedding model switch failed: the average similarity between old and new vectors is below 0.9, indicating incompatible vector spaces.", "data": data} + + async def search_datasets(tenant_id: str, req: dict): """ Search (retrieval test) across multiple datasets. diff --git a/test/testcases/test_http_api/test_dataset_management/test_update_dataset.py b/test/testcases/test_http_api/test_dataset_management/test_update_dataset.py index 58885a5395..0847a181c1 100644 --- a/test/testcases/test_http_api/test_dataset_management/test_update_dataset.py +++ b/test/testcases/test_http_api/test_dataset_management/test_update_dataset.py @@ -291,7 +291,7 @@ class TestDatasetUpdate: @pytest.mark.p1 def test_embedding_model_with_existing_chunks(self, HttpApiAuth, add_chunks): - """Guard: embedding_model cannot change when dataset has chunks (chunk_count > 0).""" + """Embedding model can be changed even when dataset has chunks (chunk_count > 0).""" dataset_id, _, _ = add_chunks res = list_datasets(HttpApiAuth, {"id": dataset_id}) @@ -306,12 +306,7 @@ class TestDatasetUpdate: payload = {"embedding_model": new_embedding} res = update_dataset(HttpApiAuth, dataset_id, payload) - assert res["code"] == 102, res - expected_message = ( - f"When chunk_num ({dataset['chunk_count']}) > 0, " - f"embedding_model must remain {current_embedding}" - ) - assert res["message"] == expected_message, res + assert res["code"] == 0, res @pytest.mark.p2 @pytest.mark.parametrize( diff --git a/test/testcases/test_web_api/test_dataset_management/test_dataset_sdk_routes_unit.py b/test/testcases/test_web_api/test_dataset_management/test_dataset_sdk_routes_unit.py index b69abb0c59..2311eb22dc 100644 --- a/test/testcases/test_web_api/test_dataset_management/test_dataset_sdk_routes_unit.py +++ b/test/testcases/test_web_api/test_dataset_management/test_dataset_sdk_routes_unit.py @@ -548,10 +548,11 @@ def test_update_route_branch_matrix_unit(monkeypatch): kb_chunked = _KB(kb_id="kb-1", name="old", chunk_num=2, embd_id="embd-1") monkeypatch.setattr(module.KnowledgebaseService, "get_or_none", lambda **kwargs: kb_chunked if kwargs.get("id") else None) + monkeypatch.setattr(module.KnowledgebaseService, "update_by_id", lambda *_args, **_kwargs: True) req_state.clear() req_state.update({"embd_id": "embd-2"}) res = _run(inspect.unwrap(module.update)("tenant-1", "kb-1")) - assert "chunk_num" in res["message"], res + assert res["code"] == module.RetCode.SUCCESS, res kb_rank = _KB(kb_id="kb-1", name="old", pagerank=0) monkeypatch.setattr(module.KnowledgebaseService, "get_or_none", lambda **kwargs: kb_rank if kwargs.get("id") else None) diff --git a/web/src/utils/api.ts b/web/src/utils/api.ts index 7c6307bc42..fbde70b7fc 100644 --- a/web/src/utils/api.ts +++ b/web/src/utils/api.ts @@ -58,7 +58,7 @@ export default { // knowledge base checkEmbedding: (datasetId: string) => - `${restAPIv1}/datasets/${datasetId}/embedding`, + `${restAPIv1}/datasets/${datasetId}/embedding/check`, kbList: `${restAPIv1}/datasets`, createKb: `${restAPIv1}/datasets`, updateKb: (datasetId: string) => `${restAPIv1}/datasets/${datasetId}`,