From 973abde9315f8eaa723ff096f03416b3871f5291 Mon Sep 17 00:00:00 2001 From: Luan Citta Date: Wed, 5 Jun 2024 10:55:29 +0100 Subject: [PATCH] small refactor --- nodes/crew_node.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/nodes/crew_node.py b/nodes/crew_node.py index a989d67..290ec41 100644 --- a/nodes/crew_node.py +++ b/nodes/crew_node.py @@ -18,11 +18,14 @@ class CrewNode(BaseNode): }, } INPUT_IS_LIST = (True, True, False, False) - RETURN_TYPES = ("STRING",) + RETURN_TYPES = ("STRING", "STRING") FUNCTION = "create_crew" - def create_crew(self, tasks, process, topic="", agents=[],): - process_type = Process.sequential if process[0] == "sequential" else Process.hierarchical + def create_crew(self, tasks, process, topic="", agents=[]): + # Determine the process type + process_type = Process.sequential if process[0] == "sequential" else Process.hierarchical + + # Create the Crew instance crew = Crew( agents=agents, tasks=tasks, @@ -30,7 +33,10 @@ class CrewNode(BaseNode): topic=topic, ) + # Kick off the process and capture the result result = crew.kickoff(inputs={'topic': topic}) - print("\n\n📎Crew AI - result", result) - return (result,) + print("\n\n📎Crew AI - result", result) + return (result, crew.usage_metrics) + +