From 46b5b7834c59eebef96b295466679a8940dc78c6 Mon Sep 17 00:00:00 2001
From: zhanglei <zhanglei@apache.org>
Date: Fri, 12 Jul 2024 16:01:54 +0800
Subject: [PATCH] Refactor data analysis workflow for better collaboration

---
 src/langchain_lab/langgraph/draw_plantuml.py  | 30 +++++++++++++
 ...etwork_operations_analysis_assistant.pluml | 44 ++++---------------
 .../network_operations_analysis_assistant.py  |  4 ++
 3 files changed, 43 insertions(+), 35 deletions(-)
 create mode 100644 src/langchain_lab/langgraph/draw_plantuml.py

diff --git a/src/langchain_lab/langgraph/draw_plantuml.py b/src/langchain_lab/langgraph/draw_plantuml.py
new file mode 100644
index 0000000..f1e1941
--- /dev/null
+++ b/src/langchain_lab/langgraph/draw_plantuml.py
@@ -0,0 +1,30 @@
+from langgraph.graph import Graph
+
+
+class PlantUMLGraph:
+    def __init__(self, graph: Graph):
+        self.graph = graph
+
+    def draw_plantuml(self) -> str:
+        plantuml = []
+        graph_json = self.graph.to_json()['graph']
+        graph_nodes = graph_json['nodes']
+        graph_edges = graph_json['edges']
+        plantuml.append("@startuml")
+        for node in graph_nodes:
+            pass
+            # plantuml.append(f"class {node['id']} {{")
+            # for attr in node['attrs']:
+            #     plantuml.append(f"  {attr['key']} : {attr['value']}")
+            # plantuml.append("}")
+
+        for edge in graph_edges:
+            source_node = edge['source'] if edge['source'] != '__start__' else '(*)'
+            target_node = edge['target'] if edge['target'] != '__end__' else '(*)'
+            plantuml.append(f"{source_node} --> {target_node}")
+
+        plantuml.append("@enduml")
+        return "\n".join(plantuml)
+
+    def print_plantuml(self):
+        print(self.draw_plantuml())
diff --git a/src/langchain_lab/langgraph/network_operations_analysis_assistant/network_operations_analysis_assistant.pluml b/src/langchain_lab/langgraph/network_operations_analysis_assistant/network_operations_analysis_assistant.pluml
index 575237d..471a0ef 100644
--- a/src/langchain_lab/langgraph/network_operations_analysis_assistant/network_operations_analysis_assistant.pluml
+++ b/src/langchain_lab/langgraph/network_operations_analysis_assistant/network_operations_analysis_assistant.pluml
@@ -1,37 +1,11 @@
 @startuml
-start
-
-:网络运营经理 (NetworkOpsManager);
-:接收统计报表并初步审核;
-:召集分析会议;
-
-:无线网络工程师 (WirelessNetworkEngineer);
-:分析无线网络性能和资源数据;
-
-
-:数据分析师 (DataAnalyst);
-:深入挖掘数据，提供建议;
-:准备数据可视化报告;
-
-:网络维护人员 (NetworkMaintenanceTech);
-:确认并解决技术问题;
-
-
-:IT运营经理 (ITOpsManager);
-:协调IT资源，支持分析和解决问题;
-:优化网络操作流程;
-
-:客户服务经理 (CustomerServiceManager);
-:整理客户投诉和反馈;
-
-
-:高层管理人员 (ExecutiveManagement);
-:审核整体报告和建议;
-:批准必要的优化项目;
-
-:质量保证（QA）团队 (QATeam);
-:评估最终网络服务质量;
-:确保符合公司标准和法规;
-
-stop
+(*) --> networkOpsManager
+networkOpsManager --> dataAnalyst
+networkOpsManager --> data_tool
+networkOpsManager --> (*)
+dataAnalyst --> networkOpsManager
+dataAnalyst --> data_tool
+dataAnalyst --> (*)
+data_tool --> networkOpsManager
+data_tool --> dataAnalyst
 @enduml
diff --git a/src/langchain_lab/langgraph/network_operations_analysis_assistant/network_operations_analysis_assistant.py b/src/langchain_lab/langgraph/network_operations_analysis_assistant/network_operations_analysis_assistant.py
index a58cc72..4a77709 100644
--- a/src/langchain_lab/langgraph/network_operations_analysis_assistant/network_operations_analysis_assistant.py
+++ b/src/langchain_lab/langgraph/network_operations_analysis_assistant/network_operations_analysis_assistant.py
@@ -16,6 +16,8 @@
 from langgraph.prebuilt.tool_executor import ToolExecutor, ToolInvocation
 from typing_extensions import TypedDict
 
+from langchain_lab.langgraph.draw_plantuml import PlantUMLGraph
+
 warnings.filterwarnings("ignore", category=UserWarning, module="openpyxl")
 
 
@@ -133,6 +135,7 @@ def __init__(self, openai_api_base: str, openai_api_key: str, model_name: str, r
         )
         workflow.set_entry_point("networkOpsManager")
         self.graph = workflow.compile()
+        PlantUMLGraph(self.graph).print_plantuml()
 
         # from IPython.display import Image
         # import matplotlib.pyplot as plt
@@ -306,3 +309,4 @@ def run(self):
                                 f.write(f"{msg.additional_kwargs}\n")
                             f.write(msg.content)
                             f.write("\n\n")  # 结束
+