joeyutong commented on code in PR #870:
URL: https://github.com/apache/flink-agents/pull/870#discussion_r3592370999
##########
python/flink_agents/api/embedding_models/embedding_model.py:
##########
@@ -120,6 +145,12 @@ def embed(
A list of floating-point numbers representing the embedding vector.
The dimension of the vector depends on the specific embedding
model used.
"""
+ return self.embed_with_usage(text, **kwargs).embeddings
+
+ def embed_with_usage(
Review Comment:
The new setup-level usage API also needs explicit cross-language bridges.
`PythonEmbeddingModelSetup` inherits the Java `embedWithUsage(...)`
implementation even though its `open()` delegates to the Python setup and never
initializes the base Java `connection`; conversely,
`JavaEmbeddingModelSetupImpl` inherits this Python method while its
`connection` remains a resource name. Those calls therefore fail instead of
crossing the language boundary. The connection wrappers currently fall back to
legacy `embed(...)` and return no usage as well. Could we bridge the methods
and result conversion in both setup and connection wrappers, and cover them in
the cross-language tests?
##########
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:
Could we handle the native DashScope response shape here?
[`dashscope.TextEmbedding.call()`
documents](https://www.alibabacloud.com/help/en/model-studio/text-embedding-synchronous-api)
top-level `usage.total_tokens` as the number of input tokens, without an
`input_tokens` or `prompt_tokens` field. The current code therefore produces
`prompt_tokens=0` for the documented response. Since that `total_tokens` value
already represents the embedding input count, could we use it as the
prompt-token fallback and test the actual shape with top-level
`usage={"total_tokens": 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]
