From 78ed8fe9a5b7ef5e07c0eb70da0db48128425c1a Mon Sep 17 00:00:00 2001 From: Kevin Hu Date: Wed, 14 Aug 2024 16:31:49 +0800 Subject: [PATCH] add doc ids to chat (#1944) ### What problem does this PR solve? ### Type of change - [x] Performance Improvement --- api/db/services/dialog_service.py | 3 +++ rag/nlp/search.py | 2 ++ 2 files changed, 5 insertions(+) diff --git a/api/db/services/dialog_service.py b/api/db/services/dialog_service.py index 2cfa8b8f83..849b931eaf 100644 --- a/api/db/services/dialog_service.py +++ b/api/db/services/dialog_service.py @@ -108,6 +108,9 @@ def chat(dialog, messages, stream=True, **kwargs): attachments = kwargs["doc_ids"].split(",") if "doc_ids" in kwargs else None if "doc_ids" in messages[-1]: attachments = messages[-1]["doc_ids"] + for m in messages[:-1]: + if "doc_ids" in m: + attachments.extend(m["doc_ids"]) embd_mdl = LLMBundle(dialog.tenant_id, LLMType.EMBEDDING, embd_nms[0]) if llm_id2llm_type(dialog.llm_id) == "image2text": diff --git a/rag/nlp/search.py b/rag/nlp/search.py index 99d0c1b96f..0ddb87ed4c 100644 --- a/rag/nlp/search.py +++ b/rag/nlp/search.py @@ -138,6 +138,8 @@ def search(self, req, idxnm, emb_mdl=None): es_logger.info("TOTAL: {}".format(self.es.getTotal(res))) if self.es.getTotal(res) == 0 and "knn" in s: bqry, _ = self.qryr.question(qst, min_match="10%") + if req.get("doc_ids"): + bqry = Q("bool", must=[]) bqry = self._add_filters(bqry, req) s["query"] = bqry.to_dict() s["knn"]["filter"] = bqry.to_dict()