[runtime][python] Manage Pemja references across actions and resources

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

@@ -22,6 +22,7 @@
 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 ChatMessage chat(
             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());
+        }
     }
 }

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

@@ -22,6 +22,7 @@
 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 ChatMessage chat(
 
         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 void setMetricGroup(FlinkAgentsMetricGroup metricGroup) {
     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());
+        }
+    }
 }

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 void setMetricGroup(FlinkAgentsMetricGroup metricGroup) {
 
     @Override
     public void close() throws Exception {
-        this.embeddingModel.close();
+        if (closed || embeddingModel == null) {
+            return;
+        }
+        closed = true;
+        try (embeddingModel) {
+            adapter.callMethod(embeddingModel, "close", Map.of());
+        }
     }
 }

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 void setMetricGroup(FlinkAgentsMetricGroup metricGroup) {
         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());
+        }
+    }
 }

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.");
+        }
+    }
+}

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

@@ -21,6 +21,7 @@
 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 @@
 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 List<Document> get(
             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 List<Document> queryEmbedding(
         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,32 +208,39 @@ public List<Document> queryNormalized(
         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
     @SuppressWarnings("unchecked")
     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 void setMetricGroup(FlinkAgentsMetricGroup metricGroup) {
         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());
+        }
+    }
 }