thestage 0.6.7__py3-none-any.whl → 0.7.0__py3-none-any.whl

thestage/{services/instance → instance/business}/instance_service.py

@@ -1,37 +1,36 @@
 from pathlib import Path
-from typing import List, Optional, Dict
+from typing import List, Optional
 
 import typer
 
-from thestage.entities.rented_instance import RentedInstanceEntity
-from thestage.entities.self_hosted_instance import SelfHostedInstanceEntity
+from thestage.instance.communication.instance_api_client import InstanceApiClient
+from thestage.instance.dto.rented_instance import RentedInstanceEntity
+from thestage.instance.dto.self_hosted_instance import SelfHostedInstanceEntity
 from thestage.i18n.translation import __
-from thestage.services.clients.thestage_api.dtos.enums.selfhosted_status import SelfhostedBusinessStatus
-from thestage.services.clients.thestage_api.dtos.enums.instance_rented_status import InstanceRentedBusinessStatus
+from thestage.instance.dto.enum.selfhosted_status import SelfhostedBusinessStatus
+from thestage.instance.dto.enum.instance_rented_status import InstanceRentedBusinessStatus
 from thestage.services.abstract_service import AbstractService
 from thestage.helpers.error_handler import error_handler
-from thestage.services.clients.thestage_api.api_client import TheStageApiClient
-from thestage.services.clients.thestage_api.dtos.instance_rented_response import InstanceRentedDto
+from thestage.instance.dto.instance_rented_response import InstanceRentedDto
 from thestage.services.clients.thestage_api.dtos.paginated_entity_list import PaginatedEntityList
-from thestage.services.clients.thestage_api.dtos.selfhosted_instance_response import SelfHostedInstanceDto
-from thestage.services.config_provider.config_provider import ConfigProvider
-from thestage.services.instance.mapper.instance_mapper import InstanceMapper
-from thestage.services.instance.mapper.selfhosted_mapper import SelfHostedMapper
-from thestage.services.remote_server_service import RemoteServerService
+from thestage.instance.dto.selfhosted_instance_response import SelfHostedInstanceDto
+from thestage.config.business.config_provider import ConfigProvider
+from thestage.instance.business.mapper.rented_instance_mapper import RentedInstanceMapper
+from thestage.instance.business.mapper.selfhosted_instance_mapper import SelfHostedInstanceMapper
+from thestage.connect.business.remote_server_service import RemoteServerService
 
 
 class InstanceService(AbstractService):
-
-    __thestage_api_client: TheStageApiClient = None
+    __instance_api_client: InstanceApiClient = None
     __config_provider: ConfigProvider = None
 
     def __init__(
             self,
-            thestage_api_client: TheStageApiClient,
+            instance_api_client: InstanceApiClient,
             config_provider: ConfigProvider,
             remote_server_service: RemoteServerService,
     ):
-        self.__thestage_api_client = thestage_api_client
+        self.__instance_api_client = instance_api_client
         self.__remote_server_service = remote_server_service
         self.__config_provider = config_provider
 
@@ -54,7 +53,7 @@ class InstanceService(AbstractService):
                 InstanceRentedBusinessStatus.TERMINATING.name,
                 InstanceRentedBusinessStatus.RENTAL_ERROR.name,
             ]:
-                typer.echo(__('Cannot connect to rented server instance: renting process failed'))
+                typer.echo(__('Cannot connect to rented server instance: renting failed'))
                 raise typer.Exit(1)
             elif instance.frontend_status.status_key in [
                 InstanceRentedBusinessStatus.STOPPED.name,
@@ -105,7 +104,7 @@ class InstanceService(AbstractService):
             instance_rented_slug: Optional[str],
             input_ssh_key_path: Optional[str]
     ):
