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 d4ceac119 [runtime][java] Preserve an existing KeyedStream's business
key when no key selector is given (#1312)
d4ceac119 is described below
commit d4ceac119453c6b6d782ca164a76f67263fd016f
Author: wank125 <[email protected]>
AuthorDate: Sat Oct 24 17:36:30 2026 +0800
[runtime][java] Preserve an existing KeyedStream's business key when no key
selector is given (#1312)
Co-authored-by: wank125 <[email protected]>
Co-authored-by: WenjinXie <[email protected]>
---
.../agents/api/AgentsExecutionEnvironment.java | 23 ++-
.../docs/development/integrate_with_flink.md | 13 +-
.../quickstart/create-agent/multiple-agents.md | 21 ++-
.../quickstart/create-agent/single-agent.md | 31 ++--
.../docs/get-started/quickstart/react_agent.md | 2 +-
.../docs/get-started/quickstart/yaml_agent.md | 22 ++-
.../integration/test/FlinkIntegrationTest.java | 6 +-
.../flink/agents/examples/ModelRoutingExample.java | 2 +-
.../agents/examples/ModelRoutingJudgeExample.java | 2 +-
.../flink/agents/examples/ReActAgentExample.java | 2 +-
.../flink/agents/examples/SkillsAgentExample.java | 5 +-
.../examples/WorkflowMultipleAgentExample.java | 16 +-
.../examples/WorkflowSingleAgentExample.java | 24 ++-
.../agents/examples/YamlWorkflowAgentExample.java | 24 ++-
.../examples/agents/ProductSuggestionAgent.java | 12 +-
.../examples/agents/ReviewAnalysisAgent.java | 10 +-
.../examples/agents/TableReviewAnalysisAgent.java | 2 +-
.../examples/openai/OpenAiModelRoutingExample.java | 2 +-
python/flink_agents/api/execution_environment.py | 5 +-
.../runtime/remote_execution_environment.py | 31 ++--
.../tests/test_remote_execution_environment.py | 50 +++++-
.../runtime/env/RemoteExecutionEnvironment.java | 23 ++-
.../env/RemoteExecutionEnvironmentTest.java | 171 +++++++++++++++++++++
.../runtime/env/TypedOutputTerminalTest.java | 4 +-
24 files changed, 394 insertions(+), 109 deletions(-)
diff --git
a/api/src/main/java/org/apache/flink/agents/api/AgentsExecutionEnvironment.java
b/api/src/main/java/org/apache/flink/agents/api/AgentsExecutionEnvironment.java
index 9095f23dd..8568f94d8 100644
---
a/api/src/main/java/org/apache/flink/agents/api/AgentsExecutionEnvironment.java
+++
b/api/src/main/java/org/apache/flink/agents/api/AgentsExecutionEnvironment.java
@@ -133,11 +133,13 @@ public abstract class AgentsExecutionEnvironment {
* Set input for agents from a DataStream. Used for remote execution.
*
* <p>This method integrates agents with Flink DataStream API, allowing
agents to process
- * streaming data with optional keying for stateful operations.
+ * streaming data keyed for stateful operations.
*
* @param input DataStream to be processed by agents.
- * @param keySelector Optional KeySelector for extracting keys from input
records. If null, the
- * stream will be processed without keying.
+ * @param keySelector Optional KeySelector for extracting keys from input
records. An explicit
+ * selector is always used, even when the input is already keyed; if
null, the input must be
+ * a KeyedStream, whose key is kept, otherwise building the pipeline
fails with a clear
+ * error rather than silently keying by the whole record.
* @param <T> Type of elements in the input DataStream.
* @param <K> Type of the key extracted by the KeySelector.
* @return AgentBuilder for configuring the agent pipeline.
@@ -146,7 +148,10 @@ public abstract class AgentsExecutionEnvironment {
DataStream<T> input, KeySelector<T, K> keySelector);
/**
- * Set input for agents from a DataStream without keying.
+ * Set input for agents from a DataStream, keeping the stream's own key.
+ *
+ * <p>The input must be a KeyedStream; its existing key is preserved. A
non-keyed input is
+ * rejected with a clear error asking for a keyed stream or a key selector.
*
* @param input DataStream to be processed by agents.
* @param <T> Type of elements in the input DataStream.
@@ -163,17 +168,23 @@ public abstract class AgentsExecutionEnvironment {
* and processing it through agents.
*
* @param input Table to be processed by agents.
- * @param keySelector Optional KeySelector for extracting keys from table
rows.
+ * @param keySelector Required KeySelector for extracting the business key
from table rows. Rows
+ * with the same key share the agent's short-term memory.
* @param <K> Type of the key extracted by the KeySelector.
* @return AgentBuilder for configuring the agent pipeline.
+ * @throws IllegalArgumentException in remote execution if keySelector is
null.
*/
public abstract <K> AgentBuilder fromTable(Table input,
KeySelector<Object, K> keySelector);
/**
- * Set input for agents from a Table without keying.
+ * Table input requires an explicit key selector.
+ *
+ * <p>Remote execution rejects this overload. Use {@link #fromTable(Table,
KeySelector)} to
+ * choose the key used for the agent's short-term memory.
*
* @param input Table to be processed by agents.
* @return AgentBuilder for configuring the agent pipeline.
+ * @throws IllegalArgumentException in remote execution because a key
selector is required.
*/
public AgentBuilder fromTable(Table input) {
return fromTable(input, null);
diff --git a/docs/content/docs/development/integrate_with_flink.md
b/docs/content/docs/development/integrate_with_flink.md
index 1bea16897..261467e21 100644
--- a/docs/content/docs/development/integrate_with_flink.md
+++ b/docs/content/docs/development/integrate_with_flink.md
@@ -119,7 +119,14 @@ outputStream.print();
{{< /tabs >}}
-The input `DataStream` must be `KeyedStream`, or user should provide
`KeySelector` to tell how to convert the input `DataStream` to `KeyedStream`.
+The input `DataStream` must be keyed. Key selection follows the same rules in
the Java and
+Python APIs:
+
+- If an explicit `KeySelector` is provided, it is always used, even when the
input is
+ already keyed.
+- Otherwise, a `KeyedStream` input keeps its existing key.
+- If neither is available, the call fails with a clear error instead of
silently keying
+ by the whole record.
For complete, runnable examples, see {{< gh_link
file="examples/src/main/java/org/apache/flink/agents/examples/WorkflowSingleAgentExample.java"
name="`WorkflowSingleAgentExample.java`" >}} (Java) and {{< gh_link
file="python/flink_agents/examples/quickstart/workflow_single_agent_example.py"
name="`workflow_single_agent_example.py`" >}} (Python).
@@ -162,6 +169,10 @@ delete earlier rows. Both `from_table()` and `fromTable()`
reject Tables that pr
updates or deletes, such as a continuous `GROUP BY` aggregation whose results
change
as new records arrive.
+Both Java and Python require an explicit key selector for Table input. The
selector
+extracts a business key from each row; rows with the same key share the Agent's
+short-term memory.
+
{{< tabs "From/To Table" >}}
{{< tab "Python" >}}
diff --git
a/docs/content/docs/get-started/quickstart/create-agent/multiple-agents.md
b/docs/content/docs/get-started/quickstart/create-agent/multiple-agents.md
index 691aab0dd..220a737a1 100644
--- a/docs/content/docs/get-started/quickstart/create-agent/multiple-agents.md
+++ b/docs/content/docs/get-started/quickstart/create-agent/multiple-agents.md
@@ -162,7 +162,7 @@ The Java table declares only the `id` and `review` columns.
A Table row reaches the Action as a `pyflink.common.Row` in Python and an
`org.apache.flink.types.Row` in Java. In Python, access fields with
`row["id"]` or call
`row.as_dict()` when dictionary methods such as `.get()` are needed. Build
Your First Agent
-gave its Agent a `ProductReview` object in Python and a JSON string in Java,
so this example
+gave its Agent a `ProductReview` object in both languages, so this example
defines a second Agent class, `TableReviewAnalysisAgent`, for the row input.
It redeclares the
prompt, tool, and chat model setup that `ReviewAnalysisAgent` declares, and
`ReviewAnalysisAgent`
itself is unchanged. The key selector that reads `id` out of a row travels
with this Agent, and
@@ -275,8 +275,8 @@ See [Integrate with Flink]({{< ref
"docs/development/integrate_with_flink#fromto
One review is too little to suggest an improvement, so the job groups the
per-review analyses by
product id in a one-minute tumbling window and reduces each group to a
`ProductReviewSummary`. That
type is declared in the shared types file and carries the product id, the
score distribution, and
-the collected unsatisfied reasons. Python returns the summary object, and Java
emits it as JSON
-text, because the next stage reads it back with Jackson. Scores outside 1–5
are excluded from
+the collected unsatisfied reasons. Both Python and Java pass the summary
object directly to
+the next Agent, keyed by product id. Scores outside 1–5 are excluded from
the distribution; if none are valid, all five percentages are `0.0%`. The
reasons are still retained.
{{< tabs "Aggregate in a Window" >}}
@@ -325,15 +325,14 @@ aggregated_analysis_res_stream = (
{{< tab "Java" >}}
```java
static class AggregateScoreDistributionAndDislikeReasons
- extends ProcessWindowFunction<ProductReviewAnalysisRes, String,
String, TimeWindow> {
+ extends ProcessWindowFunction<ProductReviewAnalysisRes,
ProductReviewSummary, String, TimeWindow> {
@Override
public void process(
String key,
Context context,
Iterable<ProductReviewAnalysisRes> elements,
- Collector<String> out)
- throws JsonProcessingException {
+ Collector<ProductReviewSummary> out) {
// Initialize rating counts for scores 1-5
int[] ratingCounts = new int[5];
@@ -371,11 +370,11 @@ static class AggregateScoreDistributionAndDislikeReasons
// Create and emit the aggregated result
ProductReviewSummary summary =
new ProductReviewSummary(key, formattedPercentages,
reasonList);
- out.collect(MAPPER.writeValueAsString(summary));
+ out.collect(summary);
}
}
-DataStream<String> aggregatedAnalysisResStream =
+DataStream<ProductReviewSummary> aggregatedAnalysisResStream =
reviewAnalysisResStream
.keyBy(ProductReviewAnalysisRes::getId)
.window(TumblingProcessingTimeWindows.of(Duration.ofMinutes(1)))
@@ -454,10 +453,8 @@ public class ProductSuggestionAgent extends Agent {
/** Process input event. */
@Action(EventType.InputEvent)
public static void processInput(Event event, RunnerContext ctx) throws
Exception {
- // The window stage emitted JSON text, so this Agent reads it back
first.
InputEvent inputEvent = (InputEvent) event;
- String input = (String) inputEvent.getInput();
- ProductReviewSummary summary = MAPPER.readValue(input,
ProductReviewSummary.class);
+ ProductReviewSummary summary = (ProductReviewSummary)
inputEvent.getInput();
// Rest of the body keeps the id and score histogram in short-term
memory and
// sends ChatRequestEvent for generateSuggestionModel.
...
@@ -495,7 +492,7 @@ agents_env.execute("Workflow Multiple Agent Example Job")
```java
DataStream<Object> productSuggestionResStream =
agentsEnv
- .fromDataStream(aggregatedAnalysisResStream)
+ .fromDataStream(aggregatedAnalysisResStream,
ProductReviewSummary::getId)
.apply(new ProductSuggestionAgent())
.toDataStream();
diff --git
a/docs/content/docs/get-started/quickstart/create-agent/single-agent.md
b/docs/content/docs/get-started/quickstart/create-agent/single-agent.md
index 24af388fa..3e305d6a7 100644
--- a/docs/content/docs/get-started/quickstart/create-agent/single-agent.md
+++ b/docs/content/docs/get-started/quickstart/create-agent/single-agent.md
@@ -263,17 +263,12 @@ public static void notifyShippingManager(
@Action(EventType.InputEvent)
public static void processInput(Event event, RunnerContext ctx) throws
Exception {
InputEvent inputEvent = (InputEvent) event;
- String input = (String) inputEvent.getInput();
- MAPPER.configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false);
CustomTypesAndResources.ProductReview inputObj =
- MAPPER.readValue(input,
CustomTypesAndResources.ProductReview.class);
+ (CustomTypesAndResources.ProductReview) inputEvent.getInput();
ctx.getShortTermMemory().set("id", inputObj.getId());
- String content =
- String.format(
- "{\n" + "\"id\": %s,\n" + "\"review\": \"%s\"\n" + "}",
- inputObj.getId(), inputObj.getReview());
+ String content = MAPPER.writeValueAsString(inputObj);
ChatMessage msg = new ChatMessage(MessageRole.USER, "");
ctx.sendEvent(
@@ -363,22 +358,28 @@ agents_env.execute("Workflow Agent Example Job")
{{< tab "Java" >}}
```java
+// Configure this mapper once on the enclosing example class:
+// private static final ObjectMapper MAPPER = new ObjectMapper()
+// .configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES,
false);
+// ProductReview is CustomTypesAndResources.ProductReview.
// Read product reviews from input_data.txt file as a streaming source.
// Each element represents a ProductReview.
File inputDataFile = copyResource("input_data.txt");
-DataStream<String> productReviewStream =
+DataStream<ProductReview> productReviewStream =
env.fromSource(
- FileSource.forRecordStreamFormat(
- new TextLineInputFormat(),
- new Path(inputDataFile.getAbsolutePath()))
- .build(),
- WatermarkStrategy.noWatermarks(),
- "streaming-agent-example");
+ FileSource.forRecordStreamFormat(
+ new TextLineInputFormat(),
+ new
Path(inputDataFile.getAbsolutePath()))
+ .build(),
+ WatermarkStrategy.noWatermarks(),
+ "streaming-agent-example")
+ .map(json -> MAPPER.readValue(json, ProductReview.class))
+ .returns(ProductReview.class);
// Use the ReviewAnalysisAgent to analyze each product review.
DataStream<Object> reviewAnalysisResStream =
agentsEnv
- .fromDataStream(productReviewStream)
+ .fromDataStream(productReviewStream, ProductReview::getId)
.apply(new ReviewAnalysisAgent())
.toDataStream();
diff --git a/docs/content/docs/get-started/quickstart/react_agent.md
b/docs/content/docs/get-started/quickstart/react_agent.md
index 325baae91..6d271d02a 100644
--- a/docs/content/docs/get-started/quickstart/react_agent.md
+++ b/docs/content/docs/get-started/quickstart/react_agent.md
@@ -255,7 +255,7 @@ DataStream<Row> productReviewStream =
// record shipping question.
DataStream<Object> reviewAnalysisResStream =
agentsEnv
- .fromDataStream(productReviewStream)
+ .fromDataStream(productReviewStream, row -> (String)
row.getField("id"))
.apply(reviewAnalysisReactAgent)
.toDataStream();
diff --git a/docs/content/docs/get-started/quickstart/yaml_agent.md
b/docs/content/docs/get-started/quickstart/yaml_agent.md
index 6f041eb92..9a6c3ecaa 100644
--- a/docs/content/docs/get-started/quickstart/yaml_agent.md
+++ b/docs/content/docs/get-started/quickstart/yaml_agent.md
@@ -318,22 +318,28 @@ agents_env.execute("YAML Agent Example Job")
File yamlFile = copyResource("yaml/yaml_review_analysis_agent.yaml");
agentsEnv.loadYaml(yamlFile.toPath());
+// Configure this mapper once on the enclosing example class:
+// private static final ObjectMapper MAPPER = new ObjectMapper()
+// .configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES,
false);
+// ProductReview is CustomTypesAndResources.ProductReview.
// Read product reviews from input_data.txt file as a streaming source.
File inputDataFile = copyResource("input_data.txt");
-DataStream<String> productReviewStream =
+DataStream<ProductReview> productReviewStream =
env.fromSource(
- FileSource.forRecordStreamFormat(
- new TextLineInputFormat(),
- new Path(inputDataFile.getAbsolutePath()))
- .build(),
- WatermarkStrategy.noWatermarks(),
- "yaml-review-analysis-example");
+ FileSource.forRecordStreamFormat(
+ new TextLineInputFormat(),
+ new
Path(inputDataFile.getAbsolutePath()))
+ .build(),
+ WatermarkStrategy.noWatermarks(),
+ "yaml-review-analysis-example")
+ .map(json -> MAPPER.readValue(json, ProductReview.class))
+ .returns(ProductReview.class);
// Apply the YAML-declared agent BY NAME — `apply` accepts either an
// Agent instance or the name of an agent registered on the environment.
DataStream<Object> reviewAnalysisResStream =
agentsEnv
- .fromDataStream(productReviewStream)
+ .fromDataStream(productReviewStream, ProductReview::getId)
.apply("review_analysis_agent")
.toDataStream();
diff --git
a/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/FlinkIntegrationTest.java
b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/FlinkIntegrationTest.java
index d7b34970a..edb51d0cc 100644
---
a/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/FlinkIntegrationTest.java
+++
b/e2e-test/flink-agents-end-to-end-tests-integration/src/test/java/org/apache/flink/agents/integration/test/FlinkIntegrationTest.java
@@ -193,7 +193,7 @@ public class FlinkIntegrationTest {
DataStream<FlinkIntegrationAgent.TypedOutput> outputStream =
agentsEnv
- .fromDataStream(inputStream)
+ .fromDataStream(inputStream, value -> value)
.apply(new FlinkIntegrationAgent.TypedOutputAgent())
.toDataStream(TypeInformation.of(FlinkIntegrationAgent.TypedOutput.class));
@@ -228,7 +228,7 @@ public class FlinkIntegrationTest {
Table outputTable =
agentsEnv
- .fromDataStream(inputStream)
+ .fromDataStream(inputStream, value -> value)
.apply(new FlinkIntegrationAgent.TypedOutputAgent())
.toTable(TypeInformation.of(FlinkIntegrationAgent.TypedOutput.class));
@@ -260,7 +260,7 @@ public class FlinkIntegrationTest {
Table outputTable =
agentsEnv
- .fromDataStream(inputStream)
+ .fromDataStream(inputStream, value -> value)
.apply(new FlinkIntegrationAgent.TypedOutputAgent())
.toTable(outputSchema);
diff --git
a/examples/src/main/java/org/apache/flink/agents/examples/ModelRoutingExample.java
b/examples/src/main/java/org/apache/flink/agents/examples/ModelRoutingExample.java
index b33dd2d3d..6ab76c87e 100644
---
a/examples/src/main/java/org/apache/flink/agents/examples/ModelRoutingExample.java
+++
b/examples/src/main/java/org/apache/flink/agents/examples/ModelRoutingExample.java
@@ -94,7 +94,7 @@ public class ModelRoutingExample {
DataStream<Object> resultStream =
agentsEnv
- .fromDataStream(requestStream)
+ .fromDataStream(requestStream, value -> value)
.apply(new ModelRoutingAgent())
.toDataStream();
diff --git
a/examples/src/main/java/org/apache/flink/agents/examples/ModelRoutingJudgeExample.java
b/examples/src/main/java/org/apache/flink/agents/examples/ModelRoutingJudgeExample.java
index 9e2a26c0b..fbc5f7313 100644
---
a/examples/src/main/java/org/apache/flink/agents/examples/ModelRoutingJudgeExample.java
+++
b/examples/src/main/java/org/apache/flink/agents/examples/ModelRoutingJudgeExample.java
@@ -96,7 +96,7 @@ public class ModelRoutingJudgeExample {
DataStream<Object> resultStream =
agentsEnv
- .fromDataStream(requestStream)
+ .fromDataStream(requestStream, value -> value)
.apply(new ModelRoutingAgent())
.toDataStream();
diff --git
a/examples/src/main/java/org/apache/flink/agents/examples/ReActAgentExample.java
b/examples/src/main/java/org/apache/flink/agents/examples/ReActAgentExample.java
index fef67f32a..f82e663ce 100644
---
a/examples/src/main/java/org/apache/flink/agents/examples/ReActAgentExample.java
+++
b/examples/src/main/java/org/apache/flink/agents/examples/ReActAgentExample.java
@@ -134,7 +134,7 @@ public class ReActAgentExample {
// record shipping question.
DataStream<Object> reviewAnalysisResStream =
agentsEnv
- .fromDataStream(productReviewStream)
+ .fromDataStream(productReviewStream, row -> (String)
row.getField("id"))
.apply(reviewAnalysisReactAgent)
.toDataStream();
diff --git
a/examples/src/main/java/org/apache/flink/agents/examples/SkillsAgentExample.java
b/examples/src/main/java/org/apache/flink/agents/examples/SkillsAgentExample.java
index 6ca22c6bf..e3758674b 100644
---
a/examples/src/main/java/org/apache/flink/agents/examples/SkillsAgentExample.java
+++
b/examples/src/main/java/org/apache/flink/agents/examples/SkillsAgentExample.java
@@ -60,7 +60,10 @@ public class SkillsAgentExample {
// Use the MathAgent to answer each question with the math-calculator
skill.
DataStream<Object> answerStream =
- agentsEnv.fromDataStream(questionStream).apply(new
MathAgent()).toDataStream();
+ agentsEnv
+ .fromDataStream(questionStream, value -> value)
+ .apply(new MathAgent())
+ .toDataStream();
// Print the answers to stdout.
answerStream.print();
diff --git
a/examples/src/main/java/org/apache/flink/agents/examples/WorkflowMultipleAgentExample.java
b/examples/src/main/java/org/apache/flink/agents/examples/WorkflowMultipleAgentExample.java
index 7218faa2d..c27639687 100644
---
a/examples/src/main/java/org/apache/flink/agents/examples/WorkflowMultipleAgentExample.java
+++
b/examples/src/main/java/org/apache/flink/agents/examples/WorkflowMultipleAgentExample.java
@@ -18,8 +18,6 @@
package org.apache.flink.agents.examples;
-import com.fasterxml.jackson.core.JsonProcessingException;
-import com.fasterxml.jackson.databind.ObjectMapper;
import org.apache.flink.agents.api.AgentsExecutionEnvironment;
import org.apache.flink.agents.api.agents.AgentExecutionOptions;
import org.apache.flink.agents.api.resource.ResourceType;
@@ -65,8 +63,6 @@ import static
org.apache.flink.agents.examples.agents.CustomTypesAndResources.Pr
*/
public class WorkflowMultipleAgentExample {
- private static final ObjectMapper MAPPER = new ObjectMapper();
-
/**
* ProcessWindowFunction to aggregate score distribution and dislike
reasons.
*
@@ -74,15 +70,15 @@ public class WorkflowMultipleAgentExample {
* calculating score distribution percentages and collecting all
unsatisfied reasons.
*/
static class AggregateScoreDistributionAndDislikeReasons
- extends ProcessWindowFunction<ProductReviewAnalysisRes, String,
String, TimeWindow> {
+ extends ProcessWindowFunction<
+ ProductReviewAnalysisRes, ProductReviewSummary, String,
TimeWindow> {
@Override
public void process(
String key,
Context context,
Iterable<ProductReviewAnalysisRes> elements,
- Collector<String> out)
- throws JsonProcessingException {
+ Collector<ProductReviewSummary> out) {
// Initialize rating counts for scores 1-5
int[] ratingCounts = new int[5];
@@ -120,7 +116,7 @@ public class WorkflowMultipleAgentExample {
// Create and emit the aggregated result
ProductReviewSummary summary =
new ProductReviewSummary(key, formattedPercentages,
reasonList);
- out.collect(MAPPER.writeValueAsString(summary));
+ out.collect(summary);
}
}
@@ -179,7 +175,7 @@ public class WorkflowMultipleAgentExample {
// Aggregate the analysis results in 1-minute tumbling windows.
// This produces a score distribution and collects all unsatisfied
reasons for
// each product.
- DataStream<String> aggregatedAnalysisResStream =
+ DataStream<ProductReviewSummary> aggregatedAnalysisResStream =
reviewAnalysisResStream
.keyBy(ProductReviewAnalysisRes::getId)
.window(TumblingProcessingTimeWindows.of(Duration.ofMinutes(1)))
@@ -189,7 +185,7 @@ public class WorkflowMultipleAgentExample {
// suggestions based on the aggregated analysis results.
DataStream<Object> productSuggestionResStream =
agentsEnv
- .fromDataStream(aggregatedAnalysisResStream)
+ .fromDataStream(aggregatedAnalysisResStream,
ProductReviewSummary::getId)
.apply(new ProductSuggestionAgent())
.toDataStream();
diff --git
a/examples/src/main/java/org/apache/flink/agents/examples/WorkflowSingleAgentExample.java
b/examples/src/main/java/org/apache/flink/agents/examples/WorkflowSingleAgentExample.java
index 62c317fe0..ddc371f67 100644
---
a/examples/src/main/java/org/apache/flink/agents/examples/WorkflowSingleAgentExample.java
+++
b/examples/src/main/java/org/apache/flink/agents/examples/WorkflowSingleAgentExample.java
@@ -17,10 +17,13 @@
*/
package org.apache.flink.agents.examples;
+import com.fasterxml.jackson.databind.DeserializationFeature;
+import com.fasterxml.jackson.databind.ObjectMapper;
import org.apache.flink.agents.api.AgentsExecutionEnvironment;
import org.apache.flink.agents.api.agents.AgentExecutionOptions;
import org.apache.flink.agents.api.resource.ResourceType;
import org.apache.flink.agents.examples.agents.CustomTypesAndResources;
+import
org.apache.flink.agents.examples.agents.CustomTypesAndResources.ProductReview;
import org.apache.flink.agents.examples.agents.ReviewAnalysisAgent;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.connector.file.src.FileSource;
@@ -48,6 +51,9 @@ import java.nio.file.StandardCopyOption;
*/
public class WorkflowSingleAgentExample {
+ private static final ObjectMapper MAPPER =
+ new
ObjectMapper().configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES,
false);
+
/** Runs the example pipeline. */
public static void main(String[] args) throws Exception {
// Set up the Flink streaming environment and the Agents execution
environment.
@@ -68,19 +74,21 @@ public class WorkflowSingleAgentExample {
// Read product reviews from input_data.txt file as a streaming source.
// Each element represents a ProductReview.
File inputDataFile = copyResource("input_data.txt");
- DataStream<String> productReviewStream =
+ DataStream<ProductReview> productReviewStream =
env.fromSource(
- FileSource.forRecordStreamFormat(
- new TextLineInputFormat(),
- new
Path(inputDataFile.getAbsolutePath()))
- .build(),
- WatermarkStrategy.noWatermarks(),
- "streaming-agent-example");
+ FileSource.forRecordStreamFormat(
+ new TextLineInputFormat(),
+ new
Path(inputDataFile.getAbsolutePath()))
+ .build(),
+ WatermarkStrategy.noWatermarks(),
+ "streaming-agent-example")
+ .map(json -> MAPPER.readValue(json,
ProductReview.class))
+ .returns(ProductReview.class);
// Use the ReviewAnalysisAgent to analyze each product review.
DataStream<Object> reviewAnalysisResStream =
agentsEnv
- .fromDataStream(productReviewStream)
+ .fromDataStream(productReviewStream,
ProductReview::getId)
.apply(new ReviewAnalysisAgent())
.toDataStream();
diff --git
a/examples/src/main/java/org/apache/flink/agents/examples/YamlWorkflowAgentExample.java
b/examples/src/main/java/org/apache/flink/agents/examples/YamlWorkflowAgentExample.java
index b425ccd91..6a7761450 100644
---
a/examples/src/main/java/org/apache/flink/agents/examples/YamlWorkflowAgentExample.java
+++
b/examples/src/main/java/org/apache/flink/agents/examples/YamlWorkflowAgentExample.java
@@ -17,8 +17,11 @@
*/
package org.apache.flink.agents.examples;
+import com.fasterxml.jackson.databind.DeserializationFeature;
+import com.fasterxml.jackson.databind.ObjectMapper;
import org.apache.flink.agents.api.AgentsExecutionEnvironment;
import org.apache.flink.agents.api.agents.AgentExecutionOptions;
+import
org.apache.flink.agents.examples.agents.CustomTypesAndResources.ProductReview;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.connector.file.src.FileSource;
import org.apache.flink.connector.file.src.reader.TextLineInputFormat;
@@ -47,6 +50,9 @@ import static
org.apache.flink.agents.examples.WorkflowSingleAgentExample.copyRe
*/
public class YamlWorkflowAgentExample {
+ private static final ObjectMapper MAPPER =
+ new
ObjectMapper().configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES,
false);
+
/** Runs the example pipeline. */
public static void main(String[] args) throws Exception {
// Set up the Flink streaming environment and the Agents execution
environment.
@@ -67,20 +73,22 @@ public class YamlWorkflowAgentExample {
// Read product reviews from input_data.txt file as a streaming source.
// Each element represents a ProductReview.
File inputDataFile = copyResource("input_data.txt");
- DataStream<String> productReviewStream =
+ DataStream<ProductReview> productReviewStream =
env.fromSource(
- FileSource.forRecordStreamFormat(
- new TextLineInputFormat(),
- new
Path(inputDataFile.getAbsolutePath()))
- .build(),
- WatermarkStrategy.noWatermarks(),
- "yaml-review-analysis-example");
+ FileSource.forRecordStreamFormat(
+ new TextLineInputFormat(),
+ new
Path(inputDataFile.getAbsolutePath()))
+ .build(),
+ WatermarkStrategy.noWatermarks(),
+ "yaml-review-analysis-example")
+ .map(json -> MAPPER.readValue(json,
ProductReview.class))
+ .returns(ProductReview.class);
// Apply the YAML-declared agent BY NAME — ``apply`` accepts either an
// Agent instance or the name of an agent registered on the
environment.
DataStream<Object> reviewAnalysisResStream =
agentsEnv
- .fromDataStream(productReviewStream)
+ .fromDataStream(productReviewStream,
ProductReview::getId)
.apply("review_analysis_agent")
.toDataStream();
diff --git
a/examples/src/main/java/org/apache/flink/agents/examples/agents/ProductSuggestionAgent.java
b/examples/src/main/java/org/apache/flink/agents/examples/agents/ProductSuggestionAgent.java
index 959ad9cfe..1c773cc07 100644
---
a/examples/src/main/java/org/apache/flink/agents/examples/agents/ProductSuggestionAgent.java
+++
b/examples/src/main/java/org/apache/flink/agents/examples/agents/ProductSuggestionAgent.java
@@ -76,17 +76,17 @@ public class ProductSuggestionAgent extends Agent {
@Action(EventType.InputEvent)
public static void processInput(Event event, RunnerContext ctx) throws
Exception {
InputEvent inputEvent = (InputEvent) event;
- String input = (String) inputEvent.getInput();
-
- ProductReviewSummary summary = MAPPER.readValue(input,
ProductReviewSummary.class);
+ ProductReviewSummary summary = (ProductReviewSummary)
inputEvent.getInput();
ctx.getShortTermMemory().set(ID, summary.getId());
ctx.getShortTermMemory().set(SCORE_HIST, summary.getScoreHist());
String content =
- String.format(
- "{\n\"id\": %s,\n\"score_histogram\":
%s,\n\"unsatisfied_reasons\": %s\n}",
- summary.getId(), summary.getScoreHist(),
summary.getUnsatisfiedReasons());
+ MAPPER.writeValueAsString(
+ Map.of(
+ "id", summary.getId(),
+ "score_histogram", summary.getScoreHist(),
+ "unsatisfied_reasons",
summary.getUnsatisfiedReasons()));
ChatMessage msg = new ChatMessage(MessageRole.USER, "");
diff --git
a/examples/src/main/java/org/apache/flink/agents/examples/agents/ReviewAnalysisAgent.java
b/examples/src/main/java/org/apache/flink/agents/examples/agents/ReviewAnalysisAgent.java
index 46e30079a..4e78da8a7 100644
---
a/examples/src/main/java/org/apache/flink/agents/examples/agents/ReviewAnalysisAgent.java
+++
b/examples/src/main/java/org/apache/flink/agents/examples/agents/ReviewAnalysisAgent.java
@@ -17,7 +17,6 @@
*/
package org.apache.flink.agents.examples.agents;
-import com.fasterxml.jackson.databind.DeserializationFeature;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.apache.flink.agents.api.Event;
@@ -91,17 +90,12 @@ public class ReviewAnalysisAgent extends Agent {
@Action(EventType.InputEvent)
public static void processInput(Event event, RunnerContext ctx) throws
Exception {
InputEvent inputEvent = (InputEvent) event;
- String input = (String) inputEvent.getInput();
- MAPPER.configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES,
false);
CustomTypesAndResources.ProductReview inputObj =
- MAPPER.readValue(input,
CustomTypesAndResources.ProductReview.class);
+ (CustomTypesAndResources.ProductReview) inputEvent.getInput();
ctx.getShortTermMemory().set("id", inputObj.getId());
- String content =
- String.format(
- "{\n" + "\"id\": %s,\n" + "\"review\": \"%s\"\n" + "}",
- inputObj.getId(), inputObj.getReview());
+ String content = MAPPER.writeValueAsString(inputObj);
ChatMessage msg = new ChatMessage(MessageRole.USER, "");
ctx.sendEvent(
diff --git
a/examples/src/main/java/org/apache/flink/agents/examples/agents/TableReviewAnalysisAgent.java
b/examples/src/main/java/org/apache/flink/agents/examples/agents/TableReviewAnalysisAgent.java
index ee9bef4e9..1c41bb62a 100644
---
a/examples/src/main/java/org/apache/flink/agents/examples/agents/TableReviewAnalysisAgent.java
+++
b/examples/src/main/java/org/apache/flink/agents/examples/agents/TableReviewAnalysisAgent.java
@@ -54,7 +54,7 @@ import static
org.apache.flink.agents.examples.agents.CustomTypesAndResources.RE
* score and reasons for dissatisfaction.
*
* <p>The main difference from {@link ReviewAnalysisAgent} is that this agent
handles {@link Row}
- * input instead of JSON strings.
+ * input instead of {@link CustomTypesAndResources.ProductReview} objects.
*/
public class TableReviewAnalysisAgent extends Agent {
diff --git
a/examples/src/main/java/org/apache/flink/agents/examples/openai/OpenAiModelRoutingExample.java
b/examples/src/main/java/org/apache/flink/agents/examples/openai/OpenAiModelRoutingExample.java
index 965d435d3..bb5f23d11 100644
---
a/examples/src/main/java/org/apache/flink/agents/examples/openai/OpenAiModelRoutingExample.java
+++
b/examples/src/main/java/org/apache/flink/agents/examples/openai/OpenAiModelRoutingExample.java
@@ -99,7 +99,7 @@ public class OpenAiModelRoutingExample {
DataStream<Object> resultStream =
agentsEnv
- .fromDataStream(requestStream)
+ .fromDataStream(requestStream, value -> value)
.apply(new ModelRoutingAgent())
.toDataStream();
diff --git a/python/flink_agents/api/execution_environment.py
b/python/flink_agents/api/execution_environment.py
index d368fd9a6..96a07186b 100644
--- a/python/flink_agents/api/execution_environment.py
+++ b/python/flink_agents/api/execution_environment.py
@@ -222,10 +222,9 @@ class AgentsExecutionEnvironment(ABC):
----------
input : Table
Receive a Table as input.
- t_env: StreamTableEnvironment
- table environment supports convert Table to/from DataStream.
key_selector : KeySelector
- Extract key from each input record.
+ Required selector extracting the business key from each table row.
+ Rows with the same key share the agent's short-term memory.
Returns:
-------
diff --git a/python/flink_agents/runtime/remote_execution_environment.py
b/python/flink_agents/runtime/remote_execution_environment.py
index 8f9e35aab..6b5478421 100644
--- a/python/flink_agents/runtime/remote_execution_environment.py
+++ b/python/flink_agents/runtime/remote_execution_environment.py
@@ -360,14 +360,17 @@ class
RemoteExecutionEnvironment(AgentsExecutionEnvironment):
def __process_input_datastream(
input: DataStream, key_selector: KeySelector | Callable | None = None
) -> KeyedStream:
+ if key_selector is not None:
+ # An explicit selector always wins, even when the input is already
keyed,
+ # so callers can re-key a keyed input without re-building the
pipeline.
+ return input.key_by(key_selector)
if isinstance(input, KeyedStream):
return input
- else:
- if key_selector is None:
- msg = "KeySelector must be provided."
- raise RuntimeError(msg)
- input = input.key_by(key_selector)
- return input
+ msg = (
+ "A key selector must be provided when the input DataStream is not
keyed; "
+ "pass a KeyedStream or a key selector."
+ )
+ raise RuntimeError(msg)
def from_datastream(
self, input: DataStream, key_selector: KeySelector | Callable | None =
None
@@ -379,8 +382,10 @@ class
RemoteExecutionEnvironment(AgentsExecutionEnvironment):
input : DataStream
Receive a DataStream as input.
key_selector : KeySelector
- Extract key from each input record, must not be None when input is
- not KeyedStream.
+ Extract key from each input record. An explicit selector always
takes
+ precedence, even over the key of an already-keyed input; otherwise
a
+ KeyedStream input keeps its key, and a non-keyed input without a
+ selector is rejected.
"""
input = self.__process_input_datastream(input, key_selector)
@@ -404,8 +409,16 @@ class
RemoteExecutionEnvironment(AgentsExecutionEnvironment):
input : Table
Receive a Table as input.
key_selector : KeySelector
- Extract key from each input record.
+ Required selector extracting the business key from each table row.
+ Rows with the same key share the agent's short-term memory.
"""
+ if key_selector is None:
+ msg = (
+ "A key selector must be provided for Table input; "
+ "pass a key selector to from_table."
+ )
+ raise RuntimeError(msg)
+
input = self.t_env.to_data_stream(table=input)
input = input.map(lambda x: x, output_type=PickledBytesTypeInfo())
diff --git
a/python/flink_agents/runtime/tests/test_remote_execution_environment.py
b/python/flink_agents/runtime/tests/test_remote_execution_environment.py
index 2f9b6f88a..6584b733f 100644
--- a/python/flink_agents/runtime/tests/test_remote_execution_environment.py
+++ b/python/flink_agents/runtime/tests/test_remote_execution_environment.py
@@ -28,7 +28,7 @@ import pytest
import yaml
from pyflink.common import Row
from pyflink.common.typeinfo import ExternalTypeInfo, Types
-from pyflink.datastream import StreamExecutionEnvironment
+from pyflink.datastream import KeyedStream, StreamExecutionEnvironment
from pyflink.table import DataTypes, Schema, StreamTableEnvironment
from flink_agents.api.agents.agent import Agent
@@ -665,3 +665,51 @@ def
test_datastream_to_table_uses_input_environment(provide_table_env, use_schem
assert
output.get_execution_environment()._j_stream_execution_environment.equals(
env._j_stream_execution_environment
)
+
+
+def test_key_selection_combinations() -> None:
+ """Key selection follows the same rules as the Java runtime.
+
+ An explicit selector always wins (even over the key of an already-keyed
+ input), a keyed input without a selector keeps its key, and an unkeyed
+ input without a selector is rejected with a clear error instead of being
+ silently keyed by the whole record.
+ """
+ env = StreamExecutionEnvironment.get_execution_environment()
+ unkeyed = env.from_collection(["same:first", "other:first"])
+ keyed = unkeyed.key_by(lambda value: value.split(":")[0])
+ process = (
+
RemoteExecutionEnvironment._RemoteExecutionEnvironment__process_input_datastream
+ )
+
+ def session(value: str) -> str:
+ return value.split(":")[0]
+
+ # An explicit selector wins even on a keyed input: the pipeline is re-keyed
+ # (a new stream) rather than silently keeping the input's own key.
+ assert process(keyed, session) is not keyed
+ # A keyed input without a selector keeps its key.
+ assert process(keyed, None) is keyed
+ # An unkeyed input with a selector is keyed by that selector.
+ assert isinstance(process(unkeyed, session), KeyedStream)
+ # An unkeyed input without a selector is rejected with a clear error.
+ with pytest.raises(RuntimeError, match="key selector"):
+ process(unkeyed, None)
+
+
[email protected]("provide_selector", [False, True])
+def test_table_requires_explicit_key_selector(provide_selector: bool) -> None:
+ """Table input must be keyed explicitly, including when None is passed."""
+ env = StreamExecutionEnvironment.get_execution_environment()
+ t_env = StreamTableEnvironment.create(env)
+ remote_env = RemoteExecutionEnvironment(env=env, t_env=t_env)
+ table = t_env.from_elements([("same", "first"), ("same", "second")],
["id", "text"])
+
+ if provide_selector:
+ builder = remote_env.from_table(table, key_selector=lambda row: row[0])
+ assert isinstance(builder._RemoteAgentBuilder__input, KeyedStream)
+ else:
+ with pytest.raises(RuntimeError, match="key selector.*Table input"):
+ remote_env.from_table(table)
+ with pytest.raises(RuntimeError, match="key selector.*Table input"):
+ remote_env.from_table(table, key_selector=None)
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/env/RemoteExecutionEnvironment.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/env/RemoteExecutionEnvironment.java
index 8cb9e599d..e2218f665 100644
---
a/runtime/src/main/java/org/apache/flink/agents/runtime/env/RemoteExecutionEnvironment.java
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/env/RemoteExecutionEnvironment.java
@@ -33,6 +33,7 @@ import org.apache.flink.api.java.typeutils.RowTypeInfo;
import org.apache.flink.configuration.ConfigConstants;
import org.apache.flink.configuration.YamlParserUtils;
import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.datastream.KeyedStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.table.api.Schema;
import org.apache.flink.table.api.Table;
@@ -94,6 +95,10 @@ public class RemoteExecutionEnvironment extends
AgentsExecutionEnvironment {
@Override
public <K> AgentBuilder fromTable(Table input, KeySelector<Object, K>
keySelector) {
+ if (keySelector == null) {
+ throw new IllegalArgumentException(
+ "A key selector must be provided for Table input; pass a
key selector to fromTable.");
+ }
return new RemoteAgentBuilder<>(
input, getTableEnvironment(), keySelector, env, config,
resources, agents);
}
@@ -220,6 +225,7 @@ public class RemoteExecutionEnvironment extends
AgentsExecutionEnvironment {
}
@Override
+ @SuppressWarnings("unchecked")
public DataStream<Object> toDataStream() {
if (agentPlan == null) {
throw new IllegalStateException("Must apply agent before
calling toDataStream");
@@ -229,10 +235,21 @@ public class RemoteExecutionEnvironment extends
AgentsExecutionEnvironment {
if (keySelector != null) {
outputDataStream =
CompileUtils.connectToAgent(inputDataStream,
keySelector, agentPlan);
- } else {
- // If no key selector provided, use a simple pass-through
key selector
+ } else if (inputDataStream instanceof KeyedStream) {
+ // An already-keyed input carries its business key;
connect through the
+ // keyed-stream path to preserve that key, instead of
replacing it with a
+ // pass-through selector that would split a session across
its messages.
outputDataStream =
- CompileUtils.connectToAgent(inputDataStream, x ->
x, agentPlan);
+ CompileUtils.connectToAgent(
+ (KeyedStream<T, ?>) inputDataStream,
agentPlan);
+ } else {
+ // Without a key selector there is no business key to key
the agent by;
+ // keying by the whole record would silently split
sessions, so require
+ // the caller to pass a keyed stream or a selector, as
documented.
+ throw new IllegalArgumentException(
+ "A key selector must be provided when the input
DataStream is not"
+ + " keyed; pass a KeyedStream or a key
selector to"
+ + " fromDataStream.");
}
}
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/env/RemoteExecutionEnvironmentTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/env/RemoteExecutionEnvironmentTest.java
index 0bb09df6f..d49e291f8 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/env/RemoteExecutionEnvironmentTest.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/env/RemoteExecutionEnvironmentTest.java
@@ -18,15 +18,32 @@
package org.apache.flink.agents.runtime.env;
import org.apache.flink.agents.api.AgentsExecutionEnvironment;
+import org.apache.flink.agents.api.Event;
+import org.apache.flink.agents.api.EventType;
+import org.apache.flink.agents.api.InputEvent;
+import org.apache.flink.agents.api.OutputEvent;
+import org.apache.flink.agents.api.agents.Agent;
+import org.apache.flink.agents.api.annotation.Action;
+import org.apache.flink.agents.api.context.MemoryObject;
+import org.apache.flink.agents.api.context.RunnerContext;
import org.apache.flink.agents.plan.AgentConfiguration;
import org.apache.flink.api.common.JobExecutionResult;
+import org.apache.flink.streaming.api.datastream.DataStream;
+import org.apache.flink.streaming.api.datastream.KeyedStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
+import org.apache.flink.table.api.DataTypes;
+import org.apache.flink.table.api.Table;
+import org.apache.flink.table.api.bridge.java.StreamTableEnvironment;
+import org.apache.flink.types.Row;
+import org.apache.flink.util.CloseableIterator;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.io.TempDir;
import java.io.File;
import java.io.FileNotFoundException;
import java.io.PrintWriter;
+import java.util.ArrayList;
+import java.util.List;
import java.util.UUID;
import static
org.apache.flink.agents.runtime.env.RemoteExecutionEnvironment.FLINK_CONF_FILENAME;
@@ -95,4 +112,158 @@ public class RemoteExecutionEnvironmentTest {
verify(flinkEnv).execute(customJobName);
}
+
+ @Test
+ void fromDataStreamPreservesExistingKeyedStreamKeyWhenNoSelectorGiven()
throws Exception {
+ StreamExecutionEnvironment env =
StreamExecutionEnvironment.getExecutionEnvironment();
+ env.setParallelism(1);
+ AgentsExecutionEnvironment agentsEnv =
+ AgentsExecutionEnvironment.getExecutionEnvironment(env);
+
+ KeyedStream<String, String> keyedBySession =
+ env.fromData("same:first", "same:second", "same:first again",
"other:first")
+ .keyBy(value -> value.split(":")[0]);
+
+ DataStream<Object> output =
+ agentsEnv.fromDataStream(keyedBySession).apply(new
CountingAgent()).toDataStream();
+
+ List<Object> results = new ArrayList<>();
+ try (CloseableIterator<Object> iterator = output.executeAndCollect()) {
+ iterator.forEachRemaining(results::add);
+ }
+
+ // Messages of the "same" session share short-term memory and count 1,
2, 3 in their
+ // arrival order, while the independent "other" session starts from 1.
A pass-through
+ // selector that replaces the session key with the whole message would
count every
+ // "same" message as 1.
+ assertThat(results).containsExactlyInAnyOrder("same:1", "same:2",
"same:3", "other:1");
+ }
+
+ @Test
+ void fromDataStreamUsesExplicitSelectorEvenOnKeyedInput() throws Exception
{
+ StreamExecutionEnvironment env =
StreamExecutionEnvironment.getExecutionEnvironment();
+ env.setParallelism(1);
+ AgentsExecutionEnvironment agentsEnv =
+ AgentsExecutionEnvironment.getExecutionEnvironment(env);
+
+ // The stream is keyed by the whole record; the explicit selector keys
by session, and
+ // the selector must win: the "same" messages share short-term memory
either way only
+ // if the session key is applied rather than the stream's own
whole-record key.
+ KeyedStream<String, String> keyedByWholeRecord =
+ env.fromData("same:first", "same:second",
"other:first").keyBy(value -> value);
+
+ DataStream<Object> output =
+ agentsEnv
+ .fromDataStream(keyedByWholeRecord, value ->
value.split(":")[0])
+ .apply(new CountingAgent())
+ .toDataStream();
+
+ List<Object> results = new ArrayList<>();
+ try (CloseableIterator<Object> iterator = output.executeAndCollect()) {
+ iterator.forEachRemaining(results::add);
+ }
+
+ assertThat(results).containsExactlyInAnyOrder("same:1", "same:2",
"other:1");
+ }
+
+ @Test
+ void fromDataStreamKeysUnkeyedInputBySelector() throws Exception {
+ StreamExecutionEnvironment env =
StreamExecutionEnvironment.getExecutionEnvironment();
+ env.setParallelism(1);
+ AgentsExecutionEnvironment agentsEnv =
+ AgentsExecutionEnvironment.getExecutionEnvironment(env);
+
+ DataStream<String> unkeyed = env.fromData("same:first", "same:second",
"other:first");
+
+ DataStream<Object> output =
+ agentsEnv
+ .fromDataStream(unkeyed, value -> value.split(":")[0])
+ .apply(new CountingAgent())
+ .toDataStream();
+
+ List<Object> results = new ArrayList<>();
+ try (CloseableIterator<Object> iterator = output.executeAndCollect()) {
+ iterator.forEachRemaining(results::add);
+ }
+
+ assertThat(results).containsExactlyInAnyOrder("same:1", "same:2",
"other:1");
+ }
+
+ @Test
+ void fromDataStreamRejectsUnkeyedInputWithoutSelector() {
+ StreamExecutionEnvironment env =
StreamExecutionEnvironment.getExecutionEnvironment();
+ env.setParallelism(1);
+ AgentsExecutionEnvironment agentsEnv =
+ AgentsExecutionEnvironment.getExecutionEnvironment(env);
+
+ DataStream<String> unkeyed = env.fromData("same:first");
+
+ assertThatThrownBy(
+ () ->
+ agentsEnv
+ .fromDataStream(unkeyed)
+ .apply(new CountingAgent())
+ .toDataStream())
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessageContaining("key selector must be provided");
+ }
+
+ @Test
+ void fromTableRejectsMissingSelector() {
+ StreamExecutionEnvironment env =
StreamExecutionEnvironment.getExecutionEnvironment();
+ StreamTableEnvironment tableEnv = StreamTableEnvironment.create(env);
+ AgentsExecutionEnvironment agentsEnv =
+ AgentsExecutionEnvironment.getExecutionEnvironment(env,
tableEnv);
+ Table input = tableEnv.fromValues(DataTypes.STRING(), "same:first");
+
+ assertThatThrownBy(() -> agentsEnv.fromTable(input))
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessageContaining("key selector")
+ .hasMessageContaining("Table input");
+ assertThatThrownBy(() -> agentsEnv.fromTable(input, null))
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessageContaining("key selector")
+ .hasMessageContaining("Table input");
+ }
+
+ @Test
+ void fromTableUsesExplicitBusinessKey() throws Exception {
+ StreamExecutionEnvironment env =
StreamExecutionEnvironment.getExecutionEnvironment();
+ env.setParallelism(1);
+ StreamTableEnvironment tableEnv = StreamTableEnvironment.create(env);
+ AgentsExecutionEnvironment agentsEnv =
+ AgentsExecutionEnvironment.getExecutionEnvironment(env,
tableEnv);
+ Table input =
+ tableEnv.fromValues(DataTypes.STRING(), "same:first",
"same:second", "other:first");
+
+ DataStream<Object> output =
+ agentsEnv
+ .fromTable(input, row -> ((String) ((Row)
row).getField(0)).split(":")[0])
+ .apply(new CountingAgent())
+ .toDataStream();
+
+ List<Object> results = new ArrayList<>();
+ try (CloseableIterator<Object> iterator = output.executeAndCollect()) {
+ iterator.forEachRemaining(results::add);
+ }
+ assertThat(results).containsExactlyInAnyOrder("same:1", "same:2",
"other:1");
+ }
+
+ /** Model-free agent counting inputs per business key in short-term
memory. */
+ public static class CountingAgent extends Agent {
+
+ @Action(EventType.InputEvent)
+ public static void countInputs(Event event, RunnerContext context)
throws Exception {
+ Object value = InputEvent.fromEvent(event).getInput();
+ String input = (String) (value instanceof Row ? ((Row)
value).getField(0) : value);
+ String session = input.split(":")[0];
+ MemoryObject memory = context.getShortTermMemory();
+ int count = 1;
+ if (memory.isExist("count")) {
+ count = ((Number) memory.get("count").getValue()).intValue() +
1;
+ }
+ memory.set("count", count);
+ context.sendEvent(new OutputEvent(session + ":" + count));
+ }
+ }
}
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/env/TypedOutputTerminalTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/env/TypedOutputTerminalTest.java
index e99c808a5..4ae759f9c 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/env/TypedOutputTerminalTest.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/env/TypedOutputTerminalTest.java
@@ -74,7 +74,9 @@ class TypedOutputTerminalTest {
AgentsExecutionEnvironment agentsEnv =
AgentsExecutionEnvironment.getExecutionEnvironment(env);
DataStream<String> inputStream = env.fromData(List.of("a", "b", "c"));
- return agentsEnv.fromDataStream(inputStream).apply(new EchoAgent());
+ // The key is irrelevant to what these tests exercise; an explicit
selector keeps the
+ // builder valid now that an unkeyed input without a selector is
rejected.
+ return agentsEnv.fromDataStream(inputStream, value -> value).apply(new
EchoAgent());
}
@Test
