Skip to content

fix(llm): preserve all vLLM stream chunks - #2234

Merged
endxxxx merged 3 commits into
MemTensor:dev-v2.0.30from
JiataiWang:fix/vllm-stream-chunks
Aug 13, 2026
Merged

fix(llm): preserve all vLLM stream chunks#2234
endxxxx merged 3 commits into
MemTensor:dev-v2.0.30from
JiataiWang:fix/vllm-stream-chunks

Conversation

@JiataiWang

@JiataiWang JiataiWang commented Aug 10, 2026

Copy link
Copy Markdown

Description

VLLMLLM.generate_stream() handled delta.content after the stream loop, so it only saw the final delta. Responses split across several content chunks were truncated, and an empty or choices-free stream could raise UnboundLocalError because delta was never assigned.

This moves content handling into the loop and closes an open <think> block when a reasoning-only stream ends. It keeps the existing remove_think_prefix behavior unchanged and adds focused regression coverage for multi-chunk content, empty/choices-free streams, reasoning-only streams, and output with think tags disabled.

No dependencies are added. I couldn't find an existing issue or open PR for this path.

Related Issue (Required): N/A

Type of change

  • Bug fix (non-breaking change which fixes an issue)

How Has This Been Tested?

  • Unit Test
  • Test Script Or Test Steps

Before the fix, the new regression file produced 2 failures and 2 errors. On the pushed commit:

  • uv run --frozen --with 'pytest==8.3.5' pytest tests/llms/test_vllm.py -q — 4 passed
  • uv run --frozen --with 'pytest==8.3.5' pytest tests/llms/ --ignore=tests/llms/test_hf.py -q — 18 passed (4 existing warnings)
  • uv run --frozen --with 'poetry>=2,<3' --with 'ruff==0.11.13' make format — Ruff passed; 617 files unchanged
  • git diff --check origin/dev-v2.0.29...HEAD — passed

The complete tests/llms/ collection was not available locally because test_hf.py imports the optional torch dependency.

Checklist

  • I have performed a self-review of my own code | 我已自行检查了自己的代码
  • I have commented my code in hard-to-understand areas (no hard-to-understand code added) | 我已在难以理解的地方对代码进行了注释
  • I have added tests that prove my fix is effective or that my feature works | 我已添加测试以证明我的修复有效或功能正常
  • I have created related documentation issue/PR in MemOS-Docs (not applicable; no documentation behavior changed) | 我已在 MemOS-Docs 中创建了相关的文档 issue/PR(如适用)
  • I have linked the issue to this PR (not applicable; no existing issue found) | 我已将 issue 链接到此 PR(如适用)
  • I have mentioned the person who will review this PR | 我已提及将审查此 PR 的人

@WeiminLee when you have a moment, could you take a look?

Reviewer Checklist

  • closes #xxxx (Replace xxxx with the GitHub issue number)
  • Made sure Checks passed
  • Tests have been provided

@Memtensor-AI Memtensor-AI added area:model llm + embedder + reranker status:in-progress Someone or AI is working on it | 人工或 AI 正在处理 labels Aug 10, 2026
@Memtensor-AI
Memtensor-AI requested a review from endxxxx August 10, 2026 08:49
@Memtensor-AI

Memtensor-AI commented Aug 10, 2026

Copy link
Copy Markdown
Collaborator

🤖 Open Code Review

Target: PR #2234
Task: b3f6befecbe91872
Base: dev-v2.0.30
Head: fix/vllm-stream-chunks

🔍 OpenCodeReview found 3 issue(s) in this PR.


1. tests/llms/test_vllm.py (L25-L27)

The test directly assigns a MagicMock to self.llm.client.chat.completions.create on a real openai.Client instance. This bypasses proper isolation: the mock is never cleaned up, and if openai.Client changes how chat.completions is exposed (e.g., via a descriptor or __slots__), this silent assignment could silently no-op, making the test call the real API and fail non-deterministically.

Prefer unittest.mock.patch.object (or unittest.mock.patch) as a context manager or decorator to ensure the mock is applied and removed correctly.

💡 Suggested Change

Before:

    def generate_chunks(self, stream_chunks):
        self.llm.client.chat.completions.create = MagicMock(return_value=iter(stream_chunks))
        return list(self.llm.generate_stream(self.messages))

After:

    def generate_chunks(self, stream_chunks):
        with unittest.mock.patch.object(
            self.llm.client.chat.completions,
            "create",
            return_value=iter(stream_chunks),
        ):
            return list(self.llm.generate_stream(self.messages))

2. tests/llms/test_vllm.py (L63-L64)

self.llm.config.remove_think_prefix is mutated directly on the shared instance without being reset afterward. While setUp is called before each test method, this mutation is not cleaned up between the subtests inside test_remove_think_prefix_omits_tags. If a future subtest is added after the mutation that expects remove_think_prefix=False, it will silently use the wrong value.

Use self.addCleanup or reset the flag explicitly at the start of the method to make the intent explicit.

💡 Suggested Change

