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 8ce3662e3..580f4be8a 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.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 messages, List tools, Map modelParams) { Map kwargs = new HashMap<>(modelParams); - List pythonMessages = new ArrayList<>(); - for (ChatMessage message : messages) { - pythonMessages.add(adapter.toPythonChatMessage(message)); - } - kwargs.put("messages", pythonMessages); + try (PythonObjectScope scope = new PythonObjectScope()) { + List pythonMessages = new ArrayList<>(); + for (ChatMessage message : messages) { + pythonMessages.add(scope.own(adapter.toPythonChatMessage(message))); + } + kwargs.put("messages", pythonMessages); - List pythonTools = new ArrayList<>(); - for (Tool tool : tools) { - pythonTools.add(adapter.convertToPythonTool(tool)); - } - kwargs.put("tools", pythonTools); + List 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 846105afb..afae1e357 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.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 kwargs = new HashMap<>(modelParams); - List pythonMessages = new ArrayList<>(); - for (ChatMessage message : messages) { - pythonMessages.add(adapter.toPythonChatMessage(message)); - } + try (PythonObjectScope scope = new PythonObjectScope()) { + List 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 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 d9bdb8d41..0e557dbcb 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 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()); + } } } 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 b6febbc58..2285bb16b 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 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()); + } + } } 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 000000000..44192ba65 --- /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. + * + *

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 ownedObjects = Collections.newSetFromMap(new IdentityHashMap<>()); + private boolean closed; + + /** Adds every {@link PyObject} reachable through the supplied bridge result to this scope. */ + public 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 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 21bb18931..4a393adc1 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 @@ 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 get( kwargs.put("limit", limit); } - Object pythonDocuments = adapter.callMethod(vectorStore, "get", kwargs); - - return adapter.fromPythonDocuments((List) pythonDocuments); + try (PythonObjectScope scope = new PythonObjectScope()) { + Object pythonDocuments = scope.own(adapter.callMethod(vectorStore, "get", kwargs)); + return adapter.fromPythonDocuments((List) pythonDocuments); + } } @Override @@ -158,8 +161,11 @@ public List queryEmbedding( if (filters != null) { kwargs.put("filters", filters); } - Object pythonDocuments = adapter.callMethod(vectorStore, "_query_embedding", kwargs); - return adapter.fromPythonDocuments((List) pythonDocuments); + try (PythonObjectScope scope = new PythonObjectScope()) { + Object pythonDocuments = + scope.own(adapter.callMethod(vectorStore, "_query_embedding", kwargs)); + return adapter.fromPythonDocuments((List) pythonDocuments); + } } /** Embed query text via the configured model (no numpy, so it stays on the async pool). */ @@ -202,8 +208,11 @@ public List queryNormalized( if (filters != null) { kwargs.put("filters", filters); } - Object pythonDocuments = adapter.callMethod(vectorStore, "_query_embedding", kwargs); - return adapter.fromPythonDocuments((List) pythonDocuments); + try (PythonObjectScope scope = new PythonObjectScope()) { + Object pythonDocuments = + scope.own(adapter.callMethod(vectorStore, "_query_embedding", kwargs)); + return adapter.fromPythonDocuments((List) pythonDocuments); + } } @Override @@ -211,23 +220,27 @@ public List queryNormalized( public List addEmbedding( List documents, @Nullable String collection, Map extraArgs) throws IOException { - Map kwargs = new HashMap<>(extraArgs); - kwargs.put("documents", adapter.toPythonDocuments(documents)); - if (collection != null) { - kwargs.put("collection_name", collection); + try (PythonObjectScope scope = new PythonObjectScope()) { + Map kwargs = new HashMap<>(extraArgs); + kwargs.put("documents", scope.own(adapter.toPythonDocuments(documents))); + if (collection != null) { + kwargs.put("collection_name", collection); + } + return (List) adapter.callMethod(vectorStore, "_add_embedding", kwargs); } - return (List) adapter.callMethod(vectorStore, "_add_embedding", kwargs); } @Override public void updateEmbedding( List documents, @Nullable String collection, Map extraArgs) { - Map kwargs = new HashMap<>(extraArgs); - kwargs.put("documents", adapter.toPythonDocuments(documents)); - if (collection != null) { - kwargs.put("collection_name", collection); + try (PythonObjectScope scope = new PythonObjectScope()) { + Map 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()); + } + } } 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 b3ccf82e1..fc063d23c 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 @@ void testGetPythonResourceWithNullChatModel() { } @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 @@ void testChat() { 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 @@ void testChat() { 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 42463b083..b08afc84b 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 @@ void testGetParameters() { } @Test - void testChat() { + void testChat() throws Exception { ChatMessage inputMessage = mock(ChatMessage.class); ChatMessage outputMessage = mock(ChatMessage.class); List messages = Collections.singletonList(inputMessage); @@ -100,8 +100,8 @@ void testChat() { 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 @@ void testChat() { 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 000000000..560aa4b51 --- /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 37b27f856..566db3fe4 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.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 @@ void testAddDocuments() throws Exception { List 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 @@ void testAddDocuments() throws Exception { assertThat(result).containsExactly("doc1", "doc2"); verify(mockAdapter).toPythonDocuments(documents); + verify(firstPythonDocument).close(); + verify(secondPythonDocument).close(); verify(mockAdapter) .callMethod( eq(mockVectorStore), @@ -246,6 +252,7 @@ void testGetDocuments() throws Exception { 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 6b91c05cd..907bfc8ee 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.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 float[] call() { } }); - final Object normalized = store.normalizeEmbedding(embedding); - - final List documents = - ctx.durableExecuteAsync( - new DurableCallable>() { - @Override - public String getId() { - return "rag-query"; - } - - @SuppressWarnings("unchecked") - @Override - public Class> getResultClass() { - return (Class>) (Class) List.class; - } - - @Override - public List 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 documents = + ctx.durableExecuteAsync( + new DurableCallable>() { + @Override + public String getId() { + return "rag-query"; + } + + @SuppressWarnings("unchecked") + @Override + public Class> getResultClass() { + return (Class>) (Class) List.class; + } + + @Override + public List 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 cb4a39351..901fa06cf 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.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 String formatString(Map kwargs) { @Override public List formatMessages(MessageRole defaultRole, Map kwargs) { Map 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 pythonMessages = (List) result; - List 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 pythonMessages = (List) result; + List 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 6bfd95d2d..08ae93f26 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 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 void setMetricGroup(FlinkAgentsMetricGroup metricGroup) { 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 8a831db5a..630f1a68a 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 @@ -41,6 +41,7 @@ public class PythonMCPTool extends Tool "python_java_utils.get_java_tool_metadata_from_tool"; private final PyObject tool; private final PythonResourceAdapter adapter; + private boolean closed; @Nullable private final String mcpServerName; /** @@ -120,4 +121,15 @@ public Map getToolExecutionMetadata(ToolParameters parameters) { } 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 000000000..47e11cb26 --- /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 000000000..9dc548242 --- /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 b9e0b48dc..eca4019e1 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 static void discoverPythonMCPResources( 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 bc4d16a19..67c0cc382 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 @@ 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 List add( MemorySet memorySet, List memoryItems, @Nullable List> metadatas) { - Map 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 kwargs = new HashMap<>(); + kwargs.put("memory_set", buildPyMemorySet(scope, memorySet)); + kwargs.put("memory_items", memoryItems); + if (metadatas != null) { + kwargs.put("metadatas", metadatas); + } + return (List) adapter.callMethod(pyMem0, "add", kwargs); } - return (List) adapter.callMethod(pyMem0, "add", kwargs); } @Override @@ -107,29 +110,33 @@ public List get( @Nullable List ids, @Nullable Map filters, @Nullable Integer limit) { - Map 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 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 ids) { - Map kwargs = new HashMap<>(); - kwargs.put("memory_set", buildPyMemorySet(memorySet)); - if (ids != null) { - kwargs.put("ids", ids); + try (PythonObjectScope scope = new PythonObjectScope()) { + Map 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 List search( int limit, @Nullable Map filters, Map extraArgs) { - Map 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 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 void close() throws Exception { } } - 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 @@ private Object buildPyMemorySet(MemorySet memorySet) { 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 e645b071a..4012bd11b 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.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 void open() throws Exception { /** * 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 String executePythonFunction(PythonFunction function, Event event) throws 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 Object getOutputFromOutputEvent(String eventJson) { * @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 31c897134..73f206aba 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.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 Object toPythonVectorStoreQuery(VectorStoreQuery query) { @Override public VectorStoreQueryResult fromPythonVectorStoreQueryResult( PyObject pythonVectorStoreQueryResult) { - List pythonDocuments = - (List) pythonVectorStoreQueryResult.getAttr("documents", List.class); - return new VectorStoreQueryResult(fromPythonDocuments(pythonDocuments)); + try (PythonObjectScope scope = new PythonObjectScope()) { + List pythonDocuments = + scope.own( + (List) + 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 000000000..36fb92177 --- /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> 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 1554be965..696b2a892 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.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 @@ void testDeleteMemorySetForwardsAndReturnsBoolean() throws Exception { } @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 @@ void testAddForwardsKwargsAndReturnsIds() throws Exception { assertThat(captureKwargs("add")) .containsKeys("memory_set", "memory_items", "metadatas") .containsEntry("memory_set", mockPyMemorySet); + verify(mockPyMemorySet).close(); } @Test @@ -131,8 +133,10 @@ void testGetOmitsNullOptionalKwargs() throws Exception { @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 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 @@ void testGetWithIdsAndFiltersConvertsItems() throws Exception { 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 5bdac5dcc..b2ac5060c 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 @@ void releasesBothPythonObjectsWhenLogicalCleanupFails() throws Exception { 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 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 f372821a5..b2225fdb3 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.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.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 @@ void testCallMethod() { .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();