| import functools | |
| from naptha_sdk.agent_service import AgentService | |
| from naptha_sdk.client.node import Node | |
| from naptha_sdk.mas import MultiAgentService | |
| from naptha_sdk.schemas import ModuleRunInput | |
| from naptha_sdk.utils import get_logger | |
| import time | |
| logger = get_logger(__name__) | |
| class App: | |
| def __init__(self, naptha): | |
| self.naptha = naptha | |
| self.agent_services = [] | |
| def agent_service(self, name, worker_node_url): | |
| def decorator(func): | |
| worker_node = Node(worker_node_url) | |
| self.agent_services.append(AgentService(name=name, fn=func, worker_node=worker_node, orchestrator_node=self.naptha.node)) | |
| return func | |
| return decorator | |
| def publish_agent_packages(self): | |
| for agent_service in self.agent_services: | |
| agent_service.publish_package(self.naptha) | |
| async def register_agent_modules(self): | |
| for agent_service in self.agent_services: | |
| await agent_service.register_module(self.naptha) | |
| async def register_agent_services(self): | |
| for agent_service in self.agent_services: | |
| await agent_service.register_service(self.naptha) | |
| self.worker_node_urls = [agent_service.worker_node.node_url for agent_service in self.agent_services] | |
| self.worker_nodes = [Node(worker_node_url) for worker_node_url in self.worker_node_urls] | |
| logger.info(f"Worker Nodes: {self.worker_nodes}") | |
| def multi_agent_service(self, name): | |
| def decorator(func): | |
| self.multi_agent_service = MultiAgentService(self.naptha, name=name, fn=func) | |
| return func | |
| return decorator | |
| def publish_multi_agent_packages(self): | |
| self.multi_agent_service.publish_package() | |
| async def register_multi_agent_modules(self): | |
| await self.multi_agent_service.register_module() | |
| async def register_multi_agent_services(self): | |
| await self.multi_agent_service.register_service() | |
| self.orchestrator_node = self.naptha.node | |
| logger.info(f"Orchestrator node: {self.orchestrator_node.node_url}") | |
| async def run(self, run_params): | |
| consumer_id = self.naptha.user["id"] | |
| flow_run_input = { | |
| "name": self.multi_agent_service.name, | |
| "type": "template", | |
| "consumer_id": consumer_id, | |
| "orchestrator_node": self.naptha.node.node_url, | |
| "worker_nodes": self.worker_node_urls, | |
| "module_name": self.multi_agent_service.module_name, | |
| "module_params": run_params, | |
| } | |
| flow_run_input = ModuleRunInput(**flow_run_input) | |
| consumer = { | |
| "public_key": consumer_id.split(':')[1], | |
| 'id': consumer_id, | |
| } | |
| logger.info(f"Starting MAS run: {flow_run_input}") | |
| logger.info(f"Checking user: {consumer}") | |
| consumer = await self.orchestrator_node.check_user(user_input=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}: {flow_run_input}") | |
| flow_run = await self.orchestrator_node.run_task(module_run_input=flow_run_input) | |
| logger.info(f"Created multi agent service run on orchestrator node {self.orchestrator_node.node_url}: {flow_run}") | |
| while True: | |
| flow_run = await self.orchestrator_node.check_task(flow_run) | |
| logger.info(flow_run.status) | |
| if flow_run.status in ["completed", "error"]: | |
| break | |
| time.sleep(3) | |
| if flow_run.status == 'completed': | |
| logger.info(flow_run.results) | |
| self.agent_service_result = flow_run.results | |
| return flow_run.results | |
| else: | |
| logger.info(flow_run.error_message) | |
| return flow_run.error_message | |