diff --git a/README.md b/README.md index 7d257af76..4d5d7d9d7 100644 --- a/README.md +++ b/README.md @@ -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: diff --git a/backend/AGENTS.md b/backend/AGENTS.md index 1c32c47bc..77662356a 100644 --- a/backend/AGENTS.md +++ b/backend/AGENTS.md @@ -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/`) diff --git a/backend/packages/harness/deerflow/models/vllm_provider.py b/backend/packages/harness/deerflow/models/vllm_provider.py index d947e1c26..ba75d8191 100644 --- a/backend/packages/harness/deerflow/models/vllm_provider.py +++ b/backend/packages/harness/deerflow/models/vllm_provider.py @@ -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" diff --git a/backend/tests/test_vllm_provider.py b/backend/tests/test_vllm_provider.py index 9e60d446f..fe64a7002 100644 --- a/backend/tests/test_vllm_provider.py +++ b/backend/tests/test_vllm_provider.py @@ -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", + ] diff --git a/config.example.yaml b/config.example.yaml index 505144eb9..a8921e16c 100644 --- a/config.example.yaml +++ b/config.example.yaml @@ -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