diff --git a/swarms/agents/utils/AgentManager.py b/swarms/agents/utils/AgentManager.py index ba1e444d..be1dd6bb 100644 --- a/swarms/agents/utils/AgentManager.py +++ b/swarms/agents/utils/AgentManager.py @@ -36,7 +36,7 @@ class AgentManager: self.memories[session] = self.create_memory() return self.memories[session] - def create_executor(self, session: str, execution: Optional[Task] = None) -> AgentExecutor: + def create_executor(self, session: str, execution: Optional[Task] = None, openai_api_key: str = None) -> AgentExecutor: try: builder = AgentBuilder(self.toolsets) builder.build_parser() @@ -45,6 +45,7 @@ class AgentManager: eval_callback = EVALCallbackHandler() eval_callback.set_parser(builder.get_parser()) callbacks.append(eval_callback) + if execution: execution_callback = ExecutionTracingCallbackHandler(execution) execution_callback.set_parser(builder.get_parser()) @@ -52,7 +53,7 @@ class AgentManager: callback_manager = CallbackManager(callbacks) - builder.build_llm(callback_manager) + builder.build_llm(callback_manager, openai_api_key) builder.build_global_tools() memory: BaseChatMemory = self.get_or_create_memory(session) diff --git a/swarms/agents/workers/WorkerUltraNode.py b/swarms/agents/workers/WorkerUltraNode.py index ab79570b..274fbf39 100644 --- a/swarms/agents/workers/WorkerUltraNode.py +++ b/swarms/agents/workers/WorkerUltraNode.py @@ -29,7 +29,9 @@ except Exception as e: logging.error(f"Failed to change directory: {e}") class WorkerUltraNode: - def __init__(self, objective: str): + def __init__(self, objective: str, openai_api_key: str): + self.openai_api_key = openai_api_key + if not isinstance(objective, str): raise TypeError("Objective must be a string") if not objective: @@ -69,7 +71,7 @@ class WorkerUltraNode: ) - self.session = self.agent_manager.create_executor(objective) + self.session = self.agent_manager.create_executor(objective, self.openai_api_key) except Exception as e: logging.error(f"Error while initializing WorkerUltraNode: {str(e)}") @@ -105,10 +107,8 @@ class WorkerUltraNode: -def WorkerUltra(objective: str): - # Instantiate WorkerUltraNode with the user's objective - worker_node = WorkerUltraNode(objective) - +def WorkerUltra(objective: str, openai_api_key: str): + worker_node = WorkerUltraNode(objective, openai_api_key) # Return the result of the execution return worker_node.result