update text2graph and alarm2graph

Author: lightislost

muagent/base_configs/prompts/simple_prompts.py

@@ -203,4 +203,103 @@
 {conversation}
 
 ## 输出
-"""
+"""
+
+
+
+text2EKG_prompt_en = '''你是一个结构化信息抽取的专家，你需要根据输入的文档，抽取其中的关键节点及节点间的连接顺序。请用json结构返回。
+
+json结构定义如下:
+{
+  "nodes": {
+    "节点序号": {
+      "type": "节点类型",
+      "content": "节点内容"
+    }
+  },
+  "edges": [
+    {
+      "start": "起始节点序号",
+      "end": "终止节点序号"
+    }
+  ]
+}
+其中 nodes 用来存放抽取的节点，每个 node 的 key 通过从0开始对递增序列表示，value 是一个字典，包含 type 和 content 两个属性， type 对应下面定义的三种节点类型，content 为抽取的节点内容。
+edges 用来存放节点间的连接顺序，它是一个列表，每个元素是一个字典，包含 start 和 end 两个属性， start 为起始 node 的 节点序号, end 为结束 node 的 节点序号。
+
+节点类型定义如下:
+Schedule:
+  表示整篇输入文档所做的事情，是对整篇输入文档的总结；
+  第一个节点永远是Schedule节点。
+Task: 
+  表示需要执行的任务。
+Phenomenon:
+  表示依据Task节点的执行结果，得到的事实结论。
+  Phenomenon节点只能连接在Task节点之后。
+Analysis:
+  表示依据Phenomenon节点的事实进行推断的过程；
+  Analysis节点只能连接在Phenomenon节点之后。
+
+以下是一个例子：
+input: 路径：排查网络问题
+1. 通过观察sofagw网关监控发现，BOLT失败数突增
+2. 且失败曲线与退保成功率曲线相关性较高，判定是网络问题。
+
+output: {
+  "nodes": {
+    "0": {
+      "type": "Schedule",
+      "content": "排查网络问题"
+    },
+    "1": {
+      "type": "Task",
+      "content": "查询sofagw网关监控BOLT失败数"
+    },
+    "2": {
+      "type": "Task",
+      "content": "查询sofagw网关监控退保成功率"
+    },
+    "3": {
+      "type": "Task",
+      "content": "判断两条时序相关性"
+    },
+    "4": {
+      "type": "Phenomenon",
+      "content": "相关性较高"
+    },
+    "5": {
+      "type": "Analysis",
+      "content": "网络问题"
+    }
+  },
+  "edges": [
+    {
+      "start": "0",
+      "end": "1"
+    },
+    {
+      "start": "1",
+      "end": "2"
+    },
+    {
+      "start": "2",
+      "end": "3"
+    },
+    {
+      "start": "3",
+      "end": "4"
+    },
+    {
+      "start": "4",
+      "end": "5"
+    }
+  ]
+}
+
+请根据上述说明和例子来对以下的输入文档抽取结构化信息:
+
+input: {text}
+
+output:'''
+
+text2EKG_prompt_zh = text2EKG_prompt_en

muagent/connector/configs/generate_prompt.py

@@ -4,7 +4,7 @@
 def replacePrompt(prompt: str, keys: list[str] = []):
     prompt = prompt.replace("{", "{{").replace("}", "}}")
     for key in keys:
-        prompt = prompt.replace(f"{{{{key}}}}", f"{{key}}")
+        prompt = prompt.replace(f"{{{{{key}}}}}", f"{{{key}}}")
     return prompt
 
 def cleanPrompt(prompt):
@@ -54,4 +54,13 @@ def createMKGPrompt(conversation, schemas, language="en", **kwargs) -> str:
     #     prompt = memory_extract_prompt_zh.format(**{"conversation": conversation, "schemas": schemas})
     # else:
     #     prompt = memory_extract_prompt_en.format(**{"conversation": conversation, "schemas": schemas})
+    return cleanPrompt(prompt)
+
+
+def createText2EKGPrompt(text, language="en", **kwargs) -> str:
+    prompt = text2EKG_prompt_zh if language == "zh" else text2EKG_prompt_en
+    prompt = replacePrompt(prompt, keys=["text"])
+    from loguru import logger
+    logger.debug(f"{prompt}")
+    prompt = prompt.format(**{"text": text,})
     return cleanPrompt(prompt)

muagent/connector/memory/hierarchical_memory_manager.py

