From 36956520ea4099e4a55e0c87e292a7be3adfc891 Mon Sep 17 00:00:00 2001 From: fengyaojie Date: Sat, 3 Feb 2024 18:01:54 +0800 Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E5=A4=8Des=20=E7=9F=A5=E8=AF=86?= =?UTF-8?q?=E5=BA=93=E6=9F=A5=E8=AF=A2bug=20(#2848)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../knowledge_base/kb_service/es_kb_service.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/server/knowledge_base/kb_service/es_kb_service.py b/server/knowledge_base/kb_service/es_kb_service.py index 2b4c8ad3..fe629a05 100644 --- a/server/knowledge_base/kb_service/es_kb_service.py +++ b/server/knowledge_base/kb_service/es_kb_service.py @@ -155,6 +155,20 @@ class ESKBService(KBService): k=top_k) return docs + def get_doc_by_ids(self, ids: List[str]) -> List[Document]: + results = [] + for doc_id in ids: + try: + response = self.es_client_python.get(index=self.index_name, id=doc_id) + source = response["_source"] + # Assuming your document has "text" and "metadata" fields + text = source.get("context", "") + metadata = source.get("metadata", {}) + results.append(Document(page_content=text, metadata=metadata)) + except Exception as e: + logger.error(f"Error retrieving document from Elasticsearch! {e}") + return results + def del_doc_by_ids(self, ids: List[str]) -> bool: for doc_id in ids: try: @@ -210,7 +224,8 @@ class ESKBService(KBService): } } } - search_results = self.es_client_python.search(body=query) + # 注意设置size，默认返回10个。 + search_results = self.es_client_python.search(body=query, size=50) if len(search_results["hits"]["hits"]) == 0: raise ValueError("召回元素个数为0") info_docs = [{"id":hit["_id"], "metadata": hit["_source"]["metadata"]} for hit in search_results["hits"]["hits"]]