update new attr and ekg_construct_test_2

Author: lightislost

muagent/db_handler/graph_db_handler/geabase_handler.py

@@ -34,6 +34,7 @@ def __init__(
     def execute(self, gql: str, option=None, return_keys: list = []) -> Dict:
         option = option or self.option
         logger.info(f"{gql}")
+        # return {"error": True}
         result = self.geabase_client.executeGQL(gql, option)
         result = json.loads(str(result.getJsonGQLResponse()))
         return result
@@ -178,12 +179,15 @@ def get_current_edge(self, src_id, dst_id, edge_type:str = None, return_keys: li
         edges = result.get("e", []) or result.get("e.attr", [])
         return [GEdge(start_id=edge["start_id"], end_id=edge["end_id"], type=edge["type"], attributes=edge) for edge in edges][0]
 
-    def get_neighbor_nodes(self, attributes: dict, node_type: str = None, return_keys: list = []) -> List[GNode]:
+    def get_neighbor_nodes(self, attributes: dict, node_type: str = None, return_keys: list = [], reverse=False) -> List[GNode]:
         # 
         extra_keys = list(set(return_keys + ["@ID", "id", "@node_type"]))
         return_str = ", ".join([f"n1.{k}" for k in extra_keys]) if return_keys else "n1"
         where_str = ' and '.join([f"n0.{k}='{v}'" for k, v in attributes.items()])
-        gql = f"MATCH (n0:{node_type})-[e]->(n1) WHERE {where_str} RETURN {return_str}"
+        if reverse:
+            gql = f"MATCH (n0:{node_type} WHERE {where_str})<-[e]-(n1) RETURN {return_str}"
+        else:
+            gql = f"MATCH (n0:{node_type})-[e]->(n1) WHERE {where_str} RETURN {return_str}"
         # 
         result = self.execute(gql, return_keys=return_keys)
         result = self.decode_result(result, gql)

muagent/db_handler/vector_db_handler/tbase_handler.py

@@ -84,7 +84,7 @@ def insert_data_hash(
                 rr = self.client.expire(key_value, expire_time or self.expire_time)
         return len(data_list)
 
-    def search(self, query, index_name: str = None, query_params: dict = {}):
+    def search(self, query, index_name: str = None, query_params: dict = {}, limit=10):
         '''
         search
         :param index_name:
@@ -96,7 +96,7 @@ def search(self, query, index_name: str = None, query_params: dict = {}):
         index = self.client.ft(index_name)
 
         if type(query) == str:
-            query = Query(query)
+            query = Query(query).paging(0, limit)
 
         res = index.search(query, query_params=query_params)
         return res

muagent/schemas/ekg/ekg_graph.py

@@ -1,6 +1,7 @@
 from pydantic import BaseModel
 from typing import List, Dict
 from enum import Enum
+import json
 
 
 SHAPE2TYPE = {
@@ -54,39 +55,77 @@ class NodeTypesEnum(Enum):
     TOOL_INSTANCE = 'opsgptkg_tool_instance'
     TEAM = 'opsgptkg_team'
     OWNER = 'opsgptkg_owner'
-    edge = 'edge'
-
+    EDGE = 'edge'
 
 # EKG Node and Edge Schemas
 class EKGNodeSchema(NodeSchema):
-    teamid: str
-    version:str # yyyy-mm-dd HH:MM:SS
+    teamids: str
+    # version:str # yyyy-mm-dd HH:MM:SS
     extra: str = ''
 
 
+
 class EKGEdgeSchema(EdgeSchema):
-    teamid: str
-    version:str # yyyy-mm-dd HH:MM:SS
+    # teamids: str
+    # version:str # yyyy-mm-dd HH:MM:SS
     extra: str = ''
 
+    def attrbutes(self, ):
+        extra_attr = json.loads(self.extra)
+        return extra_attr
+
 
 class EKGIntentNodeSchema(EKGNodeSchema):
     path: str = ''
 
+    def attrbutes(self, ):
+        extra_attr = json.loads(self.extra)
+        return {
+            **{
+                "name": self.name,
+                "description": self.description,
+                "teamids": self.teamids,
+                "path": self.path
+            }, 
+            **extra_attr
+        }
 
 class EKGScheduleNodeSchema(EKGNodeSchema):
     # do action or not
-    switch: bool
-
+    enable: bool
+
+    def attrbutes(self, ):
+        extra_attr = json.loads(self.extra)
+        return {
+            **{
+                "name": self.name,
+                "description": self.description,
+                "teamids": self.teamids,
+                "enable": self.enable
+            }, 
+            **extra_attr
+        }
 
 class EKGTaskNodeSchema(EKGNodeSchema):
-    tool: str
-    needCheck: bool
+    # tool: str
+    # needCheck: bool
     # when to access
     accessCriteria: str
     # 
-    owner: str
-
+    # owner: str
+
+    def attrbutes(self, ):
+        extra_attr = json.loads(self.extra)
+        return {
+            **{
+                "name": self.name,
+                "description": self.description,
+                "teamids": self.teamids,
+                "accessCriteria": self.accessCriteria
+            }, 
+            **extra_attr
+        }
+    
 
 class EKGAnalysisNodeSchema(EKGNodeSchema):
     # when to access
@@ -96,19 +135,38 @@ class EKGAnalysisNodeSchema(EKGNodeSchema):
     # summary template
     dslTemplate: str
 
+    def attrbutes(self, ):
+        extra_attr = json.loads(self.extra)
+        return {
+            **{
+                "name": self.name,
+                "description": self.description,
+                "teamids": self.teamids,
+                "accessCriteria": self.accessCriteria,
+                "summarySwtich": self.summarySwtich,
+                "dslTemplate": self.dslTemplate 
+            }, 
+            **extra_attr
+        }
+    
 
 class EKGPhenomenonNodeSchema(EKGNodeSchema):
-    pass
-
-
-class EKGPhenomenonNodeSchema(EKGNodeSchema):
-    pass
 
+    def attrbutes(self, ):
+        extra_attr = json.loads(self.extra)
+        return {
+            **{
+                "name": self.name,
+                "description": self.description,
+                "teamids": self.teamids,
+            }, 
+            **extra_attr
+        }
+    
 
 # Ekg Tool Schemas
 class ToolSchema(NodeSchema):
-    teamid: str
-    version:str # yyyy-mm-dd HH:MM:SS
+    # version:str # yyyy-mm-dd HH:MM:SS
     extra: str = ''
 
 
@@ -137,15 +195,21 @@ class EKGGraphSlsSchema(BaseModel):
     # {tool_id},{tool_id},{tool_id}
     tool: str = ''
     access_criteria: str = ''
-    teamid: str = ''
+    teamids: str = ''
+    extra: str = ''
+    enable: bool = False
+    dslTemplate: str = ''
 
 
 class EKGNodeTbaseSchema(BaseModel):
     node_id: str
     node_type: str
-    # node_str = 'graph_id={graph_id}'/teamid, use for searching by graph_id/teamid
+    # node_str = 'graph_id={graph_id}'/teamids, use for searching by graph_id/teamids
     node_str: str
-    node_vector: List
+    name_keyword: str
+    desc_keyword: str
+    name_vector: List
+    desc_vector: List
 
 
 class EKGEdgeTbaseSchema(BaseModel):
@@ -156,7 +220,7 @@ class EKGEdgeTbaseSchema(BaseModel):
     edge_source: str
     # end_id
     edge_target: str
-    # edge_str = 'graph_id={graph_id}'/teamid, use for searching by graph_id/teamid
+    # edge_str = 'graph_id={graph_id}'/teamids, use for searching by graph_id/teamids
     edge_str: str
 
 
@@ -178,7 +242,7 @@ class EKGSlsData(BaseModel):
     NodeTypesEnum.SCHEDULE.value: EKGScheduleNodeSchema,
     NodeTypesEnum.TOOL.value: EKGPToolTypeSchema,
     NodeTypesEnum.TOOL_INSTANCE.value: EKGPToolSchema,
-    NodeTypesEnum.edge.value: EKGEdgeSchema
+    NodeTypesEnum.EDGE.value: EKGEdgeSchema
 }