-        instance = self.__thestage_api_client.get_rented_instance(
+        instance = self.__instance_api_client.get_rented_instance(
             instance_public_id=instance_rented_public_id,
             instance_slug=instance_rented_slug,
         )
@@ -119,7 +118,7 @@ class InstanceService(AbstractService):
             if not input_ssh_key_path:
                 ssh_path_from_config = self.__config_provider.get_valid_private_key_path_by_ip_address(instance.ip_address)
                 if ssh_path_from_config:
-                    typer.echo(f"Using configured ssh key for this instance: {ssh_path_from_config}")
+                    typer.echo(f"Using configured SSH key for this instance: {ssh_path_from_config}")
 
             if not input_ssh_key_path and not ssh_path_from_config:
                 typer.echo('Using SSH agent to connect to server instance')
@@ -149,7 +148,7 @@ class InstanceService(AbstractService):
             username = 'root'
             typer.echo(__("No remote server username provided, using 'root' as username"))
 
-        instance = self.__thestage_api_client.get_selfhosted_instance(
+        instance = self.__instance_api_client.get_selfhosted_instance(
             instance_public_id=selfhosted_instance_public_id,
             instance_slug=selfhosted_instance_slug,
         )
@@ -163,7 +162,7 @@ class InstanceService(AbstractService):
             if not input_ssh_key_path:
                 ssh_path_from_config = self.__config_provider.get_valid_private_key_path_by_ip_address(instance.ip_address)
                 if ssh_path_from_config:
-                    typer.echo(f"Using configured ssh key for this instance: {ssh_path_from_config}")
+                    typer.echo(f"Using configured SSH key for this instance: {ssh_path_from_config}")
 
             if not input_ssh_key_path and not ssh_path_from_config:
                 typer.echo('Using SSH agent to connect to server instance')
@@ -187,7 +186,7 @@ class InstanceService(AbstractService):
             row: int = 5,
             page: int = 1,
     ) -> PaginatedEntityList[InstanceRentedDto]:
-        data = self.__thestage_api_client.get_rented_instance_list(
+        data = self.__instance_api_client.get_rented_instance_list(
             statuses=statuses,
             page=page,
             limit=row,
@@ -202,7 +201,7 @@ class InstanceService(AbstractService):
             row: int = 5,
             page: int = 1,
     ) -> PaginatedEntityList[SelfHostedInstanceDto]:
-        data = self.__thestage_api_client.get_selfhosted_instance_list(
+        data = self.__instance_api_client.get_selfhosted_instance_list(
             statuses=statuses,
             page=page,
             limit=row,
@@ -212,7 +211,7 @@ class InstanceService(AbstractService):
 
     @error_handler()
     def print_self_hosted_instance_list(self, statuses, row, page):
-        selfhosted_instance_status_map = self.__thestage_api_client.get_selfhosted_business_status_map()
+        selfhosted_instance_status_map = self.__instance_api_client.get_selfhosted_business_status_map()
 
         if not statuses:
             statuses = ({key: selfhosted_instance_status_map[key] for key in [
@@ -245,7 +244,7 @@ class InstanceService(AbstractService):
             func_special_params={
                 'statuses': backend_statuses,
             },
-            mapper=SelfHostedMapper(),
+            mapper=SelfHostedInstanceMapper(),
             headers=list(map(lambda x: x.alias, SelfHostedInstanceEntity.model_fields.values())),
             row=row,
             page=page,
@@ -255,7 +254,7 @@ class InstanceService(AbstractService):
 
     @error_handler()
     def print_rented_instance_list(self, statuses, row, page):
-        instance_rented_status_map = self.__thestage_api_client.get_rented_business_status_map()
+        instance_rented_status_map = self.__instance_api_client.get_rented_business_status_map()
 
         if not statuses:
             statuses = ({key: instance_rented_status_map[key] for key in [
@@ -289,7 +288,7 @@ class InstanceService(AbstractService):
             func_special_params={
                 'statuses': backend_statuses,
             },
-            mapper=InstanceMapper(),
+            mapper=RentedInstanceMapper(),
             headers=list(map(lambda x: x.alias, RentedInstanceEntity.model_fields.values())),
             row=row,
             page=page,

thestage/{services/instance/mapper/instance_mapper.py → instance/business/mapper/rented_instance_mapper.py}

@@ -1,11 +1,11 @@
 from typing import Optional
 
-from thestage.entities.rented_instance import RentedInstanceEntity
-from thestage.services.clients.thestage_api.dtos.instance_rented_response import InstanceRentedDto
+from thestage.instance.dto.rented_instance import RentedInstanceEntity
+from thestage.instance.dto.instance_rented_response import InstanceRentedDto
 from thestage.services.abstract_mapper import AbstractMapper
 
 
-class InstanceMapper(AbstractMapper):
+class RentedInstanceMapper(AbstractMapper):
 
     def build_entity(self, item: InstanceRentedDto) -> Optional[RentedInstanceEntity]:
         if not item:

thestage/{services/instance/mapper/selfhosted_mapper.py → instance/business/mapper/selfhosted_instance_mapper.py}

@@ -1,14 +1,12 @@
-from abc import ABC
-from typing import Optional, Any, Tuple
+from typing import Optional
 
-from thestage.entities.rented_instance import RentedInstanceEntity
-from thestage.entities.self_hosted_instance import SelfHostedInstanceEntity
-from thestage.services.clients.thestage_api.dtos.enums.gpu_name import InstanceGpuType
-from thestage.services.clients.thestage_api.dtos.selfhosted_instance_response import SelfHostedInstanceDto
+from thestage.instance.dto.self_hosted_instance import SelfHostedInstanceEntity
+from thestage.instance.dto.enum.gpu_name import InstanceGpuType
+from thestage.instance.dto.selfhosted_instance_response import SelfHostedInstanceDto
 from thestage.services.abstract_mapper import AbstractMapper
 
 
-class SelfHostedMapper(AbstractMapper):
+class SelfHostedInstanceMapper(AbstractMapper):
 
     def build_entity(self, item: SelfHostedInstanceDto) -> Optional[SelfHostedInstanceEntity]:
         if not item:

thestage/instance/communication/instance_api_client.py

@@ -0,0 +1,150 @@
+from typing import Optional, List, Dict
+
+from thestage.config.business.config_provider import ConfigProvider
+from thestage.services.clients.thestage_api.core.api_client_core import TheStageApiClientCore
+from thestage.services.clients.thestage_api.dtos.paginated_entity_list import PaginatedEntityList
+
+from thestage.instance.dto.instance_rented_response import (
+    InstanceRentedListResponse,
+    InstanceRentedDto,
+    InstanceRentedItemResponse,
+    InstanceRentedBusinessStatusMapperResponse
+)
+from thestage.instance.dto.selfhosted_instance_response import (
+    SelfHostedInstanceListResponse,
+    SelfHostedInstanceDto,
+    SelfHostedRentedItemResponse,
+    SelfHostedRentedRentedBusinessStatusMapperResponse
+)
+
+
+class InstanceApiClient(TheStageApiClientCore):
+    def __init__(self, config_provider: ConfigProvider):
+        super().__init__(url=config_provider.get_config().main.thestage_api_url)
+        self.__config_provider = config_provider
+
+    def get_rented_instance_list(
+            self,
+            statuses: List[str],
+            page: int = 1,
+            limit: int = 10,
+    ) -> PaginatedEntityList[InstanceRentedDto]:
+        data = {
+            #"statuses": [item.value for item in statuses],
+            "entityFilterRequest": {
+                "orderByField": "createdAt",
+                "orderByDirection": "DESC",
+                "page": page,
+                "limit": limit
+            },
+        }
+
+        if statuses:
+            data['businessStatuses'] = statuses
+
+        response = self._request(
+            method='POST',
+            url='/user-api/v3/instance-rented/list',
+            data=data,
+            token=self.__config_provider.get_config().main.thestage_auth_token,
+        )
+
+        result = InstanceRentedListResponse.model_validate(response) if response else None
+        return result.paginated_list if result and result.paginated_list else ([], None)
+
+    def get_rented_business_status_map(self) -> Optional[Dict[str, str]]:
+        response = self._request(
+            method='POST',
+            url='/user-api/v2/instance-rented/business-status-localized-map',
+            data=None,
+            token=self.__config_provider.get_config().main.thestage_auth_token,
+        )
+
+        data = InstanceRentedBusinessStatusMapperResponse.model_validate(response) if response else None
+
+        return data.instance_rented_business_status_map if data else None
+
+
+
+    def get_rented_instance(
+            self,
+            instance_public_id: Optional[str] = None,
+            instance_slug: Optional[str] = None,
+    ) -> Optional[InstanceRentedDto]:
+        if not instance_slug and not instance_public_id:
+            return None
+
+        data = {
+            "instanceRentedPublicId": instance_public_id,
+            "instanceRentedSlug": instance_slug,
+        }
+
+        response = self._request(
+            method='POST',
+            url='/user-api/v3/instance-rented/view',
+            data=data,
+            token=self.__config_provider.get_config().main.thestage_auth_token,
+        )
+
+        return InstanceRentedItemResponse.model_validate(response).instance_rented if response else None
+
+    def get_selfhosted_instance(
+            self,
+            instance_public_id: Optional[str] = None,
+            instance_slug: Optional[str] = None,
+    ) -> Optional[SelfHostedInstanceDto]:
+        if not instance_slug and not instance_public_id:
+            return None
+
+        data = {
+            "selfhostedInstancePublicId": instance_public_id,
+            "selfhostedInstanceSlug": instance_slug,
+        }
+
+        response = self._request(
+            method='POST',
+            url='/user-api/v3/self-hosted-instance/view',
+            data=data,
+            token=self.__config_provider.get_config().main.thestage_auth_token,
+        )
+
+        return SelfHostedRentedItemResponse.model_validate(response).selfhosted_instance if response else None
+
+    def get_selfhosted_instance_list(
+            self,
+            statuses: List[str],
+            page: int = 1,
+            limit: int = 10,
+    ) -> PaginatedEntityList[SelfHostedInstanceDto]:
+        data = {
+            "entityFilterRequest": {
+                "orderByField": "createdAt",
+                "orderByDirection": "DESC",
+                "page": page,
+                "limit": limit
+            }
+        }
+
+        if statuses:
+            data['businessStatuses'] = statuses
+
+        response = self._request(
+            method='POST',
+            url='/user-api/v3/self-hosted-instance/list',
+            data=data,
+            token=self.__config_provider.get_config().main.thestage_auth_token,
+        )
+
+        result = SelfHostedInstanceListResponse.model_validate(response) if response else None
+        return result.paginated_list if result and result.paginated_list else ([], None)
+
+    def get_selfhosted_business_status_map(self) -> Optional[Dict[str, str]]:
+        response = self._request(
+            method='POST',
+            url='/user-api/v2/self-hosted-instance/business-status-localized-map',
+            data=None,
+            token=self.__config_provider.get_config().main.thestage_auth_token,
+        )
+
+        data = SelfHostedRentedRentedBusinessStatusMapperResponse.model_validate(response) if response else None
+        return data.selfhosted_instance_business_status_map if data else None

thestage/{controllers/instance_controller.py → instance/communication/instance_command.py}

@@ -4,7 +4,7 @@ from typing import Optional, List
 from thestage.cli_command import CliCommand
 from thestage.cli_command_helper import get_command_group_help_panel, get_command_metadata, check_command_permission
 from thestage.helpers.logger.app_logger import app_logger
-from thestage.services.instance.instance_service import InstanceService
+from thestage.instance.business.instance_service import InstanceService
 from thestage.i18n.translation import __
 from thestage.controllers.utils_controller import \
     validate_config_and_get_service_factory, get_current_directory
@@ -91,7 +91,7 @@ def instance_connect(
     check_command_permission(command_name)
 
     if sum(v is not None for v in [public_id, slug]) != 1:
-        typer.echo("Please provide a single identifier for rented instance - ID or name.")
+        typer.echo("Provide a single identifier for rented instance - ID or name.")
         raise typer.Exit(1)
 
     if private_ssh_key_path and not Path(private_ssh_key_path).is_file():
@@ -107,7 +107,7 @@ def instance_connect(
         input_ssh_key_path=private_ssh_key_path
     )
 
-    app_logger.info(f'Stop connecting to rented server instance')
+    app_logger.info(f'Connection to rented instance closed')
     raise typer.Exit(0)
 
 
@@ -188,7 +188,7 @@ def self_hosted_connect(
     check_command_permission(command_name)
 
     if sum(v is not None for v in [public_id, slug]) != 1:
-        typer.echo("Please provide a single identifier for self-hosted instance - ID or name.")
+        typer.echo("Provide a single identifier for self-hosted instance - ID or name.")
         raise typer.Exit(1)
 
     if private_ssh_key_path and not Path(private_ssh_key_path).is_file():
@@ -205,5 +205,5 @@ def self_hosted_connect(
         input_ssh_key_path=private_ssh_key_path
     )
 
-    app_logger.info(f'Stop connecting to self-hosted instance')
+    app_logger.info(f'Connection to self-hosted instance closed')
     raise typer.Exit(0)

thestage/{services/clients/thestage_api/dtos → instance/dto}/instance_detected_gpus.py

@@ -1,9 +1,8 @@
-
 from typing import Optional, List
 
 from pydantic import Field, BaseModel, ConfigDict
 
-from thestage.services.clients.thestage_api.dtos.enums.gpu_name import InstanceGpuType
+from thestage.instance.dto.enum.gpu_name import InstanceGpuType
 
 
 class InstanceDaemonGpuDto(BaseModel):

thestage/{services/clients/thestage_api/dtos → instance/dto}/instance_rented_response.py

@@ -5,8 +5,8 @@ from pydantic import Field, BaseModel, ConfigDict
 from thestage.services.clients.thestage_api.dtos.frontend_status import FrontendStatusDto
 from thestage.services.clients.thestage_api.dtos.base_response import TheStageBasePaginatedResponse
 from thestage.services.clients.thestage_api.dtos.paginated_entity_list import PaginatedEntityList
-from thestage.services.clients.thestage_api.dtos.enums.gpu_name import InstanceGpuType
-from thestage.services.clients.thestage_api.dtos.enums.cpu_type import InstanceCpuType
+from thestage.instance.dto.enum.gpu_name import InstanceGpuType
+from thestage.instance.dto.enum.cpu_type import InstanceCpuType
 
 
 class InstanceRentedDto(BaseModel):

thestage/{services/clients/thestage_api/dtos → instance/dto}/selfhosted_instance_response.py

@@ -1,4 +1,3 @@
-from datetime import datetime
 from typing import Optional, Dict
 
 from pydantic import Field, BaseModel, ConfigDict
@@ -6,8 +5,8 @@ from pydantic import Field, BaseModel, ConfigDict
 from thestage.services.clients.thestage_api.dtos.base_response import TheStageBasePaginatedResponse
 from thestage.services.clients.thestage_api.dtos.frontend_status import FrontendStatusDto
 from thestage.services.clients.thestage_api.dtos.paginated_entity_list import PaginatedEntityList
-from thestage.services.clients.thestage_api.dtos.instance_detected_gpus import InstanceDetectedGpusDto
-from thestage.services.clients.thestage_api.dtos.enums.cpu_type import InstanceCpuType
+from thestage.instance.dto.instance_detected_gpus import InstanceDetectedGpusDto
+from thestage.instance.dto.enum.cpu_type import InstanceCpuType
 
 
 class SelfHostedInstanceDto(BaseModel):

thestage/{services/logging → logging/business}/logging_service.py

@@ -1,6 +1,6 @@
 import asyncio
 from _signal import SIGINT
-from asyncio import CancelledError, Task, StreamWriter
+from asyncio import CancelledError, StreamWriter
 from datetime import datetime
 from typing import Optional, Dict
 
@@ -8,43 +8,55 @@ import aioconsole
 import typer
 from httpx import ReadTimeout, ConnectError, ConnectTimeout
 
+from thestage.docker_container.communication.docker_container_api_client import DockerContainerApiClient
 from thestage.helpers.logger.app_logger import app_logger
-from thestage.services.clients.thestage_api.dtos.enums.container_status import DockerContainerStatus
-from thestage.services.clients.thestage_api.dtos.enums.inference_simulator_status import InferenceSimulatorStatus
-from thestage.services.clients.thestage_api.dtos.enums.task_status import TaskStatus
-from thestage.services.clients.thestage_api.dtos.inference_controller.get_inference_simulator_response import \
+from thestage.docker_container.dto.enum.container_status import DockerContainerStatus
+from thestage.inference_simulator.communication.inference_simulator_api_client import InferenceSimulatorApiClient
+from thestage.inference_simulator.dto.enum.inference_simulator_status import InferenceSimulatorStatus
+from thestage.logging.communication.logging_api_client import LoggingApiClient
+from thestage.task.communication.task_api_client import TaskApiClient
+from thestage.task.dto.enum.task_status import TaskStatus
+from thestage.inference_simulator.dto.get_inference_simulator_response import \
     GetInferenceSimulatorResponse
-from thestage.services.clients.thestage_api.dtos.task_controller.task_view_response import TaskViewResponse
-from thestage.services.logging.byte_print_style import BytePrintStyle
-from thestage.services.logging.dto.log_message import LogMessage
-from thestage.services.logging.dto.log_type import LogType
+from thestage.task.dto.view_response import TaskViewResponse
+from thestage.logging.byte_print_style import BytePrintStyle
+from thestage.logging.dto.log_message import LogMessage
+from thestage.logging.dto.log_type import LogType
 from thestage.i18n.translation import __
-from thestage.services.abstract_service import AbstractService
-from thestage.services.clients.thestage_api.dtos.container_response import DockerContainerDto
+from thestage.docker_container.dto.container_response import DockerContainerDto
 from thestage.helpers.error_handler import error_handler
-from thestage.services.clients.thestage_api.api_client import TheStageApiClient
 from rich import print
 
-from thestage.services.logging.exception.log_polling_exception import LogPollingException
-from thestage.services.logging.logging_constants import LOG_MESSAGE_CODE_TASK_FINISHED, \
+from thestage.exceptions.log_polling_exception import LogPollingException
+from thestage.logging.logging_constants import LOG_MESSAGE_CODE_TASK_FINISHED, \
     LOG_MESSAGE_CODE_INFERENCE_SIMULATOR_FAILED
 
 is_logs_streaming = False
 
 
 class LoggingService:
-    __thestage_api_client: TheStageApiClient = None
+    __logging_api_client: LoggingApiClient = None
+    __task_api_client: TaskApiClient = None
+    __inference_simulator_api_client: InferenceSimulatorApiClient = None
+    __docker_container_api_client: DockerContainerApiClient = None
 
     def __init__(
             self,
-            thestage_api_client: TheStageApiClient,
+            logging_api_client: LoggingApiClient,
+            docker_container_api_client: DockerContainerApiClient,
+            task_api_client: TaskApiClient,
+            inference_simulator_api_client: InferenceSimulatorApiClient,
+
     ):
-        self.__thestage_api_client = thestage_api_client
+        self.__logging_api_client = logging_api_client
+        self.__docker_container_api_client = docker_container_api_client
+        self.__task_api_client = task_api_client
+        self.__inference_simulator_api_client = inference_simulator_api_client
 
 
     @error_handler()
     def print_last_task_logs(self, task_public_id: str, logs_number: Optional[int]):
-        logs = self.__thestage_api_client.query_user_logs(
+        logs = self.__logging_api_client.query_user_logs(
             task_public_id=task_public_id,
             limit=logs_number
         )
@@ -54,7 +66,7 @@ class LoggingService:
 
     @error_handler()
     def print_last_inference_simulator_logs(self, inference_simulator_public_id: str, logs_number: Optional[int]):
-        logs = self.__thestage_api_client.query_user_logs(
+        logs = self.__logging_api_client.query_user_logs(
             inference_simulator_public_id=inference_simulator_public_id,
             limit=logs_number
         )
@@ -64,7 +76,7 @@ class LoggingService:
 
     @error_handler()
     def print_last_container_logs(self, container_public_id: Optional[str], container_slug: Optional[str], logs_number: Optional[int]):
-        container: Optional[DockerContainerDto] = self.__thestage_api_client.get_container(
+        container: Optional[DockerContainerDto] = self.__docker_container_api_client.get_container(
             container_public_id=container_public_id,
             container_slug=container_slug,
         )
@@ -73,7 +85,7 @@ class LoggingService:
             typer.echo("Container was not found")
             raise typer.Exit(1)
 
-        logs = self.__thestage_api_client.query_user_logs(
+        logs = self.__logging_api_client.query_user_logs(
             container_public_id=container.public_id,
             limit=logs_number
         )
@@ -90,9 +102,9 @@ class LoggingService:
 
     @error_handler()
     async def __stream_task_logs_with_controls_async(self, task_public_id: str):
-        task_view_response: Optional[TaskViewResponse] = self.__thestage_api_client.get_task(task_public_id=task_public_id,)
+        task_view_response: Optional[TaskViewResponse] = self.__task_api_client.get_task(task_public_id=task_public_id,)
 
-        task_status_map: Dict[str, str] = self.__thestage_api_client.get_task_localized_status_map()
+        task_status_map: Dict[str, str] = self.__task_api_client.get_task_localized_status_map()
 
         task = task_view_response.task
 
@@ -128,7 +140,7 @@ class LoggingService:
             print_logs_task.cancel()
             if not input_task.result():  # result is only expected if ctrl+D triggered EOFError
                 typer.echo(f"\rTask {task_public_id} will be canceled")
-                self.__thestage_api_client.cancel_task(
+                self.__task_api_client.cancel_task(
                     task_public_id=task.public_id,
                 )
 
@@ -143,12 +155,12 @@ class LoggingService:
 
     @error_handler()
     async def __stream_inference_simulator_logs_with_controls_async(self,  public_id: Optional[str], slug: Optional[str]):
-        get_inference_simulator_response: Optional[GetInferenceSimulatorResponse] = self.__thestage_api_client.get_inference_simulator(
+        get_inference_simulator_response: Optional[GetInferenceSimulatorResponse] = self.__inference_simulator_api_client.get_inference_simulator(
             public_id=public_id,
             slug=slug,
         )
 
-        inference_simulator_status_map: Dict[str, str] = self.__thestage_api_client.get_inference_simulator_business_status_map()
+        inference_simulator_status_map: Dict[str, str] = self.__inference_simulator_api_client.get_inference_simulator_business_status_map()
 
         inference_simulator = get_inference_simulator_response.inferenceSimulator
 
@@ -195,7 +207,7 @@ class LoggingService:
 
     @error_handler()
     async def __stream_container_logs_with_controls_async(self, container_public_id: Optional[str], container_slug: Optional[str]):
-        container: Optional[DockerContainerDto] = self.__thestage_api_client.get_container(
+        container: Optional[DockerContainerDto] = self.__docker_container_api_client.get_container(
             container_public_id=container_public_id,
             container_slug=container_slug,
         )
@@ -264,7 +276,7 @@ class LoggingService:
             #     stream_to_logs_diff = datetime.utcnow() - last_log_timestamp_parsed
             #     print_nonblocking(f'TDIFF {stream_to_logs_diff.total_seconds()}', writer)
             try:
-                logs_response = await self.__thestage_api_client.poll_logs_httpx(
+                logs_response = await self.__logging_api_client.poll_logs_httpx(
                     task_public_id=task_public_id,
                     inference_simulator_public_id=inference_simulator_public_id,
                     docker_container_public_id=docker_container_public_id,

thestage/logging/communication/logging_api_client.py

@@ -0,0 +1,63 @@
+from typing import Optional
+
+import httpx
+
+from thestage.config.business.config_provider import ConfigProvider
+from thestage.services.clients.thestage_api.core.api_client_core import TheStageApiClientCore
+from thestage.logging.dto.log_polling_request import LogPollingRequest
+from thestage.logging.dto.log_polling_response import LogPollingResponse
+from thestage.logging.dto.user_logs_query_request import UserLogsQueryRequest
+from thestage.logging.dto.user_logs_query_response import UserLogsQueryResponse
+
+
+class LoggingApiClient(TheStageApiClientCore):
+    def __init__(self, config_provider: ConfigProvider):
+        super().__init__(url=config_provider.get_config().main.thestage_api_url)
+        self.__config_provider = config_provider
+
+    async def poll_logs_httpx(self, docker_container_public_id: Optional[str], last_log_timestamp: str,
+                              last_log_id: str, task_public_id: Optional[str] = None,
+                              inference_simulator_public_id: Optional[str] = None) -> Optional[LogPollingResponse]:
+        request_headers = {'Content-Type': 'application/json'}
+        token = self.__config_provider.get_config().main.thestage_auth_token
+        if token: request_headers['Authorization'] = f"Bearer {token}"
+
+        request = LogPollingRequest(
+            inferenceSimulatorPublicId=inference_simulator_public_id,
+            taskPublicId=task_public_id,
+            dockerContainerPublicId=docker_container_public_id,
+            lastLogTimestamp=last_log_timestamp,
+            lastLogId=last_log_id
+        )
+
+        async with httpx.AsyncClient() as client:
+            response = await client.post(
+                url=f"{self._get_host()}/user-api/v2/logging/poll",
+                headers=request_headers,
+                json=request.model_dump(),
+                timeout=3.5
+            )
+
+            return LogPollingResponse.model_validate(response.json()) if response else None
+
+
+    def query_user_logs(self, limit: int, task_public_id: Optional[str] = None,
+                        inference_simulator_public_id: Optional[str] = None,
+                        container_public_id: Optional[str] = None) -> UserLogsQueryResponse:
+        request = UserLogsQueryRequest(
+            inferenceSimulatorPublicId=inference_simulator_public_id,
+            taskPublicId=task_public_id,
+            containerPublicId=container_public_id,
+            limit=limit,
+            ascendingOrder=False,
+        )
+
+        response = self._request(
+            method='POST',
+            url='/user-api/v2/logging/query-user-logs',
+            data=request.model_dump(),
+            token=self.__config_provider.get_config().main.thestage_auth_token,
+        )
+
+        result = UserLogsQueryResponse.model_validate(response) if response else None
+        return result