mirror of
https://github.com/NousResearch/hermes-agent.git
synced 2026-07-30 19:09:28 +00:00
fix(relay): merge Anthropic stream usage
Signed-off-by: Alex Fournier <afournier@nvidia.com>
This commit is contained in:
parent
c789bdd38a
commit
e862099dff
2 changed files with 39 additions and 1 deletions
|
|
@ -598,7 +598,12 @@ class AnthropicStreamAccumulator:
|
|||
if key in delta:
|
||||
self._message[key] = delta[key]
|
||||
if "usage" in payload:
|
||||
self._message["usage"] = payload["usage"]
|
||||
usage = payload["usage"]
|
||||
current_usage = self._message.get("usage")
|
||||
if isinstance(current_usage, dict) and isinstance(usage, dict):
|
||||
self._message["usage"] = {**current_usage, **usage}
|
||||
else:
|
||||
self._message["usage"] = usage
|
||||
|
||||
def finalize(self) -> dict[str, Any]:
|
||||
blocks = []
|
||||
|
|
|
|||
|
|
@ -247,6 +247,39 @@ def test_direct_stream_close_reaches_original_provider_resource(monkeypatch):
|
|||
assert provider_stream.closed is True
|
||||
|
||||
|
||||
def test_anthropic_stream_accumulator_merges_terminal_usage():
|
||||
accumulator = relay_llm.AnthropicStreamAccumulator()
|
||||
accumulator.observe({
|
||||
"type": "message_start",
|
||||
"message": {
|
||||
"id": "message-1",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-test",
|
||||
"usage": {
|
||||
"input_tokens": 100,
|
||||
"output_tokens": 1,
|
||||
"cache_creation_input_tokens": 20,
|
||||
"cache_read_input_tokens": 30,
|
||||
},
|
||||
},
|
||||
})
|
||||
accumulator.observe({
|
||||
"type": "message_delta",
|
||||
"delta": {"stop_reason": "end_turn", "stop_sequence": None},
|
||||
"usage": {"output_tokens": 12},
|
||||
})
|
||||
|
||||
response = accumulator.finalize()
|
||||
|
||||
assert response["usage"] == {
|
||||
"input_tokens": 100,
|
||||
"output_tokens": 12,
|
||||
"cache_creation_input_tokens": 20,
|
||||
"cache_read_input_tokens": 30,
|
||||
}
|
||||
|
||||
|
||||
def test_non_stream_preserves_raw_provider_response_identity(relay_turn):
|
||||
_relay, _turn = relay_turn
|
||||
raw_response = SimpleNamespace(model="test-model", content="raw")
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue