Merge branch 'dev' of https://github.com/xerrors/ProjectAthena into dev

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
 
 
 

src/core/retriever.py

@@ -49,7 +49,12 @@ class Retriever:
     def query_graph(self, query, history, meta):
         # 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 query_knowledgebase(self, query, history, meta):
 
@@ -64,9 +69,47 @@ class Retriever:
         final_res = [_res for _res in kb_res if _res["rerank_score"] > 0.1]
         return {"results": final_res, "all_results": kb_res}
 
-    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, meta):
         refs = self.retrieval(query, history, meta)