WorkFlow

  1import os
  2import json
  3import configparser
  4from typing import Dict, List
  5from NodeData import NodeData
  6from crewai import Agent, Task, Crew, Process
  7from langchain_community.llms import Ollama
  8from langchain.chat_models import ChatOpenAI
  9from crewai_tools import FileReadTool, BaseTool
 10import networkx as nx
 11
 12def load_nodes_from_json(filename: str) -> Dict[str, NodeData]:
 13    with open(filename, 'r') as file:
 14        data = json.load(file)
 15        node_map = {}
 16        for node_data in data["nodes"]:
 17            node = NodeData.from_dict(node_data)
 18            node_map[node.uniq_id] = node
 19        return node_map
 20
 21def find_nodes_by_type(node_map: Dict[str, NodeData], node_type: str) -> List[NodeData]:
 22    return [node for node in node_map.values() if node.type == node_type]
 23
 24def find_node_by_type(node_map: Dict[str, NodeData], node_type: str) -> NodeData:
 25    for node in node_map.values():
 26        if node.type == node_type:
 27            return node
 28    return None
 29
 30class FileWriterTool(BaseTool):
 31    name: str = "FileWriter"
 32    description: str = "Writes given content to a specified file."
 33
 34    def _run(self, filename: str, content: str) -> str:
 35        with open(filename, 'w') as file:
 36            file.write(content)
 37        return f"Content successfully written to {filename}"
 38
 39def create_agent(node: NodeData, llm) -> Agent:
 40    tools = []
 41    tools.append(FileWriterTool())
 42    tools.append(FileReadTool())
 43    
 44    return Agent(
 45        role=node.role,
 46        goal=node.goal,
 47        backstory=node.backstory,
 48        verbose=True,
 49        allow_delegation=False,
 50        llm=llm,
 51        tools=tools
 52    )
 53
 54def create_task(node: NodeData, agent: Agent, node_map: Dict[str, NodeData], task_map: Dict[str, Task]) -> Task:
 55    steps = []
 56    for step_id in node.nexts:
 57        step_node = node_map[step_id]
 58        tool_instance = None
 59        if step_node.tool == "FileWriterTool()":
 60            tool_instance = FileWriterTool()
 61        elif step_node.tool == "FileReadTool()":
 62            tool_instance = FileReadTool()
 63        step = {
 64            'tool': tool_instance,
 65            'args': step_node.arg,
 66            'output_var': step_node.output_var
 67        }
 68        steps.append(step)
 69    
 70    # Resolve dependencies with actual Task instances
 71    dependencies = [task_map[dep_id] for dep_id in node.prevs if dep_id in task_map]
 72
 73    return Task(
 74        description=node.description,
 75        expected_output=node.expected_output,
 76        agent=agent,
 77        steps=steps,
 78        dependencies=dependencies
 79    )
 80
 81def topological_sort_tasks(task_nodes: List[NodeData]) -> List[NodeData]:
 82    graph = nx.DiGraph()
 83
 84    # Add nodes to the graph
 85    for node in task_nodes:
 86        graph.add_node(node.uniq_id)
 87
 88    # Add edges to the graph
 89    for node in task_nodes:
 90        for prev_id in node.prevs:
 91            if prev_id in graph:
 92                graph.add_edge(prev_id, node.uniq_id)
 93    
 94    # Perform topological sort
 95    sorted_ids = list(nx.topological_sort(graph))
 96    
 97    # Return nodes in sorted order
 98    id_to_node = {node.uniq_id: node for node in task_nodes}
 99    sorted_tasks = [id_to_node[node_id] for node_id in sorted_ids]
100    
101    return sorted_tasks
102
103def RunWorkFlow(node: NodeData, node_map: Dict[str, NodeData], llm):
104    print(f"Start root ID: {node.uniq_id}")
105
106    # from root find team
107    sub_node_map = {next_id: node_map[next_id] for next_id in node.nexts}
108    team_node = find_node_by_type(sub_node_map, "Team")
109    if not team_node:
110        print("No Team node found")
111        return
112
113    print(f"Processing Team {team_node.name} ID: {team_node.uniq_id}")
114
115    # from team find agents
116    agent_map = {next_id: node_map[next_id] for next_id in team_node.nexts}
117    agent_nodes = find_nodes_by_type(node_map, "Agent")
118    agents = {agent_node.name: create_agent(agent_node, llm) for agent_node in agent_nodes}
119    for agent_node in agent_nodes:
120        print(f"Agent {agent_node.name} ID: {agent_node.uniq_id}")
121
122    # Use BFS to collect all task nodes
123    task_nodes = []
124    queue = find_nodes_by_type(sub_node_map, "Task")
125    
126    while queue:
127        current_node = queue.pop(0)
128        if current_node not in task_nodes:
129            print(f"Processing task_node ID: {current_node.uniq_id}")
130            task_nodes.append(current_node)
131            next_sub_node_map = {next_id: node_map[next_id] for next_id in current_node.nexts}
132            queue.extend(find_nodes_by_type(next_sub_node_map, "Task"))
133
134    # Sort tasks topologically to respect dependencies
135    sorted_task_nodes = topological_sort_tasks(task_nodes)
136
137    tasks = []
138    task_map = {}
139    
140    # Create tasks with dependencies resolved
141    for task_node in sorted_task_nodes:
142        if task_node:
143            print(f"Processing task_node ID: {task_node.description}")
144            agent = agents[task_node.agent]
145            task = create_task(task_node, agent, node_map, task_map)
146            tasks.append(task)
147            task_map[task_node.uniq_id] = task
148        else:
149            print("No task_node found")
150            return
151
152    crew = Crew(
153        agents=list(agents.values()),
154        tasks=tasks,
155        verbose=2
156    )
157    
158    result = crew.kickoff()
159    print("######################")
160    print(result)
161
162def run_workflow_from_file(filename: str, llm):
163    node_map = load_nodes_from_json(filename)
164    start_nodes = find_nodes_by_type(node_map, "Start")
165    for start_node in start_nodes:
166        RunWorkFlow(start_node, node_map, llm)