@@ -10,7 +10,7 @@
 from muagent.connector.configs.generate_prompt import *
 from muagent.connector.schema import Memory, Message
 from muagent.schemas.db import DBConfig, GBConfig, VBConfig, TBConfig
-from muagent.schemas.memory import *
+from muagent.schemas.common import *
 from muagent.db_handler import *
 from muagent.connector.memory_manager import BaseMemoryManager
 from muagent.llm_models import *

muagent/db_handler/__init__.py

@@ -6,12 +6,12 @@
 @desc:
 '''
 
-from .graph_db_handler import NebulaHandler, NetworkxHandler, AliYunSLSHandler, GeaBaseHandler
+from .graph_db_handler import NebulaHandler, NetworkxHandler, AliYunSLSHandler, GeaBaseHandler, GBHandler
 from .vector_db_handler import LocalFaissHandler, TbaseHandler, ChromaHandler
 
 
 __all__ = [
-    "NebulaHandler", "NetworkxHandler", "GeaBaseHandler",
+    "GBHandler", "NebulaHandler", "NetworkxHandler", "GeaBaseHandler", 
     "ChromaHandler", "TbaseHandler", "LocalFaissHandler", 
     "AliYunSLSHandler"
 ]

muagent/db_handler/graph_db_handler/__init__.py

@@ -5,13 +5,14 @@
 @time: 2023/11/20 下午3:07
 @desc:
 '''
+from .base_gb_handler import GBHandler
 from .nebula_handler import NebulaHandler
 from .networkx_handler import NetworkxHandler
 from .aliyun_sls_hanlder import AliYunSLSHandler
 from .geabase_handler import GeaBaseHandler
 
 
 __all__ = [
-    "NebulaHandler", "NetworkxHandler", "GeaBaseHandler",
+    "GBHandler", "NebulaHandler", "NetworkxHandler", "GeaBaseHandler",
     "AliYunSLSHandler"
 ]

muagent/db_handler/graph_db_handler/aliyun_sls_hanlder.py

@@ -3,7 +3,7 @@
 import time
 
 from aliyun.log import *
-from muagent.schemas.memory import *
+from muagent.schemas.common import *
 from muagent.schemas.db import SLSConfig
 
 
@@ -141,6 +141,6 @@ def relation_process(self, relations: List[GRelation], crud_type) -> List[List[T
         relation_list = [
             [('start_id', relation.start_id), ('end_id', relation.start_id)] +\
                 [(k, v) if k!="operation_type" else (k, crud_type) for k,v in relation.attributes.items()]
-            for node in nodes
+            for relation in relations
         ]
         return relation_list

muagent/db_handler/graph_db_handler/base_gb_handler.py

@@ -3,7 +3,7 @@
 from loguru import logger
 import json
 
-from muagent.schemas.memory import *
+from muagent.schemas.common import *
 
 
 
@@ -58,20 +58,20 @@ def search_nodes_by_attr(self, attributes: dict) -> List[GNode]:
     def search_edges_by_attr(self, attributes: dict, edge_type: str = None) -> List[GEdge]:
         pass
 
-    def get_current_node(self, attributes: dict, node_type: str = None, return_keys: list = []) -> Dict:
+    def get_current_node(self, attributes: dict, node_type: str = None, return_keys: list = []) -> GNode:
         pass
 
-    def get_current_nodes(self, attributes: dict, node_type: str = None, return_keys: list = []) -> Dict:
+    def get_current_nodes(self, attributes: dict, node_type: str = None, return_keys: list = []) -> List[GNode]:
         pass
 
-    def get_current_edge(self, src_id, dst_id, edge_type:str = None, return_keys: list = []) -> Dict:
+    def get_current_edge(self, src_id, dst_id, edge_type:str = None, return_keys: list = []) -> GEdge:
         pass
 
-    def get_neighbor_nodes(self, attributes: dict, node_type: str = None, return_keys: list = []) -> List[Dict]:
+    def get_neighbor_nodes(self, attributes: dict, node_type: str = None, return_keys: list = []) -> List[GNode]:
         pass
 
-    def get_neighbor_edges(self, attributes: dict, node_type: str = None, return_keys: list = []) -> List[Dict]:
+    def get_neighbor_edges(self, attributes: dict, node_type: str = None, return_keys: list = []) -> List[GEdge]:
         pass
 
-    def get_hop_infos(self, attributes: dict, node_type: str = None, hop: int = 2, block_attributes: dict = []):
+    def get_hop_infos(self, attributes: dict, node_type: str = None, hop: int = 2, block_attributes: dict = {}, select_attributes: dict = {}) -> Graph:
         pass