|
|
from typing import Dict, Optional |
|
|
from naptha_sdk.schemas import ModuleRun, ModuleRunInput |
|
|
from naptha_sdk.client.comms.http_client import ( |
|
|
check_user_http, register_user_http, run_task_http, check_tasks_http, check_task_http, |
|
|
create_task_run_http, update_task_run_http, read_storage_http, write_storage_http |
|
|
) |
|
|
from naptha_sdk.client.comms.ws_client import ( |
|
|
check_user_ws, register_user_ws, run_task_ws, check_task_ws, |
|
|
create_task_run_ws, update_task_run_ws, read_storage_ws, write_storage_ws |
|
|
) |
|
|
from naptha_sdk.utils import get_logger |
|
|
|
|
|
|
|
|
logger = get_logger(__name__) |
|
|
|
|
|
class Node: |
|
|
def __init__(self, node_url: Optional[str] = None, indirect_node_id: Optional[str] = None, routing_url: Optional[str] = None): |
|
|
self.node_url = node_url |
|
|
self.indirect_node_id = indirect_node_id |
|
|
self.routing_url = routing_url |
|
|
|
|
|
|
|
|
if not node_url and not indirect_node_id: |
|
|
raise ValueError("Either node_url or indirect_node_id must be set") |
|
|
|
|
|
|
|
|
if indirect_node_id and not routing_url: |
|
|
raise ValueError("routing_url must be set if indirect_node_id is set") |
|
|
|
|
|
if self.node_url: |
|
|
self.client = 'http' |
|
|
logger.info("Using http client") |
|
|
logger.info(f"Node URL: {self.node_url}") |
|
|
else: |
|
|
self.client = 'ws' |
|
|
logger.info("Using ws client") |
|
|
logger.info(f"Routing URL: {self.routing_url}") |
|
|
logger.info(f"Indirect Node ID: {self.indirect_node_id}") |
|
|
|
|
|
self.access_token = None |
|
|
|
|
|
|
|
|
async def check_user(self, user_input): |
|
|
if self.client == 'http': |
|
|
return await check_user_http(self.node_url, user_input) |
|
|
else: |
|
|
return await check_user_ws(self.routing_url, self.indirect_node_id, user_input) |
|
|
|
|
|
|
|
|
async def register_user(self, user_input): |
|
|
if self.client == 'http': |
|
|
return await register_user_http(self.node_url, user_input) |
|
|
else: |
|
|
return await register_user_ws(self.routing_url, self.indirect_node_id, user_input) |
|
|
|
|
|
async def run_task(self, module_run_input: ModuleRunInput) -> ModuleRun: |
|
|
if self.client == 'http': |
|
|
return await run_task_http( |
|
|
node_url=self.node_url, |
|
|
module_run_input=module_run_input, |
|
|
access_token=self.access_token |
|
|
) |
|
|
else: |
|
|
return await run_task_ws( |
|
|
routing_url=self.routing_url, |
|
|
indirect_node_id=self.indirect_node_id, |
|
|
module_run_input=module_run_input |
|
|
) |
|
|
|
|
|
async def check_tasks(self): |
|
|
if self.client == 'http': |
|
|
return await check_tasks_http(self.node_url) |
|
|
else: |
|
|
raise NotImplementedError("check_tasks is not implemented for ws client") |
|
|
|
|
|
async def check_task(self, module_run: ModuleRun) -> ModuleRun: |
|
|
if self.client == 'http': |
|
|
return await check_task_http(self.node_url, module_run) |
|
|
else: |
|
|
return await check_task_ws(self.routing_url, self.indirect_node_id, module_run) |
|
|
|
|
|
async def create_task_run(self, module_run_input: ModuleRunInput) -> ModuleRun: |
|
|
if self.client == 'http': |
|
|
logger.info(f"Creating task run with input: {module_run_input}") |
|
|
logger.info(f"Node URL: {self.node_url}") |
|
|
return await create_task_run_http(self.node_url, module_run_input) |
|
|
else: |
|
|
return await create_task_run_ws(self.routing_url, self.indirect_node_id, module_run_input) |
|
|
|
|
|
async def update_task_run(self, module_run: ModuleRun): |
|
|
if self.client == 'http': |
|
|
return await update_task_run_http(self.node_url, module_run) |
|
|
else: |
|
|
return await update_task_run_ws(self.routing_url, self.indirect_node_id, module_run) |
|
|
|
|
|
async def read_storage(self, module_run_id, output_dir, ipfs=False): |
|
|
if self.client == 'http': |
|
|
return await read_storage_http(self.node_url, module_run_id, output_dir, ipfs) |
|
|
else: |
|
|
return await read_storage_ws(self.routing_url, self.indirect_node_id, module_run_id, output_dir, ipfs) |
|
|
|
|
|
async def write_storage(self, storage_input: str, ipfs: bool = False, publish_to_ipns: bool = False, update_ipns_name: Optional[str] = None): |
|
|
if self.client == 'http': |
|
|
return await write_storage_http(self.node_url, storage_input, ipfs, publish_to_ipns, update_ipns_name) |
|
|
else: |
|
|
return await write_storage_ws(self.routing_url, self.indirect_node_id, storage_input, ipfs, publish_to_ipns, update_ipns_name) |
