File size: 4,116 Bytes
4ae0b03 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 |
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
|