feat(models): support cumulative vLLM stream usage (#4537)

* feat(models): support cumulative vLLM stream usage

* fix(models): preserve active cumulative usage streams
This commit is contained in:
阿泽 2026-07-28 19:56:40 +08:00 committed by GitHub
parent 1bb93f9517
commit 94003c1f47
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 436 additions and 2 deletions

View File

@ -181,7 +181,7 @@ That prompt is intended for coding agents. It tells the agent to clone the repo
To route OpenAI models through `/v1/responses`, keep using `langchain_openai:ChatOpenAI` and set `use_responses_api: true` with `output_version: responses/v1`.
For vLLM 0.19.0, use `deerflow.models.vllm_provider:VllmChatModel`. For Qwen-style reasoning models, DeerFlow toggles reasoning with `extra_body.chat_template_kwargs.enable_thinking` and preserves vLLM's non-standard `reasoning` field across multi-turn tool-call conversations. Legacy `thinking` configs are normalized automatically for backward compatibility. Reasoning models may also require the server to be started with `--reasoning-parser ...`. If your local vLLM deployment accepts any non-empty API key, you can still set `VLLM_API_KEY` to a placeholder value.
For vLLM 0.19.0, use `deerflow.models.vllm_provider:VllmChatModel`. For Qwen-style reasoning models, DeerFlow toggles reasoning with `extra_body.chat_template_kwargs.enable_thinking` and preserves vLLM's non-standard `reasoning` field across multi-turn tool-call conversations. Legacy `thinking` configs are normalized automatically for backward compatibility. If the endpoint reports a cumulative usage snapshot on every streaming chunk, set `cumulative_stream_usage: true` so DeerFlow converts those snapshots into per-chunk deltas; the option is disabled by default and leaves usage unchanged when a stable completion id is unavailable. Reasoning models may also require the server to be started with `--reasoning-parser ...`. If your local vLLM deployment accepts any non-empty API key, you can still set `VLLM_API_KEY` to a placeholder value.
CLI-backed provider examples:

View File

@ -727,6 +727,7 @@ Lets a caller pass per-request, short-lived end-user credentials (e.g. an ERP to
- `VllmChatModel` subclasses `langchain_openai:ChatOpenAI` for vLLM 0.19.0 OpenAI-compatible endpoints
- Preserves vLLM's non-standard assistant `reasoning` field on full responses, streaming deltas, and follow-up tool-call turns
- Designed for configs that enable thinking through `extra_body.chat_template_kwargs.enable_thinking` on vLLM 0.19.0 Qwen reasoning models, while accepting the older `thinking` alias
- `cumulative_stream_usage` is an opt-in model setting (default `false`) for endpoints that repeat cumulative token totals on each streaming chunk. The provider converts snapshots to deltas only when a stable completion id is present, isolates interleaved streams by id, and leaves the original usage untouched otherwise. Per-model tracking is lock-protected and cleared on the trailing empty-`choices` frame whether or not that frame carries usage. A soft cap of 1024 ids evicts only entries idle for at least one hour; active streams may temporarily exceed the cap so eviction cannot corrupt their deltas. Regression coverage lives in `tests/test_vllm_provider.py`.
### IM Channels System (`app/channels/`)

View File

@ -15,6 +15,9 @@ This provider preserves ``reasoning`` on:
from __future__ import annotations
import json
import threading
import time
from collections import OrderedDict
from collections.abc import Mapping
from typing import Any, cast
@ -30,10 +33,15 @@ from langchain_core.messages import (
SystemMessageChunk,
ToolMessageChunk,
)
from langchain_core.messages.ai import UsageMetadata, subtract_usage
from langchain_core.messages.tool import tool_call_chunk
from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
from langchain_openai import ChatOpenAI
from langchain_openai.chat_models.base import _create_usage_metadata
from pydantic import Field, PrivateAttr
_CUMULATIVE_USAGE_TRACKER_CAPACITY = 1024
_CUMULATIVE_USAGE_TRACKER_IDLE_SECONDS = 60 * 60
def _normalize_vllm_chat_template_kwargs(payload: dict[str, Any]) -> None:
@ -156,15 +164,59 @@ def _restore_reasoning_field(payload_msg: dict[str, Any], orig_msg: AIMessage) -
payload_msg["reasoning"] = reasoning
def _get_completion_id(chunk: Mapping[str, Any]) -> str | None:
"""Return a stable completion id when the provider supplied one."""
candidates = [chunk.get("id")]
nested_chunk = chunk.get("chunk")
if isinstance(nested_chunk, Mapping):
candidates.append(nested_chunk.get("id"))
for candidate in candidates:
if isinstance(candidate, str) and candidate.strip():
return candidate
return None
class VllmChatModel(ChatOpenAI):
"""ChatOpenAI variant that preserves vLLM reasoning fields across turns."""
model_config = {"arbitrary_types_allowed": True}
cumulative_stream_usage: bool = Field(
default=False,
description=("Treat streaming usage snapshots from vLLM as cumulative per completion and convert them to per-chunk deltas."),
)
_cumulative_usage_by_completion: OrderedDict[str, tuple[UsageMetadata, float]] = PrivateAttr(default_factory=OrderedDict)
_cumulative_usage_lock: Any = PrivateAttr(default_factory=threading.Lock)
@property
def _llm_type(self) -> str:
return "vllm-openai-compatible"
def _usage_delta(self, completion_id: str, usage: UsageMetadata, *, terminal: bool) -> UsageMetadata:
"""Convert a completion's cumulative usage snapshot into a delta."""
with self._cumulative_usage_lock:
previous_snapshot = self._cumulative_usage_by_completion.get(completion_id)
previous = previous_snapshot[0] if previous_snapshot is not None else None
delta = subtract_usage(usage, previous)
if terminal:
self._cumulative_usage_by_completion.pop(completion_id, None)
else:
now = time.monotonic()
self._cumulative_usage_by_completion[completion_id] = (usage, now)
self._cumulative_usage_by_completion.move_to_end(completion_id)
while len(self._cumulative_usage_by_completion) > _CUMULATIVE_USAGE_TRACKER_CAPACITY:
oldest_completion_id, (_, updated_at) = next(iter(self._cumulative_usage_by_completion.items()))
if now - updated_at < _CUMULATIVE_USAGE_TRACKER_IDLE_SECONDS:
break
self._cumulative_usage_by_completion.pop(oldest_completion_id, None)
return delta
def _clear_usage_snapshot(self, completion_id: str) -> None:
"""Forget a completed stream even when its terminal frame has no usage."""
with self._cumulative_usage_lock:
self._cumulative_usage_by_completion.pop(completion_id, None)
def _get_request_payload(
self,
input_: LanguageModelInput,
@ -224,9 +276,15 @@ class VllmChatModel(ChatOpenAI):
token_usage = chunk.get("usage")
choices = chunk.get("choices", []) or chunk.get("chunk", {}).get("choices", [])
usage_metadata = _create_usage_metadata(token_usage, chunk.get("service_tier")) if token_usage else None
completion_id = _get_completion_id(chunk) if self.cumulative_stream_usage else None
if len(choices) == 0:
generation_chunk = ChatGenerationChunk(message=default_chunk_class(content="", usage_metadata=usage_metadata), generation_info=base_generation_info)
if completion_id is not None:
if usage_metadata is not None and isinstance(generation_chunk.message, AIMessageChunk):
generation_chunk.message.usage_metadata = self._usage_delta(completion_id, usage_metadata, terminal=True)
else:
self._clear_usage_snapshot(completion_id)
if self.output_version == "v1":
generation_chunk.message.content = []
generation_chunk.message.response_metadata["output_version"] = "v1"
@ -252,6 +310,8 @@ class VllmChatModel(ChatOpenAI):
generation_info["logprobs"] = logprobs
if usage_metadata and isinstance(message_chunk, AIMessageChunk):
if completion_id is not None:
usage_metadata = self._usage_delta(completion_id, usage_metadata, terminal=False)
message_chunk.usage_metadata = usage_metadata
message_chunk.response_metadata["model_provider"] = "openai"

View File

@ -5,14 +5,58 @@ from langchain_core.messages import AIMessage, AIMessageChunk, HumanMessage
from deerflow.models.vllm_provider import VllmChatModel
def _make_model() -> VllmChatModel:
def _make_model(*, cumulative_stream_usage: bool = False) -> VllmChatModel:
return VllmChatModel(
model="Qwen/QwQ-32B",
api_key="dummy",
base_url="http://localhost:8000/v1",
cumulative_stream_usage=cumulative_stream_usage,
)
def _stream_chunk(
*,
completion_id: str | None,
prompt_tokens: int,
completion_tokens: int,
content: str = "",
reasoning: str | None = None,
finish_reason: str | None = None,
) -> dict:
delta = {"role": "assistant", "content": content}
if reasoning is not None:
delta["reasoning"] = reasoning
chunk = {
"model": "Qwen/QwQ-32B",
"choices": [
{
"delta": delta,
"finish_reason": finish_reason,
}
],
"usage": {
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
"total_tokens": prompt_tokens + completion_tokens,
},
}
if completion_id is not None:
chunk["id"] = completion_id
return chunk
def _convert_stream_chunk(model: VllmChatModel, chunk: dict):
return model._convert_chunk_to_generation_chunk(chunk, AIMessageChunk, {})
def _assert_usage(message, *, input_tokens: int, output_tokens: int, total_tokens: int) -> None:
assert message.usage_metadata is not None
assert message.usage_metadata["input_tokens"] == input_tokens
assert message.usage_metadata["output_tokens"] == output_tokens
assert message.usage_metadata["total_tokens"] == total_tokens
def test_vllm_provider_restores_reasoning_in_request_payload():
model = _make_model()
payload = model._get_request_payload(
@ -136,3 +180,330 @@ def test_vllm_provider_preserves_empty_reasoning_values_in_streaming_chunks():
assert chunk.message.additional_kwargs["reasoning"] == ""
assert "reasoning_content" not in chunk.message.additional_kwargs
assert chunk.message.content == "Still replying..."
def test_vllm_provider_converts_cumulative_stream_usage_to_deltas():
model = _make_model(cumulative_stream_usage=True)
first = _convert_stream_chunk(
model,
_stream_chunk(
completion_id="chatcmpl-1",
prompt_tokens=10,
completion_tokens=1,
content="A",
reasoning="Inspect the evidence.",
),
)
second = _convert_stream_chunk(
model,
_stream_chunk(
completion_id="chatcmpl-1",
prompt_tokens=10,
completion_tokens=3,
content="B",
finish_reason="stop",
),
)
terminal = _convert_stream_chunk(
model,
{
"id": "chatcmpl-1",
"model": "Qwen/QwQ-32B",
"choices": [],
"usage": {
"prompt_tokens": 10,
"completion_tokens": 3,
"total_tokens": 13,
},
},
)
assert first is not None
_assert_usage(first.message, input_tokens=10, output_tokens=1, total_tokens=11)
assert second is not None
_assert_usage(second.message, input_tokens=0, output_tokens=2, total_tokens=2)
assert terminal is not None
_assert_usage(terminal.message, input_tokens=0, output_tokens=0, total_tokens=0)
combined = first + second + terminal
_assert_usage(combined.message, input_tokens=10, output_tokens=3, total_tokens=13)
assert combined.message.content == "AB"
assert combined.message.additional_kwargs["reasoning"] == "Inspect the evidence."
assert not model._cumulative_usage_by_completion
def test_vllm_provider_leaves_standard_usage_only_terminal_frame_unchanged():
model = _make_model(cumulative_stream_usage=True)
terminal = _convert_stream_chunk(
model,
{
"id": "chatcmpl-terminal",
"model": "Qwen/QwQ-32B",
"choices": [],
"usage": {
"prompt_tokens": 10,
"completion_tokens": 4,
"total_tokens": 14,
},
},
)
assert terminal is not None
_assert_usage(terminal.message, input_tokens=10, output_tokens=4, total_tokens=14)
assert not model._cumulative_usage_by_completion
def test_vllm_provider_clears_snapshot_on_terminal_frame_without_usage():
model = _make_model(cumulative_stream_usage=True)
_convert_stream_chunk(
model,
_stream_chunk(
completion_id="chatcmpl-no-terminal-usage",
prompt_tokens=10,
completion_tokens=2,
),
)
terminal = _convert_stream_chunk(
model,
{
"id": "chatcmpl-no-terminal-usage",
"model": "Qwen/QwQ-32B",
"choices": [],
},
)
assert terminal is not None
assert terminal.message.usage_metadata is None
assert not model._cumulative_usage_by_completion
def test_vllm_provider_does_not_advance_usage_for_discarded_null_delta():
model = _make_model(cumulative_stream_usage=True)
first = _convert_stream_chunk(
model,
_stream_chunk(
completion_id="chatcmpl-null-delta",
prompt_tokens=10,
completion_tokens=1,
),
)
discarded = _convert_stream_chunk(
model,
{
"id": "chatcmpl-null-delta",
"model": "Qwen/QwQ-32B",
"choices": [
{
"delta": None,
"finish_reason": None,
}
],
"usage": {
"prompt_tokens": 10,
"completion_tokens": 2,
"total_tokens": 12,
},
},
)
next_chunk = _convert_stream_chunk(
model,
_stream_chunk(
completion_id="chatcmpl-null-delta",
prompt_tokens=10,
completion_tokens=3,
),
)
assert first is not None
assert discarded is None
assert next_chunk is not None
_assert_usage(next_chunk.message, input_tokens=0, output_tokens=2, total_tokens=2)
def test_vllm_provider_tracks_concurrent_streams_by_completion_id():
model = _make_model(cumulative_stream_usage=True)
stream_a_first = _convert_stream_chunk(
model,
_stream_chunk(
completion_id="chatcmpl-a",
prompt_tokens=10,
completion_tokens=1,
),
)
stream_b_first = _convert_stream_chunk(
model,
_stream_chunk(
completion_id="chatcmpl-b",
prompt_tokens=20,
completion_tokens=2,
),
)
stream_a_second = _convert_stream_chunk(
model,
_stream_chunk(
completion_id="chatcmpl-a",
prompt_tokens=10,
completion_tokens=5,
),
)
assert stream_a_first is not None
assert stream_a_first.message.usage_metadata["total_tokens"] == 11
assert stream_b_first is not None
assert stream_b_first.message.usage_metadata["total_tokens"] == 22
assert stream_a_second is not None
_assert_usage(stream_a_second.message, input_tokens=0, output_tokens=4, total_tokens=4)
def test_vllm_provider_preserves_reasoning_when_converting_cumulative_usage():
model = _make_model(cumulative_stream_usage=True)
chunk = _convert_stream_chunk(
model,
_stream_chunk(
completion_id="chatcmpl-reasoning",
prompt_tokens=12,
completion_tokens=2,
content="Answer",
reasoning="Check the evidence first.",
),
)
assert chunk is not None
assert chunk.message.additional_kwargs["reasoning"] == "Check the evidence first."
assert chunk.message.additional_kwargs["reasoning_content"] == "Check the evidence first."
assert chunk.message.content == "Answer"
_assert_usage(chunk.message, input_tokens=12, output_tokens=2, total_tokens=14)
def test_vllm_provider_leaves_cumulative_usage_unchanged_by_default():
model = _make_model()
_convert_stream_chunk(
model,
_stream_chunk(
completion_id="chatcmpl-default",
prompt_tokens=10,
completion_tokens=1,
),
)
second = _convert_stream_chunk(
model,
_stream_chunk(
completion_id="chatcmpl-default",
prompt_tokens=10,
completion_tokens=3,
),
)
assert second is not None
_assert_usage(second.message, input_tokens=10, output_tokens=3, total_tokens=13)
assert model.cumulative_stream_usage is False
assert not model._cumulative_usage_by_completion
def test_vllm_provider_leaves_usage_unchanged_without_stable_completion_id():
model = _make_model(cumulative_stream_usage=True)
_convert_stream_chunk(
model,
_stream_chunk(
completion_id=None,
prompt_tokens=10,
completion_tokens=1,
),
)
second = _convert_stream_chunk(
model,
_stream_chunk(
completion_id=None,
prompt_tokens=10,
completion_tokens=3,
),
)
assert second is not None
_assert_usage(second.message, input_tokens=10, output_tokens=3, total_tokens=13)
assert not model._cumulative_usage_by_completion
def test_vllm_provider_does_not_evict_active_streams_at_soft_capacity(monkeypatch):
monkeypatch.setattr(
"deerflow.models.vllm_provider._CUMULATIVE_USAGE_TRACKER_CAPACITY",
2,
)
now = [0.0]
monkeypatch.setattr("deerflow.models.vllm_provider.time.monotonic", lambda: now[0])
model = _make_model(cumulative_stream_usage=True)
for index in range(3):
_convert_stream_chunk(
model,
_stream_chunk(
completion_id=f"chatcmpl-{index}",
prompt_tokens=10,
completion_tokens=1,
),
)
assert list(model._cumulative_usage_by_completion) == [
"chatcmpl-0",
"chatcmpl-1",
"chatcmpl-2",
]
now[0] = 1.0
next_chunk = _convert_stream_chunk(
model,
_stream_chunk(
completion_id="chatcmpl-0",
prompt_tokens=10,
completion_tokens=5,
),
)
assert next_chunk is not None
_assert_usage(next_chunk.message, input_tokens=0, output_tokens=4, total_tokens=4)
def test_vllm_provider_evicts_only_idle_streams_above_soft_capacity(monkeypatch):
monkeypatch.setattr(
"deerflow.models.vllm_provider._CUMULATIVE_USAGE_TRACKER_CAPACITY",
2,
)
monkeypatch.setattr(
"deerflow.models.vllm_provider._CUMULATIVE_USAGE_TRACKER_IDLE_SECONDS",
10,
)
now = [0.0]
monkeypatch.setattr("deerflow.models.vllm_provider.time.monotonic", lambda: now[0])
model = _make_model(cumulative_stream_usage=True)
for index in range(3):
_convert_stream_chunk(
model,
_stream_chunk(
completion_id=f"chatcmpl-{index}",
prompt_tokens=10,
completion_tokens=1,
),
)
now[0] = 11.0
_convert_stream_chunk(
model,
_stream_chunk(
completion_id="chatcmpl-3",
prompt_tokens=10,
completion_tokens=1,
),
)
assert list(model._cumulative_usage_by_completion) == [
"chatcmpl-2",
"chatcmpl-3",
]

View File

@ -607,6 +607,8 @@ models:
# model: Qwen/Qwen3-32B
# api_key: $VLLM_API_KEY
# base_url: http://localhost:8000/v1
# # Enable only when the endpoint reports cumulative usage on every stream chunk.
# cumulative_stream_usage: true
# request_timeout: 600.0
# max_retries: 2
# max_tokens: 8192