From e7f0a1a8f677bee6137b94c82e28eee4f1f739af Mon Sep 17 00:00:00 2001 From: peilong Date: Fri, 27 Dec 2024 11:48:02 +0800 Subject: [PATCH] =?UTF-8?q?=E5=A2=9E=E5=8A=A0spo=E6=A3=80=E7=B4=A2?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- kag/solver/execute/default_lf_executor.py | 2 +- kag/solver/logic/core_modules/common/one_hop_graph.py | 7 ++++++- kag/solver/retriever/impl/default_chunk_retrieval.py | 10 ++++------ 3 files changed, 11 insertions(+), 8 deletions(-) diff --git a/kag/solver/execute/default_lf_executor.py b/kag/solver/execute/default_lf_executor.py index 11277abf..e41bc8c1 100644 --- a/kag/solver/execute/default_lf_executor.py +++ b/kag/solver/execute/default_lf_executor.py @@ -112,7 +112,7 @@ def _execute_chunk_answer(self, req_id: str, query: str, lf: LFPlan, process_inf # force chunk retriever, so we clear kg solved answer process_info['kg_solved_answer'] = [] # chunk retriever - all_related_entities = kg_graph.get_all_entity() + all_related_entities = kg_graph.get_all_spo() all_related_entities = list(set(all_related_entities)) sub_query = self._generate_sub_query_with_history_qa(history, lf.query) doc_retrieved = self.chunk_retriever.recall_docs(queries=[query, sub_query], diff --git a/kag/solver/logic/core_modules/common/one_hop_graph.py b/kag/solver/logic/core_modules/common/one_hop_graph.py index fb02eb9e..ed0bb252 100644 --- a/kag/solver/logic/core_modules/common/one_hop_graph.py +++ b/kag/solver/logic/core_modules/common/one_hop_graph.py @@ -730,7 +730,12 @@ def get_all_entity(self): for d in self.entity_map[k]: all_entity.append(d) return list(set(all_entity)) - + def get_all_spo(self): + all_spo = [] + for k in self.edge_map.keys(): + for d in self.edge_map[k]: + all_spo.append(d) + return all_spo def _graph_to_json(self): total_entity_map = {} edge_dict = {} diff --git a/kag/solver/retriever/impl/default_chunk_retrieval.py b/kag/solver/retriever/impl/default_chunk_retrieval.py index 90ea0c6f..216729bd 100644 --- a/kag/solver/retriever/impl/default_chunk_retrieval.py +++ b/kag/solver/retriever/impl/default_chunk_retrieval.py @@ -325,12 +325,10 @@ def _add_extra_entity_from_spo(self, matched_entities: Dict, retrieved_spo: List all_related_entities = [] if retrieved_spo: for spo in retrieved_spo: - if spo.type not in ['Text', 'attribute']: - all_related_entities.append(spo) - # if spo.from_entity.type not in ['Text', 'attribute']: - # all_related_entities.append(spo.from_entity) - # if spo.end_entity.type not in ['Text', 'attribute']: - # all_related_entities.append(spo.end_entity) + if spo.from_entity.type not in ['Text', 'attribute']: + all_related_entities.append(spo.from_entity) + if spo.end_entity.type not in ['Text', 'attribute']: + all_related_entities.append(spo.end_entity) all_related_entities = list(set(all_related_entities)) if len(all_related_entities) == 0: