unit test

Signed-off-by: Andrew Xia <axia@fb.com>

Author: Andrew Xia

tests/entrypoints/openai/test_response_api_parsable_context.py

@@ -165,6 +165,7 @@ async def test_mcp_tool_call(client: OpenAI, model_name: str):
         model=model_name,
         input="What is 13 * 24? Use python to calculate the result.",
         tools=[{"type": "code_interpreter", "container": {"type": "auto"}}],
+        extra_body={"enable_response_messages": True},
         temperature=0.0,
     )
 
@@ -178,3 +179,8 @@ async def test_mcp_tool_call(client: OpenAI, model_name: str):
     # make sure the correct math is in the final output
     assert response.output[3].type == "message"
     assert "312" in response.output[3].content[0].text
+
+    # test raw input_messages / output_messages
+    assert len(response.input_messages) == 1
+    assert len(response.output_messages) == 3
+    assert "312" in response.output_messages[2]["message"]

vllm/entrypoints/context.py

@@ -182,21 +182,23 @@ def append_output(self, output) -> None:
         self.num_cached_tokens = output.num_cached_tokens or 0
         self.num_output_tokens += len(output.outputs[0].token_ids or [])
 
-        if len(self.input_messages) == 0:
-            output_prompt = output.prompt or ""
-            output_prompt_token_ids = output.prompt_token_ids or []
-            self.input_messages.append(
+        # only store if enable_response_messages is True, save memory
+        if self.request.enable_response_messages:
+            if len(self.input_messages) == 0:
+                output_prompt = output.prompt or ""
+                output_prompt_token_ids = output.prompt_token_ids or []
+                self.input_messages.append(
+                    ResponseRawMessageAndToken(
+                        message=output_prompt,
+                        tokens=output_prompt_token_ids,
+                    )
+                )
+            self.output_messages.append(
                 ResponseRawMessageAndToken(
-                    message=output_prompt,
-                    tokens=output_prompt_token_ids,
+                    message=output.outputs[0].text,
+                    tokens=output.outputs[0].token_ids,
                 )
             )
-        self.output_messages.append(
-            ResponseRawMessageAndToken(
-                message=output.outputs[0].text,
-                tokens=output.outputs[0].token_ids,
-            )
-        )
 
     def append_tool_output(self, output) -> None:
         raise NotImplementedError("Should not be called.")
@@ -274,30 +276,31 @@ def append_output(self, output: RequestOutput) -> None:
         self.num_cached_tokens = output.num_cached_tokens or 0
         self.num_output_tokens += len(output.outputs[0].token_ids or [])
         self.parser.process(output.outputs[0])
-        output_prompt = output.prompt or ""
-        output_prompt_token_ids = output.prompt_token_ids or []
-        if len(self.input_messages) == 0:
-            self.input_messages.append(
-                ResponseRawMessageAndToken(
-                    message=output_prompt,
-                    tokens=output_prompt_token_ids,
+
+        # only store if enable_response_messages is True, save memory
+        if self.request.enable_response_messages:
+            output_prompt = output.prompt or ""
+            output_prompt_token_ids = output.prompt_token_ids or []
+            if len(self.input_messages) == 0:
+                self.input_messages.append(
+                    ResponseRawMessageAndToken(
+                        message=output_prompt,
+                        tokens=output_prompt_token_ids,
+                    )
+                )
+            else:
+                self.output_messages.append(
+                    ResponseRawMessageAndToken(
+                        message=output_prompt,
+                        tokens=output_prompt_token_ids,
+                    )
                 )
-            )
-        else:
-            # TODO: merge them in properly together
-            # TODO: responsesParser doesn't parse kimi k2 sentences correctly
             self.output_messages.append(
                 ResponseRawMessageAndToken(
-                    message=output_prompt,
-                    tokens=output_prompt_token_ids,
+                    message=output.outputs[0].text,
+                    tokens=output.outputs[0].token_ids,
                 )
             )
-        self.output_messages.append(
-            ResponseRawMessageAndToken(
-                message=output.outputs[0].text,
-                tokens=output.outputs[0].token_ids,
-            )
-        )
 
     def append_tool_output(self, output: list[ResponseInputOutputItem]) -> None:
         self.parser.response_messages.extend(output)

vllm/entrypoints/openai/serving_responses.py

@@ -318,8 +318,6 @@ async def create_responses(
         if maybe_validation_error is not None:
             return maybe_validation_error
 
-        fbvscode.set_trace()
-
         # If the engine is dead, raise the engine's DEAD_ERROR.
         # This is required for the streaming case, where we return a
         # success status before we actually start generating text :).