Merge pull request #79 from codefuse-ai/ekg2_wyp_dev

LewisLiu-code · web-flow · commit 85221a3d5bfc · 2024-11-13T11:37:00.000+08:00
feat: add two func for process biz datas in httpapis
diff --git a/examples/ekg_examples/who_is_spy_game.py b/examples/ekg_examples/who_is_spy_game.py
@@ -5,6 +5,7 @@
 
 from muagent.schemas.ekg.ekg_graph import TYPE2SCHEMA
 from muagent.schemas.common import GNode, GEdge
+from muagent.service.utils import decode_biznodes, encode_biznodes
 
 
 import math
@@ -234,7 +235,8 @@ def autofill_nodes(nodes: List[GNode]):
 
 
 def add_nodes(ekg_service, nodes: list[GNode]):
-    newnodes = autofill_nodes(nodes)
+    # newnodes = autofill_nodes(nodes)
+    newnodes, newedges = decode_biznodes(nodes)
     logger.info('尝试查插入节点')
     for one_node  in newnodes:   
         one_node.attributes['description']  = one_node.attributes['description']
@@ -250,8 +252,9 @@ def add_nodes(ekg_service, nodes: list[GNode]):
             one_node.attributes['enable'] = True
 
         ekg_service.add_nodes([one_node], teamid=teamid)
-        # ekg_service.gb.add_node(one_node)
     
+    # add task-tool edge or task-agent edge
+    add_edges(ekg_service, newedges)
 
 def add_edges(ekg_service, edges):
 
diff --git a/muagent/httpapis/ekg_construct/api.py b/muagent/httpapis/ekg_construct/api.py
@@ -11,6 +11,7 @@
 from muagent.service.ekg_construct.ekg_construct_base import EKGConstructService
 from muagent.schemas.apis.ekg_api_schema import *
 from muagent.schemas.ekg import *
+from muagent.service.utils import decode_biznodes, encode_biznodes
 from muagent.service.ekg_reasoning.src.graph_search.graph_search_main import main
 
 
@@ -221,15 +222,20 @@ async def update_graph(request: EKGFeaturesRequest):
         origin_edges = request.features.query.get('originEdges', [])
         edges = request.features.query.get('edges', [])
 
-        # 构建 nodeid 到 type 的字典，方便后续查找
-        nodeid2type_dict = {n["id"]: n["type"] for n in origin_nodes + nodes}
-
         # 将 origin_nodes 和 nodes 转换为 GNode 对象
         origin_nodes = [GNode(**n) for n in origin_nodes]
-        origin_nodes = autofill_nodes(origin_nodes)
+        # origin_nodes = autofill_nodes(origin_nodes)
+        origin_nodes, origin_new_edges = decode_biznodes(origin_nodes)
         nodes = [GNode(**n) for n in nodes]
-        nodes = autofill_nodes(nodes)
+        # nodes = autofill_nodes(nodes)
+        nodes, new_edges = decode_biznodes(nodes)
+
+        origin_edges.extend([e.dict() for e in origin_new_edges])
+        edges.extend([e.dict() for e in new_edges])
                                       
+        # 构建 nodeid 到 type 的字典，方便后续查找
+        nodeid2type_dict = {n.id: n.type for n in origin_nodes + nodes}
+
         # 处理 origin_edges，给每个 edge 设置 type 字段
         origin_edges = [
             GEdge(
@@ -293,7 +299,7 @@ def get_node(request: EKGFeaturesRequest):
             node = ekg_construct_service.get_node_by_id(
                 query.nodeid, query.nodeType
             )
-            # node = node.dict()
+            # might lost agents and tools
         except Exception as e:
             errorMessage = str(e)
             successCode = False
@@ -325,10 +331,11 @@ def get_graph(request: EKGFeaturesRequest):
                 )
             nodes = graph.nodes
             edges = graph.edges
+            nodes, edges = encode_biznodes(nodes, edges)
         except Exception as e:
             errorMessage = str(e)
             successCode = False
-            nodes, edges = {}, {}
+            nodes, edges = [], []
             
         result = EKGGraphResponse(
             successCode=successCode, errorMessage=errorMessage,
@@ -377,14 +384,13 @@ def get_ancestor(request: EKGFeaturesRequest):
                 nodeid=query.nodeid, node_type=query.nodeType, 
                 rootid=query.rootid
             )
-            # nodes = graph.nodes.dict()
-            # edges = graph.edges.dict()
             nodes = graph.nodes
             edges = graph.edges
+            nodes, edges = encode_biznodes(nodes, edges)
         except Exception as e:
             errorMessage = str(e)
             successCode = False
-            nodes, edges = {}, {}
+            nodes, edges = [], []
             
             
         result = EKGGraphResponse(
diff --git a/muagent/schemas/ekg/ekg_graph.py b/muagent/schemas/ekg/ekg_graph.py
@@ -185,6 +185,12 @@ class EKGToolSchema(NodeSchema):
 class EKGAgentSchema(NodeSchema):
     modelname: str = Field("", description='llm model name')
 
+
+class EKGTaskBizNodeSchema(EKGTaskNodeSchema):
+    tools: List['EKGToolSchema'] = []
+    agents: List['EKGAgentSchema'] = []
+
+
 # SLS / Tbase 
 class EKGGraphSlsSchema(BaseModel):
     # node_{NodeTypesEnum}
@@ -254,6 +260,8 @@ class EKGSlsData(BaseModel):
     NodeTypesEnum.EDGE.value: EKGEdgeSchema
 }
 
+TYPE2SCHEMA_BIZ = copy.deepcopy(TYPE2SCHEMA)
+TYPE2SCHEMA_BIZ[NodeTypesEnum.TASK.value] = EKGTaskBizNodeSchema
 
 #####################
 ##### yuque dsl #####
diff --git a/muagent/service/ekg_construct/ekg_construct_base.py b/muagent/service/ekg_construct/ekg_construct_base.py
@@ -922,6 +922,11 @@ def search_nodes_by_text(
                 node.id, node.type, f"ekg_team_{teamid}"
             ).paths) > 0
         ]
