[api][runtime][python] Propagate metric groups to cross-language reso…

api/src/main/java/org/apache/flink/agents/api/chat/model/python/PythonChatModelConnection.java

@@ -65,6 +65,11 @@ public Object getPythonResource() {
         return chatModel;
     }
 
+    @Override
+    public PythonResourceAdapter getPythonResourceAdapter() {
+        return adapter;
+    }
+
     @Override
     public ChatMessage chat(
             List<ChatMessage> messages, List<Tool> tools, Map<String, Object> modelParams) {

api/src/main/java/org/apache/flink/agents/api/chat/model/python/PythonChatModelSetup.java

@@ -89,6 +89,11 @@ public Object getPythonResource() {
         return chatModelSetup;
     }
 
+    @Override
+    public PythonResourceAdapter getPythonResourceAdapter() {
+        return adapter;
+    }
+
     @Override
     public Map<String, Object> getParameters() {
         return Map.of();

api/src/main/java/org/apache/flink/agents/api/embedding/model/python/PythonEmbeddingModelConnection.java

@@ -123,6 +123,11 @@ public Object getPythonResource() {
         return embeddingModel;
     }
 
+    @Override
+    public PythonResourceAdapter getPythonResourceAdapter() {
+        return adapter;
+    }
+
     @Override
     public void close() throws Exception {
         this.embeddingModel.close();

api/src/main/java/org/apache/flink/agents/api/embedding/model/python/PythonEmbeddingModelSetup.java

@@ -132,4 +132,9 @@ public Map<String, Object> getParameters() {
     public Object getPythonResource() {
         return embeddingModelSetup;
     }
+
+    @Override
+    public PythonResourceAdapter getPythonResourceAdapter() {
+        return adapter;
+    }
 }

api/src/main/java/org/apache/flink/agents/api/resource/Resource.java

@@ -19,6 +19,7 @@
 package org.apache.flink.agents.api.resource;
 
 import org.apache.flink.agents.api.metrics.FlinkAgentsMetricGroup;
+import org.apache.flink.agents.api.resource.python.PythonResourceWrapper;
 
 /**
  * Base interface for all kinds of resources, including chat models, tools, prompts and so on.
@@ -60,6 +61,9 @@ public ResourceContext getResourceContext() {
      */
     public void setMetricGroup(FlinkAgentsMetricGroup metricGroup) {
         this.metricGroup = metricGroup;
+        if (this instanceof PythonResourceWrapper) {
+            ((PythonResourceWrapper) this).setPythonResourceMetricGroup(metricGroup);
+        }
     }
 
     /**

api/src/main/java/org/apache/flink/agents/api/resource/python/PythonResourceAdapter.java

@@ -19,6 +19,7 @@
 package org.apache.flink.agents.api.resource.python;
 
 import org.apache.flink.agents.api.chat.messages.ChatMessage;
+import org.apache.flink.agents.api.metrics.FlinkAgentsMetricGroup;
 import org.apache.flink.agents.api.tools.Tool;
 import org.apache.flink.agents.api.vectorstores.Document;
 import org.apache.flink.agents.api.vectorstores.VectorStoreQuery;
@@ -120,6 +121,14 @@ public interface PythonResourceAdapter {
      */
     Object callMethod(Object obj, String methodName, Map<String, Object> kwargs);
 
+    /**
+     * Binds a Java metric group to a Python resource.
+     *
+     * @param pythonResource the Python resource object
+     * @param metricGroup the Java metric group to expose through Python's metric group API
+     */
+    default void setMetricGroup(Object pythonResource, FlinkAgentsMetricGroup metricGroup) {}
+
     /**
      * Invokes a method with the specified name and arguments.
      *

api/src/main/java/org/apache/flink/agents/api/resource/python/PythonResourceWrapper.java

@@ -17,6 +17,8 @@
  */
 package org.apache.flink.agents.api.resource.python;
 
+import org.apache.flink.agents.api.metrics.FlinkAgentsMetricGroup;
+
 /**
  * Wrapper interface for Python resource objects. This interface provides a unified way to access
  * the underlying Python resource from Java objects that encapsulate Python functionality.
@@ -29,4 +31,26 @@ public interface PythonResourceWrapper {
      * @return the wrapped Python resource object
      */
     Object getPythonResource();
+
+    /**
+     * Returns the adapter that owns the wrapped Python resource.
+     *
+     * @return the Python resource adapter, or null if metric forwarding is unsupported
+     */
+    default PythonResourceAdapter getPythonResourceAdapter() {
+        return null;
+    }
+
+    /**
+     * Binds the current Java metric group to the wrapped Python resource.
+     *
+     * @param metricGroup the metric group to bind
+     */
+    default void setPythonResourceMetricGroup(FlinkAgentsMetricGroup metricGroup) {
+        PythonResourceAdapter adapter = getPythonResourceAdapter();
+        Object pythonResource = getPythonResource();
+        if (adapter != null && pythonResource != null) {
+            adapter.setMetricGroup(pythonResource, metricGroup);
+        }
+    }
 }

api/src/main/java/org/apache/flink/agents/api/vectorstores/python/PythonVectorStore.java

@@ -233,4 +233,9 @@ public void updateEmbedding(
     public Object getPythonResource() {
         return vectorStore;
     }
+
+    @Override
+    public PythonResourceAdapter getPythonResourceAdapter() {
+        return adapter;
+    }
 }

api/src/test/java/org/apache/flink/agents/api/chat/model/python/PythonChatModelSetupTest.java

@@ -18,6 +18,7 @@
 package org.apache.flink.agents.api.chat.model.python;
 
 import org.apache.flink.agents.api.chat.messages.ChatMessage;
+import org.apache.flink.agents.api.metrics.FlinkAgentsMetricGroup;
 import org.apache.flink.agents.api.resource.ResourceContext;
 import org.apache.flink.agents.api.resource.ResourceDescriptor;
 import org.apache.flink.agents.api.resource.python.PythonResourceAdapter;
@@ -161,4 +162,13 @@ void testImplementsPythonResourceWrapper() {
                 .isInstanceOf(
                         org.apache.flink.agents.api.resource.python.PythonResourceWrapper.class);
     }
+
+    @Test
+    void testSetMetricGroupPropagatesToPythonResource() {
+        FlinkAgentsMetricGroup metricGroup = mock(FlinkAgentsMetricGroup.class);
+
+        pythonChatModelSetup.setMetricGroup(metricGroup);
+
+        verify(mockAdapter).setMetricGroup(mockChatModelSetup, metricGroup);
+    }
 }

plan/src/main/java/org/apache/flink/agents/plan/resource/python/PythonMCPPrompt.java

@@ -47,6 +47,11 @@ public Object getPythonResource() {
         return prompt;
     }
 
+    @Override
+    public PythonResourceAdapter getPythonResourceAdapter() {
+        return adapter;
+    }
+
     public String getName() {
         if (name == null) {
             name = prompt.getAttr("name").toString();

plan/src/main/java/org/apache/flink/agents/plan/resource/python/PythonMCPServer.java

@@ -84,6 +84,11 @@ public Object getPythonResource() {
         return server;
     }
 
+    @Override
+    public PythonResourceAdapter getPythonResourceAdapter() {
+        return adapter;
+    }
+
     @Override
     public ResourceType getResourceType() {
         return ResourceType.MCP_SERVER;

plan/src/main/java/org/apache/flink/agents/plan/resource/python/PythonMCPTool.java

@@ -75,6 +75,11 @@ public Object getPythonResource() {
         return tool;
     }
 
+    @Override
+    public PythonResourceAdapter getPythonResourceAdapter() {
+        return adapter;
+    }
+
     @Override
     public ToolType getToolType() {
         return ToolType.MCP;

python/flink_agents/runtime/java/java_chat_model.py

@@ -26,6 +26,9 @@
 )
 from flink_agents.api.resource import ResourceType
 from flink_agents.api.tools.tool import Tool
+from flink_agents.runtime.java.java_resource_wrapper import (
+    set_java_resource_metric_group,
+)
 
 
 class JavaChatModelConnectionImpl(JavaChatModelConnection):
@@ -51,6 +54,11 @@ def __init__(self, j_resource: Any, j_resource_adapter: Any, **kwargs: Any) -> N
         self._j_resource = j_resource
         self._j_resource_adapter = j_resource_adapter
 
+    @override
+    def set_metric_group(self, metric_group: Any) -> None:
+        super().set_metric_group(metric_group)
+        set_java_resource_metric_group(self._j_resource, metric_group)
+
     @override
     def chat(
         self,
@@ -114,6 +122,11 @@ def __init__(self, j_resource: Any, j_resource_adapter: Any, **kwargs: Any) -> N
         self._j_resource = j_resource
         self._j_resource_adapter = j_resource_adapter
 
+    @override
+    def set_metric_group(self, metric_group: Any) -> None:
+        super().set_metric_group(metric_group)
+        set_java_resource_metric_group(self._j_resource, metric_group)
+
     @property
     @override
     def model_kwargs(self) -> Dict[str, Any]:

python/flink_agents/runtime/java/java_embedding_model.py

@@ -23,6 +23,9 @@
     JavaEmbeddingModelConnection,
     JavaEmbeddingModelSetup,
 )
+from flink_agents.runtime.java.java_resource_wrapper import (
+    set_java_resource_metric_group,
+)
 
 
 class JavaEmbeddingModelConnectionImpl(JavaEmbeddingModelConnection):
@@ -48,6 +51,11 @@ def __init__(self, j_resource: Any, j_resource_adapter: Any, **kwargs: Any) -> N
         self._j_resource = j_resource
         self._j_resource_adapter = j_resource_adapter
 
+    @override
+    def set_metric_group(self, metric_group: Any) -> None:
+        super().set_metric_group(metric_group)
+        set_java_resource_metric_group(self._j_resource, metric_group)
+
     def embed(
         self, text: str | Sequence[str], **kwargs: Any
     ) -> list[float] | list[list[float]]:
@@ -92,6 +100,11 @@ def __init__(self, j_resource: Any, j_resource_adapter: Any, **kwargs: Any) -> N
         self._j_resource = j_resource
         self._j_resource_adapter = j_resource_adapter
 
+    @override
+    def set_metric_group(self, metric_group: Any) -> None:
+        super().set_metric_group(metric_group)
+        set_java_resource_metric_group(self._j_resource, metric_group)
+
     @property
     def model_kwargs(self) -> Dict[str, Any]:
         """Return embedding model settings.

python/flink_agents/runtime/java/java_resource_wrapper.py

@@ -27,6 +27,14 @@
 from flink_agents.api.tools.tool import Tool, ToolMetadata, ToolType
 
 
+def set_java_resource_metric_group(j_resource: Any, metric_group: Any) -> None:
+    """Bind the underlying Java metric group to a wrapped Java resource."""
+    if j_resource is None:
+        return
+    j_metric_group = getattr(metric_group, "_j_metric_group", metric_group)
+    j_resource.setMetricGroup(j_metric_group)
+
+
 class JavaTool(Tool):
     """Java Tool that carries tool metadata and can be recognized by PythonChatModel."""
 

python/flink_agents/runtime/java/java_vector_store.py

@@ -27,6 +27,9 @@
     Document,
     _maybe_cast_to_list,
 )
+from flink_agents.runtime.java.java_resource_wrapper import (
+    set_java_resource_metric_group,
+)
 from flink_agents.runtime.python_java_utils import from_java_document
 
 
@@ -60,6 +63,11 @@ def __init__(self, j_resource: Any, j_resource_adapter: Any, **kwargs: Any) -> N
         self._j_resource = j_resource
         self._j_resource_adapter = j_resource_adapter
 
+    @override
+    def set_metric_group(self, metric_group: Any) -> None:
+        super().set_metric_group(metric_group)
+        set_java_resource_metric_group(self._j_resource, metric_group)
+
     @property
     @override
     def store_kwargs(self) -> Dict[str, Any]:

python/flink_agents/runtime/python_java_utils.py

@@ -377,3 +377,13 @@ def call_method(obj: Any, method_name: str, kwargs: Dict[str, Any]) -> Any:
 
     method = getattr(obj, method_name)
     return method(**kwargs)
+
+
+def set_metric_group(obj: Resource, j_metric_group: Any) -> None:
+    """Bind a Java metric group to a Python resource."""
+    from flink_agents.runtime.flink_metric_group import FlinkMetricGroup
+
+    metric_group = (
+        FlinkMetricGroup(j_metric_group) if j_metric_group is not None else None
+    )
+    obj.set_metric_group(metric_group)

python/flink_agents/runtime/tests/test_cross_language_metric_group.py

@@ -0,0 +1,86 @@
+################################################################################
+#  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 and
+# limitations under the License.
+#################################################################################
+from typing import Any
+
+import pytest
+
+from flink_agents.runtime.java.java_chat_model import (
+    JavaChatModelConnectionImpl,
+    JavaChatModelSetupImpl,
+)
+from flink_agents.runtime.java.java_embedding_model import (
+    JavaEmbeddingModelConnectionImpl,
+    JavaEmbeddingModelSetupImpl,
+)
+from flink_agents.runtime.java.java_resource_wrapper import (
+    set_java_resource_metric_group,
+)
+
+
+class _JavaResource:
+    def __init__(self) -> None:
+        self.metric_group: Any = None
+
+    def setMetricGroup(self, metric_group: Any) -> None:
+        self.metric_group = metric_group
+
+
+class _MetricGroup:
+    def __init__(self) -> None:
+        self._j_metric_group = object()
+
+
+@pytest.mark.parametrize(
+    "resource",
+    [
+        JavaChatModelConnectionImpl(
+            j_resource=_JavaResource(), j_resource_adapter=None
+        ),
+        JavaChatModelSetupImpl(
+            j_resource=_JavaResource(),
+            j_resource_adapter=None,
+            connection="connection",
+            model="model",
+        ),
+        JavaEmbeddingModelConnectionImpl(
+            j_resource=_JavaResource(), j_resource_adapter=None
+        ),
+        JavaEmbeddingModelSetupImpl(
+            j_resource=_JavaResource(),
+            j_resource_adapter=None,
+            connection="connection",
+            model="model",
+        ),
+    ],
+)
+def test_java_resource_wrappers_forward_metric_group(resource):
+    metric_group = _MetricGroup()
+
+    resource.set_metric_group(metric_group)
+
+    assert resource.metric_group is metric_group
+    assert resource._j_resource.metric_group is metric_group._j_metric_group
+
+
+def test_set_java_resource_metric_group_unwraps_flink_metric_group():
+    java_resource = _JavaResource()
+    metric_group = _MetricGroup()
+
+    set_java_resource_metric_group(java_resource, metric_group)
+
+    assert java_resource.metric_group is metric_group._j_metric_group