zxs1633079383 commented on code in PR #870:
URL: https://github.com/apache/flink-agents/pull/870#discussion_r3618994795
##########
python/flink_agents/integrations/embedding_models/tongyi_embedding_model.py:
##########
@@ -103,8 +124,25 @@ def embed(
msg = f"DashScope TextEmbedding call failed: {response.message}"
raise RuntimeError(msg)
+ usage = getattr(response, "usage", None)
+ if usage is None and isinstance(response.output, dict):
+ usage = response.output.get("usage")
+ prompt_tokens = _get_usage_value(usage, "input_tokens",
"prompt_tokens")
+ total_tokens = _get_usage_value(usage, "total_tokens")
+ token_usage = None
+ if prompt_tokens is not None or total_tokens is not None:
+ token_usage = EmbeddingTokenUsage(
+ prompt_tokens=int(prompt_tokens or 0),
Review Comment:
Updated in 0dc947f.
Tongyi now falls back from `usage.total_tokens` to `prompt_tokens` when
neither `input_tokens` nor `prompt_tokens` is present. This matches the
documented DashScope top-level response shape, where `total_tokens` is the
embedding input-token count.
The regression test now uses the native shape: top-level
`usage={"total_tokens": 6}` and verifies both prompt and total tokens are
reported as 6.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
