weiqingy commented on code in PR #1047:
URL: https://github.com/apache/flink-agents/pull/1047#discussion_r3891277408
##########
api/src/main/java/org/apache/flink/agents/api/embedding/model/BaseEmbeddingModelSetup.java:
##########
@@ -93,9 +94,52 @@ public String getModel() {
return model;
}
+ /**
+ * Record embedding token usage metrics for the given model on this
setup's bound metric group.
+ *
+ * <p>Mirrors {@code BaseChatModelSetup#recordTokenMetrics} but records
input-side tokens only,
+ * since embeddings have no completion tokens. Counters are placed under
the same {@code model}
+ * key-value group used by chat metrics, so embedding and chat usage for a
model share one
+ * dimension.
+ *
+ * <p>Unlike the chat path, embedding calls do not run inside a plan
action that hands in a
+ * request-scoped metric group (vector-store, RAG, and direct calls reach
this setup directly),
Review Comment:
This line says vector-store and RAG calls reach this setup, but they do not
reach `embedWithUsage`, which is where the recording happens. All four call
sites use `embed`: `BaseVectorStore.java:179` (RAG query),
`BaseVectorStore.java:337` (auto-embed on add and update),
`vector_store.py:290`, and `vector_store.py:349`.
I grepped `embedWithUsage` and `embed_with_usage` on `main` and found no
vector-store or RAG call site at all. That matches your own Javadoc on
`embed(String)`, which says usage is discarded there. Those are the paths #858
asks for, so today they would still record nothing.
Is extending to the `embed` paths in scope for this PR? Or would you rather
land the direct-call case first and reword this sentence, and the matching
claim in the PR body, to match what it covers?
One more thing on the line above: it says embedding calls do not run inside
a plan action, but RAG does. `ContextRetrievalAction` is registered as
`context_retrieval_action` at `ContextRetrievalAction.java:44`. And while you
are in the PR body, the sentence about `RowTypeInfo` and "a schema that
renders" looks like it came from a different change. Worth dropping?
##########
api/src/test/java/org/apache/flink/agents/api/embedding/model/BaseEmbeddingModelSetupTokenMetricsTest.java:
##########
@@ -0,0 +1,318 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions of
+ * limitations under the License.
+ */
+
+package org.apache.flink.agents.api.embedding.model;
+
+import org.apache.flink.agents.api.metrics.FlinkAgentsMetricGroup;
+import org.apache.flink.agents.api.metrics.UpdatableGauge;
+import org.apache.flink.agents.api.resource.ResourceContext;
+import org.apache.flink.agents.api.resource.ResourceDescriptor;
+import org.apache.flink.metrics.Counter;
+import org.apache.flink.metrics.Histogram;
+import org.apache.flink.metrics.Meter;
+import org.apache.flink.metrics.SimpleCounter;
+import org.junit.jupiter.api.DisplayName;
+import org.junit.jupiter.api.Test;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.mockito.Mockito.mock;
+
+/**
+ * Test cases for embedding token usage metrics recorded by {@link
BaseEmbeddingModelSetup}. Mirrors
+ * {@code BaseChatModelSetupTokenMetricsTest}: embedding providers already
populate {@link
+ * EmbeddingTokenUsage} on the returned {@link EmbeddingResult}, but nothing
records it until this
+ * setup reads it back at the {@code embedWithUsage} chokepoint.
+ */
+class BaseEmbeddingModelSetupTokenMetricsTest {
+
+ /** Value-based metric group that mirrors the one in the chat
token-metrics test. */
+ private static class TestMetricGroup implements FlinkAgentsMetricGroup {
+ final Map<String, TestMetricGroup> subGroups = new HashMap<>();
+ final Map<String, SimpleCounter> counters = new HashMap<>();
+
+ @Override
+ public FlinkAgentsMetricGroup getSubGroup(String name) {
+ return subGroups.computeIfAbsent(name, k -> new TestMetricGroup());
+ }
+
+ @Override
+ public FlinkAgentsMetricGroup getSubGroup(String key, String value) {
+ return subGroups.computeIfAbsent(key + "=" + value, k -> new
TestMetricGroup());
+ }
+
+ @Override
+ public Counter getCounter(String name) {
+ return counters.computeIfAbsent(name, k -> new SimpleCounter());
+ }
+
+ @Override
+ public UpdatableGauge getGauge(String name) {
+ return null;
+ }
+
+ @Override
+ public Meter getMeter(String name) {
+ return null;
+ }
+
+ @Override
+ public Meter getMeter(String name, Counter counter) {
+ return null;
+ }
+
+ @Override
+ public Histogram getHistogram(String name) {
+ return null;
+ }
+
+ @Override
+ public Histogram getHistogram(String name, int windowSize) {
+ return null;
+ }
+ }
+
+ private static final float[] VEC = new float[] {0.1f, 0.2f};
+
+ /**
+ * Builds a setup bound to a connection that reports the given usage on
single-text embed, with
+ * the given model name in its descriptor (may be {@code null} to exercise
the guard).
+ */
+ private static BaseEmbeddingModelSetup setupWithSingleUsageAndModel(
+ EmbeddingTokenUsage usage, String model) {
+ BaseEmbeddingModelSetup setup =
+ new BaseEmbeddingModelSetup(
+ new ResourceDescriptor("test", descriptorArgs(model)),
+ mock(ResourceContext.class)) {
+ @Override
+ public Map<String, Object> getParameters() {
+ return new HashMap<>();
+ }
+ };
+ setup.connection =
+ new BaseEmbeddingModelConnection(
+ new ResourceDescriptor("conn", Collections.emptyMap()),
+ mock(ResourceContext.class)) {
+ @Override
+ public float[] embed(String text, Map<String, Object>
parameters) {
+ return VEC;
+ }
+
+ @Override
+ public List<float[]> embed(List<String> texts, Map<String,
Object> parameters) {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public EmbeddingResult<float[]> embedWithUsage(
+ String text, Map<String, Object> parameters) {
+ return new EmbeddingResult<>(VEC, usage);
+ }
+ };
+ return setup;
+ }
+
+ /** Builds a setup bound to a connection that reports the given usage on
single-text embed. */
+ private static BaseEmbeddingModelSetup
setupWithSingleUsage(EmbeddingTokenUsage usage) {
+ return setupWithSingleUsageAndModel(usage, "bedrock-text");
+ }
+
+ /** Descriptor args with an optional model (omitted when null/blank so it
stays unset). */
+ private static Map<String, String> descriptorArgs(String model) {
Review Comment:
This helper returns `Map<String, String>`, but `ResourceDescriptor`'s
constructor takes `Map<String, Object>`. Java generics are invariant, so the
call at `:106` does not compile.
That is what all 18 red CI checks are hitting. Every one of them stops at
`flink-agents-api` testCompile, including the Elasticsearch job, which builds
this module through `-am`:
```
BaseEmbeddingModelSetupTokenMetricsTest.java:[106,70] incompatible types:
java.util.Map<java.lang.String,java.lang.String> cannot be converted to
java.util.Map<java.lang.String,java.lang.Object>
```
Would changing this helper and its local `HashMap` to `Map<String, Object>`
be enough? I tried just that change locally, and the file compiles, all 9 new
tests pass, and `spotless:check` stays green. The two sibling call sites
already work because `Map.of(...)` at `:156` and `Collections.emptyMap()` at
`:115` take their type from the target on the spot, so only the explicitly
typed helper trips.
One thing worth knowing, since you mentioned leaning on CI: a green `Code
Style Check` does not tell you the Java side builds. Spotless formats test
sources without compiling them, so it stays green straight through a compile
error.
##########
python/flink_agents/api/embedding_models/embedding_model.py:
##########
@@ -155,4 +155,50 @@ def embed_with_usage(
"""Generate embeddings and return provider token usage when
available."""
merged_kwargs = self.model_kwargs.copy()
merged_kwargs.update(kwargs)
- return self._get_connection().embed_with_usage(text, **merged_kwargs)
+ result = self._get_connection().embed_with_usage(text, **merged_kwargs)
+ self._record_token_usage(result.token_usage)
+ return result
+
+ def _record_token_metrics(
+ self, model_name: str, prompt_tokens: int, total_tokens: int
+ ) -> None:
+ """Record embedding token usage metrics for the given model.
+
+ Mirrors ``BaseChatModelSetup._record_token_metrics`` but records
input-side
+ tokens only, since embeddings have no completion tokens. Counters are
placed
+ under the same ``model`` key-value group used by chat metrics, so
embedding
+ and chat usage for a model share one dimension.
+
+ Unlike the chat path, embedding calls do not run inside a plan action
that
+ hands in a request-scoped metric group (vector-store, RAG, and direct
calls
+ reach this setup directly), so the resource-bound metric group
injected via
+ ``set_metric_group`` is used instead.
+
+ Parameters
+ ----------
+ model_name : str
+ The name of the model used
+ prompt_tokens : int
+ The number of prompt tokens
+ total_tokens : int
+ The total number of tokens reported by the provider
+ """
+ metric_group = self.metric_group
+ if metric_group is None:
+ return
+
+ model_group = metric_group.get_sub_group("model", model_name)
+ model_group.get_counter("promptTokens").inc(prompt_tokens)
+ model_group.get_counter("totalTokens").inc(total_tokens)
+
+ def _record_token_usage(self, token_usage: EmbeddingTokenUsage | None) ->
None:
+ """Record the provider-reported embedding token usage, if any.
+
+ Called from ``embed_with_usage`` so direct calls and vector-store/RAG
paths
+ are both covered without each provider repeating the recording.
+ """
+ if token_usage is None or not self.model:
Review Comment:
`not self.model` and Java's `model == null || model.isBlank()`
(`BaseEmbeddingModelSetup.java:131`) do not agree when the model name is only
spaces. I ran the PR's own code to check: with `model=' '` Python records
under a group literally named `model= `, while Java skips it. With `model=''`
both skip.
`AGENTS.md` asks for the Java, Python and YAML APIs to stay semantically
aligned.
There is no Python test on this guard right now. Removing `or not
self.model` leaves all 10 tests passing, while Java has
`testEmbedWithUsageNullModelRecordsNothing` covering the same thing.
Is the whitespace case worth folding in? Something like this, in case it
helps:
```python
if token_usage is None or not (self.model or "").strip():
```
##########
api/src/main/java/org/apache/flink/agents/api/embedding/model/python/PythonEmbeddingModelSetup.java:
##########
@@ -137,7 +137,10 @@ public EmbeddingResult<float[]> embedWithUsage(String
text, Map<String, Object>
Map<String, Object> kwargs = new HashMap<>(parameters);
kwargs.put("text", text);
Object result = adapter.invoke(CALL_EMBED_WITH_USAGE,
embeddingModelSetup, kwargs);
- return EmbeddingModelUtils.toSingleEmbeddingResult(result);
+ EmbeddingResult<float[]> embeddingResult =
+ EmbeddingModelUtils.toSingleEmbeddingResult(result);
+ recordTokenUsage(embeddingResult.getTokenUsage());
Review Comment:
I think this records the same tokens twice, for a Java agent using a Python
embedding model.
`adapter.invoke(CALL_EMBED_WITH_USAGE, ...)` at `:139` calls into the Python
`BaseEmbeddingModelSetup.embed_with_usage`, and this PR makes that method
record too, at `embedding_model.py:159`. Then this line records again.
Both writes land on the same Java `SimpleCounter`. `setMetricGroup` here
(`:177-181`) binds the Java field and forwards the same
`FlinkAgentsMetricGroup` to the Python side, and `FlinkMetricGroup` and
`FlinkCounter` just pass through to that Java object. Neither guard stops it,
because both `model` fields read the same descriptor argument
(`EmbeddingCrossLanguageAgent.java:68` sets it).
The other direction stays at one write:
`JavaEmbeddingModelSetupImpl.embed_with_usage`
(`java_embedding_model.py:174-182`) calls the Java resource without chaining to
`super()`.
Nothing would catch it today either. I deleted both `recordTokenUsage(...)`
lines and the whole `api` module still passed 385/385.
`PythonEmbeddingModelSetupTest` builds the setup from a `@Mock
ResourceDescriptor` (`:50`), so `getArgument("model")` comes back null and the
recording returns early, and no metric group is ever bound.
So: which side should own the write for a cross-language resource? Dropping
these two lines and letting the Python setup own it would match what
`JavaEmbeddingModelSetupImpl` already does, unless there is a reason the Java
wrapper needs its own. Either way, would a test with a real descriptor carrying
`model` plus a bound group be worth adding, so these lines are covered?
##########
api/src/main/java/org/apache/flink/agents/api/embedding/model/BaseEmbeddingModelSetup.java:
##########
@@ -93,9 +94,52 @@ public String getModel() {
return model;
}
+ /**
+ * Record embedding token usage metrics for the given model on this
setup's bound metric group.
+ *
+ * <p>Mirrors {@code BaseChatModelSetup#recordTokenMetrics} but records
input-side tokens only,
+ * since embeddings have no completion tokens. Counters are placed under
the same {@code model}
+ * key-value group used by chat metrics, so embedding and chat usage for a
model share one
+ * dimension.
+ *
+ * <p>Unlike the chat path, embedding calls do not run inside a plan
action that hands in a
+ * request-scoped metric group (vector-store, RAG, and direct calls reach
this setup directly),
+ * so the resource-bound metric group injected via {@link #setMetricGroup}
is used instead.
+ *
+ * @param modelName the name of the model used
+ * @param promptTokens the number of prompt tokens
+ * @param totalTokens the total number of tokens reported by the provider
+ */
+ public void recordTokenMetrics(String modelName, long promptTokens, long
totalTokens) {
+ Preconditions.checkArgument(
+ modelName != null && !modelName.isBlank(), "Model name must
not be null or blank.");
+ FlinkAgentsMetricGroup metricGroup = getMetricGroup();
Review Comment:
Reading the group bound to the resource here, and at
`embedding_model.py:186`, is the pattern #859 reported and #861 moved chat off.
Here is what worries me. `Resource.metricGroup` is a single mutable field on
an object that is cached and shared across the whole subtask.
`RunnerContextImpl.getResource:485-495` rewrites it to the current action's
group on every fetch. Actions yield to other keys while they run. So another
action can rebind that field in between a `ctx.getResource(...)` call and the
later `embedWithUsage`.
Chat takes a different route: it captures the group up front and passes it
in (`ChatModelInvoker.java:121`, used at `:176`). There is a test guarding
exactly that, `BaseChatModelSetupTokenMetricsTest.java:95-111`, which asserts
the bound group is not used.
The vector-store path looks worse than racy. A vector store resolves its
model through `ResourceContext` (`BaseVectorStore.java:96-101`), and that path
never sets `metricGroup` at all, so it stays null there.
Given RAG does run inside an action, what would make passing the group in
the way chat does hard here? And whichever way it settles, could a test pin it?
None of the new tests bind two groups, so nothing would fail today if the
choice were reversed.
##########
api/src/test/java/org/apache/flink/agents/api/embedding/model/BaseEmbeddingModelSetupTokenMetricsTest.java:
##########
@@ -0,0 +1,318 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions of
Review Comment:
nit: `permissions of` looks like it wants to be `permissions and`. RAT
accepts the file either way, so nothing is broken. It is the only Java file in
the repo with that wording though. Worth folding into the compile fix?
--
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]
