diff --git a/api/apps/restful_apis/chunk_api.py b/api/apps/restful_apis/chunk_api.py index 049ed60e04..e5a78cdf96 100644 --- a/api/apps/restful_apis/chunk_api.py +++ b/api/apps/restful_apis/chunk_api.py @@ -105,6 +105,7 @@ class Chunk(BaseModel): questions: list = Field(default_factory=list) question_tks: str = "" image_id: str = "" + doc_type_kwd: str = "" available: bool = True positions: list[list[int]] = Field(default_factory=list) @@ -489,6 +490,7 @@ async def list_chunks(tenant_id, dataset_id, document_id): "questions": chunk.get("question_kwd", []), "dataset_id": chunk.get("kb_id", chunk.get("dataset_id")), "image_id": chunk.get("img_id", ""), + "doc_type_kwd": chunk.get("doc_type_kwd") if isinstance(chunk.get("doc_type_kwd"), str) else "text", "available": bool(chunk.get("available_int", 1)), "positions": chunk.get("position_int", []), "tag_kwd": chunk.get("tag_kwd", []), @@ -516,6 +518,7 @@ async def list_chunks(tenant_id, dataset_id, document_id): "questions": sres.field[chunk_id].get("question_kwd", []), "dataset_id": sres.field[chunk_id].get("kb_id", sres.field[chunk_id].get("dataset_id")), "image_id": sres.field[chunk_id].get("img_id", ""), + "doc_type_kwd": sres.field[chunk_id].get("doc_type_kwd") if isinstance(sres.field[chunk_id].get("doc_type_kwd"), str) else "text", "available": bool(int(sres.field[chunk_id].get("available_int", "1"))), "positions": sres.field[chunk_id].get("position_int", []), } @@ -969,6 +972,7 @@ async def add_chunk(tenant_id, dataset_id, document_id): "id": chunk_id, "content_ltks": rag_tokenizer.tokenize(req["content"]), "content_with_weight": req["content"], + "doc_type_kwd": "text", } d["content_sm_ltks"] = rag_tokenizer.fine_grained_tokenize(d["content_ltks"]) d["important_kwd"] = req.get("important_keywords", []) @@ -1024,6 +1028,7 @@ async def add_chunk(tenant_id, dataset_id, document_id): "create_time": "create_time", "document_keyword": "document", "img_id": "image_id", + "doc_type_kwd": "doc_type_kwd", } renamed_chunk = {new_key: d[key] for key, new_key in key_mapping.items() if key in d} _ = Chunk(**renamed_chunk) diff --git a/test/testcases/test_web_api/test_chunk_app/test_chunk_routes_unit.py b/test/testcases/test_web_api/test_chunk_app/test_chunk_routes_unit.py index 5a2f2aee4b..3708c76d43 100644 --- a/test/testcases/test_web_api/test_chunk_app/test_chunk_routes_unit.py +++ b/test/testcases/test_web_api/test_chunk_app/test_chunk_routes_unit.py @@ -521,6 +521,13 @@ def test_restful_chunk_list_get_and_delete_unit(monkeypatch): assert res["data"]["total"] == 1, res assert res["data"]["chunks"][0]["id"] == "chunk-1", res assert res["data"]["chunks"][0]["available"] is True, res + assert res["data"]["chunks"][0]["doc_type_kwd"] == "text", res + + module.request = SimpleNamespace(args={"id": "chunk-1"}, headers={}) + module.settings.docStoreConn.chunk["doc_type_kwd"] = "image" + res = _run(_route_core(module.list_chunks)("tenant-1", "kb-1", "doc-1")) + assert res["code"] == 0, res + assert res["data"]["chunks"][0]["doc_type_kwd"] == "image", res res = _run(_route_core(module.get_chunk)("tenant-1", "kb-1", "doc-1", "chunk-1")) assert res["code"] == 0, res @@ -715,3 +722,4 @@ def test_restful_add_chunk_valid_image_base64_stores_before_insert(monkeypatch): inserted = module.settings.docStoreConn.inserted[-1] assert inserted.get("img_id"), inserted assert inserted.get("doc_type_kwd") == "image", inserted + assert res["data"]["chunk"]["doc_type_kwd"] == "image", res