From b3c923da6be1f5b5d36878917ab00a6597857d43 Mon Sep 17 00:00:00 2001 From: KevinHuSh Date: Fri, 17 May 2024 17:51:54 +0800 Subject: [PATCH] add doc ids in API: completion (#827) ### What problem does this PR solve? #808 ### Type of change - [x] New Feature (non-breaking change which adds functionality) --- api/db/services/dialog_service.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/api/db/services/dialog_service.py b/api/db/services/dialog_service.py index 8d928afab..94c2285c7 100644 --- a/api/db/services/dialog_service.py +++ b/api/db/services/dialog_service.py @@ -118,7 +118,9 @@ def chat(dialog, messages, stream=True, **kwargs): else: kbinfos = retrievaler.retrieval(" ".join(questions), embd_mdl, dialog.tenant_id, dialog.kb_ids, 1, dialog.top_n, dialog.similarity_threshold, - dialog.vector_similarity_weight, top=1024, aggs=False) + dialog.vector_similarity_weight, + doc_ids=kwargs.get("doc_ids", "").split(","), + top=1024, aggs=False) knowledges = [ck["content_with_weight"] for ck in kbinfos["chunks"]] chat_logger.info( "{}->{}".format(" ".join(questions), "\n->".join(knowledges)))