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 9b5d223f [api][java][python] Reject non-positive max_results in
ContextRetrievalRequestEvent (#1157)
9b5d223f is described below
commit 9b5d223f29d438bf079b17e892cef561a7bfddb9
Author: william <[email protected]>
AuthorDate: Tue Sep 29 15:30:45 2026 +0800
[api][java][python] Reject non-positive max_results in
ContextRetrievalRequestEvent (#1157)
---
.../api/event/ContextRetrievalRequestEvent.java | 31 ++++++
.../event/ContextRetrievalRequestEventTest.java | 112 +++++++++++++++++++++
.../api/events/context_retrieval_event.py | 16 ++-
.../api/tests/test_context_retrieval_event.py | 52 ++++++++++
4 files changed, 210 insertions(+), 1 deletion(-)
diff --git
a/api/src/main/java/org/apache/flink/agents/api/event/ContextRetrievalRequestEvent.java
b/api/src/main/java/org/apache/flink/agents/api/event/ContextRetrievalRequestEvent.java
index 25580308..fb336137 100644
---
a/api/src/main/java/org/apache/flink/agents/api/event/ContextRetrievalRequestEvent.java
+++
b/api/src/main/java/org/apache/flink/agents/api/event/ContextRetrievalRequestEvent.java
@@ -37,18 +37,49 @@ public class ContextRetrievalRequestEvent extends Event {
this(query, vectorStore, DEFAULT_MAX_RESULTS);
}
+ /**
+ * Creates a new context retrieval request event.
+ *
+ * @param query the search query text
+ * @param vectorStore the name of the vector store to query
+ * @param maxResults the maximum number of results to return; must be
positive
+ * @throws IllegalArgumentException if {@code maxResults} is non-positive
+ */
public ContextRetrievalRequestEvent(String query, String vectorStore, int
maxResults) {
super(EVENT_TYPE);
+ validateMaxResults(maxResults);
setAttr("query", query);
setAttr("vector_store", vectorStore);
setAttr("max_results", maxResults);
}
+ /**
+ * Reconstructs a context retrieval request event from its serialized form.
+ *
+ * <p>The {@code max_results} attribute is validated here as well, so that
reconstructed events
+ * are rejected the same way as directly constructed ones and stay
consistent with the Python
+ * implementation.
+ *
+ * @param id the event id
+ * @param attributes the event attributes
+ * @throws IllegalArgumentException if {@code max_results} is present and
non-positive
+ */
@JsonCreator
public ContextRetrievalRequestEvent(
@JsonProperty("id") UUID id,
@JsonProperty("attributes") Map<String, Object> attributes) {
super(id, EVENT_TYPE, attributes);
+ Object maxResults = attributes == null ? null :
attributes.get("max_results");
+ if (maxResults instanceof Number) {
+ validateMaxResults(((Number) maxResults).intValue());
+ }
+ }
+
+ private static void validateMaxResults(int maxResults) {
+ if (maxResults <= 0) {
+ throw new IllegalArgumentException(
+ String.format("`max_results` must be positive, but was
%d.", maxResults));
+ }
}
/**
diff --git
a/api/src/test/java/org/apache/flink/agents/api/event/ContextRetrievalRequestEventTest.java
b/api/src/test/java/org/apache/flink/agents/api/event/ContextRetrievalRequestEventTest.java
new file mode 100644
index 00000000..79dc8d85
--- /dev/null
+++
b/api/src/test/java/org/apache/flink/agents/api/event/ContextRetrievalRequestEventTest.java
@@ -0,0 +1,112 @@
+/*
+ * 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.event;
+
+import org.apache.flink.agents.api.Event;
+import org.junit.jupiter.api.DisplayName;
+import org.junit.jupiter.api.Test;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.UUID;
+
+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 for the max results validation of {@link
ContextRetrievalRequestEvent}. */
+class ContextRetrievalRequestEventTest {
+
+ @Test
+ @DisplayName("Zero max results is rejected with a message naming the
parameter")
+ void testZeroMaxResultsRejected() {
+ assertThatThrownBy(() -> new ContextRetrievalRequestEvent("query",
"store", 0))
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessageContaining("max_results");
+ }
+
+ @Test
+ @DisplayName(
+ "Negative max results is rejected with a message naming the
parameter and its value")
+ void testNegativeMaxResultsRejected() {
+ assertThatThrownBy(() -> new ContextRetrievalRequestEvent("query",
"store", -1))
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessageContaining("max_results")
+ .hasMessageContaining("-1");
+ }
+
+ @Test
+ @DisplayName("Positive max results is accepted")
+ void testPositiveMaxResultsAccepted() {
+ assertThat(new ContextRetrievalRequestEvent("query", "store",
5).getMaxResults())
+ .isEqualTo(5);
+ }
+
+ @Test
+ @DisplayName("Default max results is unchanged")
+ void testDefaultMaxResultsUnchanged() {
+ assertThat(new ContextRetrievalRequestEvent("query",
"store").getMaxResults()).isEqualTo(3);
+ }
+
+ @Test
+ @DisplayName("Reconstruction from attributes rejects non-positive max
results")
+ void testReconstructedNonPositiveMaxResultsRejected() {
+ assertThatThrownBy(() -> new
ContextRetrievalRequestEvent(UUID.randomUUID(), attributes(0)))
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessageContaining("max_results");
+ }
+
+ @Test
+ @DisplayName("Reconstruction from a base event rejects non-positive max
results")
+ void testFromEventNonPositiveMaxResultsRejected() {
+ assertThatThrownBy(() ->
ContextRetrievalRequestEvent.fromEvent(baseEvent(-2)))
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessageContaining("max_results")
+ .hasMessageContaining("-2");
+ }
+
+ @Test
+ @DisplayName("Reconstruction of a valid event keeps max results")
+ void testFromEventValidMaxResultsAccepted() {
+
assertThat(ContextRetrievalRequestEvent.fromEvent(baseEvent(5)).getMaxResults())
+ .isEqualTo(5);
+ }
+
+ @Test
+ @DisplayName("Reconstruction tolerates a missing max results attribute")
+ void testReconstructionWithoutMaxResultsAttribute() {
+ Map<String, Object> attrs = attributes(5);
+ attrs.remove("max_results");
+ assertThatCode(() -> new
ContextRetrievalRequestEvent(UUID.randomUUID(), attrs))
+ .doesNotThrowAnyException();
+ }
+
+ private static Map<String, Object> attributes(int maxResults) {
+ Map<String, Object> attrs = new HashMap<>();
+ attrs.put("query", "query");
+ attrs.put("vector_store", "store");
+ attrs.put("max_results", maxResults);
+ return attrs;
+ }
+
+ private static Event baseEvent(int maxResults) {
+ return new Event(
+ UUID.randomUUID(), ContextRetrievalRequestEvent.EVENT_TYPE,
attributes(maxResults));
+ }
+}
diff --git a/python/flink_agents/api/events/context_retrieval_event.py
b/python/flink_agents/api/events/context_retrieval_event.py
index d014e185..7dde0ded 100644
--- a/python/flink_agents/api/events/context_retrieval_event.py
+++ b/python/flink_agents/api/events/context_retrieval_event.py
@@ -43,7 +43,21 @@ class ContextRetrievalRequestEvent(Event):
EVENT_TYPE: ClassVar[str] = "_context_retrieval_request_event"
def __init__(self, query: str, vector_store: str, max_results: int = 3) ->
None:
- """Create a ContextRetrievalRequestEvent."""
+ """Create a ContextRetrievalRequestEvent.
+
+ Args:
+ ----
+ query: The search query text to find relevant context for.
+ vector_store: Name of the vector store setup resource to use.
+ max_results: Maximum number of results to return (default: 3). Must be
positive.
+
+ Raises:
+ ------
+ ValueError: If `max_results` is not positive.
+ """
+ if max_results <= 0:
+ msg = f"`max_results` must be positive, but was {max_results}."
+ raise ValueError(msg)
super().__init__(
type=ContextRetrievalRequestEvent.EVENT_TYPE,
attributes={
diff --git a/python/flink_agents/api/tests/test_context_retrieval_event.py
b/python/flink_agents/api/tests/test_context_retrieval_event.py
new file mode 100644
index 00000000..19db8417
--- /dev/null
+++ b/python/flink_agents/api/tests/test_context_retrieval_event.py
@@ -0,0 +1,52 @@
+################################################################################
+# 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.
+#################################################################################
+"""Tests for the max_results validation of
:class:`ContextRetrievalRequestEvent`."""
+
+from __future__ import annotations
+
+import pytest
+
+from flink_agents.api.events.context_retrieval_event import
ContextRetrievalRequestEvent
+
+
+def test_zero_max_results_rejected() -> None:
+ """A zero max_results should be rejected."""
+ with pytest.raises(ValueError, match="max_results"):
+ ContextRetrievalRequestEvent(query="flink", vector_store="store",
max_results=0)
+
+
+def test_negative_max_results_rejected() -> None:
+ """A negative max_results should be rejected."""
+ with pytest.raises(ValueError, match="max_results"):
+ ContextRetrievalRequestEvent(
+ query="flink", vector_store="store", max_results=-1
+ )
+
+
+def test_positive_max_results_accepted() -> None:
+ """A positive max_results should be accepted."""
+ event = ContextRetrievalRequestEvent(
+ query="flink", vector_store="store", max_results=5
+ )
+ assert event.max_results == 5
+
+
+def test_default_max_results_unchanged() -> None:
+ """The default max_results should stay unchanged."""
+ event = ContextRetrievalRequestEvent(query="flink", vector_store="store")
+ assert event.max_results == 3
