mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-10-26 21:26:39 +00:00
fix(client): accumulate incremental streaming token usage (#6939)
* fix(client): accumulate incremental streaming token usage * test(client): cover OpenAI and vLLM stream usage * fix(client): clamp completed usage reconciliation deltas --------- Co-authored-by: haotianqi666 <250384991+haotianqi666@users.noreply.github.com> Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
This commit is contained in:
parent
2464566942
commit
9c74b1529b
@ -2902,6 +2902,13 @@ client.clear_goal("thread-1")
|
||||
|
||||
When continuing a thread, `values` still contains the full conversation state, while `messages-tuple` and `end.usage` cover only the current turn.
|
||||
|
||||
The embedded client's `end.usage` sums usage deltas across streamed chunks,
|
||||
including providers such as native Gemini that report usage more than once per
|
||||
reply. The completed message's usage is reconciled without counting those
|
||||
chunks again; a lower completed-message count does not subtract tokens already
|
||||
reported by chunks. Signed provider corrections in chunks remain supported, as
|
||||
do providers that report usage only at completion.
|
||||
|
||||
Concurrent `upload_files()` calls preserve filenames absent at the start of each
|
||||
call by assigning numeric suffixes and returning the actual saved names. The same
|
||||
rule applies to new Markdown companions. Exact names already present retain
|
||||
|
||||
@ -8,7 +8,7 @@
|
||||
|
||||
- DeerFlow 有**两条并行**的流式路径:**Gateway 路径**(async / HTTP SSE / JSON 序列化)服务浏览器和 IM 渠道;**DeerFlowClient 路径**(sync / in-process / 原生 LangChain 对象)服务 Jupyter、脚本、测试。它们**无法合并**——消费者模型不同。
|
||||
- 两条路径都从 `create_agent()` 工厂出发,核心都是订阅 LangGraph 的 `stream_mode=["values", "messages", "custom"]`。`values` 是节点级 state 快照,`messages` 是 LLM token 级 delta,`custom` 是显式 `StreamWriter` 事件。DeerFlow 内置 custom 事件同时通过 callback dispatch 暴露为 `astream_events(version="v2")` 的 `on_custom_event`,供 AG-UI 等 callback 型消费者使用。**这些接口不是详细程度的梯度,而是独立事件源**,消费者必须订阅自己需要的接口。
|
||||
- 嵌入式 client 为每个 `stream()` 调用维护四个 `set[str]`:`seen_ids` / `streamed_ids` / `counted_usage_ids` / `historical_message_ids`。前三者分别负责去重与用量幂等;最后一个隔离续聊时的历史消息。
|
||||
- 嵌入式 client 按消息 id 分别跟踪 values 去重、已流式发送的文本、已核算的完整用量和历史消息;`streamed_usage_by_id` 累加尚未核算完整消息的 chunk 用量。
|
||||
|
||||
---
|
||||
|
||||
@ -226,8 +226,9 @@ return "".join(chunks.get(last_id, ()))
|
||||
```python
|
||||
seen_ids: set[str] = set() # values 路径内部 dedup
|
||||
streamed_ids: set[str] = set() # messages → values 跨模式 dedup
|
||||
counted_usage_ids: set[str] = set() # usage_metadata 幂等计数
|
||||
counted_usage_ids: set[str] = set() # 完整消息的 usage 已核算
|
||||
historical_message_ids: set[str] = set() # 续聊时隔离先前回合
|
||||
streamed_usage_by_id: dict[str, dict[str, int]] = {} # 同 id 的 chunk 用量之和
|
||||
```
|
||||
|
||||
乍看用途相近,实际每个管**不同的不变式**。
|
||||
@ -236,7 +237,7 @@ historical_message_ids: set[str] = set() # 续聊时隔离先前回合
|
||||
|---|---|---|---|
|
||||
| `seen_ids` | 连续两个 `values` 快照里同一条 message 只生成一个 `messages-tuple` 事件 | values 分支每处理一条消息就加入 | values 分支处理下一条消息前检查 |
|
||||
| `streamed_ids` | 如果一条消息已经通过 `messages` 模式 token 级流过,values 快照到达时**不要**再合成一次完整 `messages-tuple` | messages 分支每发一个 AI/tool 事件就加入 | values 分支看到消息时检查 |
|
||||
| `counted_usage_ids` | 同一个 `usage_metadata` 在 messages 末尾 chunk 和 values 快照的 final AIMessage 里各带一份,**累计总量只算一次** | `_account_usage()` 每次接受 usage 就加入 | `_account_usage()` 每次调用时检查 |
|
||||
| `counted_usage_ids` | 完整 AIMessage 的 usage 只核算一次,已累加的 chunk 部分不再重复计入 | `_account_usage()` 接受非 chunk 的 usage 时加入 | `_account_usage()` 每次调用时检查 |
|
||||
| `historical_message_ids` | 续聊时旧消息不再作为本轮增量事件,也不计入本轮 `end.usage` | values 快照里当前用户消息之前的 id | messages 和 values 分支发射前检查 |
|
||||
|
||||
### 为什么不能只用一个 set
|
||||
@ -253,19 +254,20 @@ sequenceDiagram
|
||||
|
||||
Note over M: 第一个 AI text chunk 到达
|
||||
M->>SS: add(msg_id)
|
||||
Note over M: 最后一个 chunk 带 usage
|
||||
M->>SU: add(msg_id)
|
||||
Note over M: 每个带 usage 的 chunk 都累加到 streamed_usage_by_id
|
||||
Note over V: snapshot 到达,包含同一条 AI message
|
||||
V->>SE: add(msg_id)
|
||||
V->>SS: 查询 → 已存在,跳过文本合成
|
||||
V->>SU: 查询 → 已存在,不重复计数
|
||||
V->>SU: 核算完整 usage 与 chunk 累计之差,add(msg_id)
|
||||
```
|
||||
|
||||
- `seen_ids` **永远在 values 快照到达时**加入,所以它是 "values 已处理" 的标记。一条只出现在 messages 流里的消息(罕见但可能),`seen_ids` 里永远没有它。
|
||||
- `streamed_ids` **在 messages 流的第一个有效事件时**加入。一条只通过 values 快照到达的非 AI 消息(HumanMessage、被 truncate 的 tool 消息),`streamed_ids` 里永远没有它。
|
||||
- `counted_usage_ids` **只在看到非空 `usage_metadata` 时**加入。一条完全没有 usage 的消息(tool message、错误消息)永远不会进去。
|
||||
- `counted_usage_ids` 只在看到完整消息的非空 `usage_metadata` 时加入。带用量的 `AIMessageChunk` 是增量,会逐段累加,不能收到第一段就将该 id 标为完成。
|
||||
|
||||
**集合包含关系**:`counted_usage_ids ⊆ (streamed_ids ∪ seen_ids)` 大致成立,但**不是严格子集**,因为一条消息可以在 messages 模式流完 text 但**在最后那个带 usage 的 chunk 之前**就被 values snapshot 赶上——此时它已经在 `streamed_ids` 里,但还不在 `counted_usage_ids` 里。把它们合并成一个 dict-of-flags 会让这个微妙的时序依赖**从类型系统里消失**,变成注释里的一句话。三个独立的 set 把不变式显式化了:每个 set 名对应一个可以口头回答的问题。
|
||||
`streamed_usage_by_id` 保存每条消息已计入总量的 chunk 用量。完整消息到达时,各字段只计入 `max(完整用量 - chunk 累计, 0)`,再清除该累计值并标记 `counted_usage_ids`。中间件若调低完整消息的用量,不会扣回已经报告的 chunk 用量;这只限制完整消息的核算差值,chunk 内的有符号供应商修正仍按增量累加。重复的完整消息或晚到的同 id chunk 不会再次计数;没有完整快照时,已收到的 chunk 用量仍保留在总量中。没有 id 的消息无法跨模式匹配,维持逐事件计数。
|
||||
|
||||
例如原生 Gemini 适配器把服务端累计值转换成 chunk 增量:先 `10/1/11`,再 `0/3/3`(输入/输出/总量),最终 AIMessage 为 `10/4/14`。客户端必须累加两段并避免快照重复计数,而不是在第一段后忽略该 id。纯用量 chunk 即使没有文本也要计入;适配器返回的输入 token 修正值同样按增量累加。`tests/test_client_stream_usage.py` 用真实 Gemini、vLLM 和 OpenAI 适配器及 LangGraph、仅替换 SDK 网络调用,覆盖增量用量和仅在末尾汇总用量的路径。
|
||||
|
||||
---
|
||||
|
||||
@ -294,11 +296,11 @@ sequenceDiagram
|
||||
|
||||
Note over A: LLM finish_reason=stop,最后一个 chunk 带 usage
|
||||
A-->>C: ("messages", (AIMessageChunk(content="", usage_metadata={...}), meta))
|
||||
C->>C: counted_usage_ids.add(ai-1)<br/>(无文本,不 yield)
|
||||
C->>C: 累加 chunk usage<br/>(无文本,不 yield)
|
||||
|
||||
A-->>C: ("values", {messages: [..., AIMessage(complete)]})
|
||||
C->>C: ai-1 in streamed_ids → 跳过合成
|
||||
C->>C: 捕获 usage (已在 counted_usage_ids,no-op)
|
||||
C->>C: 核算完整 usage 与 chunk 累计之差 (0)<br/>counted_usage_ids.add(ai-1)
|
||||
C-->>U: StreamEvent(type="values", ...)
|
||||
|
||||
C-->>U: StreamEvent(type="end", data={usage:{...}})
|
||||
@ -308,7 +310,7 @@ sequenceDiagram
|
||||
|
||||
1. 用户看到 **35 个 messages-tuple 事件**,跨越约 476ms,每个事件带一个 token delta 和同一个 `id=ai-1`。
|
||||
2. 最后一个 `values` 快照里的 `AIMessage` **不会**再触发一个完整的 `messages-tuple` 事件——因为 `ai-1 in streamed_ids` 跳过了合成。
|
||||
3. `end` 事件里的 `usage` 正好等于那一份 cumulative usage,**不是它的两倍**——`counted_usage_ids` 在 messages 末尾 chunk 上已经吸收了,values 分支的重复访问是 no-op。
|
||||
3. `end` 事件里的 `usage` 正好等于那一份 cumulative usage,**不是它的两倍**:values 分支减去已累加的 chunk 部分,之后再通过 `counted_usage_ids` 排除重复完整消息。
|
||||
4. 消费者拿到的 `content` 是**增量**:"ele" 只包含 3 个字符,不是 "one\ntwo\n...ele"。想要完整文本要按 `id` 累加,`chat()` 已经帮你做了。
|
||||
|
||||
---
|
||||
|
||||
@ -72,7 +72,7 @@ drift.
|
||||
- `"values"` — state snapshot (title, messages, artifacts, summary_text). Always forward `summary_text` (current summary or `None`), including unchanged values/resets. Never re-emit AI text delivered via `messages`; serialized `ToolMessage` entries retain non-`None` native `artifact`
|
||||
- `"messages-tuple"` — current-turn AI text **deltas** by `id` and each tool call/result once; excludes resumed history and preserves result `artifact`
|
||||
- `"custom"` — forwarded from `StreamWriter`; DeerFlow-built-in custom events are dual-emitted through `deerflow.utils.custom_events`, so `astream_events(version="v2")` consumers also receive one `on_custom_event` with `name=payload["type"]` and the unchanged payload as `data`
|
||||
- `"end"` — current-turn cumulative `usage`, counted once per message id
|
||||
- `"end"` — turn `usage`: signed chunks + positive full-message remainders, deduped by id
|
||||
- **Custom-event invariant** — use `emit_custom_event` / `aemit_custom_event`, never `StreamWriter` alone. Built-in payloads require a non-empty string `type`; typeless payloads stay writer-only, absent from `astream_events`. The writer runs first and is authoritative for Gateway/Web UI/embedded clients; best-effort callbacks must not break it. Async graph hooks must await the async helper, never dispatch synchronously on a running event loop.
|
||||
- Lazy graph creation uses `create_agent()` + `build_middlewares()`.
|
||||
- Cache graphs by storage `user_id` and the unordered set of named-agent `mcp_plugins`. `stream()` materializes `user_id` before worker/loop boundaries in every auth mode.
|
||||
|
||||
@ -1088,39 +1088,41 @@ class DeerFlowClient:
|
||||
# arguments only parse once the message is complete, so their
|
||||
# tool_calls event is emitted from the values snapshot instead.
|
||||
pending_tool_call_ids: set[str] = set()
|
||||
# The same message id carries identical cumulative ``usage_metadata``
|
||||
# in both the final ``messages`` chunk and the values snapshot —
|
||||
# count it only on whichever arrives first.
|
||||
# Chunks carry additive usage; complete messages carry the total.
|
||||
# Reconcile that total once, without counting the streamed part again.
|
||||
streamed_usage_by_id: dict[str, dict[str, int]] = {}
|
||||
counted_usage_ids: set[str] = set()
|
||||
sent_additional_kwargs_by_id: dict[str, dict[str, Any]] = {}
|
||||
cumulative_usage: dict[str, int] = {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0}
|
||||
|
||||
def _account_usage(msg_id: str | None, usage: Any) -> dict | None:
|
||||
"""Add *usage* to cumulative totals if this id has not been counted.
|
||||
def _account_usage(msg_id: str | None, usage: Any, *, is_chunk: bool = False) -> dict | None:
|
||||
"""Accumulate chunk deltas or reconcile a complete message once.
|
||||
|
||||
``usage`` is a ``langchain_core.messages.UsageMetadata`` TypedDict
|
||||
or ``None``; typed as ``Any`` because TypedDicts are not
|
||||
structurally assignable to plain ``dict`` under strict type
|
||||
checking. Returns the normalized usage dict (for attaching
|
||||
checking. Returns the newly accounted usage (for attaching
|
||||
to an event) when we accepted it, otherwise ``None``.
|
||||
"""
|
||||
if not usage:
|
||||
return None
|
||||
if msg_id and msg_id in counted_usage_ids:
|
||||
return None
|
||||
delta = {key: usage.get(key, 0) or 0 for key in cumulative_usage}
|
||||
if msg_id:
|
||||
counted_usage_ids.add(msg_id)
|
||||
input_tokens = usage.get("input_tokens", 0) or 0
|
||||
output_tokens = usage.get("output_tokens", 0) or 0
|
||||
total_tokens = usage.get("total_tokens", 0) or 0
|
||||
cumulative_usage["input_tokens"] += input_tokens
|
||||
cumulative_usage["output_tokens"] += output_tokens
|
||||
cumulative_usage["total_tokens"] += total_tokens
|
||||
return {
|
||||
"input_tokens": input_tokens,
|
||||
"output_tokens": output_tokens,
|
||||
"total_tokens": total_tokens,
|
||||
}
|
||||
if is_chunk:
|
||||
streamed = streamed_usage_by_id.setdefault(msg_id, {})
|
||||
for key, value in delta.items():
|
||||
streamed[key] = streamed.get(key, 0) + value
|
||||
else:
|
||||
counted_usage_ids.add(msg_id)
|
||||
streamed = streamed_usage_by_id.pop(msg_id, {})
|
||||
# State rewrites may lower totals; do not retract streamed usage.
|
||||
# Signed provider corrections remain valid in the chunk path.
|
||||
delta = {key: max(value - streamed.get(key, 0), 0) for key, value in delta.items()}
|
||||
for key, value in delta.items():
|
||||
cumulative_usage[key] += value
|
||||
return delta
|
||||
|
||||
def _unsent_additional_kwargs(msg_id: str | None, additional_kwargs: dict[str, Any] | None) -> dict[str, Any] | None:
|
||||
if not additional_kwargs:
|
||||
@ -1167,7 +1169,7 @@ class DeerFlowClient:
|
||||
if isinstance(msg_chunk, AIMessage):
|
||||
text = self._extract_text(msg_chunk.content)
|
||||
additional_kwargs = self._serialize_additional_kwargs(msg_chunk)
|
||||
counted_usage = _account_usage(msg_id, msg_chunk.usage_metadata)
|
||||
counted_usage = _account_usage(msg_id, msg_chunk.usage_metadata, is_chunk=isinstance(msg_chunk, AIMessageChunk))
|
||||
sent_additional_kwargs = False
|
||||
|
||||
if text:
|
||||
|
||||
220
backend/tests/test_client_stream_usage.py
Normal file
220
backend/tests/test_client_stream_usage.py
Normal file
@ -0,0 +1,220 @@
|
||||
"""Embedded stream usage must match LangChain's assembled message usage."""
|
||||
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from langchain_core.messages import AIMessage, AIMessageChunk, HumanMessage
|
||||
from test_client import client as client
|
||||
from test_client import mock_app_config as mock_app_config
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("last_text", "initial_input"), [(" world", 10), ("", 10), (" world", 12)])
|
||||
def test_stream_accounts_native_gemini_usage(client, last_text, initial_input):
|
||||
from google.genai.types import Candidate, Content, GenerateContentResponse, GenerateContentResponseUsageMetadata, Part
|
||||
from langchain.agents import create_agent
|
||||
from langchain_google_genai import ChatGoogleGenerativeAI
|
||||
|
||||
model = ChatGoogleGenerativeAI(model="gemini-2.5-flash", api_key="offline-placeholder")
|
||||
responses = [
|
||||
GenerateContentResponse(
|
||||
candidates=[Candidate(content=Content(role="model", parts=[Part(text=text)]), finish_reason="STOP" if index else None)],
|
||||
usage_metadata=GenerateContentResponseUsageMetadata(prompt_token_count=inputs, candidates_token_count=outputs, total_token_count=inputs + outputs),
|
||||
)
|
||||
for index, (text, inputs, outputs) in enumerate([("Hello", initial_input, 1), (last_text, 10, 4)])
|
||||
]
|
||||
try:
|
||||
# Only replace the SDK transport; retain the real adapter and graph.
|
||||
with (
|
||||
patch.object(model.client.models, "generate_content_stream", return_value=iter(responses)) as transport,
|
||||
patch.object(client, "_ensure_agent"),
|
||||
patch.object(client, "_agent", create_agent(model, tools=[])),
|
||||
):
|
||||
events = list(client.stream("Say hello", thread_id="gemini-usage"))
|
||||
transport.assert_called_once()
|
||||
finally:
|
||||
model.client.close()
|
||||
|
||||
expected = {"input_tokens": 10, "output_tokens": 4, "total_tokens": 14}
|
||||
final_values = [event for event in events if event.type == "values"][-1]
|
||||
assembled = final_values.data["messages"][-1]
|
||||
assert {key: assembled["usage_metadata"][key] for key in expected} == expected
|
||||
assert assembled["content"] == "Hello" + last_text
|
||||
assert "".join(event.data["content"] for event in events if event.type == "messages-tuple") == assembled["content"]
|
||||
assert events[-1].type == "end"
|
||||
assert events[-1].data["usage"] == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize("provider", ["openai", "vllm"])
|
||||
def test_stream_accounts_openai_compatible_usage(client, provider):
|
||||
from langchain.agents import create_agent
|
||||
from langchain_openai import ChatOpenAI
|
||||
|
||||
from deerflow.models.vllm_provider import VllmChatModel
|
||||
|
||||
model_class = VllmChatModel if provider == "vllm" else ChatOpenAI
|
||||
model = model_class(
|
||||
model="offline-model",
|
||||
api_key="offline-placeholder",
|
||||
use_responses_api=False,
|
||||
stream_usage=True,
|
||||
**({"cumulative_stream_usage": True} if provider == "vllm" else {}),
|
||||
)
|
||||
responses = [
|
||||
{
|
||||
"id": "chatcmpl-usage",
|
||||
"model": "offline-model",
|
||||
"choices": [{"index": 0, "delta": {"role": "assistant", "content": text}, "finish_reason": "stop" if index else None}],
|
||||
"usage": {"prompt_tokens": 10, "completion_tokens": output, "total_tokens": 10 + output} if provider == "vllm" else None,
|
||||
}
|
||||
for index, (text, output) in enumerate([("Hello", 1), (" world", 4)])
|
||||
]
|
||||
responses.append(
|
||||
{
|
||||
"id": "chatcmpl-usage",
|
||||
"model": "offline-model",
|
||||
"choices": [],
|
||||
"usage": {"prompt_tokens": 10, "completion_tokens": 4, "total_tokens": 14},
|
||||
}
|
||||
)
|
||||
response_stream = MagicMock()
|
||||
response_stream.__enter__.return_value = iter(responses)
|
||||
try:
|
||||
with (
|
||||
patch.object(model.client, "create", return_value=response_stream) as transport,
|
||||
patch.object(client, "_ensure_agent"),
|
||||
patch.object(client, "_agent", create_agent(model, tools=[])),
|
||||
):
|
||||
events = list(client.stream("Say hello", thread_id=f"{provider}-usage"))
|
||||
transport.assert_called_once()
|
||||
finally:
|
||||
model.root_client.close()
|
||||
|
||||
expected = {"input_tokens": 10, "output_tokens": 4, "total_tokens": 14}
|
||||
assembled = [event for event in events if event.type == "values"][-1].data["messages"][-1]
|
||||
assert {key: assembled["usage_metadata"][key] for key in expected} == expected
|
||||
assert assembled["content"] == "Hello world"
|
||||
assert "".join(event.data["content"] for event in events if event.type == "messages-tuple") == "Hello world"
|
||||
assert events[-1].type == "end"
|
||||
assert events[-1].data["usage"] == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize("include_snapshots", [False, True])
|
||||
def test_stream_sums_usage_deltas_per_message(client, include_snapshots):
|
||||
chunks = [
|
||||
AIMessageChunk(id="a", content="A", usage_metadata={"input_tokens": 10, "output_tokens": 1, "total_tokens": 11}),
|
||||
AIMessageChunk(id="b", content="B", usage_metadata={"input_tokens": 20, "output_tokens": 2, "total_tokens": 22}),
|
||||
AIMessageChunk(id="a", content="C", usage_metadata={"input_tokens": 0, "output_tokens": 3, "total_tokens": 3}),
|
||||
AIMessageChunk(id="b", content="", usage_metadata={"input_tokens": 0, "output_tokens": 4, "total_tokens": 4}),
|
||||
]
|
||||
items = [("messages", (chunk, {})) for chunk in chunks]
|
||||
if include_snapshots:
|
||||
messages = [
|
||||
AIMessage(id="a", content="AC", usage_metadata={"input_tokens": 10, "output_tokens": 4, "total_tokens": 14}),
|
||||
AIMessage(id="b", content="B", usage_metadata={"input_tokens": 20, "output_tokens": 6, "total_tokens": 26}),
|
||||
]
|
||||
items.extend([("values", {"messages": messages}), ("values", {"messages": [message.model_copy() for message in messages]})])
|
||||
agent = MagicMock()
|
||||
agent.stream.return_value = iter(items)
|
||||
with patch.object(client, "_ensure_agent"), patch.object(client, "_agent", agent):
|
||||
events = list(client.stream("hi"))
|
||||
|
||||
assert events[-1].data["usage"] == {"input_tokens": 30, "output_tokens": 10, "total_tokens": 40}
|
||||
text_events = [event for event in events if event.type == "messages-tuple"]
|
||||
assert [event.data["content"] for event in text_events] == ["A", "B", "C"]
|
||||
assert [event.data["usage_metadata"] for event in text_events] == [chunk.usage_metadata for chunk in chunks[:3]]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("snapshot_mode", ["messages", "values"])
|
||||
@pytest.mark.parametrize("snapshot_first", [False, True])
|
||||
def test_stream_reconciles_full_usage_without_double_counting(client, snapshot_mode, snapshot_first):
|
||||
full = AIMessage(id="a", content="Hello", usage_metadata={"input_tokens": 10, "output_tokens": 4, "total_tokens": 14})
|
||||
partial = AIMessageChunk(id="a", content="Hello", usage_metadata={"input_tokens": 10, "output_tokens": 1, "total_tokens": 11})
|
||||
snapshot = (snapshot_mode, (full, {}) if snapshot_mode == "messages" else {"messages": [full]})
|
||||
items = [("messages", (partial, {}))]
|
||||
items.insert(0 if snapshot_first else 1, snapshot)
|
||||
items.append(snapshot)
|
||||
agent = MagicMock()
|
||||
agent.stream.return_value = iter(items)
|
||||
with patch.object(client, "_ensure_agent"), patch.object(client, "_agent", agent):
|
||||
events = list(client.stream("hi"))
|
||||
|
||||
assert events[-1].data["usage"] == full.usage_metadata
|
||||
|
||||
|
||||
def test_resumed_stream_sums_only_current_usage_deltas(client):
|
||||
old = AIMessage(id="old", content="Previous reply", usage_metadata={"input_tokens": 100, "output_tokens": 40, "total_tokens": 140})
|
||||
first = AIMessageChunk(id="new", content="Hello", usage_metadata={"input_tokens": 10, "output_tokens": 1, "total_tokens": 11})
|
||||
last = AIMessageChunk(id="new", content=" again", usage_metadata={"input_tokens": 0, "output_tokens": 3, "total_tokens": 3})
|
||||
full = AIMessage(id="new", content="Hello again", usage_metadata={"input_tokens": 10, "output_tokens": 4, "total_tokens": 14})
|
||||
|
||||
def stream(state, **kwargs):
|
||||
current = state["messages"][0]
|
||||
messages = [HumanMessage(id="previous-user", content="Previous request"), old, current]
|
||||
yield "values", {"messages": messages}
|
||||
yield "messages", (old, {})
|
||||
yield "messages", (first, {})
|
||||
yield "messages", (last, {})
|
||||
yield "values", {"messages": [*messages, full]}
|
||||
|
||||
agent = MagicMock()
|
||||
agent.stream.side_effect = stream
|
||||
with patch.object(client, "_ensure_agent"), patch.object(client, "_agent", agent):
|
||||
events = list(client.stream("hi", thread_id="existing-thread"))
|
||||
|
||||
assert events[-1].data["usage"] == full.usage_metadata
|
||||
assert [event.data["content"] for event in events if event.type == "messages-tuple"] == ["Hello", " again"]
|
||||
|
||||
|
||||
def test_stream_keeps_no_id_usage_deltas(client):
|
||||
agent = MagicMock()
|
||||
agent.stream.return_value = iter(
|
||||
[
|
||||
("messages", (AIMessageChunk(content="Hi", usage_metadata={"input_tokens": 10, "output_tokens": 1, "total_tokens": 11}), {})),
|
||||
("messages", (AIMessageChunk(content="", usage_metadata={"input_tokens": 0, "output_tokens": 3, "total_tokens": 3}), {})),
|
||||
]
|
||||
)
|
||||
with patch.object(client, "_ensure_agent"), patch.object(client, "_agent", agent):
|
||||
events = list(client.stream("hi"))
|
||||
|
||||
assert events[-1].data["usage"] == {"input_tokens": 10, "output_tokens": 4, "total_tokens": 14}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("snapshot_mode", ["messages", "values"])
|
||||
@pytest.mark.parametrize(
|
||||
("final_usage", "expected"),
|
||||
[
|
||||
({"input_tokens": 5, "output_tokens": 2, "total_tokens": 7}, {"input_tokens": 10, "output_tokens": 4, "total_tokens": 14}),
|
||||
({"input_tokens": 0, "output_tokens": 0, "total_tokens": 0}, {"input_tokens": 10, "output_tokens": 4, "total_tokens": 14}),
|
||||
({"input_tokens": 5, "output_tokens": 12, "total_tokens": 17}, {"input_tokens": 10, "output_tokens": 12, "total_tokens": 17}),
|
||||
],
|
||||
ids=["lower", "zero", "mixed"],
|
||||
)
|
||||
def test_stream_snapshot_cannot_subtract_already_accounted_usage(client, snapshot_mode, final_usage, expected):
|
||||
# A middleware may rewrite usage in the completed message. Reconciliation
|
||||
# may add missing tokens, but must not retract usage already reported.
|
||||
partial = AIMessageChunk(id="a", content="Hello", usage_metadata={"input_tokens": 10, "output_tokens": 4, "total_tokens": 14})
|
||||
full = AIMessage(id="a", content="Hello again", usage_metadata=final_usage)
|
||||
snapshot = (snapshot_mode, (full, {}) if snapshot_mode == "messages" else {"messages": [full]})
|
||||
duplicate = full.model_copy()
|
||||
duplicate_snapshot = (snapshot_mode, (duplicate, {}) if snapshot_mode == "messages" else {"messages": [duplicate]})
|
||||
other = AIMessage(id="b", content="Another reply", usage_metadata={"input_tokens": 3, "output_tokens": 1, "total_tokens": 4})
|
||||
agent = MagicMock()
|
||||
agent.stream.return_value = iter(
|
||||
[
|
||||
("messages", (partial, {})),
|
||||
snapshot,
|
||||
duplicate_snapshot,
|
||||
("messages", (partial.model_copy(update={"content": ""}), {})),
|
||||
("messages", (other, {})),
|
||||
]
|
||||
)
|
||||
with patch.object(client, "_ensure_agent"), patch.object(client, "_agent", agent):
|
||||
events = list(client.stream("hi"))
|
||||
|
||||
assert events[-1].data["usage"] == {key: value + other.usage_metadata[key] for key, value in expected.items()}
|
||||
# Raw values remain faithful to state; only accounting is clamped.
|
||||
if snapshot_mode == "values":
|
||||
assert [event for event in events if event.type == "values"][-1].data["messages"][0]["usage_metadata"] == final_usage
|
||||
else:
|
||||
reconciled = next(event for event in events if event.type == "messages-tuple" and event.data["content"] == "Hello again")
|
||||
assert reconciled.data["usage_metadata"] == {key: max(final_usage[key] - partial.usage_metadata[key], 0) for key in expected}
|
||||
Loading…
x
Reference in New Issue
Block a user