Fix typing in generate function for AsyncLLMEngine

vllm-project · Dec 14, 2023 · f09c990 · f09c990
1 parent f1c8520
commit f09c990
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 6 deletions.
diff --git a/requirements-dev.txt b/requirements-dev.txt
@@ -1,5 +1,6 @@
 # formatting
 yapf==0.32.0
+toml==0.10.2
 ruff==0.1.5
 
 # type checking

diff --git a/vllm/engine/async_llm_engine.py b/vllm/engine/async_llm_engine.py
@@ -2,7 +2,7 @@
 import time
 from functools import partial
 from typing import (Any, Dict, Iterable, List, Optional, Set, Tuple, Type,
-                    Union)
+                    Union, AsyncIterator)
 
 from vllm.config import ModelConfig
 from vllm.engine.arg_utils import AsyncEngineArgs
@@ -401,11 +401,12 @@ async def add_request(
         return stream
 
     async def generate(
-            self,
-            prompt: Optional[str],
-            sampling_params: SamplingParams,
-            request_id: str,
-            prompt_token_ids: Optional[List[int]] = None) -> RequestOutput:
+        self,
+        prompt: Optional[str],
+        sampling_params: SamplingParams,
+        request_id: str,
+        prompt_token_ids: Optional[List[int]] = None
+    ) -> AsyncIterator[RequestOutput]:
         """Generate outputs for a request.
 
         Generate outputs for a request. This method is a coroutine. It adds the