fix: Fix mypy errors on trtllm examples (#1277)

Author: tanmayv25

examples/tensorrt_llm/common/utils.py

@@ -21,7 +21,7 @@
 import weakref
 from enum import Enum
 from queue import Queue
-from typing import Callable, Optional, TypedDict, Union
+from typing import Any, Callable, Coroutine, Optional, TypedDict, Union
 
 logger = logging.getLogger(__name__)
 
@@ -54,7 +54,9 @@ class ConversationMessage(TypedDict):
 class ManagedThread(threading.Thread):
     def __init__(
         self,
-        task: Optional[Union[Callable[..., bool], weakref.WeakMethod]],
+        task: Optional[
+            Union[Callable[..., Coroutine[Any, Any, bool]], weakref.WeakMethod]
+        ],
         error_queue: Optional[Queue] = None,
         name: Optional[str] = None,
         loop: Optional[asyncio.AbstractEventLoop] = None,
@@ -74,7 +76,9 @@ def set_loop(self, loop: asyncio.AbstractEventLoop):
 
     def run(self):
         while not self.stop_event.is_set():
-            task: Optional[Union[Callable[..., bool], weakref.WeakMethod]] = self.task
+            task: Optional[
+                Union[Callable[..., Coroutine[Any, Any, bool]], weakref.WeakMethod]
+            ] = self.task
             if isinstance(task, weakref.WeakMethod):
                 task = task()
                 if task is None: