This is an automated email from the ASF dual-hosted git repository.
yunfengzhou-hub 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 b6702839c [hotfix][integrations][python][docs] Align vector stores
with their declared contracts (#1280)
b6702839c is described below
commit b6702839c03d5298530d8fe029b34f7c629b62b9
Author: yunfengzhou-hub <[email protected]>
AuthorDate: Tue Oct 20 15:39:31 2026 +0800
[hotfix][integrations][python][docs] Align vector stores with their
declared contracts (#1280)
---
.../vector-search/vector_stores/_index.md | 2 +-
.../elasticsearch/ElasticsearchVectorStore.java | 10 ++-
...ElasticsearchVectorStoreKNumCandidatesTest.java | 72 ++++++++++++++++++
.../ElasticsearchVectorStoreTest.java | 88 ++++++++++++++++++++++
python/pyproject.toml | 2 +-
5 files changed, 168 insertions(+), 6 deletions(-)
diff --git
a/docs/content/docs/development/vector-search/vector_stores/_index.md
b/docs/content/docs/development/vector-search/vector_stores/_index.md
index 6e773e0a7..e82c19fe3 100644
--- a/docs/content/docs/development/vector-search/vector_stores/_index.md
+++ b/docs/content/docs/development/vector-search/vector_stores/_index.md
@@ -328,7 +328,7 @@ vectorStore.delete(null, "my_collection", null, Map.of());
{"user_id": "u1", "run_id": "r1"}
```
-`None` / `null` means "no filter". Richer operators (ranges, set membership,
OR, NOT, etc.) are out of scope here. Callers needing backend-specific
operators should pass them through `extra_args` (Python
`VectorStoreQuery.extra_args` or `**kwargs`) or `extraArgs` (Java) — for
example, ChromaDB's native `where` dict. Implementations that receive an
unsupported operator via `filters` raise `NotImplementedError` (Python) or
`UnsupportedOperationException` (Java).
+`None` / `null` means "no filter". Richer operators (ranges, set membership,
OR, NOT, etc.) are out of scope here. Callers needing backend-specific
operators should pass them through `extra_args` (Python
`VectorStoreQuery.extra_args` or `**kwargs`) or `extraArgs` (Java) — for
example, ChromaDB's native `where` dict. An unsupported operator is not
portable. Chroma and Milvus reject some operator forms locally with
`NotImplementedError` (Python) or `UnsupportedOperationException` (Java); o
[...]
### Usage Example
diff --git
a/integrations/vector-stores/elasticsearch/src/main/java/org/apache/flink/agents/integrations/vectorstores/elasticsearch/ElasticsearchVectorStore.java
b/integrations/vector-stores/elasticsearch/src/main/java/org/apache/flink/agents/integrations/vectorstores/elasticsearch/ElasticsearchVectorStore.java
index 4be3d6be1..c04c0781b 100644
---
a/integrations/vector-stores/elasticsearch/src/main/java/org/apache/flink/agents/integrations/vectorstores/elasticsearch/ElasticsearchVectorStore.java
+++
b/integrations/vector-stores/elasticsearch/src/main/java/org/apache/flink/agents/integrations/vectorstores/elasticsearch/ElasticsearchVectorStore.java
@@ -181,9 +181,9 @@ public class ElasticsearchVectorStore extends
BaseVectorStore
this.numCandidates = descriptor.getArgument("num_candidates");
if (this.k != null && this.numCandidates != null) {
- if (this.k < this.numCandidates) {
+ if (this.numCandidates < this.k) {
throw new IllegalArgumentException(
- "'k' should be greater or equals than
'num_candidates'");
+ "'num_candidates' should be greater than or equal to
'k'");
}
}
@@ -318,8 +318,8 @@ public class ElasticsearchVectorStore extends
BaseVectorStore
m.put("vector_field", this.vectorField);
m.put("dims", this.dims);
if (this.k != null) {
- // Expose as top_k (preferred name); keep legacy key for
compatibility if referenced
- m.put("top_k", this.k);
+ // queryEmbedding() reads this under "k"; expose it under the same
key
+ m.put("k", this.k);
}
if (this.numCandidates != null) {
m.put("num_candidates", this.numCandidates);
@@ -626,6 +626,8 @@ public class ElasticsearchVectorStore extends
BaseVectorStore
SearchRequest.Builder builder =
new SearchRequest.Builder()
.index(index)
+ // k selects neighbours; size bounds hits and
defaults to 10.
+ .size(Math.min(limit, k))
.knn(
kb -> {
kb.field(this.vectorField)
diff --git
a/integrations/vector-stores/elasticsearch/src/test/java/org/apache/flink/agents/integrations/vectorstores/elasticsearch/ElasticsearchVectorStoreKNumCandidatesTest.java
b/integrations/vector-stores/elasticsearch/src/test/java/org/apache/flink/agents/integrations/vectorstores/elasticsearch/ElasticsearchVectorStoreKNumCandidatesTest.java
new file mode 100644
index 000000000..fb2c3f908
--- /dev/null
+++
b/integrations/vector-stores/elasticsearch/src/test/java/org/apache/flink/agents/integrations/vectorstores/elasticsearch/ElasticsearchVectorStoreKNumCandidatesTest.java
@@ -0,0 +1,72 @@
+/*
+ * 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.integrations.vectorstores.elasticsearch;
+
+import org.apache.flink.agents.api.resource.ResourceContext;
+import org.apache.flink.agents.api.resource.ResourceDescriptor;
+import org.junit.jupiter.api.Test;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatCode;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/**
+ * Tests the {@code k} / {@code num_candidates} constructor validation of
{@link
+ * ElasticsearchVectorStore}.
+ *
+ * <p>Elasticsearch requires {@code num_candidates >= k} for a KNN search, so
that combination is
+ * the one the constructor must accept.
+ */
+class ElasticsearchVectorStoreKNumCandidatesTest {
+
+ private static ElasticsearchVectorStore store(int k, int numCandidates) {
+ ResourceDescriptor descriptor =
+
ResourceDescriptor.Builder.newBuilder(ElasticsearchVectorStore.class.getName())
+ .addInitialArgument("embedding_model", "emb")
+ .addInitialArgument("index", "idx")
+ .addInitialArgument("dims", 5)
+ .addInitialArgument("k", k)
+ .addInitialArgument("num_candidates", numCandidates)
+ .build();
+ return new ElasticsearchVectorStore(
+ descriptor, ResourceContext.fromGetResource((name, type) ->
null));
+ }
+
+ @Test
+ void acceptsNumCandidatesGreaterThanK() {
+ assertThatCode(() -> store(10, 20)).doesNotThrowAnyException();
+ }
+
+ @Test
+ void acceptsNumCandidatesEqualToK() {
+ assertThatCode(() -> store(10, 10)).doesNotThrowAnyException();
+ }
+
+ @Test
+ void rejectsNumCandidatesLessThanK() {
+ assertThatThrownBy(() -> store(20, 10))
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessageContaining("num_candidates");
+ }
+
+ @Test
+ void exposesStoreLevelKUnderTheKeyQueryReads() {
+ assertThat(store(7, 20).getStoreKwargs()).containsEntry("k", 7);
+ }
+}
diff --git
a/integrations/vector-stores/elasticsearch/src/test/java/org/apache/flink/agents/integrations/vectorstores/elasticsearch/ElasticsearchVectorStoreTest.java
b/integrations/vector-stores/elasticsearch/src/test/java/org/apache/flink/agents/integrations/vectorstores/elasticsearch/ElasticsearchVectorStoreTest.java
index 9c04b0c14..21bb0f916 100644
---
a/integrations/vector-stores/elasticsearch/src/test/java/org/apache/flink/agents/integrations/vectorstores/elasticsearch/ElasticsearchVectorStoreTest.java
+++
b/integrations/vector-stores/elasticsearch/src/test/java/org/apache/flink/agents/integrations/vectorstores/elasticsearch/ElasticsearchVectorStoreTest.java
@@ -310,4 +310,92 @@ public class ElasticsearchVectorStoreTest {
((CollectionManageableVectorStore) store).deleteCollection(name);
}
+
+ /**
+ * Builds a store whose descriptor sets a store-level {@code k}, mirroring
{@link
+ * #initialize()}.
+ */
+ private static BaseVectorStore storeWithK(int k) throws Exception {
+ ResourceDescriptor descriptor =
+
ResourceDescriptor.Builder.newBuilder(ElasticsearchVectorStore.class.getName())
+ .addInitialArgument("embedding_model",
"embeddingModel")
+ .addInitialArgument(
+ "host",
+ System.getenv().getOrDefault("ES_HOST",
"http://localhost:9200"))
+ .addInitialArgument("dims", 5)
+ .addInitialArgument("username", "elastic")
+ .addInitialArgument("password",
System.getenv("ES_PASSWORD"))
+ .addInitialArgument("k", k)
+ .build();
+ BaseVectorStore kStore =
+ new ElasticsearchVectorStore(
+ descriptor,
+
ResourceContext.fromGetResource(ElasticsearchVectorStoreTest::getResource));
+ kStore.open();
+ return kStore;
+ }
+
+ /** Documents sharing one direction, with explicit embeddings so no embed
stub is needed. */
+ private static List<Document> sameDirectionDocs(int count, String prefix) {
+ List<Document> docs = new ArrayList<>();
+ for (int i = 0; i < count; i++) {
+ Document doc =
+ new Document(
+ prefix + " document " + i, Map.of("src", "test"),
prefix + "_" + i);
+ doc.setEmbedding(new float[] {1.0f, 0.0f, 0.0f, 0.0f, 0.0f});
+ docs.add(doc);
+ }
+ return docs;
+ }
+
+ @Test
+ public void testStoreLevelKDoesNotExceedQueryLimit() throws Exception {
+ // Contract: VectorStoreQuery.limit is the maximum number of documents
returned. A
+ // store-level k chooses how many neighbours the KNN clause considers;
reaching
+ // queryEmbedding through the merged store kwargs must not widen the
result past limit.
+ String name = "store_k_bounded_by_limit";
+ BaseVectorStore kStore = storeWithK(8);
+ try {
+ ((CollectionManageableVectorStore)
kStore).createCollectionIfNotExists(name, Map.of());
+ kStore.addEmbedding(sameDirectionDocs(12, "k_bounded"), name,
Collections.emptyMap());
+ Thread.sleep(1000);
+
+ List<Document> hits =
+ kStore.query(
+ new VectorStoreQuery(
+ "Elasticsearch is a search engine",
+ 3,
+ name,
+ Collections.emptyMap()))
+ .getDocuments();
+
+ Assertions.assertEquals(3, hits.size());
+ } finally {
+ ((CollectionManageableVectorStore) kStore).deleteCollection(name);
+ kStore.close();
+ }
+ }
+
+ @Test
+ public void testQueryLimitAboveElasticsearchDefaultSizeIsHonored() throws
Exception {
+ // Elasticsearch defaults size to 10, so a larger limit has to be sent
explicitly or the
+ // result is silently truncated below what the caller asked for.
+ String name = "limit_above_default_size";
+ ((CollectionManageableVectorStore)
store).createCollectionIfNotExists(name, Map.of());
+ store.addEmbedding(sameDirectionDocs(12, "limit_12"), name,
Collections.emptyMap());
+ Thread.sleep(1000);
+
+ List<Document> hits =
+ store.query(
+ new VectorStoreQuery(
+ "Elasticsearch is a search engine",
+ 12,
+ name,
+ Collections.emptyMap()))
+ .getDocuments();
+
+ Assertions.assertEquals(12, hits.size());
+
+ ((CollectionManageableVectorStore) store).deleteCollection(name);
+ }
}
diff --git a/python/pyproject.toml b/python/pyproject.toml
index 9913ddd26..0cbc66166 100644
--- a/python/pyproject.toml
+++ b/python/pyproject.toml
@@ -64,7 +64,7 @@ dependencies = [
# incompatible with the protobuf runtimes selected by supported PyFlink
releases.
"googleapis-common-protos<1.72.1",
"chromadb==1.0.21",
- "mem0ai>=0.1.43,<2.0.0",
+ "mem0ai>=0.1.74,<2.0.0",
"onnxruntime<1.24.1;python_version<'3.11'",
]