This is an automated email from the ASF dual-hosted git repository.
wenjin272 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/flink-agents.git
The following commit(s) were added to refs/heads/main by this push:
new e81638a8 [runtime][python] Manage Pemja references across actions and
resources (#1058)
e81638a8 is described below
commit e81638a831c785e0b84b78635cfa217a0b7dc4a9
Author: Yu Tang <[email protected]>
AuthorDate: Tue Sep 15 14:23:43 2026 +0800
[runtime][python] Manage Pemja references across actions and resources
(#1058)
Generated-by: Codex 0.154.0-alpha.6.2 (GPT-6 Astra)
Co-Authored-By: Codex <[email protected]>
---
.../model/python/PythonChatModelConnection.java | 36 ++++---
.../chat/model/python/PythonChatModelSetup.java | 32 ++++--
.../python/PythonEmbeddingModelConnection.java | 9 +-
.../model/python/PythonEmbeddingModelSetup.java | 12 +++
.../api/resource/python/PythonObjectScope.java | 109 +++++++++++++++++++++
.../api/vectorstores/python/PythonVectorStore.java | 58 +++++++----
.../python/PythonChatModelConnectionTest.java | 11 ++-
.../model/python/PythonChatModelSetupTest.java | 8 +-
.../api/resource/python/PythonObjectScopeTest.java | 45 +++++++++
.../PythonCollectionManageableVectorStoreTest.java | 9 +-
.../plan/actions/ContextRetrievalAction.java | 59 +++++------
.../plan/resource/python/PythonMCPPrompt.java | 35 +++++--
.../plan/resource/python/PythonMCPServer.java | 13 +++
.../agents/plan/resource/python/PythonMCPTool.java | 12 +++
.../resource/python/PythonMCPResourceTest.java | 106 ++++++++++++++++++++
.../PythonResourceProviderTest.java | 65 ++++++++++++
.../agents/runtime/PythonMCPResourceDiscovery.java | 1 +
.../agents/runtime/memory/Mem0LongTermMemory.java | 86 +++++++++-------
.../runtime/python/utils/PythonActionExecutor.java | 41 +++++---
.../python/utils/PythonResourceAdapterImpl.java | 11 ++-
.../runtime/PythonMCPResourceDiscoveryTest.java | 73 ++++++++++++++
.../runtime/memory/Mem0LongTermMemoryTest.java | 11 ++-
.../python/utils/PythonActionExecutorTest.java | 100 ++++++++++++++++++-
.../utils/PythonResourceAdapterImplTest.java | 18 ++++
24 files changed, 814 insertions(+), 146 deletions(-)
diff --git
a/api/src/main/java/org/apache/flink/agents/api/chat/model/python/PythonChatModelConnection.java
b/api/src/main/java/org/apache/flink/agents/api/chat/model/python/PythonChatModelConnection.java
index 8ce3662e..580f4be8 100644
---
a/api/src/main/java/org/apache/flink/agents/api/chat/model/python/PythonChatModelConnection.java
+++
b/api/src/main/java/org/apache/flink/agents/api/chat/model/python/PythonChatModelConnection.java
@@ -22,6 +22,7 @@ import
org.apache.flink.agents.api.chat.model.BaseChatModelConnection;
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.PythonObjectScope;
import org.apache.flink.agents.api.resource.python.PythonResourceAdapter;
import org.apache.flink.agents.api.resource.python.PythonResourceWrapper;
import org.apache.flink.agents.api.tools.Tool;
@@ -41,6 +42,7 @@ public class PythonChatModelConnection extends
BaseChatModelConnection
implements PythonResourceWrapper {
private final PyObject chatModel;
private final PythonResourceAdapter adapter;
+ private boolean closed;
/**
* Creates a new PythonChatModelConnection.
@@ -82,24 +84,32 @@ public class PythonChatModelConnection extends
BaseChatModelConnection
List<ChatMessage> messages, List<Tool> tools, Map<String, Object>
modelParams) {
Map<String, Object> kwargs = new HashMap<>(modelParams);
- List<Object> pythonMessages = new ArrayList<>();
- for (ChatMessage message : messages) {
- pythonMessages.add(adapter.toPythonChatMessage(message));
- }
- kwargs.put("messages", pythonMessages);
+ try (PythonObjectScope scope = new PythonObjectScope()) {
+ List<Object> pythonMessages = new ArrayList<>();
+ for (ChatMessage message : messages) {
+
pythonMessages.add(scope.own(adapter.toPythonChatMessage(message)));
+ }
+ kwargs.put("messages", pythonMessages);
- List<Object> pythonTools = new ArrayList<>();
- for (Tool tool : tools) {
- pythonTools.add(adapter.convertToPythonTool(tool));
- }
- kwargs.put("tools", pythonTools);
+ List<Object> pythonTools = new ArrayList<>();
+ for (Tool tool : tools) {
+ pythonTools.add(scope.own(adapter.convertToPythonTool(tool)));
+ }
+ kwargs.put("tools", pythonTools);
- Object pythonMessageResponse = adapter.callMethod(chatModel, "chat",
kwargs);
- return adapter.fromPythonChatMessage(pythonMessageResponse);
+ Object pythonMessageResponse =
scope.own(adapter.callMethod(chatModel, "chat", kwargs));
+ return adapter.fromPythonChatMessage(pythonMessageResponse);
+ }
}
@Override
public void close() throws Exception {
- this.chatModel.close();
+ if (closed || chatModel == null) {
+ return;
+ }
+ closed = true;
+ try (chatModel) {
+ adapter.callMethod(chatModel, "close", Map.of());
+ }
}
}
diff --git
a/api/src/main/java/org/apache/flink/agents/api/chat/model/python/PythonChatModelSetup.java
b/api/src/main/java/org/apache/flink/agents/api/chat/model/python/PythonChatModelSetup.java
index 846105af..afae1e35 100644
---
a/api/src/main/java/org/apache/flink/agents/api/chat/model/python/PythonChatModelSetup.java
+++
b/api/src/main/java/org/apache/flink/agents/api/chat/model/python/PythonChatModelSetup.java
@@ -22,6 +22,7 @@ import
org.apache.flink.agents.api.chat.model.BaseChatModelSetup;
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.PythonObjectScope;
import org.apache.flink.agents.api.resource.python.PythonResourceAdapter;
import org.apache.flink.agents.api.resource.python.PythonResourceWrapper;
import pemja.core.object.PyObject;
@@ -46,6 +47,7 @@ public class PythonChatModelSetup extends BaseChatModelSetup
implements PythonRe
private final PyObject chatModelSetup;
private final PythonResourceAdapter adapter;
+ private boolean closed;
public PythonChatModelSetup(
PythonResourceAdapter adapter,
@@ -73,16 +75,19 @@ public class PythonChatModelSetup extends
BaseChatModelSetup implements PythonRe
Map<String, Object> kwargs = new HashMap<>(modelParams);
- List<Object> pythonMessages = new ArrayList<>();
- for (ChatMessage message : messages) {
- pythonMessages.add(adapter.toPythonChatMessage(message));
- }
+ try (PythonObjectScope scope = new PythonObjectScope()) {
+ List<Object> pythonMessages = new ArrayList<>();
+ for (ChatMessage message : messages) {
+
pythonMessages.add(scope.own(adapter.toPythonChatMessage(message)));
+ }
- kwargs.put("messages", pythonMessages);
- kwargs.put("prompt_args", promptArgs != null ? promptArgs :
Collections.emptyMap());
+ kwargs.put("messages", pythonMessages);
+ kwargs.put("prompt_args", promptArgs != null ? promptArgs :
Collections.emptyMap());
- Object pythonMessageResponse = adapter.callMethod(chatModelSetup,
"chat", kwargs);
- return adapter.fromPythonChatMessage(pythonMessageResponse);
+ Object pythonMessageResponse =
+ scope.own(adapter.callMethod(chatModelSetup, "chat",
kwargs));
+ return adapter.fromPythonChatMessage(pythonMessageResponse);
+ }
}
@Override
@@ -105,4 +110,15 @@ public class PythonChatModelSetup extends
BaseChatModelSetup implements PythonRe
public Map<String, Object> getParameters() {
return Map.of();
}
+
+ @Override
+ public void close() throws Exception {
+ if (closed || chatModelSetup == null) {
+ return;
+ }
+ closed = true;
+ try (chatModelSetup) {
+ adapter.callMethod(chatModelSetup, "close", Map.of());
+ }
+ }
}
diff --git
a/api/src/main/java/org/apache/flink/agents/api/embedding/model/python/PythonEmbeddingModelConnection.java
b/api/src/main/java/org/apache/flink/agents/api/embedding/model/python/PythonEmbeddingModelConnection.java
index d9bdb8d4..0e557dbc 100644
---
a/api/src/main/java/org/apache/flink/agents/api/embedding/model/python/PythonEmbeddingModelConnection.java
+++
b/api/src/main/java/org/apache/flink/agents/api/embedding/model/python/PythonEmbeddingModelConnection.java
@@ -48,6 +48,7 @@ public class PythonEmbeddingModelConnection extends
BaseEmbeddingModelConnection
private final PyObject embeddingModel;
private final PythonResourceAdapter adapter;
+ private boolean closed;
/**
* Creates a new PythonEmbeddingModelConnection.
@@ -166,6 +167,12 @@ public class PythonEmbeddingModelConnection extends
BaseEmbeddingModelConnection
@Override
public void close() throws Exception {
- this.embeddingModel.close();
+ if (closed || embeddingModel == null) {
+ return;
+ }
+ closed = true;
+ try (embeddingModel) {
+ adapter.callMethod(embeddingModel, "close", Map.of());
+ }
}
}
diff --git
a/api/src/main/java/org/apache/flink/agents/api/embedding/model/python/PythonEmbeddingModelSetup.java
b/api/src/main/java/org/apache/flink/agents/api/embedding/model/python/PythonEmbeddingModelSetup.java
index b6febbc5..2285bb16 100644
---
a/api/src/main/java/org/apache/flink/agents/api/embedding/model/python/PythonEmbeddingModelSetup.java
+++
b/api/src/main/java/org/apache/flink/agents/api/embedding/model/python/PythonEmbeddingModelSetup.java
@@ -48,6 +48,7 @@ public class PythonEmbeddingModelSetup extends
BaseEmbeddingModelSetup
private final PyObject embeddingModelSetup;
private final PythonResourceAdapter adapter;
+ private boolean closed;
/**
* Creates a new PythonEmbeddingModelSetup.
@@ -173,4 +174,15 @@ public class PythonEmbeddingModelSetup extends
BaseEmbeddingModelSetup
super.setMetricGroup(metricGroup);
setPythonResourceMetricGroup(metricGroup);
}
+
+ @Override
+ public void close() throws Exception {
+ if (closed || embeddingModelSetup == null) {
+ return;
+ }
+ closed = true;
+ try (embeddingModelSetup) {
+ adapter.callMethod(embeddingModelSetup, "close", Map.of());
+ }
+ }
}
diff --git
a/api/src/main/java/org/apache/flink/agents/api/resource/python/PythonObjectScope.java
b/api/src/main/java/org/apache/flink/agents/api/resource/python/PythonObjectScope.java
new file mode 100644
index 00000000..44192ba6
--- /dev/null
+++
b/api/src/main/java/org/apache/flink/agents/api/resource/python/PythonObjectScope.java
@@ -0,0 +1,109 @@
+/*
+ * 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.
+ */
+package org.apache.flink.agents.api.resource.python;
+
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.util.ExceptionUtils;
+import org.apache.flink.util.LambdaUtil;
+import pemja.core.object.PyObject;
+
+import java.util.Collections;
+import java.util.IdentityHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+/**
+ * Owns temporary Pemja references consumed within a Java-to-Python bridge
operation.
+ *
+ * <p>Callers may register a {@link PyObject} directly or nested in a Java
map, list, or object
+ * array. Every registered handle must be fully consumed before the scope
closes. Values returned to
+ * callers must not contain handles owned by this scope.
+ */
+@Internal
+public final class PythonObjectScope implements AutoCloseable {
+
+ private final Set<PyObject> ownedObjects = Collections.newSetFromMap(new
IdentityHashMap<>());
+ private boolean closed;
+
+ /** Adds every {@link PyObject} reachable through the supplied bridge
result to this scope. */
+ public <T> T own(T value) {
+ ensureOpen();
+ visit(value, Collections.newSetFromMap(new IdentityHashMap<>()));
+ return value;
+ }
+
+ @Override
+ public void close() {
+ if (closed) {
+ return;
+ }
+ closed = true;
+
+ try {
+ LambdaUtil.applyToAllWhileSuppressingExceptions(ownedObjects,
PyObject::close);
+ } catch (Exception e) {
+ ExceptionUtils.rethrow(e);
+ } finally {
+ ownedObjects.clear();
+ }
+ }
+
+ private void visit(Object value, Set<Object> visitedContainers) {
+ if (value == null) {
+ return;
+ }
+ if (value instanceof PyObject) {
+ ownedObjects.add((PyObject) value);
+ return;
+ }
+ if (value instanceof Map) {
+ if (!visitedContainers.add(value)) {
+ return;
+ }
+ for (Map.Entry<?, ?> entry : ((Map<?, ?>) value).entrySet()) {
+ visit(entry.getKey(), visitedContainers);
+ visit(entry.getValue(), visitedContainers);
+ }
+ return;
+ }
+ if (value instanceof List) {
+ if (!visitedContainers.add(value)) {
+ return;
+ }
+ for (Object element : (List<?>) value) {
+ visit(element, visitedContainers);
+ }
+ return;
+ }
+ if (value instanceof Object[]) {
+ if (!visitedContainers.add(value)) {
+ return;
+ }
+ for (Object element : (Object[]) value) {
+ visit(element, visitedContainers);
+ }
+ }
+ }
+
+ private void ensureOpen() {
+ if (closed) {
+ throw new IllegalStateException("PythonObjectScope is already
closed.");
+ }
+ }
+}
diff --git
a/api/src/main/java/org/apache/flink/agents/api/vectorstores/python/PythonVectorStore.java
b/api/src/main/java/org/apache/flink/agents/api/vectorstores/python/PythonVectorStore.java
index 21bb1893..4a393adc 100644
---
a/api/src/main/java/org/apache/flink/agents/api/vectorstores/python/PythonVectorStore.java
+++
b/api/src/main/java/org/apache/flink/agents/api/vectorstores/python/PythonVectorStore.java
@@ -21,6 +21,7 @@ package org.apache.flink.agents.api.vectorstores.python;
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.PythonObjectScope;
import org.apache.flink.agents.api.resource.python.PythonResourceAdapter;
import org.apache.flink.agents.api.resource.python.PythonResourceWrapper;
import org.apache.flink.agents.api.vectorstores.BaseVectorStore;
@@ -54,6 +55,7 @@ import java.util.Map;
public class PythonVectorStore extends BaseVectorStore implements
PythonResourceWrapper {
protected final PyObject vectorStore;
protected final PythonResourceAdapter adapter;
+ private boolean closed;
/**
* Creates a new PythonVectorStore.
@@ -106,9 +108,10 @@ public class PythonVectorStore extends BaseVectorStore
implements PythonResource
kwargs.put("limit", limit);
}
- Object pythonDocuments = adapter.callMethod(vectorStore, "get",
kwargs);
-
- return adapter.fromPythonDocuments((List<PyObject>) pythonDocuments);
+ try (PythonObjectScope scope = new PythonObjectScope()) {
+ Object pythonDocuments = scope.own(adapter.callMethod(vectorStore,
"get", kwargs));
+ return adapter.fromPythonDocuments((List<PyObject>)
pythonDocuments);
+ }
}
@Override
@@ -158,8 +161,11 @@ public class PythonVectorStore extends BaseVectorStore
implements PythonResource
if (filters != null) {
kwargs.put("filters", filters);
}
- Object pythonDocuments = adapter.callMethod(vectorStore,
"_query_embedding", kwargs);
- return adapter.fromPythonDocuments((List<PyObject>) pythonDocuments);
+ try (PythonObjectScope scope = new PythonObjectScope()) {
+ Object pythonDocuments =
+ scope.own(adapter.callMethod(vectorStore,
"_query_embedding", kwargs));
+ return adapter.fromPythonDocuments((List<PyObject>)
pythonDocuments);
+ }
}
/** Embed query text via the configured model (no numpy, so it stays on
the async pool). */
@@ -202,8 +208,11 @@ public class PythonVectorStore extends BaseVectorStore
implements PythonResource
if (filters != null) {
kwargs.put("filters", filters);
}
- Object pythonDocuments = adapter.callMethod(vectorStore,
"_query_embedding", kwargs);
- return adapter.fromPythonDocuments((List<PyObject>) pythonDocuments);
+ try (PythonObjectScope scope = new PythonObjectScope()) {
+ Object pythonDocuments =
+ scope.own(adapter.callMethod(vectorStore,
"_query_embedding", kwargs));
+ return adapter.fromPythonDocuments((List<PyObject>)
pythonDocuments);
+ }
}
@Override
@@ -211,23 +220,27 @@ public class PythonVectorStore extends BaseVectorStore
implements PythonResource
public List<String> addEmbedding(
List<Document> documents, @Nullable String collection, Map<String,
Object> extraArgs)
throws IOException {
- Map<String, Object> kwargs = new HashMap<>(extraArgs);
- kwargs.put("documents", adapter.toPythonDocuments(documents));
- if (collection != null) {
- kwargs.put("collection_name", collection);
+ try (PythonObjectScope scope = new PythonObjectScope()) {
+ Map<String, Object> kwargs = new HashMap<>(extraArgs);
+ kwargs.put("documents",
scope.own(adapter.toPythonDocuments(documents)));
+ if (collection != null) {
+ kwargs.put("collection_name", collection);
+ }
+ return (List<String>) adapter.callMethod(vectorStore,
"_add_embedding", kwargs);
}
- return (List<String>) adapter.callMethod(vectorStore,
"_add_embedding", kwargs);
}
@Override
public void updateEmbedding(
List<Document> documents, @Nullable String collection, Map<String,
Object> extraArgs) {
- Map<String, Object> kwargs = new HashMap<>(extraArgs);
- kwargs.put("documents", adapter.toPythonDocuments(documents));
- if (collection != null) {
- kwargs.put("collection_name", collection);
+ try (PythonObjectScope scope = new PythonObjectScope()) {
+ Map<String, Object> kwargs = new HashMap<>(extraArgs);
+ kwargs.put("documents",
scope.own(adapter.toPythonDocuments(documents)));
+ if (collection != null) {
+ kwargs.put("collection_name", collection);
+ }
+ adapter.callMethod(vectorStore, "_update_embedding", kwargs);
}
- adapter.callMethod(vectorStore, "_update_embedding", kwargs);
}
@Override
@@ -245,4 +258,15 @@ public class PythonVectorStore extends BaseVectorStore
implements PythonResource
super.setMetricGroup(metricGroup);
setPythonResourceMetricGroup(metricGroup);
}
+
+ @Override
+ public void close() throws Exception {
+ if (closed || vectorStore == null) {
+ return;
+ }
+ closed = true;
+ try (vectorStore) {
+ adapter.callMethod(vectorStore, "close", Map.of());
+ }
+ }
}
diff --git
a/api/src/test/java/org/apache/flink/agents/api/chat/model/python/PythonChatModelConnectionTest.java
b/api/src/test/java/org/apache/flink/agents/api/chat/model/python/PythonChatModelConnectionTest.java
index b3ccf82e..fc063d23 100644
---
a/api/src/test/java/org/apache/flink/agents/api/chat/model/python/PythonChatModelConnectionTest.java
+++
b/api/src/test/java/org/apache/flink/agents/api/chat/model/python/PythonChatModelConnectionTest.java
@@ -83,7 +83,7 @@ public class PythonChatModelConnectionTest {
}
@Test
- void testChat() {
+ void testChat() throws Exception {
ChatMessage inputMessage = mock(ChatMessage.class);
ChatMessage outputMessage = mock(ChatMessage.class);
Tool mockTool = mock(Tool.class);
@@ -93,9 +93,9 @@ public class PythonChatModelConnectionTest {
modelParams.put("temperature", 0.7);
modelParams.put("max_tokens", 100);
- Object pythonInputMessage = new Object();
- Object pythonOutputMessage = new Object();
- Object pythonTool = new Object();
+ PyObject pythonInputMessage = mock(PyObject.class);
+ PyObject pythonOutputMessage = mock(PyObject.class);
+ PyObject pythonTool = mock(PyObject.class);
when(mockAdapter.toPythonChatMessage(inputMessage)).thenReturn(pythonInputMessage);
when(mockAdapter.convertToPythonTool(mockTool)).thenReturn(pythonTool);
@@ -133,6 +133,9 @@ public class PythonChatModelConnectionTest {
return true;
}));
verify(mockAdapter).fromPythonChatMessage(pythonOutputMessage);
+ verify(pythonInputMessage).close();
+ verify(pythonTool).close();
+ verify(pythonOutputMessage).close();
}
@Test
diff --git
a/api/src/test/java/org/apache/flink/agents/api/chat/model/python/PythonChatModelSetupTest.java
b/api/src/test/java/org/apache/flink/agents/api/chat/model/python/PythonChatModelSetupTest.java
index 42463b08..b08afc84 100644
---
a/api/src/test/java/org/apache/flink/agents/api/chat/model/python/PythonChatModelSetupTest.java
+++
b/api/src/test/java/org/apache/flink/agents/api/chat/model/python/PythonChatModelSetupTest.java
@@ -90,7 +90,7 @@ public class PythonChatModelSetupTest {
}
@Test
- void testChat() {
+ void testChat() throws Exception {
ChatMessage inputMessage = mock(ChatMessage.class);
ChatMessage outputMessage = mock(ChatMessage.class);
List<ChatMessage> messages = Collections.singletonList(inputMessage);
@@ -100,8 +100,8 @@ public class PythonChatModelSetupTest {
modelParams.put("temperature", 0.7);
modelParams.put("max_tokens", 100);
- Object pythonInputMessage = new Object();
- Object pythonOutputMessage = new Object();
+ PyObject pythonInputMessage = mock(PyObject.class);
+ PyObject pythonOutputMessage = mock(PyObject.class);
when(mockAdapter.toPythonChatMessage(inputMessage)).thenReturn(pythonInputMessage);
when(mockAdapter.callMethod(eq(mockChatModelSetup), eq("chat"),
any(Map.class)))
@@ -132,6 +132,8 @@ public class PythonChatModelSetupTest {
return true;
}));
verify(mockAdapter).fromPythonChatMessage(pythonOutputMessage);
+ verify(pythonInputMessage).close();
+ verify(pythonOutputMessage).close();
}
@Test
diff --git
a/api/src/test/java/org/apache/flink/agents/api/resource/python/PythonObjectScopeTest.java
b/api/src/test/java/org/apache/flink/agents/api/resource/python/PythonObjectScopeTest.java
new file mode 100644
index 00000000..560aa4b5
--- /dev/null
+++
b/api/src/test/java/org/apache/flink/agents/api/resource/python/PythonObjectScopeTest.java
@@ -0,0 +1,45 @@
+/*
+ * 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.
+ */
+package org.apache.flink.agents.api.resource.python;
+
+import org.junit.jupiter.api.Test;
+import pemja.core.object.PyObject;
+
+import java.util.List;
+import java.util.Map;
+
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+
+class PythonObjectScopeTest {
+
+ @Test
+ void closesNestedReferencesOnce() throws Exception {
+ PyObject first = mock(PyObject.class);
+ PyObject second = mock(PyObject.class);
+
+ PythonObjectScope scope = new PythonObjectScope();
+ scope.own(Map.of("values", List.of(first, second, first)));
+ scope.close();
+ scope.close();
+
+ verify(first, times(1)).close();
+ verify(second, times(1)).close();
+ }
+}
diff --git
a/api/src/test/java/org/apache/flink/agents/api/vectorstores/python/PythonCollectionManageableVectorStoreTest.java
b/api/src/test/java/org/apache/flink/agents/api/vectorstores/python/PythonCollectionManageableVectorStoreTest.java
index 37b27f85..566db3fe 100644
---
a/api/src/test/java/org/apache/flink/agents/api/vectorstores/python/PythonCollectionManageableVectorStoreTest.java
+++
b/api/src/test/java/org/apache/flink/agents/api/vectorstores/python/PythonCollectionManageableVectorStoreTest.java
@@ -39,6 +39,7 @@ import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.argThat;
import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
@@ -172,7 +173,10 @@ public class PythonCollectionManageableVectorStoreTest {
List<String> expectedIds = Arrays.asList("doc1", "doc2");
- when(mockAdapter.toPythonDocuments(documents)).thenReturn(new
Object());
+ PyObject firstPythonDocument = mock(PyObject.class);
+ PyObject secondPythonDocument = mock(PyObject.class);
+ when(mockAdapter.toPythonDocuments(documents))
+ .thenReturn(List.of(firstPythonDocument,
secondPythonDocument));
when(mockAdapter.callMethod(eq(mockVectorStore), eq("_add_embedding"),
any(Map.class)))
.thenReturn(expectedIds);
@@ -183,6 +187,8 @@ public class PythonCollectionManageableVectorStoreTest {
assertThat(result).containsExactly("doc1", "doc2");
verify(mockAdapter).toPythonDocuments(documents);
+ verify(firstPythonDocument).close();
+ verify(secondPythonDocument).close();
verify(mockAdapter)
.callMethod(
eq(mockVectorStore),
@@ -246,6 +252,7 @@ public class PythonCollectionManageableVectorStoreTest {
assertThat(result).isNotNull();
assertThat(result).hasSize(2);
+ verify(mockPythonDocument).close();
verify(mockAdapter)
.callMethod(
diff --git
a/plan/src/main/java/org/apache/flink/agents/plan/actions/ContextRetrievalAction.java
b/plan/src/main/java/org/apache/flink/agents/plan/actions/ContextRetrievalAction.java
index 6b91c05c..907bfc8e 100644
---
a/plan/src/main/java/org/apache/flink/agents/plan/actions/ContextRetrievalAction.java
+++
b/plan/src/main/java/org/apache/flink/agents/plan/actions/ContextRetrievalAction.java
@@ -25,6 +25,7 @@ import org.apache.flink.agents.api.context.RunnerContext;
import org.apache.flink.agents.api.event.ContextRetrievalRequestEvent;
import org.apache.flink.agents.api.event.ContextRetrievalResponseEvent;
import org.apache.flink.agents.api.resource.ResourceType;
+import org.apache.flink.agents.api.resource.python.PythonObjectScope;
import org.apache.flink.agents.api.vectorstores.BaseVectorStore;
import org.apache.flink.agents.api.vectorstores.Document;
import org.apache.flink.agents.api.vectorstores.VectorStoreQuery;
@@ -136,33 +137,35 @@ public class ContextRetrievalAction {
}
});
- final Object normalized = store.normalizeEmbedding(embedding);
-
- final List<Document> documents =
- ctx.durableExecuteAsync(
- new DurableCallable<List<Document>>() {
- @Override
- public String getId() {
- return "rag-query";
- }
-
- @SuppressWarnings("unchecked")
- @Override
- public Class<List<Document>> getResultClass() {
- return (Class<List<Document>>) (Class<?>)
List.class;
- }
-
- @Override
- public List<Document> call() {
- return store.queryNormalized(
- normalized,
- query.getLimit(),
- query.getCollection(),
- query.getFilters(),
- store.getStoreKwargs());
- }
- });
-
- return new VectorStoreQueryResult(documents);
+ try (PythonObjectScope scope = new PythonObjectScope()) {
+ final Object normalized =
scope.own(store.normalizeEmbedding(embedding));
+
+ final List<Document> documents =
+ ctx.durableExecuteAsync(
+ new DurableCallable<List<Document>>() {
+ @Override
+ public String getId() {
+ return "rag-query";
+ }
+
+ @SuppressWarnings("unchecked")
+ @Override
+ public Class<List<Document>> getResultClass() {
+ return (Class<List<Document>>) (Class<?>)
List.class;
+ }
+
+ @Override
+ public List<Document> call() {
+ return store.queryNormalized(
+ normalized,
+ query.getLimit(),
+ query.getCollection(),
+ query.getFilters(),
+ store.getStoreKwargs());
+ }
+ });
+
+ return new VectorStoreQueryResult(documents);
+ }
}
}
diff --git
a/plan/src/main/java/org/apache/flink/agents/plan/resource/python/PythonMCPPrompt.java
b/plan/src/main/java/org/apache/flink/agents/plan/resource/python/PythonMCPPrompt.java
index cb4a3935..901fa06c 100644
---
a/plan/src/main/java/org/apache/flink/agents/plan/resource/python/PythonMCPPrompt.java
+++
b/plan/src/main/java/org/apache/flink/agents/plan/resource/python/PythonMCPPrompt.java
@@ -21,6 +21,7 @@ import org.apache.flink.agents.api.chat.messages.ChatMessage;
import org.apache.flink.agents.api.chat.messages.MessageRole;
import org.apache.flink.agents.api.metrics.FlinkAgentsMetricGroup;
import org.apache.flink.agents.api.prompt.Prompt;
+import org.apache.flink.agents.api.resource.python.PythonObjectScope;
import org.apache.flink.agents.api.resource.python.PythonResourceAdapter;
import org.apache.flink.agents.api.resource.python.PythonResourceWrapper;
import pemja.core.object.PyObject;
@@ -36,6 +37,7 @@ public class PythonMCPPrompt extends Prompt implements
PythonResourceWrapper {
private final PyObject prompt;
private final PythonResourceAdapter adapter;
+ private boolean closed;
private String name;
public PythonMCPPrompt(PythonResourceAdapter adapter, PyObject prompt) {
@@ -75,18 +77,31 @@ public class PythonMCPPrompt extends Prompt implements
PythonResourceWrapper {
@Override
public List<ChatMessage> formatMessages(MessageRole defaultRole,
Map<String, String> kwargs) {
Map<String, Object> parameters = new HashMap<>(kwargs);
- Object pythonRole = adapter.invoke(FROM_JAVA_MESSAGE_ROLE,
defaultRole);
- parameters.put("role", pythonRole);
+ try (PythonObjectScope scope = new PythonObjectScope()) {
+ Object pythonRole = adapter.invoke(FROM_JAVA_MESSAGE_ROLE,
defaultRole);
+ parameters.put("role", pythonRole);
- Object result = adapter.callMethod(prompt, "format_messages",
parameters);
- if (result instanceof List) {
- List<Object> pythonMessages = (List<Object>) result;
- List<ChatMessage> messages = new
ArrayList<>(pythonMessages.size());
- for (Object pythonMessage : pythonMessages) {
- messages.add(adapter.fromPythonChatMessage(pythonMessage));
+ Object result = scope.own(adapter.callMethod(prompt,
"format_messages", parameters));
+ if (result instanceof List) {
+ List<Object> pythonMessages = (List<Object>) result;
+ List<ChatMessage> messages = new
ArrayList<>(pythonMessages.size());
+ for (Object pythonMessage : pythonMessages) {
+ messages.add(adapter.fromPythonChatMessage(pythonMessage));
+ }
+ return messages;
}
- return messages;
+ return Collections.emptyList();
+ }
+ }
+
+ @Override
+ public void close() throws Exception {
+ if (closed || prompt == null) {
+ return;
+ }
+ closed = true;
+ try (prompt) {
+ adapter.callMethod(prompt, "close", Map.of());
}
- return Collections.emptyList();
}
}
diff --git
a/plan/src/main/java/org/apache/flink/agents/plan/resource/python/PythonMCPServer.java
b/plan/src/main/java/org/apache/flink/agents/plan/resource/python/PythonMCPServer.java
index 6bfd95d2..08ae93f2 100644
---
a/plan/src/main/java/org/apache/flink/agents/plan/resource/python/PythonMCPServer.java
+++
b/plan/src/main/java/org/apache/flink/agents/plan/resource/python/PythonMCPServer.java
@@ -31,10 +31,12 @@ import javax.annotation.Nullable;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
+import java.util.Map;
public class PythonMCPServer extends Resource implements PythonResourceWrapper
{
private final PyObject server;
private final PythonResourceAdapter adapter;
+ private boolean closed;
/**
* Creates a new PythonMCPServer.
@@ -107,4 +109,15 @@ public class PythonMCPServer extends Resource implements
PythonResourceWrapper {
public ResourceType getResourceType() {
return ResourceType.MCP_SERVER;
}
+
+ @Override
+ public void close() throws Exception {
+ if (closed || server == null) {
+ return;
+ }
+ closed = true;
+ try (server) {
+ adapter.callMethod(server, "close", Map.of());
+ }
+ }
}
diff --git
a/plan/src/main/java/org/apache/flink/agents/plan/resource/python/PythonMCPTool.java
b/plan/src/main/java/org/apache/flink/agents/plan/resource/python/PythonMCPTool.java
index ecd16286..859d67ce 100644
---
a/plan/src/main/java/org/apache/flink/agents/plan/resource/python/PythonMCPTool.java
+++
b/plan/src/main/java/org/apache/flink/agents/plan/resource/python/PythonMCPTool.java
@@ -43,6 +43,7 @@ public class PythonMCPTool extends Tool
"python_java_utils.invoke_python_tool_instance";
private final PyObject tool;
private final PythonResourceAdapter adapter;
+ private boolean closed;
@Nullable private final String mcpServerName;
/**
@@ -122,4 +123,15 @@ public class PythonMCPTool extends Tool
}
return metadata;
}
+
+ @Override
+ public void close() throws Exception {
+ if (closed || tool == null) {
+ return;
+ }
+ closed = true;
+ try (tool) {
+ adapter.callMethod(tool, "close", Map.of());
+ }
+ }
}
diff --git
a/plan/src/test/java/org/apache/flink/agents/plan/resource/python/PythonMCPResourceTest.java
b/plan/src/test/java/org/apache/flink/agents/plan/resource/python/PythonMCPResourceTest.java
new file mode 100644
index 00000000..47e11cb2
--- /dev/null
+++
b/plan/src/test/java/org/apache/flink/agents/plan/resource/python/PythonMCPResourceTest.java
@@ -0,0 +1,106 @@
+/*
+ * 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.
+ */
+package org.apache.flink.agents.plan.resource.python;
+
+import org.apache.flink.agents.api.chat.messages.ChatMessage;
+import org.apache.flink.agents.api.chat.messages.MessageRole;
+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;
+import org.junit.jupiter.api.Test;
+import pemja.core.object.PyObject;
+
+import java.util.List;
+import java.util.Map;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+class PythonMCPResourceTest {
+
+ @Test
+ void closesDiscoveredResourcesOnce() throws Exception {
+ PythonResourceAdapter adapter = mock(PythonResourceAdapter.class);
+ PyObject serverObject = mock(PyObject.class);
+ PyObject toolObject = mock(PyObject.class);
+ PyObject promptObject = mock(PyObject.class);
+ PythonMCPServer server =
+ new PythonMCPServer(
+ adapter,
+ serverObject,
+ mock(ResourceDescriptor.class),
+ mock(ResourceContext.class));
+
+ when(adapter.callMethod(serverObject, "list_tools", Map.of()))
+ .thenReturn(List.of(toolObject));
+
when(adapter.invoke("python_java_utils.get_java_tool_metadata_from_tool",
toolObject))
+ .thenReturn(
+ Map.of(
+ "name", "tool",
+ "description", "description",
+ "inputSchema", "{}"));
+ when(adapter.callMethod(serverObject, "list_prompts", Map.of()))
+ .thenReturn(List.of(promptObject));
+
+ PythonMCPTool tool = server.listTools("server").get(0);
+ PythonMCPPrompt prompt = server.listPrompts().get(0);
+
+ verify(toolObject, never()).close();
+ verify(promptObject, never()).close();
+
+ tool.close();
+ prompt.close();
+ server.close();
+ tool.close();
+ prompt.close();
+ server.close();
+
+ verify(adapter).callMethod(toolObject, "close", Map.of());
+ verify(adapter).callMethod(promptObject, "close", Map.of());
+ verify(adapter).callMethod(serverObject, "close", Map.of());
+ verify(toolObject).close();
+ verify(promptObject).close();
+ verify(serverObject).close();
+ }
+
+ @Test
+ void releasesPromptBridgeValuesAfterConversion() throws Exception {
+ PythonResourceAdapter adapter = mock(PythonResourceAdapter.class);
+ PyObject promptObject = mock(PyObject.class);
+ PyObject messageObject = mock(PyObject.class);
+ ChatMessage message = mock(ChatMessage.class);
+ PythonMCPPrompt prompt = new PythonMCPPrompt(adapter, promptObject);
+
+ when(adapter.invoke("python_java_utils.from_java_message_role",
MessageRole.USER))
+ .thenReturn("user");
+ when(adapter.callMethod(eq(promptObject), eq("format_messages"),
any(Map.class)))
+ .thenReturn(List.of(messageObject));
+ when(adapter.fromPythonChatMessage(messageObject)).thenReturn(message);
+
+ assertThat(prompt.formatMessages(MessageRole.USER,
Map.of())).containsExactly(message);
+
+ verify(messageObject).close();
+
+ prompt.close();
+ }
+}
diff --git
a/plan/src/test/java/org/apache/flink/agents/plan/resourceprovider/PythonResourceProviderTest.java
b/plan/src/test/java/org/apache/flink/agents/plan/resourceprovider/PythonResourceProviderTest.java
new file mode 100644
index 00000000..9dc54824
--- /dev/null
+++
b/plan/src/test/java/org/apache/flink/agents/plan/resourceprovider/PythonResourceProviderTest.java
@@ -0,0 +1,65 @@
+/*
+ * 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.
+ */
+package org.apache.flink.agents.plan.resourceprovider;
+
+import org.apache.flink.agents.api.resource.Resource;
+import org.apache.flink.agents.api.resource.ResourceContext;
+import org.apache.flink.agents.api.resource.ResourceDescriptor;
+import org.apache.flink.agents.api.resource.ResourceType;
+import org.apache.flink.agents.api.resource.python.PythonResourceAdapter;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.EnumSource;
+import pemja.core.object.PyObject;
+
+import java.util.Map;
+
+import static org.mockito.ArgumentMatchers.anyMap;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+class PythonResourceProviderTest {
+
+ @ParameterizedTest
+ @EnumSource(
+ value = ResourceType.class,
+ names = {
+ "CHAT_MODEL", "CHAT_MODEL_CONNECTION", "EMBEDDING_MODEL",
+ "EMBEDDING_MODEL_CONNECTION", "VECTOR_STORE", "MCP_SERVER"
+ })
+ void providedResourceOwnsAndClosesPythonHandleOnce(ResourceType type)
throws Exception {
+ PythonResourceAdapter adapter = mock(PythonResourceAdapter.class);
+ PyObject pythonResource = mock(PyObject.class);
+ ResourceDescriptor descriptor =
+ new ResourceDescriptor("example.module", "ExampleModel",
Map.of());
+ PythonResourceProvider provider = new PythonResourceProvider("model",
type, descriptor);
+ provider.setPythonResourceAdapter(adapter);
+ when(adapter.initPythonResource(anyString(), anyString(), anyMap()))
+ .thenReturn(pythonResource);
+
+ Resource resource = provider.provide(mock(ResourceContext.class));
+
+ verify(pythonResource, never()).close();
+ resource.close();
+ resource.close();
+ verify(adapter).callMethod(pythonResource, "close", Map.of());
+ verify(pythonResource).close();
+ }
+}
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/PythonMCPResourceDiscovery.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/PythonMCPResourceDiscovery.java
index b9e0b48d..eca4019e 100644
---
a/runtime/src/main/java/org/apache/flink/agents/runtime/PythonMCPResourceDiscovery.java
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/PythonMCPResourceDiscovery.java
@@ -71,6 +71,7 @@ public class PythonMCPResourceDiscovery {
provider.setPythonResourceAdapter(adapter);
PythonMCPServer server = (PythonMCPServer)
provider.provide(cache.getResourceContext());
+ cache.put(provider.getName(), MCP_SERVER, server);
for (PythonMCPTool tool : server.listTools(provider.getName())) {
cache.put(tool.getName(), TOOL, tool);
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/memory/Mem0LongTermMemory.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/memory/Mem0LongTermMemory.java
index bc4d16a1..67c0cc38 100644
---
a/runtime/src/main/java/org/apache/flink/agents/runtime/memory/Mem0LongTermMemory.java
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/memory/Mem0LongTermMemory.java
@@ -19,6 +19,7 @@ package org.apache.flink.agents.runtime.memory;
import org.apache.flink.agents.api.memory.MemorySet;
import org.apache.flink.agents.api.memory.MemorySetItem;
+import org.apache.flink.agents.api.resource.python.PythonObjectScope;
import org.apache.flink.agents.api.resource.python.PythonResourceAdapter;
import pemja.core.object.PyObject;
@@ -92,13 +93,15 @@ public class Mem0LongTermMemory implements
InteranlBaseLongTermMemory {
MemorySet memorySet,
List<String> memoryItems,
@Nullable List<Map<String, Object>> metadatas) {
- Map<String, Object> kwargs = new HashMap<>();
- kwargs.put("memory_set", buildPyMemorySet(memorySet));
- kwargs.put("memory_items", memoryItems);
- if (metadatas != null) {
- kwargs.put("metadatas", metadatas);
+ try (PythonObjectScope scope = new PythonObjectScope()) {
+ Map<String, Object> kwargs = new HashMap<>();
+ kwargs.put("memory_set", buildPyMemorySet(scope, memorySet));
+ kwargs.put("memory_items", memoryItems);
+ if (metadatas != null) {
+ kwargs.put("metadatas", metadatas);
+ }
+ return (List<String>) adapter.callMethod(pyMem0, "add", kwargs);
}
- return (List<String>) adapter.callMethod(pyMem0, "add", kwargs);
}
@Override
@@ -107,29 +110,33 @@ public class Mem0LongTermMemory implements
InteranlBaseLongTermMemory {
@Nullable List<String> ids,
@Nullable Map<String, Object> filters,
@Nullable Integer limit) {
- Map<String, Object> kwargs = new HashMap<>();
- kwargs.put("memory_set", buildPyMemorySet(memorySet));
- if (ids != null) {
- kwargs.put("ids", ids);
- }
- if (filters != null) {
- kwargs.put("filters", filters);
- }
- if (limit != null) {
- kwargs.put("limit", limit);
+ try (PythonObjectScope scope = new PythonObjectScope()) {
+ Map<String, Object> kwargs = new HashMap<>();
+ kwargs.put("memory_set", buildPyMemorySet(scope, memorySet));
+ if (ids != null) {
+ kwargs.put("ids", ids);
+ }
+ if (filters != null) {
+ kwargs.put("filters", filters);
+ }
+ if (limit != null) {
+ kwargs.put("limit", limit);
+ }
+ Object pyItems = scope.own(adapter.callMethod(pyMem0, "get",
kwargs));
+ return convertItems(pyItems);
}
- Object pyItems = adapter.callMethod(pyMem0, "get", kwargs);
- return convertItems(pyItems);
}
@Override
public void delete(MemorySet memorySet, @Nullable List<String> ids) {
- Map<String, Object> kwargs = new HashMap<>();
- kwargs.put("memory_set", buildPyMemorySet(memorySet));
- if (ids != null) {
- kwargs.put("ids", ids);
+ try (PythonObjectScope scope = new PythonObjectScope()) {
+ Map<String, Object> kwargs = new HashMap<>();
+ kwargs.put("memory_set", buildPyMemorySet(scope, memorySet));
+ if (ids != null) {
+ kwargs.put("ids", ids);
+ }
+ adapter.callMethod(pyMem0, "delete", kwargs);
}
- adapter.callMethod(pyMem0, "delete", kwargs);
}
@Override
@@ -139,15 +146,17 @@ public class Mem0LongTermMemory implements
InteranlBaseLongTermMemory {
int limit,
@Nullable Map<String, Object> filters,
Map<String, Object> extraArgs) {
- Map<String, Object> kwargs = new HashMap<>(extraArgs);
- kwargs.put("memory_set", buildPyMemorySet(memorySet));
- kwargs.put("query", query);
- kwargs.put("limit", limit);
- if (filters != null) {
- kwargs.put("filters", filters);
+ try (PythonObjectScope scope = new PythonObjectScope()) {
+ Map<String, Object> kwargs = new HashMap<>(extraArgs);
+ kwargs.put("memory_set", buildPyMemorySet(scope, memorySet));
+ kwargs.put("query", query);
+ kwargs.put("limit", limit);
+ if (filters != null) {
+ kwargs.put("filters", filters);
+ }
+ Object pyItems = scope.own(adapter.callMethod(pyMem0, "search",
kwargs));
+ return convertItems(pyItems);
}
- Object pyItems = adapter.callMethod(pyMem0, "search", kwargs);
- return convertItems(pyItems);
}
@Override
@@ -203,7 +212,7 @@ public class Mem0LongTermMemory implements
InteranlBaseLongTermMemory {
}
}
- private Object buildPyMemorySet(MemorySet memorySet) {
+ private Object buildPyMemorySet(PythonObjectScope scope, MemorySet
memorySet) {
// Mem0 ignores a falsy agent_id rather than matching on it, so
forwarding an
// unbound or empty-keyed set would widen the operation to every key
sharing the job
// id and set name, which for a delete means deleting another key's
items.
@@ -216,12 +225,13 @@ public class Mem0LongTermMemory implements
InteranlBaseLongTermMemory {
memorySet.getName()));
}
requireNonEmptyPartitionKey(memorySet.getPartitionKey());
- return adapter.invoke(
- TO_PYTHON_MEMORY_SET,
- memorySet.getName(),
- memorySet.getPartitionKey(),
- memorySet.getObservationId(),
- memorySet.isObservationSuppressed());
+ return scope.own(
+ adapter.invoke(
+ TO_PYTHON_MEMORY_SET,
+ memorySet.getName(),
+ memorySet.getPartitionKey(),
+ memorySet.getObservationId(),
+ memorySet.isObservationSuppressed()));
}
/** Returns the partition key in scope, refusing what Mem0 cannot scope an
operation to. */
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutor.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutor.java
index e645b071..4012bd11 100644
---
a/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutor.java
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutor.java
@@ -23,6 +23,7 @@ import org.apache.flink.agents.api.Event;
import org.apache.flink.agents.api.agents.AgentExecutionOptions;
import org.apache.flink.agents.api.resource.Resource;
import org.apache.flink.agents.api.resource.ResourceType;
+import org.apache.flink.agents.api.resource.python.PythonObjectScope;
import org.apache.flink.agents.plan.AgentPlan;
import org.apache.flink.agents.plan.PythonFunction;
import org.apache.flink.agents.runtime.operator.ActionTask;
@@ -230,9 +231,9 @@ public class PythonActionExecutor implements AutoCloseable {
/**
* Execute the Python function, which may return a Python coroutine
(awaitable) that needs to be
* processed in the future. Due to an issue in Pemja regarding incorrect
object reference
- * counting, this may lead to garbage collection of the object. To prevent
this, we use the set
- * and get methods to manually increment the object's reference count,
then return the name of
- * the Python awaitable variable.
+ * counting, this may lead to garbage collection of the object. To prevent
this, we store the
+ * awaitable in the interpreter globals, then return the name of that
variable. The temporary
+ * Java wrapper can be closed after the interpreter takes ownership of its
own reference.
*
* @return The name of the Python awaitable variable. It may be null if
the Python function does
* not return a coroutine.
@@ -242,10 +243,10 @@ public class PythonActionExecutor implements
AutoCloseable {
function.setInterpreter(interpreter);
String eventJson = new ObjectMapper().writeValueAsString(event);
- Object pythonEventObject =
interpreter.invoke(CONVERT_JSON_TO_PYTHON_EVENT, eventJson);
-
- try {
- Object calledResult = function.call(pythonEventObject,
pythonRunnerContext);
+ try (PyObject pythonEventObject =
+ (PyObject)
interpreter.invoke(CONVERT_JSON_TO_PYTHON_EVENT, eventJson);
+ PyObject calledResult =
+ (PyObject) function.call(pythonEventObject,
pythonRunnerContext)) {
if (calledResult == null) {
return null;
} else {
@@ -301,15 +302,23 @@ public class PythonActionExecutor implements
AutoCloseable {
* @return true if the awaitable has completed; false otherwise
*/
public boolean callPythonAwaitable(String pythonAwaitableRef) {
- // Calling awaitable.send(None) in Python returns a tuple of
(finished, output).
- Object pythonAwaitable = interpreter.get(pythonAwaitableRef);
- checkState(
- pythonAwaitable != null,
- "Python awaitable '%s' not found in interpreter. ",
- pythonAwaitableRef);
- Object invokeResult = interpreter.invoke(CALL_PYTHON_AWAITABLE,
pythonAwaitable);
- checkState(invokeResult.getClass().isArray() && ((Object[])
invokeResult).length == 2);
- return (boolean) ((Object[]) invokeResult)[0];
+ try (PythonObjectScope scope = new PythonObjectScope()) {
+ PyObject pythonAwaitable = scope.own((PyObject)
interpreter.get(pythonAwaitableRef));
+ checkState(
+ pythonAwaitable != null,
+ "Python awaitable '%s' not found in interpreter.",
+ pythonAwaitableRef);
+ // Actions communicate through Events, so this caller consumes
only the completion flag.
+ Object invokeResult =
+ scope.own(interpreter.invoke(CALL_PYTHON_AWAITABLE,
pythonAwaitable));
+ checkState(invokeResult instanceof Object[] && ((Object[])
invokeResult).length == 2);
+ Object[] result = (Object[]) invokeResult;
+ boolean finished = (boolean) result[0];
+ if (finished) {
+ interpreter.exec("del " + pythonAwaitableRef);
+ }
+ return finished;
+ }
}
@Override
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonResourceAdapterImpl.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonResourceAdapterImpl.java
index 31c89713..73f206ab 100644
---
a/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonResourceAdapterImpl.java
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/python/utils/PythonResourceAdapterImpl.java
@@ -24,6 +24,7 @@ import org.apache.flink.agents.api.prompt.Prompt;
import org.apache.flink.agents.api.resource.Resource;
import org.apache.flink.agents.api.resource.ResourceContext;
import org.apache.flink.agents.api.resource.ResourceType;
+import org.apache.flink.agents.api.resource.python.PythonObjectScope;
import org.apache.flink.agents.api.resource.python.PythonResourceAdapter;
import org.apache.flink.agents.api.resource.python.PythonResourceWrapper;
import org.apache.flink.agents.api.tools.Tool;
@@ -194,9 +195,13 @@ public class PythonResourceAdapterImpl implements
PythonResourceAdapter, AutoClo
@Override
public VectorStoreQueryResult fromPythonVectorStoreQueryResult(
PyObject pythonVectorStoreQueryResult) {
- List<PyObject> pythonDocuments =
- (List<PyObject>)
pythonVectorStoreQueryResult.getAttr("documents", List.class);
- return new
VectorStoreQueryResult(fromPythonDocuments(pythonDocuments));
+ try (PythonObjectScope scope = new PythonObjectScope()) {
+ List<PyObject> pythonDocuments =
+ scope.own(
+ (List<PyObject>)
+
pythonVectorStoreQueryResult.getAttr("documents", List.class));
+ return new
VectorStoreQueryResult(fromPythonDocuments(pythonDocuments));
+ }
}
@Override
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/PythonMCPResourceDiscoveryTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/PythonMCPResourceDiscoveryTest.java
new file mode 100644
index 00000000..36fb9217
--- /dev/null
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/PythonMCPResourceDiscoveryTest.java
@@ -0,0 +1,73 @@
+/*
+ * 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.
+ */
+package org.apache.flink.agents.runtime;
+
+import org.apache.flink.agents.api.resource.ResourceType;
+import org.apache.flink.agents.api.resource.python.PythonResourceAdapter;
+import org.apache.flink.agents.plan.resource.python.PythonMCPPrompt;
+import org.apache.flink.agents.plan.resource.python.PythonMCPServer;
+import org.apache.flink.agents.plan.resource.python.PythonMCPTool;
+import org.apache.flink.agents.plan.resourceprovider.PythonResourceProvider;
+import org.apache.flink.agents.plan.resourceprovider.ResourceProvider;
+import org.junit.jupiter.api.Test;
+
+import java.util.List;
+import java.util.Map;
+
+import static org.apache.flink.agents.api.resource.ResourceType.MCP_SERVER;
+import static org.apache.flink.agents.api.resource.ResourceType.PROMPT;
+import static org.apache.flink.agents.api.resource.ResourceType.TOOL;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+class PythonMCPResourceDiscoveryTest {
+
+ @Test
+ void cachesServerAndDiscoveredResourcesForNormalShutdown() throws
Exception {
+ PythonResourceAdapter adapter = mock(PythonResourceAdapter.class);
+ PythonResourceProvider serverProvider =
mock(PythonResourceProvider.class);
+ PythonMCPServer server = mock(PythonMCPServer.class);
+ PythonMCPTool tool = mock(PythonMCPTool.class);
+ PythonMCPPrompt prompt = mock(PythonMCPPrompt.class);
+ Map<ResourceType, Map<String, ResourceProvider>> providers =
+ Map.of(MCP_SERVER, Map.of("server", serverProvider));
+ ResourceCache cache = new ResourceCache(providers);
+
+ when(serverProvider.getName()).thenReturn("server");
+ when(serverProvider.provide(any())).thenReturn(server);
+ when(server.listTools("server")).thenReturn(List.of(tool));
+ when(server.listPrompts()).thenReturn(List.of(prompt));
+ when(tool.getName()).thenReturn("tool");
+ when(prompt.getName()).thenReturn("prompt");
+
+ PythonMCPResourceDiscovery.discoverPythonMCPResources(providers,
adapter, cache);
+
+ assertThat(cache.getResource("server", MCP_SERVER)).isSameAs(server);
+ assertThat(cache.getResource("tool", TOOL)).isSameAs(tool);
+ assertThat(cache.getResource("prompt", PROMPT)).isSameAs(prompt);
+
+ cache.close();
+
+ verify(server).close();
+ verify(tool).close();
+ verify(prompt).close();
+ }
+}
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/memory/Mem0LongTermMemoryTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/memory/Mem0LongTermMemoryTest.java
index 1554be96..696b2a89 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/memory/Mem0LongTermMemoryTest.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/memory/Mem0LongTermMemoryTest.java
@@ -37,6 +37,7 @@ import static
org.assertj.core.api.Assertions.assertThatThrownBy;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
@@ -102,7 +103,7 @@ public class Mem0LongTermMemoryTest {
}
@Test
- void testAddForwardsKwargsAndReturnsIds() throws Exception {
+ void testAddForwardsKwargsAndReleasesTemporaryMemorySet() throws Exception
{
MemorySet ms = ltm.getMemorySet("notes");
when(mockAdapter.callMethod(eq(mockPyMem0), eq("add"), any()))
.thenReturn(List.of("a", "b"));
@@ -114,6 +115,7 @@ public class Mem0LongTermMemoryTest {
assertThat(captureKwargs("add"))
.containsKeys("memory_set", "memory_items", "metadatas")
.containsEntry("memory_set", mockPyMemorySet);
+ verify(mockPyMemorySet).close();
}
@Test
@@ -131,8 +133,10 @@ public class Mem0LongTermMemoryTest {
@Test
void testGetWithIdsAndFiltersConvertsItems() throws Exception {
MemorySet ms = ltm.getMemorySet("notes");
- when(mockAdapter.callMethod(eq(mockPyMem0), eq("get"),
any())).thenReturn("py_items");
- when(mockAdapter.invoke(eq("python_java_utils.mem0_items_to_java"),
eq("py_items")))
+ PyObject pythonItem = mock(PyObject.class);
+ List<PyObject> pythonItems = List.of(pythonItem);
+ when(mockAdapter.callMethod(eq(mockPyMem0), eq("get"),
any())).thenReturn(pythonItems);
+ when(mockAdapter.invoke(eq("python_java_utils.mem0_items_to_java"),
eq(pythonItems)))
.thenReturn(
List.of(
Map.of(
@@ -149,6 +153,7 @@ public class Mem0LongTermMemoryTest {
assertThat(item.getId()).isEqualTo("id1");
assertThat(item.getValue()).isEqualTo("hello");
assertThat(item.getAdditionalMetadata()).containsEntry("k", "v");
+ verify(pythonItem).close();
assertThat(item.getCreatedAt()).isNull();
assertThat(captureKwargs("get"))
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutorTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutorTest.java
index 5bdac5dc..b2ac5060 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutorTest.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/PythonActionExecutorTest.java
@@ -18,24 +18,33 @@
package org.apache.flink.agents.runtime.python.utils;
import com.fasterxml.jackson.databind.ObjectMapper;
+import org.apache.flink.agents.api.InputEvent;
import org.apache.flink.agents.api.agents.AgentExecutionOptions;
import org.apache.flink.agents.plan.AgentPlan;
+import org.apache.flink.agents.plan.PythonFunction;
import org.apache.flink.agents.runtime.python.context.PythonRunnerContextImpl;
import org.apache.flink.types.Row;
import org.junit.jupiter.api.Test;
+import org.mockito.ArgumentCaptor;
import org.mockito.InOrder;
import pemja.core.PythonInterpreter;
import pemja.core.object.PyObject;
import java.lang.reflect.Field;
import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.ArgumentMatchers.isNull;
+import static org.mockito.ArgumentMatchers.same;
import static org.mockito.Mockito.clearInvocations;
import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.inOrder;
import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.verifyNoInteractions;
import static org.mockito.Mockito.when;
@@ -51,6 +60,9 @@ class PythonActionExecutorTest {
"flink_runner_context.create_flink_runner_context";
private static final String CLOSE_FLINK_RUNNER_CONTEXT =
"flink_runner_context.close_flink_runner_context";
+ private static final String CONVERT_JSON_TO_PYTHON_EVENT =
+ "python_java_utils.convert_json_to_python_event";
+ private static final String CALL_PYTHON_AWAITABLE =
"function.call_python_awaitable";
@Test
void resolvesPickledPythonKeyTextFromPyFlinkKeyRow() throws Exception {
@@ -243,9 +255,95 @@ class PythonActionExecutorTest {
fixture.interpreter, fixture.asyncThreadPool,
fixture.runnerContextObject);
}
+ @Test
+ void closesPythonEventAfterSynchronousAction() throws Exception {
+ PythonInterpreter interpreter = mock(PythonInterpreter.class);
+ PythonRunnerContextImpl runnerContext =
mock(PythonRunnerContextImpl.class);
+ PythonActionExecutor executor = newExecutor(interpreter,
runnerContext);
+ PythonFunction function = mock(PythonFunction.class);
+ PyObject pythonEvent = mock(PyObject.class);
+ when(interpreter.invoke(same(CONVERT_JSON_TO_PYTHON_EVENT),
anyString()))
+ .thenReturn(pythonEvent);
+ when(function.call(same(pythonEvent), isNull())).thenReturn(null);
+
+ assertThat(executor.executePythonFunction(function, new
InputEvent(1L))).isNull();
+
+ verify(function).setInterpreter(interpreter);
+ verify(pythonEvent).close();
+ }
+
+ @Test
+ void closesTemporaryWrappersAfterStoringAwaitable() throws Exception {
+ PythonInterpreter interpreter = mock(PythonInterpreter.class);
+ PythonRunnerContextImpl runnerContext =
mock(PythonRunnerContextImpl.class);
+ PythonActionExecutor executor = newExecutor(interpreter,
runnerContext);
+ PythonFunction function = mock(PythonFunction.class);
+ PyObject pythonEvent = mock(PyObject.class);
+ PyObject pythonAwaitable = mock(PyObject.class);
+ when(interpreter.invoke(same(CONVERT_JSON_TO_PYTHON_EVENT),
anyString()))
+ .thenReturn(pythonEvent);
+ when(function.call(same(pythonEvent),
isNull())).thenReturn(pythonAwaitable);
+
+ String pythonAwaitableRef = executor.executePythonFunction(function,
new InputEvent(1L));
+
+ ArgumentCaptor<String> refCaptor =
ArgumentCaptor.forClass(String.class);
+ verify(interpreter).set(refCaptor.capture(), same(pythonAwaitable));
+ assertThat(pythonAwaitableRef)
+ .isEqualTo(refCaptor.getValue())
+ .startsWith("python_awaitable_");
+ InOrder closeOrder = inOrder(interpreter, pythonAwaitable,
pythonEvent);
+ closeOrder.verify(interpreter).set(pythonAwaitableRef,
pythonAwaitable);
+ closeOrder.verify(pythonAwaitable).close();
+ closeOrder.verify(pythonEvent).close();
+ }
+
+ @Test
+ void closesRetrievedAwaitableWhileItIsPending() throws Exception {
+ PythonInterpreter interpreter = mock(PythonInterpreter.class);
+ PythonActionExecutor executor = newExecutor(interpreter);
+ PyObject pythonAwaitable = mock(PyObject.class);
+ PyObject yieldedValue = mock(PyObject.class);
+ String pythonAwaitableRef = "python_awaitable_1";
+ when(interpreter.get(pythonAwaitableRef)).thenReturn(pythonAwaitable);
+ when(interpreter.invoke(CALL_PYTHON_AWAITABLE, pythonAwaitable))
+ .thenReturn(
+ new Object[] {false, Map.of("items",
List.of(yieldedValue, yieldedValue))});
+
+ assertThat(executor.callPythonAwaitable(pythonAwaitableRef)).isFalse();
+
+ verify(pythonAwaitable).close();
+ verify(yieldedValue).close();
+ verify(interpreter, never()).exec(anyString());
+ }
+
+ @Test
+ void deletesCompletedAwaitableAndClosesRetrievedWrapper() throws Exception
{
+ PythonInterpreter interpreter = mock(PythonInterpreter.class);
+ PythonActionExecutor executor = newExecutor(interpreter);
+ PyObject pythonAwaitable = mock(PyObject.class);
+ PyObject returnedValue = mock(PyObject.class);
+ String pythonAwaitableRef = "python_awaitable_1";
+ when(interpreter.get(pythonAwaitableRef)).thenReturn(pythonAwaitable);
+ when(interpreter.invoke(CALL_PYTHON_AWAITABLE, pythonAwaitable))
+ .thenReturn(new Object[] {true, new Object[] {returnedValue}});
+
+ assertThat(executor.callPythonAwaitable(pythonAwaitableRef)).isTrue();
+
+ InOrder interpreterOrder = inOrder(interpreter);
+ interpreterOrder.verify(interpreter).invoke(CALL_PYTHON_AWAITABLE,
pythonAwaitable);
+ interpreterOrder.verify(interpreter).exec("del " + pythonAwaitableRef);
+ verify(returnedValue).close();
+ verify(pythonAwaitable).close();
+ }
+
private static PythonActionExecutor newExecutor(PythonInterpreter
interpreter)
throws Exception {
- return new PythonActionExecutor(interpreter, null, null, null,
"test-job");
+ return newExecutor(interpreter, null);
+ }
+
+ private static PythonActionExecutor newExecutor(
+ PythonInterpreter interpreter, PythonRunnerContextImpl
runnerContext) throws Exception {
+ return new PythonActionExecutor(interpreter, null, null,
runnerContext, "test-job");
}
private static TestFixture createOpenedExecutor() throws Exception {
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/PythonResourceAdapterImplTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/PythonResourceAdapterImplTest.java
index f372821a..b2225fdb 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/PythonResourceAdapterImplTest.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/python/utils/PythonResourceAdapterImplTest.java
@@ -25,6 +25,7 @@ import org.apache.flink.agents.api.resource.ResourceContext;
import org.apache.flink.agents.api.resource.ResourceType;
import org.apache.flink.agents.api.resource.python.PythonResourceWrapper;
import org.apache.flink.agents.api.tools.Tool;
+import org.apache.flink.agents.api.vectorstores.VectorStoreQueryResult;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
@@ -34,6 +35,7 @@ import pemja.core.PythonInterpreter;
import pemja.core.object.PyObject;
import java.util.HashMap;
+import java.util.List;
import java.util.Map;
import static org.assertj.core.api.Assertions.assertThat;
@@ -196,6 +198,22 @@ public class PythonResourceAdapterImplTest {
.invoke(PythonResourceAdapterImpl.CALL_METHOD, obj,
methodName, kwargs);
}
+ @Test
+ void closesRetrievedDocumentsAfterQueryResultConversion() throws Exception
{
+ PyObject pythonResult = mock(PyObject.class);
+ PyObject pythonDocument = mock(PyObject.class);
+ when(pythonResult.getAttr("documents",
List.class)).thenReturn(List.of(pythonDocument));
+ when(pythonDocument.getAttr("content")).thenReturn("content");
+ when(pythonDocument.getAttr("metadata",
Map.class)).thenReturn(Map.of("source", "test"));
+ when(pythonDocument.getAttr("id")).thenReturn("doc-1");
+
+ VectorStoreQueryResult result =
+
pythonResourceAdapter.fromPythonVectorStoreQueryResult(pythonResult);
+
+
assertThat(result.getDocuments().get(0).getContent()).isEqualTo("content");
+ verify(pythonDocument).close();
+ }
+
@Test
void testSetMetricGroup() {
Object pythonResource = new Object();
