fix(vllm-serve): allow null logprobs in responses

Author: LeonEricsson

tests/test_vllm_client_server.py

@@ -179,6 +179,22 @@ def test_chat(self):
         for seq in completion_ids:
             assert all(isinstance(tok, int) for tok in seq)
 
+    def test_generate_with_logprobs_none(self):
+        outputs = self.client.generate(["Hello, AI!"], logprobs=None)
+
+        assert isinstance(outputs["prompt_ids"], list)
+        assert isinstance(outputs["completion_ids"], list)
+        assert outputs["logprobs"] is None
+        assert outputs["logprob_token_ids"] is None
+
+    def test_chat_with_logprobs_none(self):
+        outputs = self.client.chat([[{"role": "user", "content": "Hello, AI!"}]], logprobs=None)
+
+        assert isinstance(outputs["prompt_ids"], list)
+        assert isinstance(outputs["completion_ids"], list)
+        assert outputs["logprobs"] is None
+        assert outputs["logprob_token_ids"] is None
+
     def test_generate_with_params(self):
         prompts = ["Hello, AI!", "Tell me a joke"]
         completion_ids = self.client.generate(prompts, n=2, repetition_penalty=0.9, temperature=0.8, max_tokens=32)[

trl/scripts/vllm_serve.py

@@ -506,8 +506,8 @@ class GenerateRequest(BaseModel):
     class GenerateResponse(BaseModel):
         prompt_ids: list[list[int]]
         completion_ids: list[list[int]]
-        logprobs: list[list[list[float]]]
-        logprob_token_ids: list[list[list[int]]]
+        logprobs: list[list[list[float | None]]] | None
+        logprob_token_ids: list[list[list[int]]] | None
 
     @app.post("/generate/", response_model=GenerateResponse)
     async def generate(request: GenerateRequest):
@@ -672,8 +672,8 @@ class ChatRequest(BaseModel):
     class ChatResponse(BaseModel):
         prompt_ids: list[list[int]]
         completion_ids: list[list[int]]
-        logprobs: list[list[list[float]]]
-        logprob_token_ids: list[list[list[int]]]
+        logprobs: list[list[list[float | None]]] | None
+        logprob_token_ids: list[list[list[int]]] | None
 
     @app.post("/chat/", response_model=ChatResponse)
     async def chat(request: ChatRequest):