| import inspect | |
| import os | |
| from naptha_sdk.agent_service_engine import run_agent_service | |
| from naptha_sdk.code_extraction import create_poetry_package, extract_packages, install_poetry_packages, publish_hf_package, transform_code_as | |
| from naptha_sdk.schemas import ModuleRunInput | |
| from naptha_sdk.utils import get_logger, AsyncMixin | |
| logger = get_logger(__name__) | |
| class AgentService(AsyncMixin): | |
| def __init__(self, name, fn, worker_node, orchestrator_node, flow_run=None): | |
| self.name = name | |
| self.fn = fn | |
| if isinstance(self.fn, str): | |
| self.module_name = self.fn | |
| else: | |
| self.module_name = self.fn.__name__ | |
| self.worker_node = worker_node | |
| self.orchestrator_node = orchestrator_node | |
| self.flow_run = flow_run | |
| self.repo_id = f"as_{self.module_name}" | |
| super().__init__() | |
| async def __ainit__(self): | |
| self.publish_package() | |
| await self.register_module() | |
| await self.register_service() | |
| def publish_package(self, naptha): | |
| logger.info(f"Publishing Package...") | |
| as_code = inspect.getsource(self.fn) | |
| as_code = transform_code_as(as_code) | |
| packages = extract_packages(as_code) | |
| logger.info(f"PACKAGES: {packages}") | |
| create_poetry_package(self.module_name) | |
| install_poetry_packages(self.module_name, packages) | |
| publish_hf_package(naptha.hf, self.module_name, self.repo_id, as_code, naptha.hf_username) | |
| async def register_module(self, naptha): | |
| module_config = { | |
| "name": self.module_name, | |
| "description": self.module_name, | |
| "author": f"user:{naptha.hf_username}", | |
| "url": f"https://huggingface.co/{naptha.hf_username}/{self.repo_id}", | |
| "type": "template" | |
| } | |
| logger.info(f"Registering Agent Module {module_config}") | |
| module = await naptha.hub.create_module(module_config) | |
| async def register_service(self, naptha): | |
| agent_service_name = self.name | |
| agent_service_config = { | |
| "name": agent_service_name, | |
| "description": agent_service_name, | |
| "module_name": self.module_name, | |
| "worker_node_url": self.worker_node.node_url, | |
| } | |
| logger.info(f"Registering Agent Service {agent_service_config}") | |
| service = await naptha.hub.create_service(agent_service_config) | |
| async def __call__(self, *args, **kwargs): | |
| return await run_agent_service(agent_service=self, flow_run=self.flow_run, parameters=kwargs) | |