1313from aignostics .platform import NotFoundException , RunState
1414from aignostics .utils import console , get_logger , get_user_data_directory , sanitize_path
1515
16- from ._service import DownloadProgress , DownloadProgressState , Service
16+ from ._models import DownloadProgress , DownloadProgressState
17+ from ._service import Service
1718from ._utils import (
1819 application_run_status_to_str ,
1920 get_mime_type_for_artifact ,
@@ -903,7 +904,7 @@ def result_download( # noqa: C901, PLR0913, PLR0915, PLR0917
903904 with Live (panel ):
904905 main_task = main_download_progress_ui .add_task (description = "" , total = None , extra_description = "" )
905906
906- def update_progress (progress : DownloadProgress ) -> None :
907+ def update_progress (progress : DownloadProgress ) -> None : # noqa: C901
907908 """Update progress bar for file downloads."""
908909 if progress .run :
909910 panel .title = (
@@ -920,18 +921,43 @@ def update_progress(progress: DownloadProgress) -> None:
920921 panel .subtitle += f", status: { status_text } ({ progress .run .termination_reason } )."
921922 else :
922923 panel .subtitle += f", status: { application_run_status_to_str (progress .run .state )} ."
924+ # Determine the status message based on progress state
925+ if progress .status is DownloadProgressState .DOWNLOADING_INPUT :
926+ status_message = (
927+ f"Downloading input slide { progress .item_index + 1 } of { progress .item_count } "
928+ if progress .item_index is not None and progress .item_count
929+ else "Downloading input slide ..."
930+ )
931+ elif progress .status is DownloadProgressState .DOWNLOADING and progress .total_artifact_index is not None :
932+ status_message = (
933+ f"Downloading artifact { progress .total_artifact_index + 1 } of { progress .total_artifact_count } "
934+ )
935+ else :
936+ status_message = progress .status
937+
923938 main_download_progress_ui .update (
924939 main_task ,
925- description = (
926- progress .status
927- if progress .status is not DownloadProgressState .DOWNLOADING or not progress .total_artifact_index
928- else (
929- f"Downloading artifact { progress .total_artifact_index + 1 } "
930- f"of { progress .total_artifact_count } "
931- )
932- ).ljust (50 ),
940+ description = status_message .ljust (50 ),
933941 )
934- if progress .artifact_path :
942+ # Handle input slide download progress
943+ if progress .status is DownloadProgressState .DOWNLOADING_INPUT and progress .input_slide_path :
944+ task_key = str (progress .input_slide_path .absolute ())
945+ if task_key not in download_tasks :
946+ download_tasks [task_key ] = artifact_download_progress_ui .add_task (
947+ f"{ progress .input_slide_path .name } " .ljust (50 ),
948+ total = progress .input_slide_size ,
949+ extra_description = f"Input from { progress .input_slide_url or 'gs://' } "
950+ if progress .input_slide_url
951+ else "Input slide" ,
952+ )
953+
954+ artifact_download_progress_ui .update (
955+ download_tasks [task_key ],
956+ total = progress .input_slide_size ,
957+ advance = progress .input_slide_downloaded_chunk_size ,
958+ )
959+ # Handle artifact download progress
960+ elif progress .artifact_path :
935961 task_key = str (progress .artifact_path .absolute ())
936962 if task_key not in download_tasks :
937963 download_tasks [task_key ] = artifact_download_progress_ui .add_task (
0 commit comments