| import asyncio | |
| from datetime import datetime | |
| import json | |
| from naptha_sdk.client.node import Node | |
| from naptha_sdk.utils import get_logger | |
| import os | |
| import pytz | |
| import requests | |
| import time | |
| import traceback | |
| logger = get_logger(__name__) | |
| async def run_mas(multi_agent_service, mas_run) -> None: | |
| mas_engine = MASEngine(multi_agent_service, mas_run) | |
| await mas_engine.start_run() | |
| class MASEngine: | |
| def __init__(self, multi_agent_service, mas_run): | |
| self.mas = multi_agent_service | |
| self.mas_run = mas_run | |
| self.mas_name = multi_agent_service.module_name | |
| self.parameters = mas_run.module_params | |
| self.orchestrator_node = Node(self.mas_run.orchestrator_node) | |
| logger.info(f"Orchestrator node: {self.orchestrator_node.node_url}") | |
| if mas_run.worker_nodes is not None: | |
| self.worker_nodes = [Node(worker_node) for worker_node in mas_run.worker_nodes] | |
| else: | |
| self.worker_nodes = None | |
| logger.info(f"Worker Nodes: {self.worker_nodes}") | |
| self.consumer = { | |
| "public_key": mas_run.consumer_id.split(':')[1], | |
| 'id': mas_run.consumer_id, | |
| } | |
| async def start_run(self): | |
| logger.info(f"Starting MAS run: {self.mas_run}") | |
| logger.info(f"Checking user: {self.consumer}") | |
| consumer = await self.orchestrator_node.check_user(user_input=self.consumer) | |
| if consumer["is_registered"] == True: | |
| logger.info("Found user...", consumer) | |
| elif consumer["is_registered"] == False: | |
| logger.info("No user found. Registering user...") | |
| consumer = await self.orchestrator_node.register_user(user_input=consumer) | |
| logger.info(f"User registered: {consumer}.") | |
| logger.info(f"Running multi agent service on orchestrator node {self.orchestrator_node.node_url}: {self.mas_run}") | |
| mas_run = await self.orchestrator_node.run_task(module_run_input=self.mas_run) | |
| logger.info(f"Created multi agent service run on orchestrator node {self.orchestrator_node.node_url}: {mas_run}") | |
| while True: | |
| mas_run = await self.orchestrator_node.check_task(mas_run) | |
| logger.info(mas_run.status) | |
| if mas_run.status in ["completed", "error"]: | |
| break | |
| time.sleep(3) | |
| if mas_run.status == 'completed': | |
| logger.info(mas_run.results) | |
| self.agent_service_result = mas_run.results | |
| return mas_run.results | |
| else: | |
| logger.info(mas_run.error_message) | |
| return mas_run.error_message | |