[BugFix] Typing fixes to RequestOutput.prompt and beam search (#9473)

vllm/engine/protocol.py

@@ -59,7 +59,7 @@ class EngineClient(ABC):
 
     async def beam_search(
         self,
-        prompt: Union[PromptType, List[int]],
+        prompt: Union[str, List[int]],
         request_id: str,
         params: BeamSearchParams,
     ) -> AsyncGenerator[RequestOutput, None]:
@@ -71,9 +71,13 @@ class EngineClient(ABC):
         length_penalty = params.length_penalty
 
         tokenizer = await self.get_tokenizer(lora_request=None)
-        tokenizedPrompt = prompt if isinstance(
-            prompt, list) else tokenizer.encode(prompt)
-        tokenizedLength = len(tokenizedPrompt)
+        if isinstance(prompt, str):
+            tokenized_prompt = tokenizer.encode(prompt)
+            prompt_text = prompt
+        else:
+            tokenized_prompt = prompt
+            prompt_text = None
+        tokenized_length = len(tokenized_prompt)
 
         sort_beams_key = create_sort_beams_key_function(
             tokenizer.eos_token_id, length_penalty)
@@ -81,7 +85,11 @@ class EngineClient(ABC):
         beam_search_params = SamplingParams(logprobs=2 * beam_width,
                                             max_tokens=1,
                                             temperature=temperature)
-        all_beams = [BeamSearchSequence(tokens=tokenizedPrompt, cum_logprob=0)]
+        all_beams = [
+            BeamSearchSequence(tokens=tokenized_prompt,
+                               logprobs=[],
+                               cum_logprob=0)
+        ]
         completed = []
 
         for _ in range(max_tokens):
@@ -114,6 +122,7 @@ class EngineClient(ABC):
                     for token_id, logprob_obj in logprobs.items():
                         new_beam = BeamSearchSequence(
                             tokens=current_beam.tokens + [token_id],
+                            logprobs=current_beam.logprobs + [logprobs],
                             cum_logprob=current_beam.cum_logprob +
                             logprob_obj.logprob)
 
@@ -131,22 +140,22 @@ class EngineClient(ABC):
         best_beams = sorted_completed[:beam_width]
 
         for beam in best_beams:
-            beam.text = tokenizer.decode(beam.tokens[tokenizedLength:])
+            beam.text = tokenizer.decode(beam.tokens[tokenized_length:])
 
         beam_search_output = RequestOutput(
             request_id=request_id,
-            prompt=prompt,
+            prompt=prompt_text,
             outputs=[
                 CompletionOutput(
                     text=beam.text,
                     cumulative_logprob=beam.cum_logprob,
-                    token_ids=beam.tokens,
+                    token_ids=beam.tokens[tokenized_length:],
                     index=i,
-                    logprobs=beam.cum_logprob,
+                    logprobs=beam.logprobs,
                 ) for (i, beam) in enumerate(best_beams)
             ],
             finished=True,
-            prompt_token_ids=tokenizedPrompt,
+            prompt_token_ids=tokenized_prompt,
             prompt_logprobs=None)
 
         yield beam_search_output