A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 8, 15:15 UTC
[responsesAPI] fix simpleContext streaming output_messages (#34188)
Signed-off-by: Andrew Xia <axia@meta.com> Signed-off-by: Andrew Xia <axia@fb.com> Co-authored-by: Andrew Xia <axia@fb.com>
Andrew Xia committed
Feb 9, 2026 at 22:53 UTC
9608844f96e0e739bead72520b7710f1b6f82b65
3 files changed
+265
-5
tests/entrypoints/test_context.py
+246
@@ -8,6 +8,7 @@ from openai_harmony import Author, Message, Role, StreamState, TextContent
8
9
from vllm.entrypoints.openai.responses.context import (
10
HarmonyContext,
11
+ SimpleContext,
12
StreamingHarmonyContext,
13
TurnMetrics,
14
)
@@ -597,3 +598,248 @@ def test_turn_metrics_copy_and_reset():
598
assert copied_metrics.output_tokens == 20
599
assert copied_metrics.cached_input_tokens == 5
600
assert copied_metrics.tool_output_tokens == 3
601
+
602
+
603
+# ==================== SimpleContext Tests ====================
604
+
605
+
606
+def create_simple_context_output(
607
+ text="",
608
+ token_ids=None,
609
+ prompt="Test prompt",
610
+ prompt_token_ids=None,
611
+ num_cached_tokens=0,
612
+ logprobs=None,
613
+ finished=True,
614
+):
615
+ """Helper to create a RequestOutput with customizable text for
616
+ SimpleContext tests."""
617
+ if token_ids is None:
618
+ token_ids = []
619
+ return RequestOutput(
620
+ request_id="test-id",
621
+ prompt=prompt,
622
+ prompt_token_ids=prompt_token_ids,
623
+ prompt_logprobs=None,
624
+ outputs=[
625
+ CompletionOutput(
626
+ index=0,
627
+ text=text,
628
+ token_ids=token_ids,
629
+ cumulative_logprob=0.0,
630
+ logprobs=logprobs,
631
+ finish_reason=None,
632
+ stop_reason=None,
633
+ )
634
+ ],
635
+ finished=finished,
636
+ num_cached_tokens=num_cached_tokens,
637
+ )
638
+
639
+
640
+def test_simple_context_output_messages_empty():
641
+ """output_messages should be empty before any output is appended."""
642
+ context = SimpleContext()
643
+ assert context.output_messages == []
644
+
645
+
646
+def test_simple_context_output_messages_single_call():
647
+ """Non-streaming: single append_output produces a single output message."""
648
+ context = SimpleContext()
649
+ output = create_simple_context_output(
650
+ text="Hello world",
651
+ token_ids=[10, 20, 30],
652
+ prompt_token_ids=[1, 2, 3],
653
+ )
654
+ context.append_output(output)
655
+
656
+ messages = context.output_messages
657
+ assert len(messages) == 1
658
+ assert messages[0].message == "Hello world"
659
+ assert messages[0].tokens == [10, 20, 30]
660
+ assert messages[0].type == "raw_message_tokens"
661
+
662
+
663
+def test_simple_context_output_messages_streaming_consolidation():
664
+ """Streaming: multiple append_output calls consolidate into one message."""
665
+ context = SimpleContext()
666
+
667
+ # Simulate 3 streaming deltas
668
+ context.append_output(
669
+ create_simple_context_output(
670
+ text="Hello",
671
+ token_ids=[10],
672
+ prompt_token_ids=[1, 2, 3],
673
+ )
674
+ )
675
+ context.append_output(
676
+ create_simple_context_output(
677
+ text=" world",
678
+ token_ids=[20],
679
+ prompt_token_ids=[1, 2, 3],
680
+ )
681
+ )
682
+ context.append_output(
683
+ create_simple_context_output(
684
+ text="!",
685
+ token_ids=[30],
686
+ prompt_token_ids=[1, 2, 3],
687
+ )
688
+ )
689
+
690
+ messages = context.output_messages
691
+ assert len(messages) == 1
692
+ assert messages[0].message == "Hello world!"
693
+ assert messages[0].tokens == [10, 20, 30]
694
+
695
+
696
+def test_simple_context_output_messages_many_deltas():
697
+ """Streaming with many small deltas still produces a single message."""
698
+ context = SimpleContext()
699
+
700
+ words = ["The", " quick", " brown", " fox", " jumps"]
701
+ for i, word in enumerate(words):
702
+ context.append_output(
703
+ create_simple_context_output(
704
+ text=word,
705
+ token_ids=[100 + i],
706
+ prompt_token_ids=[1, 2],
707
+ )
708
+ )
709
+
710
+ messages = context.output_messages
711
+ assert len(messages) == 1
712
+ assert messages[0].message == "The quick brown fox jumps"
713
+ assert messages[0].tokens == [100, 101, 102, 103, 104]
714
+
715
+
716
+def test_simple_context_input_messages():
717
+ """input_messages is populated on the first append_output call."""
718
+ context = SimpleContext()
719
+ assert context.input_messages == []
720
+
721
+ context.append_output(
722
+ create_simple_context_output(
723
+ text="Hi",
724
+ token_ids=[10],
725
+ prompt="My prompt text",
726
+ prompt_token_ids=[1, 2, 3],
727
+ )
728
+ )
729
+
730
+ assert len(context.input_messages) == 1
731
+ assert context.input_messages[0].message == "My prompt text"
732
+ assert context.input_messages[0].tokens == [1, 2, 3]
733
+
734
+ # Second call should not add another input message
735
+ context.append_output(
736
+ create_simple_context_output(
737
+ text=" there",
738
+ token_ids=[20],
739
+ prompt="My prompt text",
740
+ prompt_token_ids=[1, 2, 3],
741
+ )
742
+ )
743
+
744
+ assert len(context.input_messages) == 1
745
+
746
+
747
+def test_simple_context_token_counting():
748
+ """Token counting accumulates across streaming deltas."""
749
+ context = SimpleContext()
750
+
751
+ context.append_output(
752
+ create_simple_context_output(
753
+ text="a",
754
+ token_ids=[10, 11],
755
+ prompt_token_ids=[1, 2, 3, 4, 5],
756
+ num_cached_tokens=2,
757
+ )
758
+ )
759
+ context.append_output(
760
+ create_simple_context_output(
761
+ text="b",
762
+ token_ids=[12],
763
+ prompt_token_ids=[1, 2, 3, 4, 5],
764
+ num_cached_tokens=2,
765
+ )
766
+ )
767
+
768
+ assert context.num_prompt_tokens == 5
769
+ assert context.num_output_tokens == 3 # 2 + 1
770
+ assert context.num_cached_tokens == 2
771
+
772
+
773
+def test_simple_context_final_output():
774
+ """final_output reconstructs accumulated text and token_ids."""
775
+ context = SimpleContext()
776
+
777
+ context.append_output(
778
+ create_simple_context_output(
779
+ text="foo",
780
+ token_ids=[1, 2],
781
+ prompt_token_ids=[10],
782
+ )
783
+ )
784
+ context.append_output(
785
+ create_simple_context_output(
786
+ text="bar",
787
+ token_ids=[3],
788
+ prompt_token_ids=[10],
789
+ )
790
+ )
791
+
792
+ final = context.final_output
793
+ assert final is not None
794
+ assert final.outputs[0].text == "foobar"
795
+ assert final.outputs[0].token_ids == (1, 2, 3)
796
+
797
+
798
+def test_simple_context_output_messages_empty_text_with_tokens():
799
+ """output_messages should be returned when tokens exist even if text is
800
+ empty (e.g. special tokens)."""
801
+ context = SimpleContext()
802
+ context.append_output(
803
+ create_simple_context_output(
804
+ text="",
805
+ token_ids=[99],
806
+ prompt_token_ids=[1],
807
+ )
808
+ )
809
+
810
+ messages = context.output_messages
811
+ assert len(messages) == 1
812
+ assert messages[0].message == ""
813
+ assert messages[0].tokens == [99]
814
+
815
+
816
+def test_simple_context_output_messages_no_mutation():
817
+ """Each call to output_messages returns a fresh list; callers can't
818
+ corrupt internal state."""
819
+ context = SimpleContext()
820
+ context.append_output(
821
+ create_simple_context_output(
822
+ text="hello",
823
+ token_ids=[1],
824
+ prompt_token_ids=[10],
825
+ )
826
+ )
827
+
828
+ msgs1 = context.output_messages
829
+ msgs2 = context.output_messages
830
+ assert msgs1 is not msgs2
831
+ assert msgs1[0].message == msgs2[0].message
832
+
833
+ # Appending more output updates the property
834
+ context.append_output(
835
+ create_simple_context_output(
836
+ text=" world",
837
+ token_ids=[2],
838
+ prompt_token_ids=[10],
839
+ )
840
+ )
841
+
842
+ msgs3 = context.output_messages
843
+ assert len(msgs3) == 1
844
+ assert msgs3[0].message == "hello world"
845
+ assert msgs3[0].tokens == [1, 2]
vllm/benchmarks/datasets.py
+5
@@ -1379,6 +1379,11 @@ def add_dataset_parser(parser: FlexibleArgumentParser):
1379
action="store_true",
1380
help="Disable shuffling of dataset samples for deterministic ordering.",
1381
)
1382
+ parser.add_argument(
1383
+ "--trust-remote-code",
1384
+ action="store_true",
1385
+ help="Trust remote code from HuggingFace.",
1386
+ )
1387
1388
# group for dataset specific arguments
1389
custom_group = parser.add_argument_group("custom dataset options")
vllm/entrypoints/openai/responses/context.py
+14
-5
@@ -182,7 +182,6 @@ class SimpleContext(ConversationContext):
182
self.all_turn_metrics = []
183
184
self.input_messages: list[ResponseRawMessageAndToken] = []
185
- self.output_messages: list[ResponseRawMessageAndToken] = []
185
186
def append_output(self, output) -> None:
187
self.last_output = output
@@ -208,12 +207,22 @@ class SimpleContext(ConversationContext):
207
tokens=output_prompt_token_ids,
208
)
209
)
211
- self.output_messages.append(
210
+
211
+ @property
212
+ def output_messages(self) -> list[ResponseRawMessageAndToken]:
213
+ """Return consolidated output as a single message.
214
+
215
+ In streaming mode, text and tokens are accumulated across many deltas.
216
+ This property returns them as a single entry rather than one per delta.
217
+ """
218
+ if not self._accumulated_text and not self._accumulated_token_ids:
219
+ return []
220
+ return [
221
ResponseRawMessageAndToken(
213
- message=delta_output.text,
214
- tokens=delta_output.token_ids,
222
+ message=self._accumulated_text,
223
+ tokens=list(self._accumulated_token_ids),
224
)
216
- )
225
+ ]
226
227
@property
228
def final_output(self) -> RequestOutput | None: