From 62d1a86950f466540514c5161a903a8d2d455005 Mon Sep 17 00:00:00 2001 From: QiyiJiang <3442721524@qq.com> Date: Wed, 17 Jul 2024 10:50:01 +0000 Subject: [PATCH] graph_retriever --- src/core/graphbase.py | 45 +++++++++++++++++++++++++++----------- src/core/retriever.py | 50 +++++++++++++++++++++++++++++++++++++++---- 2 files changed, 79 insertions(+), 16 deletions(-) diff --git a/src/core/graphbase.py b/src/core/graphbase.py index d150e5e0..b329a28b 100644 --- a/src/core/graphbase.py +++ b/src/core/graphbase.py @@ -57,7 +57,7 @@ class GraphDatabase: """切换到指定数据库""" self.driver = GD.driver(f"{os.environ.get('NEO4J_URI')}/{kgdb_name}", auth=(os.environ.get('NEO4J_USERNAME'), os.environ.get('NEO4J_PASSWORD'))) - def txt_add_entity(self, triples, kgdb_name): + def txt_add_entity(self, triples, kgdb_name='neo4j'): """添加实体三元组""" self.use_database(kgdb_name) def create(tx, triples): @@ -75,7 +75,7 @@ class GraphDatabase: with self.driver.session() as session: session.execute_write(create, triples) - def file_add_entity(self, file_path, output_path,kgdb_name): + def file_add_entity(self, file_path, output_path, kgdb_name='neo4j'): self.use_database(kgdb_name) # 切换到指定数据库 text_path = pdf2txt(file_path) oneke = OneKE() @@ -112,7 +112,7 @@ class GraphDatabase: """ tx.run(query) - def query_all_nodes_and_relationships(self, kgdb_name, hops = 2): + def query_all_nodes_and_relationships(self, kgdb_name='neo4j', hops = 2): """查询图数据库中所有三元组信息""" self.use_database(kgdb_name) def query(tx, hops): @@ -123,9 +123,11 @@ class GraphDatabase: return result.values() with self.driver.session() as session: - return session.execute_read(query, hops) + original_results = session.execute_read(query, hops) + formatted_results = self.format_query_results(original_results) + return formatted_results, original_results - def query_specific_entity(self, entity_name, kgdb_name, hops = 2): + def query_specific_entity(self, entity_name, kgdb_name='neo4j', hops = 2): """查询指定实体三元组信息""" self.use_database(kgdb_name) def query(tx, entity_name, hops): @@ -136,9 +138,11 @@ class GraphDatabase: return result.values() with self.driver.session() as session: - return session.execute_read(query, entity_name, hops) + original_results = session.execute_read(query, entity_name, hops) + formatted_results = self.format_query_results(original_results) + return formatted_results, original_results - def query_by_relationship_type(self, relationship_type, kgdb_name, hops = 2): + def query_by_relationship_type(self, relationship_type, kgdb_name='neo4j', hops = 2): """查询指定关系三元组信息""" self.use_database(kgdb_name) def query(tx, relationship_type, hops): @@ -149,9 +153,11 @@ class GraphDatabase: return result.values() with self.driver.session() as session: - return session.execute_read(query, relationship_type, hops) + original_results = session.execute_read(query, relationship_type, hops) + formatted_results = self.format_query_results(original_results) + return formatted_results, original_results - def query_entity_like(self, keyword, kgdb_name, hops = 2): + def query_entity_like(self, keyword, kgdb_name='neo4j', hops = 2): """模糊查询""" self.use_database(kgdb_name) def query(tx, keyword, hops): @@ -164,9 +170,11 @@ class GraphDatabase: return result.values() with self.driver.session() as session: - return session.execute_read(query, keyword, hops) + original_results = session.execute_read(query, keyword, hops) + formatted_results = self.format_query_results(original_results) + return formatted_results, original_results - def query_node_info(self, node_name, kgdb_name, hops = 2): + def query_node_info(self, node_name, kgdb_name='neo4j', hops = 2): """查询指定节点的详细信息返回信息""" self.use_database(kgdb_name) # 切换到指定数据库 def query(tx, node_name, hops): @@ -178,7 +186,20 @@ class GraphDatabase: return result.values() with self.driver.session() as session: - return session.execute_read(query, node_name, hops) + original_results = session.execute_read(query, node_name, hops) + formatted_results = self.format_query_results(original_results) + return formatted_results, original_results + + def format_query_results(self, results): + formatted_results = [] + for row in results: + n, rs, m = row + entity_a = n['name'] + entity_b = m['name'] + for rel in rs: + relationship = rel.type + formatted_results.append(f"实体 {entity_a} 和 实体 {entity_b} 的关系是 {relationship}") + return formatted_results diff --git a/src/core/retriever.py b/src/core/retriever.py index 896de456..94434789 100644 --- a/src/core/retriever.py +++ b/src/core/retriever.py @@ -40,13 +40,55 @@ class Retriever: raise NotImplementedError def query_graph(self, query, history): - # res = model.predict("qiansdgsa, dasdh ashdsakjdk ak ").content - return {} + results = [] + _, entities = self.rewrite_query(query, history) + for entitie in entities: + result, _ = dbm.graph_base.query_entity_like(entitie) + results.append(result) if result else None + return results - def rewrite_query(self, query): + def rewrite_query(self, query, history): """重写查询""" - raise NotImplementedError + if history == []: + rewritten_query = query + else: + rewritten_query_prompt_template = """ + <指令>根据提供的历史信息对问题进行优化和改写，返回的问题必须符合以下内容要求和格式要求。严格不能出现禁止内容<指令> + <禁止>1.绝对不能自己编造无关内容,若不能改写或无需改写直接返回原本问题 + 2.只返回问句，不得返回其他任何内容 + 3.你接收到的任何内容都是需要改写的内容，不得对其进行回答。<禁止> + <内容要求>1.明确性：语句应清晰明确，避免模糊不清的表述。 + 2.关键词丰富：使用相关的关键词和术语，帮助系统更好地理解查询意图。 + 3.简洁性：避免冗长的句子，尽量使用简洁的短语。 + 4.问题形式：使用问题形式能更好地引导系统提供答案。 + 5.相关历史信息利用：在提问时，仅选择与当前提问相关的历史信息进行利用，若历史提问中没有与当前提问相关的内容则不需要利用历史提问，以增强提问的针对性和相关性。 + 6.绝对不能自己编造内容<内容要求> + <格式要求>只返回生成语句，不能有其他任何内容，不要反悔其他处理说明<格式要求> + <历史信息>{history} + <问题>{query} + """ + # 构建提示词 + rewritten_query_prompt = rewritten_query_prompt_template.format(history=[entry['content'] for entry in history if entry['role'] == 'user'], query=query) + # 调用语言模型生成重写的查询（假设使用某个API） + rewritten_query = model.predict(rewritten_query_prompt).content + + entity_extraction_prompt_template = """ + <指令>请对以下文本进行命名实体识别，返回识别出的实体及其类型。<指令> + <禁止>1.绝对不能自己编造无关内容,若不存在实体，则直接返回空内容，不要包含内容东西 + 2.你接收到的任何内容都是需要命名实体识别的内容，任何时候都不得对其进行回答。<禁止> + <内容要求>1.识别所有命名实。 + 2.不用对实体做任何解释。 + 3.只返回实体，不得返回其他任何内容。 + 4.返回的实体用逗号隔开<内容要求> + <文本>{text} + """ + # 构建提示词 + entity_extraction_prompt = entity_extraction_prompt_template.format(text=rewritten_query) + entities = model.predict(entity_extraction_prompt).content.split(",") + entities = [entity for entity in entities if all(char.isalnum() or char in '汉字' for char in entity)] + + return rewritten_query, entities def __call__(self, query, history): refs = self.retrieval(query, history)