feat: refine token usage display modes (#2329)

* 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>
2026-05-21 07:26:50 +00:00 · 2026-05-04 09:56:16 +08:00
parent 82e7936d36
commit d02f762ab0
20 changed files with 2346 additions and 222 deletions
@@ -0,0 +1,53 @@
+"""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"}]},
+    }