Before:

    def test_remove_think_prefix_omits_tags(self):
        self.llm.config.remove_think_prefix = True

After:

    def test_remove_think_prefix_omits_tags(self):
        self.llm.config.remove_think_prefix = True
        self.addCleanup(setattr, self.llm.config, "remove_think_prefix", False)

3. tests/llms/test_vllm.py (L42-L46)

There is no test case for a stream chunk whose delta has neither a reasoning nor a content attribute. In the production code (generate_stream), such a chunk is silently skipped (both hasattr checks return False). While this is currently correct behavior, the absence of a test means any future regression (e.g., an AttributeError on an unexpected attribute access) would go undetected. Consider adding a subtest to test_empty_or_choiceless_stream_returns_no_chunks to cover this path.

💡 Suggested Change

Before:

    def test_empty_or_choiceless_stream_returns_no_chunks(self):
        streams = {
            "empty": [],
            "choices_empty": [SimpleNamespace(choices=[])],
        }

After:

    def test_empty_or_choiceless_stream_returns_no_chunks(self):
        streams = {
            "empty": [],
            "choices_empty": [SimpleNamespace(choices=[])],
            "delta_no_fields": [self.make_chunk()],  # delta has neither reasoning nor content
        }

Generated by cloud-assistant via Open Code Review.

@Memtensor-AI

Copy link
Copy Markdown
Collaborator

✅ Automated Test Results: PASSED

All tests passed (3/3 executed). memos_python_core/changed-repo-python: 3/3. Duration: 5s [advisory, non-gating] AI-generated tests on branch test/auto-gen-d68ac4a0e84e54cc-20260810165338: 34/34 passed — these do NOT affect the PR verdict; review the branch manually.

Branch: fix/vllm-stream-chunks

@Memtensor-AI Memtensor-AI added status:ready Ready for implementation; waiting for assignee or AI dispatch | 可进入实现,等待认领或派发 status:in-progress Someone or AI is working on it | 人工或 AI 正在处理 and removed status:in-progress Someone or AI is working on it | 人工或 AI 正在处理 status:ready Ready for implementation; waiting for assignee or AI dispatch | 可进入实现,等待认领或派发 labels Aug 10, 2026
@JiataiWang

Copy link
Copy Markdown
Author

Good catch — added both cases in 9ee2711: reasoning plus content with tags disabled, and a reasoning-only stream to make sure no closing tag leaks. The focused file is now 4/4, and the LLM suite (excluding the optional torch test) is 18/18.

@Memtensor-AI

Copy link
Copy Markdown
Collaborator

✅ Automated Test Results: PASSED

All tests passed (4/4 executed). memos_python_core/changed-repo-python: 4/4. Duration: 5s [advisory, non-gating] AI-generated tests on branch test/auto-gen-8a14dde71c250f5a-20260810170421: 37/37 passed — these do NOT affect the PR verdict; review the branch manually.

Branch: fix/vllm-stream-chunks

@Memtensor-AI Memtensor-AI added status:ready Ready for implementation; waiting for assignee or AI dispatch | 可进入实现,等待认领或派发 and removed status:in-progress Someone or AI is working on it | 人工或 AI 正在处理 labels Aug 10, 2026
@JiataiWang

Copy link
Copy Markdown
Author

I double-checked the pushed SHA: the reasoning check on line 183 and the content check on line 189 both have 16 spaces, so they are already sibling branches. The first regression test also sends three content-only chunks after a separate reasoning chunk and asserts that all three are yielded; the new autotest run passed 4/4. This looks like an OCR false positive, so I'm leaving the implementation unchanged.

@endxxxx
endxxxx changed the base branch from dev-v2.0.29 to dev-v2.0.30 August 13, 2026 08:38
@Memtensor-AI Memtensor-AI added status:in-progress Someone or AI is working on it | 人工或 AI 正在处理 and removed status:ready Ready for implementation; waiting for assignee or AI dispatch | 可进入实现,等待认领或派发 labels Aug 13, 2026
@Memtensor-AI

Copy link
Copy Markdown
Collaborator

✅ Automated Test Results: PASSED

All tests passed (4/4 executed). memos_python_core/changed-repo-python: 4/4. Duration: 7s [advisory, non-gating] AI-generated tests on branch test/auto-gen-b3f6befecbe91872-20260813172203: 32/33 passed, 1 failed — these do NOT affect the PR verdict; review the branch manually.

Branch: fix/vllm-stream-chunks

@Memtensor-AI Memtensor-AI added status:ready Ready for implementation; waiting for assignee or AI dispatch | 可进入实现,等待认领或派发 and removed status:in-progress Someone or AI is working on it | 人工或 AI 正在处理 labels Aug 13, 2026
@endxxxx
endxxxx merged commit 15645e9 into MemTensor:dev-v2.0.30 Aug 13, 2026
18 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:model llm + embedder + reranker status:ready Ready for implementation; waiting for assignee or AI dispatch | 可进入实现,等待认领或派发

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants