YuJitang
d02f762ab0
* feat: refine token usage display modes * docs: clarify token usage accounting semantics * fix: avoid duplicate subtask debug keys * style: format token usage tests * chore: address token attribution review feedback * Update test_token_usage_middleware.py * Update test_token_usage_middleware.py * chore: simplify token attribution fallback * fix token usage metadata follow-up handling --------- Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
54 lines
1.5 KiB
Python
54 lines
1.5 KiB
Python
"""Tests for DeerFlowClient message serialization helpers."""
|
|
|
|
from langchain_core.messages import AIMessage, HumanMessage
|
|
|
|
from deerflow.client import DeerFlowClient
|
|
|
|
|
|
def test_serialize_ai_message_preserves_additional_kwargs():
|
|
message = AIMessage(
|
|
content="done",
|
|
additional_kwargs={
|
|
"token_usage_attribution": {
|
|
"version": 1,
|
|
"kind": "final_answer",
|
|
"shared_attribution": False,
|
|
"actions": [],
|
|
}
|
|
},
|
|
usage_metadata={"input_tokens": 12, "output_tokens": 3, "total_tokens": 15},
|
|
)
|
|
|
|
serialized = DeerFlowClient._serialize_message(message)
|
|
|
|
assert serialized["type"] == "ai"
|
|
assert serialized["usage_metadata"] == {
|
|
"input_tokens": 12,
|
|
"output_tokens": 3,
|
|
"total_tokens": 15,
|
|
}
|
|
assert serialized["additional_kwargs"] == {
|
|
"token_usage_attribution": {
|
|
"version": 1,
|
|
"kind": "final_answer",
|
|
"shared_attribution": False,
|
|
"actions": [],
|
|
}
|
|
}
|
|
|
|
|
|
def test_serialize_human_message_preserves_additional_kwargs():
|
|
message = HumanMessage(
|
|
content="hello",
|
|
additional_kwargs={"files": [{"name": "diagram.png"}]},
|
|
)
|
|
|
|
serialized = DeerFlowClient._serialize_message(message)
|
|
|
|
assert serialized == {
|
|
"type": "human",
|
|
"content": "hello",
|
|
"id": None,
|
|
"additional_kwargs": {"files": [{"name": "diagram.png"}]},
|
|
}
|