@@ -365,7 +365,9 @@ def get_raft_unique_request_id(self) -> str:
365365 return RaftIdGenerator .new_id ()
366366
367367
368- class GetOngoingTaskInfoOperation (MaintenanceOperation [OngoingTask ]):
368+ class GetOngoingTaskInfoOperation (
369+ MaintenanceOperation [Union [OngoingTask , OngoingTaskGenAi , OngoingTaskEmbeddingsGeneration ]]
370+ ):
369371 """
370372 Operation to retrieve detailed information about a specific ongoing task.
371373 Ongoing tasks include various types of tasks such as replication, ETL, backup, and subscriptions.
@@ -398,14 +400,16 @@ def __init__(self, task_id_or_name: Union[int, str], task_type: OngoingTaskType)
398400
399401 self ._task_type = task_type
400402
401- def get_command (self , conventions : "DocumentConventions" ) -> RavenCommand [OngoingTask ]:
403+ def get_command (
404+ self , conventions : "DocumentConventions"
405+ ) -> RavenCommand [OngoingTask | OngoingTaskGenAi | OngoingTaskEmbeddingsGeneration ]:
402406 if self ._task_name is not None :
403407 return GetOngoingTaskInfoOperation ._GetOngoingTaskInfoCommand (
404408 task_name = self ._task_name , task_type = self ._task_type
405409 )
406410 return GetOngoingTaskInfoOperation ._GetOngoingTaskInfoCommand (task_id = self ._task_id , task_type = self ._task_type )
407411
408- class _GetOngoingTaskInfoCommand (RavenCommand [OngoingTask ]):
412+ class _GetOngoingTaskInfoCommand (RavenCommand [OngoingTask | OngoingTaskGenAi | OngoingTaskEmbeddingsGeneration ]):
409413 def __init__ (
410414 self ,
411415 task_type : OngoingTaskType ,
@@ -432,7 +436,9 @@ def set_response(self, response: Optional[str], from_cache: bool) -> None:
432436 json_dict = json .loads (response )
433437 self .result = self ._deserialize_task (json_dict )
434438
435- def _deserialize_task (self , json_dict : dict ) -> OngoingTask :
439+ def _deserialize_task (
440+ self , json_dict : dict
441+ ) -> OngoingTask | OngoingTaskGenAi | OngoingTaskEmbeddingsGeneration :
436442 """Deserialize the task based on its type."""
437443 if self ._task_type == OngoingTaskType .GEN_AI :
438444 return OngoingTaskGenAi .from_json (json_dict )
0 commit comments