From 3235e9d502d60e041d6fb02aa7b6c510eb296f70 Mon Sep 17 00:00:00 2001
From: Luan Citta <dev.luan.citta@gmail.com>
Date: Wed, 5 Jun 2024 09:00:11 +0100
Subject: [PATCH] Create nodes

---
 __init__.py        | 12 ++++++++++--
 agent_node.py      | 32 ++++++++++++++++++++++++++++++++
 crew_node.py       | 33 +++++++++++++++++++++++++++++++++
 list_agent_node.py | 24 ++++++++++++++++++++++++
 list_task_node.py  | 25 +++++++++++++++++++++++++
 task_node.py       | 28 ++++++++++++++++++++++++++++
 6 files changed, 152 insertions(+), 2 deletions(-)
 create mode 100644 agent_node.py
 create mode 100644 crew_node.py
 create mode 100644 list_agent_node.py
 create mode 100644 list_task_node.py
 create mode 100644 task_node.py

diff --git a/__init__.py b/__init__.py
index a51cc18..bce10b0 100644
--- a/__init__.py
+++ b/__init__.py
@@ -1,8 +1,16 @@
 import os
-
+from .agent_node import AgentNode
+from .task_node import TaskNode
+from .crew_node import CrewNode
+from .list_task_node import TaskList
+from .list_agent_node import AgentList
 
 NODE_CLASS_MAPPINGS = {
-
+    "CrewAI Agent Node": AgentNode,
+    "CrewAI Task Node": TaskNode,
+    "CrewAI Crew Node": CrewNode,
+    "CrewAI TaskList Node": TaskList,
+    "CrewAI AgentList Node": AgentList,
 }
 
 WEB_DIRECTORY = "./js"
diff --git a/agent_node.py b/agent_node.py
new file mode 100644
index 0000000..9af2d65
--- /dev/null
+++ b/agent_node.py
@@ -0,0 +1,32 @@
+from crewai import Agent
+
+class AgentNode:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "role": ("STRING", {"default": ""}),
+                "goal": ("STRING", {"default": ""}),
+                "backstory": ("STRING", {"default": ""}),
+                "verbose": ("BOOL", {"default": False}),
+                "allow_delegation": ("BOOL", {"default": False}),
+            },
+        }
+
+    RETURN_TYPES = ("AGENT",)
+    FUNCTION = "create_agent"
+    CATEGORY = "📎CrewAi"
+
+    def create_agent(self, role, goal, backstory, verbose, allow_delegation):
+       
+        agent = Agent(
+            role=role,
+            goal=goal,
+            backstory=backstory,
+            verbose=verbose,
+            allow_delegation=allow_delegation
+        )
+        return (agent,)
diff --git a/crew_node.py b/crew_node.py
new file mode 100644
index 0000000..e001cba
--- /dev/null
+++ b/crew_node.py
@@ -0,0 +1,33 @@
+from crewai import Crew, Process
+
+class CrewNode:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "agents": ("AGENT", {"default": []}),
+                "tasks": ("TASK", {"default": []}),
+                "process": ("STRING", {"default": "sequential", "choices": ["sequential", "hierarchical"]}),
+                "topic": ("STRING", {"default": ""}),
+            },
+        }
+    INPUT_IS_LIST = (True, True, False, False)
+    RETURN_TYPES = ("STRING",)
+    FUNCTION = "create_crew"
+    CATEGORY = "📎CrewAi"
+
+    def create_crew(self, agents, tasks, process, topic):
+        process_type = Process.sequential if process == "sequential" else Process.hierarchical
+        print("\n\n📎Crew AI - input", agents, tasks, process, topic)
+        crew = Crew(
+            agents=agents,
+            tasks=tasks,
+            process=process_type
+        )
+        result = crew.kickoff(inputs={'topic': topic})
+        print("\n\n📎Crew AI - result", result)       
+        return (result,)
+
diff --git a/list_agent_node.py b/list_agent_node.py
new file mode 100644
index 0000000..795322d
--- /dev/null
+++ b/list_agent_node.py
@@ -0,0 +1,24 @@
+class AgentList:
+    @classmethod
+    def INPUT_TYPES(cls):
+        # Generate a dictionary for up to n agents dynamically
+        inputs = {"required": {"agent-1": ("AGENT",)}}
+        # for i in range(2, 101):  # Adjust the range as needed
+        #     inputs["optional"] = inputs.get("optional", {})
+        #     inputs["optional"][f"agent{i}"] = ("AGENT",)
+        return inputs
+
+    OUTPUT_IS_LIST = (True,)
+    RETURN_TYPES = ("AGENT",)
+    FUNCTION = "create_list"
+    CATEGORY = "📎CrewAi"
+
+    def create_list(self, **kwargs):
+        agents = []
+
+        # Collect all agents from kwargs dynamically
+        for key, value in kwargs.items():
+            if value is not None:
+                agents.append(value)
+
+        return (agents,)
\ No newline at end of file
diff --git a/list_task_node.py b/list_task_node.py
new file mode 100644
index 0000000..083157a
--- /dev/null
+++ b/list_task_node.py
@@ -0,0 +1,25 @@
+from crewai import Crew, Process
+
+class TaskList:
+    @classmethod
+    def INPUT_TYPES(s):
+        # Generate a dictionary for up to n agents dynamically
+        inputs = {"required": {"task-1": ("TASK",)}}
+        # for i in range(2, 101):  # Adjust the range as needed
+        #     inputs["optional"] = inputs.get("optional", {})
+        #     inputs["optional"][f"agent{i}"] = ("AGENT",)
+        return inputs
+
+
+    OUTPUT_IS_LIST = (True,)
+    RETURN_TYPES = ("TASK",)
+    FUNCTION = "create_list"
+    CATEGORY = "📎CrewAi"
+
+    def create_list(self, **kwargs):
+        task = []
+
+        for k, v in kwargs.items():
+            task.append(v)
+
+        return (task, )
diff --git a/task_node.py b/task_node.py
new file mode 100644
index 0000000..cea0b90
--- /dev/null
+++ b/task_node.py
@@ -0,0 +1,28 @@
+from crewai import Task
+
+class TaskNode:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "description": ("STRING", {"default": ""}),
+                "expected_output": ("STRING", {"default": ""}),
+                "agent": ("AGENT",),
+            },
+        }
+
+    RETURN_TYPES = ("TASK",)
+    FUNCTION = "create_task"
+    CATEGORY = "📎CrewAi"
+
+    def create_task(self, description, expected_output, agent):
+        task = Task(
+            description=description,
+            expected_output=expected_output,
+            agent=agent
+        )
+        return (task,)
+