+        # 
+        nodes = [
+            node for node in nodes
+            if node.type not in ["opsgptkg_tool", "opsgptkg_agent"]
+        ]
         return nodes
 
     def search_rootpath_by_nodeid(
diff --git a/muagent/service/utils.py b/muagent/service/utils.py
@@ -0,0 +1,125 @@
+import json
+import hashlib
+from loguru import logger
+from typing import List, Tuple, Dict
+
+from muagent.schemas.common import GNode, GEdge
+from muagent.schemas.ekg import NodeSchema, TYPE2SCHEMA_BIZ
+
+
+
+def decode_biznodes(
+        nodes: List[GNode]
+    ) -> Tuple[List[GNode], List[GEdge]]:
+    '''
+    compatible biz data : 
+        - node's attributes "extra"
+        - task node with tools and agents
+    '''
+    new_nodes = []
+    new_edges = []
+    for node in nodes:
+        schema = TYPE2SCHEMA_BIZ.get(node.type,)
+        extra = node.attributes.pop("extra", {})
+        if extra and isinstance(extra, str):
+            extra = json.loads(extra)
+        elif extra and isinstance(extra, dict):
+            pass
+        else:
+            extra = {}
+
+        node.attributes.update(extra)
+        node_data = schema(
+            **{**{"id": node.id, "type": node.type}, **node.attributes}
+        )
+
+        if node.type == "opsgptkg_task":
+            logger.debug(f"schema:{ schema}")
+            logger.debug(f"node_data:{ type(node_data)}")
+            logger.debug(f"node_data:{ node_data}")
+            
+        node_data = {
+            k:v
+            for k, v in node_data.dict().items()
+            if k not in ["type", "ID", "id", "extra"]
+        }
+
+        if node.type == "opsgptkg_task":
+            logger.debug(f"node_data:{ node_data}")
+
+        # update agent/tool nodes and edges
+        agents = node_data.pop("agents", [])
+        tools = node_data.pop("tools", [])
+
+        for agent_or_tool in agents + tools:
+
+            if agent_or_tool["type"] not in ["opsgptkg_agent", "opsgptkg_tool"]:
+                continue
+            schema = TYPE2SCHEMA_BIZ.get(agent_or_tool["type"])
+            agent_or_tool: NodeSchema = schema(**agent_or_tool)
+            new_nodes.append(GNode(**{
+                "id": agent_or_tool.id, 
+                "type": agent_or_tool.type,
+                "attributes": agent_or_tool.attributes()
+            }))
+            # may lost edge attributes
+            new_edges.append(GEdge(
+                start_id=node.id,
+                end_id=agent_or_tool.id,
+                type='_route_'.join(['opsgptkg_task', agent_or_tool.type]),
+                attributes={}
+            ))
+
+        if node.type == "opsgptkg_task":
+            logger.debug(f"node_data:{ node_data}")
+            logger.debug(f"node.attributes:{ node.attributes}")
+
+        new_nodes.append(GNode(**{
+            "id": node.id, 
+            "type": node.type,
+            "attributes": {**node_data, **node.attributes}
+        }))
+    return new_nodes, new_edges
+
+
+
+def encode_biznodes(
+        nodes: List[GNode], 
+        edges: List[GEdge]
+    ) -> Tuple[List[GNode], List[GEdge]]:
+    '''
+    compatible biz data: 
+        - node's attributes "extra"
+        - task node with tools and agents
+    '''
+    task_nodes_by_id: Dict[str, GNode] = {
+        n.id: n for n in nodes if n.type=="opsgptkg_task"
+    }
+    agent_nodes_by_id: Dict[str, GNode] = {
+        n.id: n for n in nodes if n.type in ["opsgptkg_tool"]
+    }
+    tool_nodes_by_id: Dict[str, GNode] = {
+        n.id: n for n in nodes if n.type in ["opsgptkg_agent"]
+    }
+
+    new_edges = []
+    for edge in edges:
+        if edge.start_id in task_nodes_by_id:
+            task_nodes_by_id[edge.start_id].attributes.setdefault("agents", [])
+            task_nodes_by_id[edge.start_id].attributes.setdefault("tools", [])
+
+            if edge.end_id in agent_nodes_by_id:
+                task_nodes_by_id[edge.end_id].attributes["agents"].append(
+                    agent_nodes_by_id[edge.end_id].dict()
+                )
+                continue
+
+            if edge.end_id in tool_nodes_by_id:
+                task_nodes_by_id[edge.end_id].attributes["tools"].append(
+                    tool_nodes_by_id[edge.end_id].dict()
+                )
+                continue
+        new_edges.append(edge)
+
+    
+    return nodes, new_edges
diff --git a/tests/httpapis/api_func_test2.py b/tests/httpapis/api_func_test2.py
