mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-05-03 07:18:25 +00:00
33 lines
894 B
Python
33 lines
894 B
Python
from unittest.mock import MagicMock, patch
|
|
|
|
from langchain_core.messages import AIMessage
|
|
|
|
from deerflow.agents.middlewares.token_usage_middleware import TokenUsageMiddleware
|
|
|
|
|
|
def test_after_model_logs_usage_metadata_counts():
|
|
middleware = TokenUsageMiddleware()
|
|
state = {
|
|
"messages": [
|
|
AIMessage(
|
|
content="done",
|
|
usage_metadata={
|
|
"input_tokens": 10,
|
|
"output_tokens": 5,
|
|
"total_tokens": 15,
|
|
},
|
|
)
|
|
]
|
|
}
|
|
|
|
with patch("deerflow.agents.middlewares.token_usage_middleware.logger.info") as info_mock:
|
|
result = middleware.after_model(state=state, runtime=MagicMock())
|
|
|
|
assert result is None
|
|
info_mock.assert_called_once_with(
|
|
"LLM token usage: input=%s output=%s total=%s",
|
|
10,
|
|
5,
|
|
15,
|
|
)
|