Repository: LangStream/langstream (language: python)

Issue:
Add agent that can write to one of many possible topics
Can we create an agent that needs to run some logic or call an external API and based on the results, it decides which topic to write to? For example an agent that need to call a REST API for a multiclass classification model, and based on model prediction, the agent decides which topic it will write to from a list of available topics.


Discussion / hints:
I think that a good approach to provide this feature is to add some "topic-sink" that allows to dispatch the incoming message to a topic depending on some properties or values.

So in your case the custom agent may simply add a "property" (or "header") to the message with the destination and then the "topic-sink" will dispatch the message.

I am referring to a sink because a simple "processor" is required to have a single "input" and a single "output", because that's the model in the ExecutionPlan and having a processor that is able to by pass the pipeline configuration may lead to tricky situations.



That's a great idea. A dispatcher will solve this problem if it can route an incoming message to a topic based on some values or properties in the message itself.

Fix patch:
diff --git a/langstream-agents/langstream-agents-commons/src/main/java/ai/langstream/ai/agents/commons/TransformContext.java b/langstream-agents/langstream-agents-commons/src/main/java/ai/langstream/ai/agents/commons/TransformContext.java
index a7e1c02a6..25a1ceb6c 100644
--- a/langstream-agents/langstream-agents-commons/src/main/java/ai/langstream/ai/agents/commons/TransformContext.java
+++ b/langstream-agents/langstream-agents-commons/src/main/java/ai/langstream/ai/agents/commons/TransformContext.java
@@ -15,6 +15,9 @@
  */
 package ai.langstream.ai.agents.commons;
 
+import ai.langstream.api.runner.code.Header;
+import ai.langstream.api.runner.code.Record;
+import ai.langstream.api.runner.code.SimpleRecord;
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.core.type.TypeReference;
 import com.fasterxml.jackson.databind.JsonNode;
@@ -22,13 +25,23 @@
 import com.fasterxml.jackson.databind.node.ObjectNode;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.sql.Time;
+import java.sql.Timestamp;
+import java.time.Instant;
+import java.time.LocalDate;
+import java.time.LocalDateTime;
+import java.time.LocalTime;
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.Date;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
+import java.util.stream.Collectors;
 import lombok.Data;
 import lombok.SneakyThrows;
 import lombok.extern.slf4j.Slf4j;
@@ -358,4 +371,172 @@ public static Object safeClone(Object object) {
         }
         throw new UnsupportedOperationException("Cannot copy a value of " + object.getClass());
     }
+
+    public static TransformContext recordToTransformContext(
+            Record record, boolean attemptJsonConversion) {
+        TransformContext context = new TransformContext();
+        context.setKeyObject(record.key());
+        context.setKeySchemaType(
+                record.key() == null ? null : getSchemaType(record.key().getClass()));
+        // TODO: temporary hack. We should be able to get the schema from the record
+        if (record.key() instanceof GenericRecord) {
+            context.setKeyNativeSchema(((GenericRecord) record.key()).getSchema());
+        }
+        context.setValueObject(record.value());
+        context.setValueSchemaType(
+                record.value() == null ? null : getSchemaType(record.value().getClass()));
+        // TODO: temporary hack. We should be able to get the schema from the record
+        if (record.value() instanceof GenericRecord) {
+            context.setKeyNativeSchema(((GenericRecord) record.value()).getSchema());
+        }
+        context.setInputTopic(record.origin());
+        context.setEventTime(record.timestamp());
+        if (attemptJsonConversion) {
+            context.setKeyObject(attemptJsonConversion(context.getKeyObject()));
+            context.setValueObject(attemptJsonConversion(context.getValueObject()));
+        }
+        // the headers must be Strings, this is a tentative conversion
+        // in the future we need a better way to handle headers
+        context.setProperties(
+                record.headers().stream()
+                        .filter(h -> h.key() != null && h.value() != null)
+                        .collect(
+                                Collectors.toMap(
+                                        Header::key,
+                                        (h -> {
+                                            if (h.value() == null) {
+                                                return null;
+                                            }
+                                            if (h.value() instanceof byte[]) {
+                                                return new String(
+                                                        (byte[]) h.value(), StandardCharsets.UTF_8);
+                                            } else {
+                                                return h.value().toString();
+                                            }
+                                        }))));
+        return context;
+    }
+
+    public static Optional<Record> transformContextToRecord(TransformContext context) {
+        if (context.isDropCurrentRecord()) {
+            return Optional.empty();
+        }
+        List<Header> headers = new ArrayList<>();
+        context.getProperties()
+                .forEach(
+                        (key, value) -> {
+                            SimpleRecord.SimpleHeader header =
+                                    new SimpleRecord.SimpleHeader(key, value);
+                            headers.add(header);
+                        });
+        return Optional.of(new TransformRecord(context, headers));
+    }
+
+    private record TransformRecord(TransformContext context, Collection<Header> headers)
+            implements Record {
+        private TransformRecord(TransformContext context, Collection<Header> headers) {
+            this.context = context;
+            this.headers = new ArrayList<>(headers);
+        }
+
+        @Override
+        public Object key() {
+            return context.getKeyObject();
+        }
+
+        @Override
+        public Object value() {
+            return context.getValueObject();
+        }
+
+        @Override
+        public String origin() {
+            return context.getInputTopic();
+        }
+
+        @Override
+        public Long timestamp() {
+            return context.getEventTime();
+        }
+    }
+
+    private static TransformSchemaType getSchemaType(Class<?> javaType) {
+        if (String.class.isAssignableFrom(javaType)) {
+            return TransformSchemaType.STRING;
+        }
+        if (Byte.class.isAssignableFrom(javaType)) {
+            return TransformSchemaType.INT8;
+        }
+        if (Short.class.isAssignableFrom(javaType)) {
+            return TransformSchemaType.INT16;
+        }
+        if (Integer.class.isAssignableFrom(javaType)) {
+            return TransformSchemaType.INT32;
+        }
+        if (Long.class.isAssignableFrom(javaType)) {
+            return TransformSchemaType.INT64;
+        }
+        if (Double.class.isAssignableFrom(javaType)) {
+            return TransformSchemaType.DOUBLE;
+        }
+        if (Float.class.isAssignableFrom(javaType)) {
+            return TransformSchemaType.FLOAT;
+        }
+        if (Boolean.class.isAssignableFrom(javaType)) {
+            return TransformSchemaType.BOOLEAN;
+        }
+        if (byte[].class.isAssignableFrom(javaType)) {
+            return TransformSchemaType.BYTES;
+        }
+        // Must be before DATE
+        if (Time.class.isAssignableFrom(javaType)) {
+            return TransformSchemaType.TIME;
+        }
+        // Must be before DATE
+        if (Timestamp.class.isAssignableFrom(javaType)) {
+            return TransformSchemaType.TIMESTAMP;
+        }
+        if (Date.class.isAssignableFrom(javaType)) {
+            return TransformSchemaType.DATE;
+        }
+        if (Instant.class.isAssignableFrom(javaType)) {
+            return TransformSchemaType.INSTANT;
+        }
+        if (LocalDate.class.isAssignableFrom(javaType)) {
+            return TransformSchemaType.LOCAL_DATE;
+        }
+        if (LocalTime.class.isAssignableFrom(javaType)) {
+            return TransformSchemaType.LOCAL_TIME;
+        }
+        if (LocalDateTime.class.isAssignableFrom(javaType)) {
+            return TransformSchemaType.LOCAL_DATE_TIME;
+        }
+        if (GenericRecord.class.isAssignableFrom(javaType)) {
+            return TransformSchemaType.AVRO;
+        }
+        if (JsonNode.class.isAssignableFrom(javaType)) {
+            return TransformSchemaType.JSON;
+        }
+        if (Map.class.isAssignableFrom(javaType)) {
+            return TransformSchemaType.MAP;
+        }
+        throw new IllegalArgumentException("Unsupported data type: " + javaType);
+    }
+
+    public static Object attemptJsonConversion(Object value) {
+        try {
+            if (value instanceof String) {
+                return OBJECT_MAPPER.readValue(
+                        (String) value, new TypeReference<Map<String, Object>>() {});
+            } else if (value instanceof byte[]) {
+                return OBJECT_MAPPER.readValue(
+                        (byte[]) value, new TypeReference<Map<String, Object>>() {});
+            }
+        } catch (IOException e) {
+            if (log.isDebugEnabled()) {
+                log.debug("Cannot convert value to json", e);
+            }
+        }
+        return value;
+    }
 }
diff --git a/langstream-agents/langstream-agents-flow-control/pom.xml b/langstream-agents/langstream-agents-flow-control/pom.xml
new file mode 100644
index 000000000..a4daeb85d
--- /dev/null
+++ b/langstream-agents/langstream-agents-flow-control/pom.xml
@@ -0,0 +1,91 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!--
+
+    Copyright DataStax, Inc.
+
+    Licensed 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.
+
+-->
+<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+  <parent>
+    <groupId>ai.langstream</groupId>
+    <artifactId>langstream-agents</artifactId>
+    <version>0.1.1-SNAPSHOT</version>
+  </parent>
+  <modelVersion>4.0.0</modelVersion>
+  <artifactId>langstream-agents-flow-control</artifactId>
+  <packaging>jar</packaging>
+  <name>LangStream - Flow control agents</name>
+  <properties>
+    <maven.compiler.source>17</maven.compiler.source>
+    <maven.compiler.target>17</maven.compiler.target>
+  </properties>
+  <dependencies>
+    <dependency>
+      <groupId>${project.groupId}</groupId>
+      <artifactId>langstream-api</artifactId>
+      <version>${project.version}</version>
+      <scope>provided</scope>
+    </dependency>
+    <dependency>
+      <groupId>${project.groupId}</groupId>
+      <artifactId>langstream-agents-commons</artifactId>
+      <version>${project.version}</version>
+    </dependency>
+    <dependency>
+      <groupId>org.projectlombok</groupId>
+      <artifactId>lombok</artifactId>
+      <scope>provided</scope>
+    </dependency>
+    <dependency>
+      <groupId>org.slf4j</groupId>
+      <artifactId>slf4j-api</artifactId>
+    </dependency>
+    <dependency>
+      <groupId>ch.qos.logback</groupId>
+      <artifactId>logback-core</artifactId>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>ch.qos.logback</groupId>
+      <artifactId>logback-classic</artifactId>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>org.junit.jupiter</groupId>
+      <artifactId>junit-jupiter</artifactId>
+      <scope>test</scope>
+    </dependency>
+  </dependencies>
+  <build>
+    <plugins>
+      <plugin>
+        <groupId>org.apache.nifi</groupId>
+        <artifactId>nifi-nar-maven-plugin</artifactId>
+        <extensions>true</extensions>
+        <configuration>
+          <classifier>nar</classifier>
+        </configuration>
+        <executions>
+          <execution>
+            <id>default-nar</id>
+            <phase>package</phase>
+            <goals>
+              <goal>nar</goal>
+            </goals>
+          </execution>
+        </executions>
+      </plugin>
+    </plugins>
+  </build>
+</project>
diff --git a/langstream-agents/langstream-agents-flow-control/src/main/java/ai/langstream/agents/flow/DispatchAgent.java b/langstream-agents/langstream-agents-flow-control/src/main/java/ai/langstream/agents/flow/DispatchAgent.java
new file mode 100644
index 000000000..217982644
--- /dev/null
+++ b/langstream-agents/langstream-agents-flow-control/src/main/java/ai/langstream/agents/flow/DispatchAgent.java
@@ -0,0 +1,178 @@
+/*
+ * Copyright DataStax, Inc.
+ *
+ * Licensed 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 ai.langstream.agents.flow;
+
+import ai.langstream.ai.agents.commons.TransformContext;
+import ai.langstream.ai.agents.commons.jstl.predicate.JstlPredicate;
+import ai.langstream.api.runner.code.AbstractAgentCode;
+import ai.langstream.api.runner.code.AgentContext;
+import ai.langstream.api.runner.code.AgentProcessor;
+import ai.langstream.api.runner.code.Record;
+import ai.langstream.api.runner.code.RecordSink;
+import ai.langstream.api.runner.topics.TopicProducer;
+import ai.langstream.api.runtime.ComponentType;
+import ai.langstream.api.util.ConfigurationUtils;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import lombok.extern.slf4j.Slf4j;
+
+@Slf4j
+public class DispatchAgent extends AbstractAgentCode implements AgentProcessor {
+
+    record Route(String destination, boolean drop, JstlPredicate predicate) {}
+
+    private final List<Route> routes = new ArrayList<>();
+    private final Map<String, TopicProducer> producers = new HashMap<>();
+    private AgentContext agentContext;
+
+    @SuppressWarnings("unchecked")
+    @Override
+    public void init(Map<String, Object> configuration) {
+        List<Map<String, Object>> routes =
+                (List<Map<String, Object>>) configuration.getOrDefault("routes", List.of());
+        routes.forEach(
+                r -> {
+                    String when = ConfigurationUtils.getString("when", "", r);
+                    String destination = ConfigurationUtils.getString("destination", "", r);
+                    String action = ConfigurationUtils.getString("action", "dispatch", r);
+                    final boolean drop;
+                    switch (action) {
+                        case "dispatch":
+                            drop = false;
+                            break;
+                        case "drop":
+                            drop = true;
+                            break;
+                        default:
+                            throw new IllegalArgumentException();
+                    }
+                    if (drop) {
+                        log.info("Condition: \"{}\", action = drop", when);
+                        if (!destination.isEmpty()) {
+                            throw new IllegalStateException(
+                                    "drop action cannot have a destination");
+                        }
+                    } else {
+                        log.info(
+                                "Condition: \"{}\", action = dispatch, destination: {}",
+                                when,
+                                destination);
+                    }
+                    this.routes.add(new Route(destination, drop, new JstlPredicate(when)));
+                });
+    }
+
+    @Override
+    public void setContext(AgentContext context) throws Exception {
+        this.agentContext = context;
+    }
+
+    @Override
+    public void start() throws Exception {
+        routes.forEach(
+                r -> {
+                    String topic = r.destination;
+                    if (topic != null && !topic.isEmpty()) {
+                        TopicProducer producer =
+                                agentContext
+                                        .getTopicConnectionProvider()
+                                        .createProducer(
+                                                agentContext.getGlobalAgentId(), topic, Map.of());
+                        producer.start();
+                        producers.put(topic, producer);
+                    }
+                });
+    }
+
+    @Override
+    public void process(List<Record> records, RecordSink recordSink) {
+        for (Record record : records) {
+            processRecord(record, recordSink);
+        }
+    }
+
+    @Override
+    public ComponentType componentType() {
+        return ComponentType.PROCESSOR;
+    }
+
+    public void processRecord(Record record, RecordSink recordSink) {
+        try {
+            TransformContext context = TransformContext.recordToTransformContext(record, true);
+
+            for (Route r : routes) {
+                boolean test = r.predicate.test(context);
+                if (test) {
+                    if (r.drop) {
+                        if (log.isDebugEnabled()) {
+                            log.debug("Discarding record {} - action=drop", record);
+                        }
+                        recordSink.emit(new SourceRecordAndResult(record, List.of(), null));
+                    } else {
+                        if (log.isDebugEnabled()) {
+                            log.debug("Dispatching record {} to topic {}", record, r.destination);
+                        }
+                        TopicProducer topicProducer = producers.get(r.destination);
+                        topicProducer
+                                .write(record)
+                                .whenComplete(
+                                        (__, e) -> {
+                                            if (e != null) {
+                                                log.error(
+                                                        "Error writing record to topic {}",
+                                                        r.destination,
+                                                        e);
+                                                recordSink.emit(
+                                                        new SourceRecordAndResult(record, null, e));
+                                            } else {
+                                                // the record is to be marked as processed, but not
+                                                // emitted to the
+                                                // next agent
+                                                recordSink.emit(
+                                                        new SourceRecordAndResult(
+                                                                record, List.of(), null));
+                                            }
+                                        });
+                    }
+                    return;
+                }
+            }
+
+            if (log.isDebugEnabled()) {
+                log.debug("Sending record {} to the default destination", record);
+            }
+            recordSink.emit(new SourceRecordAndResult(record, List.of(record), null));
+        } catch (Throwable error) {
+            log.error("Error processing record: {}", record, error);
+            recordSink.emit(new SourceRecordAndResult(record, null, error));
+        }
+    }
+
+    @Override
+    public void close() throws Exception {
+        producers.forEach(
+                (destination, producer) -> {
+                    log.info("Closing producer for destination {}", destination);
+                    try {
+                        producer.close();
+                    } catch (Exception e) {
+                        log.error("Error closing producer for destination {}", destination, e);
+                    }
+                });
+    }
+}
diff --git a/langstream-agents/langstream-agents-flow-control/src/main/java/ai/langstream/agents/flow/FlowControlAgentsCodeProvider.java b/langstream-agents/langstream-agents-flow-control/src/main/java/ai/langstream/agents/flow/FlowControlAgentsCodeProvider.java
new file mode 100644
index 000000000..47cb232ec
--- /dev/null
+++ b/langstream-agents/langstream-agents-flow-control/src/main/java/ai/langstream/agents/flow/FlowControlAgentsCodeProvider.java
@@ -0,0 +1,37 @@
+/*
+ * Copyright DataStax, Inc.
+ *
+ * Licensed 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 ai.langstream.agents.flow;
+
+import ai.langstream.api.runner.code.AgentCodeProvider;
+import ai.langstream.api.runner.code.AgentProcessor;
+import java.util.Map;
+import java.util.function.Supplier;
+
+public class FlowControlAgentsCodeProvider implements AgentCodeProvider {
+
+    private static final Map<String, Supplier<AgentProcessor>> FACTORIES =
+            Map.of("dispatch", DispatchAgent::new);
+
+    @Override
+    public boolean supports(String agentType) {
+        return FACTORIES.containsKey(agentType);
+    }
+
+    @Override
+    public AgentProcessor createInstance(String agentType) {
+        return FACTORIES.get(agentType).get();
+    }
+}
diff --git a/langstream-agents/langstream-agents-flow-control/src/main/resources/META-INF/ai.langstream.agents.index b/langstream-agents/langstream-agents-flow-control/src/main/resources/META-INF/ai.langstream.agents.index
new file mode 100644
index 000000000..554505981
--- /dev/null
+++ b/langstream-agents/langstream-agents-flow-control/src/main/resources/META-INF/ai.langstream.agents.index
@@ -0,0 +1,1 @@
+dispatch
\ No newline at end of file
diff --git a/langstream-agents/langstream-agents-flow-control/src/main/resources/META-INF/services/ai.langstream.api.runner.code.AgentCodeProvider b/langstream-agents/langstream-agents-flow-control/src/main/resources/META-INF/services/ai.langstream.api.runner.code.AgentCodeProvider
new file mode 100644
index 000000000..32cef2bdd
--- /dev/null
+++ b/langstream-agents/langstream-agents-flow-control/src/main/resources/META-INF/services/ai.langstream.api.runner.code.AgentCodeProvider
@@ -0,0 +1,1 @@
+ai.langstream.agents.flow.FlowControlAgentsCodeProvider
\ No newline at end of file
diff --git a/langstream-agents/langstream-ai-agents/src/main/java/ai/langstream/ai/agents/GenAIToolKitAgent.java b/langstream-agents/langstream-ai-agents/src/main/java/ai/langstream/ai/agents/GenAIToolKitAgent.java
index 8053584ae..4516f7e7d 100644
--- a/langstream-agents/langstream-ai-agents/src/main/java/ai/langstream/ai/agents/GenAIToolKitAgent.java
+++ b/langstream-agents/langstream-ai-agents/src/main/java/ai/langstream/ai/agents/GenAIToolKitAgent.java
@@ -15,17 +15,17 @@
  */
 package ai.langstream.ai.agents;
 
+import static ai.langstream.ai.agents.commons.TransformContext.recordToTransformContext;
+import static ai.langstream.ai.agents.commons.TransformContext.transformContextToRecord;
+
 import ai.langstream.ai.agents.commons.TransformContext;
-import ai.langstream.ai.agents.commons.TransformSchemaType;
 import ai.langstream.ai.agents.datasource.DataSourceProviderRegistry;
 import ai.langstream.ai.agents.services.ServiceProviderRegistry;
 import ai.langstream.api.runner.code.AbstractAgentCode;
 import ai.langstream.api.runner.code.AgentContext;
 import ai.langstream.api.runner.code.AgentProcessor;
-import ai.langstream.api.runner.code.Header;
 import ai.langstream.api.runner.code.Record;
 import ai.langstream.api.runner.code.RecordSink;
-import ai.langstream.api.runner.code.SimpleRecord;
 import ai.langstream.api.runner.topics.TopicProducer;
 import ai.langstream.api.runtime.ComponentType;
 import com.datastax.oss.streaming.ai.StepPredicatePair;
@@ -37,18 +37,7 @@
 import com.datastax.oss.streaming.ai.streaming.StreamingAnswersConsumer;
 import com.datastax.oss.streaming.ai.streaming.StreamingAnswersConsumerFactory;
 import com.datastax.oss.streaming.ai.util.TransformFunctionUtil;
-import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
-import java.nio.charset.StandardCharsets;
-import java.sql.Time;
-import java.sql.Timestamp;
-import java.time.Instant;
-import java.time.LocalDate;
-import java.time.LocalDateTime;
-import java.time.LocalTime;
-import java.util.ArrayList;
-import java.util.Collection;
-import java.util.Date;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -56,10 +45,8 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
 import java.util.function.Predicate;
-import java.util.stream.Collectors;
 import lombok.SneakyThrows;
 import lombok.extern.slf4j.Slf4j;
-import org.apache.avro.generic.GenericRecord;
 
 @Slf4j
 public class GenAIToolKitAgent extends AbstractAgentCode implements AgentProcessor {
@@ -191,159 +178,6 @@ public void close() throws Exception {
         }
     }
 
-    public static TransformContext recordToTransformContext(
-            Record record, boolean attemptJsonConversion) {
-        TransformContext context = new TransformContext();
-        context.setKeyObject(record.key());
-        context.setKeySchemaType(
-                record.key() == null ? null : getSchemaType(record.key().getClass()));
-        // TODO: temporary hack. We should be able to get the schema from the record
-        if (record.key() instanceof GenericRecord) {
-            context.setKeyNativeSchema(((GenericRecord) record.key()).getSchema());
-        }
-        context.setValueObject(record.value());
-        context.setValueSchemaType(
-                record.value() == null ? null : getSchemaType(record.value().getClass()));
-        // TODO: temporary hack. We should be able to get the schema from the record
-        if (record.value() instanceof GenericRecord) {
-            context.setKeyNativeSchema(((GenericRecord) record.value()).getSchema());
-        }
-        context.setInputTopic(record.origin());
-        context.setEventTime(record.timestamp());
-        if (attemptJsonConversion) {
-            context.setKeyObject(
-                    TransformFunctionUtil.attemptJsonConversion(context.getKeyObject()));
-            context.setValueObject(
-                    TransformFunctionUtil.attemptJsonConversion(context.getValueObject()));
-        }
-        // the headers must be Strings, this is a tentative conversion
-        // in the future we need a better way to handle headers
-        context.setProperties(
-                record.headers().stream()
-                        .filter(h -> h.key() != null && h.value() != null)
-                        .collect(
-                                Collectors.toMap(
-                                        Header::key,
-                                        (h -> {
-                                            if (h.value() == null) {
-                                                return null;
-                                            }
-                                            if (h.value() instanceof byte[]) {
-                                                return new String(
-                                                        (byte[]) h.value(), StandardCharsets.UTF_8);
-                                            } else {
-                                                return h.value().toString();
-                                            }
-                                        }))));
-        return context;
-    }
-
-    public static Optional<Record> transformContextToRecord(TransformContext context) {
-        if (context.isDropCurrentRecord()) {
-            return Optional.empty();
-        }
-        List<Header> headers = new ArrayList<>();
-        context.getProperties()
-                .forEach(
-                        (key, value) -> {
-                            SimpleRecord.SimpleHeader header =
-                                    new SimpleRecord.SimpleHeader(key, value);
-                            headers.add(header);
-                        });
-        return Optional.of(new TransformRecord(context, headers));
-    }
-
-    private record TransformRecord(TransformContext context, Collection<Header> headers)
-            implements Record {
-        private TransformRecord(TransformContext context, Collection<Header> headers) {
-            this.context = context;
-            this.headers = new ArrayList<>(headers);
-        }
-
-        @Override
-        public Object key() {
-            return context.getKeyObject();
-        }
-
-        @Override
-        public Object value() {
-            return context.getValueObject();
-        }
-
-        @Override
-        public String origin() {
-            return context.getInputTopic();
-        }
-
-        @Override
-        public Long timestamp() {
-            return context.getEventTime();
-        }
-    }
-
-    private static TransformSchemaType getSchemaType(Class<?> javaType) {
-        if (String.class.isAssignableFrom(javaType)) {
-            return TransformSchemaType.STRING;
-        }
-        if (Byte.class.isAssignableFrom(javaType)) {
-            return TransformSchemaType.INT8;
-        }
-        if (Short.class.isAssignableFrom(javaType)) {
-            return TransformSchemaType.INT16;
-        }
-        if (Integer.class.isAssignableFrom(javaType)) {
-            return TransformSchemaType.INT32;
-        }
-        if (Long.class.isAssignableFrom(javaType)) {
-            return TransformSchemaType.INT64;
-        }
-        if (Double.class.isAssignableFrom(javaType)) {
-            return TransformSchemaType.DOUBLE;
-        }
-        if (Float.class.isAssignableFrom(javaType)) {
-            return TransformSchemaType.FLOAT;
-        }
-        if (Boolean.class.isAssignableFrom(javaType)) {
-            return TransformSchemaType.BOOLEAN;
-        }
-        if (byte[].class.isAssignableFrom(javaType)) {
-            return TransformSchemaType.BYTES;
-        }
-        // Must be before DATE
-        if (Time.class.isAssignableFrom(javaType)) {
-            return TransformSchemaType.TIME;
-        }
-        // Must be before DATE
-        if (Timestamp.class.isAssignableFrom(javaType)) {
-            return TransformSchemaType.TIMESTAMP;
-        }
-        if (Date.class.isAssignableFrom(javaType)) {
-            return TransformSchemaType.DATE;
-        }
-        if (Instant.class.isAssignableFrom(javaType)) {
-            return TransformSchemaType.INSTANT;
-        }
-        if (LocalDate.class.isAssignableFrom(javaType)) {
-            return TransformSchemaType.LOCAL_DATE;
-        }
-        if (LocalTime.class.isAssignableFrom(javaType)) {
-            return TransformSchemaType.LOCAL_TIME;
-        }
-        if (LocalDateTime.class.isAssignableFrom(javaType)) {
-            return TransformSchemaType.LOCAL_DATE_TIME;
-        }
-        if (GenericRecord.class.isAssignableFrom(javaType)) {
-            return TransformSchemaType.AVRO;
-        }
-        if (JsonNode.class.isAssignableFrom(javaType)) {
-            return TransformSchemaType.JSON;
-        }
-        if (Map.class.isAssignableFrom(javaType)) {
-            return TransformSchemaType.MAP;
-        }
-        throw new IllegalArgumentException("Unsupported data type: " + javaType);
-    }
-
     private static class TopicProducerStreamingAnswersConsumerFactory
             implements StreamingAnswersConsumerFactory {
         private AgentContext agentContext;
@@ -355,8 +189,7 @@ public StreamingAnswersConsumer create(String topicName) {
             TopicProducer topicProducer =
                     agentContext
                             .getTopicConnectionProvider()
-                            .createProducer(
-                                    agentContext.getGlobalAgentId(), Map.of("topic", topicName));
+                            .createProducer(agentContext.getGlobalAgentId(), topicName, Map.of());
             topicProducer.start();
             return new TopicStreamingAnswersConsumer(topicProducer);
         }
diff --git a/langstream-agents/langstream-ai-agents/src/main/java/ai/langstream/ai/agents/rerank/ReRankAgent.java b/langstream-agents/langstream-ai-agents/src/main/java/ai/langstream/ai/agents/rerank/ReRankAgent.java
index cf47d4c88..e14524902 100644
--- a/langstream-agents/langstream-ai-agents/src/main/java/ai/langstream/ai/agents/rerank/ReRankAgent.java
+++ b/langstream-agents/langstream-ai-agents/src/main/java/ai/langstream/ai/agents/rerank/ReRankAgent.java
@@ -15,9 +15,9 @@
  */
 package ai.langstream.ai.agents.rerank;
 
-import static ai.langstream.ai.agents.GenAIToolKitAgent.transformContextToRecord;
+import static ai.langstream.ai.agents.commons.TransformContext.recordToTransformContext;
+import static ai.langstream.ai.agents.commons.TransformContext.transformContextToRecord;
 
-import ai.langstream.ai.agents.GenAIToolKitAgent;
 import ai.langstream.ai.agents.commons.TransformContext;
 import ai.langstream.ai.agents.commons.jstl.JstlEvaluator;
 import ai.langstream.api.runner.code.Record;
@@ -101,8 +101,7 @@ public List<Record> processRecord(Record record) throws Exception {
         if (record == null) {
             return List.of();
         }
-        TransformContext transformContext =
-                GenAIToolKitAgent.recordToTransformContext(record, true).copy();
+        TransformContext transformContext = recordToTransformContext(record, true).copy();
 
         List<Object> currentList = (List<Object>) fieldAccessor.evaluate(transformContext);
 
diff --git a/langstream-agents/langstream-ai-agents/src/main/java/com/datastax/oss/streaming/ai/CastStep.java b/langstream-agents/langstream-ai-agents/src/main/java/com/datastax/oss/streaming/ai/CastStep.java
index 53455b37c..23d97e5b8 100644
--- a/langstream-agents/langstream-ai-agents/src/main/java/com/datastax/oss/streaming/ai/CastStep.java
+++ b/langstream-agents/langstream-ai-agents/src/main/java/com/datastax/oss/streaming/ai/CastStep.java
@@ -15,7 +15,7 @@
  */
 package com.datastax.oss.streaming.ai;
 
-import static com.datastax.oss.streaming.ai.util.TransformFunctionUtil.attemptJsonConversion;
+import static ai.langstream.ai.agents.commons.TransformContext.attemptJsonConversion;
 
 import ai.langstream.ai.agents.commons.TransformContext;
 import ai.langstream.ai.agents.commons.TransformSchemaType;
diff --git a/langstream-agents/langstream-ai-agents/src/main/java/com/datastax/oss/streaming/ai/util/TransformFunctionUtil.java b/langstream-agents/langstream-ai-agents/src/main/java/com/datastax/oss/streaming/ai/util/TransformFunctionUtil.java
index fd3ec3712..dfbc48274 100644
--- a/langstream-agents/langstream-ai-agents/src/main/java/com/datastax/oss/streaming/ai/util/TransformFunctionUtil.java
+++ b/langstream-agents/langstream-ai-agents/src/main/java/com/datastax/oss/streaming/ai/util/TransformFunctionUtil.java
@@ -63,9 +63,7 @@
 import com.datastax.oss.streaming.ai.model.config.UnwrapKeyValueConfig;
 import com.datastax.oss.streaming.ai.services.ServiceProvider;
 import com.datastax.oss.streaming.ai.streaming.StreamingAnswersConsumerFactory;
-import com.fasterxml.jackson.core.type.TypeReference;
 import com.fasterxml.jackson.databind.ObjectMapper;
-import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.nio.charset.Charset;
 import java.nio.charset.StandardCharsets;
@@ -388,23 +386,6 @@ public static void processStep(TransformContext transformContext, StepPredicateP
         }
     }
 
-    public static Object attemptJsonConversion(Object value) {
-        try {
-            if (value instanceof String) {
-                return OBJECT_MAPPER.readValue(
-                        (String) value, new TypeReference<Map<String, Object>>() {});
-            } else if (value instanceof byte[]) {
-                return OBJECT_MAPPER.readValue(
-                        (byte[]) value, new TypeReference<Map<String, Object>>() {});
-            }
-        } catch (IOException e) {
-            if (log.isDebugEnabled()) {
-                log.debug("Cannot convert value to json", e);
-            }
-        }
-        return value;
-    }
-
     public static byte[] getBytes(ByteBuffer byteBuffer) {
         if (byteBuffer == null) {
             return null;
diff --git a/langstream-agents/langstream-vector-agents/src/main/java/ai/langstream/agents/vector/QueryVectorDBAgent.java b/langstream-agents/langstream-vector-agents/src/main/java/ai/langstream/agents/vector/QueryVectorDBAgent.java
index b156af21f..798cdcdf1 100644
--- a/langstream-agents/langstream-vector-agents/src/main/java/ai/langstream/agents/vector/QueryVectorDBAgent.java
+++ b/langstream-agents/langstream-vector-agents/src/main/java/ai/langstream/agents/vector/QueryVectorDBAgent.java
@@ -15,7 +15,9 @@
  */
 package ai.langstream.agents.vector;
 
-import ai.langstream.ai.agents.GenAIToolKitAgent;
+import static ai.langstream.ai.agents.commons.TransformContext.recordToTransformContext;
+import static ai.langstream.ai.agents.commons.TransformContext.transformContextToRecord;
+
 import ai.langstream.ai.agents.commons.TransformContext;
 import ai.langstream.ai.agents.datasource.DataSourceProviderRegistry;
 import ai.langstream.api.runner.code.Record;
@@ -63,10 +65,10 @@ public List<Record> processRecord(Record record) throws Exception {
         if (log.isDebugEnabled()) {
             log.debug("Processing record {}", record);
         }
-        TransformContext context = GenAIToolKitAgent.recordToTransformContext(record, true);
+        TransformContext context = recordToTransformContext(record, true);
         TransformFunctionUtil.processTransformSteps(context, steps);
         context.convertMapToStringOrBytes();
-        Optional<Record> recordResult = GenAIToolKitAgent.transformContextToRecord(context);
+        Optional<Record> recordResult = transformContextToRecord(context);
         if (log.isDebugEnabled()) {
             log.debug("recordResult {}", recordResult);
         }
diff --git a/langstream-agents/langstream-vector-agents/src/main/java/ai/langstream/agents/vector/jdbc/JdbcWriter.java b/langstream-agents/langstream-vector-agents/src/main/java/ai/langstream/agents/vector/jdbc/JdbcWriter.java
index 51813659e..0756b17f1 100644
--- a/langstream-agents/langstream-vector-agents/src/main/java/ai/langstream/agents/vector/jdbc/JdbcWriter.java
+++ b/langstream-agents/langstream-vector-agents/src/main/java/ai/langstream/agents/vector/jdbc/JdbcWriter.java
@@ -15,7 +15,8 @@
  */
 package ai.langstream.agents.vector.jdbc;
 
-import ai.langstream.ai.agents.GenAIToolKitAgent;
+import static ai.langstream.ai.agents.commons.TransformContext.recordToTransformContext;
+
 import ai.langstream.ai.agents.commons.TransformContext;
 import ai.langstream.ai.agents.commons.jstl.JstlEvaluator;
 import ai.langstream.ai.agents.datasource.impl.JdbcDataSourceProvider;
@@ -145,8 +146,7 @@ public synchronized CompletableFuture<?> upsert(
                 Record record, Map<String, Object> context) {
             CompletableFuture<?> handle = new CompletableFuture<>();
             try {
-                TransformContext transformContext =
-                        GenAIToolKitAgent.recordToTransformContext(record, true);
+                TransformContext transformContext = recordToTransformContext(record, true);
 
                 List<Object> primaryKeyValues = prepareValueList(transformContext, primaryKey);
                 List<Object> otherValues = prepareValueList(transformContext, columns);
diff --git a/langstream-agents/langstream-vector-agents/src/main/java/ai/langstream/agents/vector/milvus/MilvusWriter.java b/langstream-agents/langstream-vector-agents/src/main/java/ai/langstream/agents/vector/milvus/MilvusWriter.java
index 9c64521f1..66663c4ec 100644
--- a/langstream-agents/langstream-vector-agents/src/main/java/ai/langstream/agents/vector/milvus/MilvusWriter.java
+++ b/langstream-agents/langstream-vector-agents/src/main/java/ai/langstream/agents/vector/milvus/MilvusWriter.java
@@ -15,7 +15,8 @@
  */
 package ai.langstream.agents.vector.milvus;
 
-import ai.langstream.ai.agents.GenAIToolKitAgent;
+import static ai.langstream.ai.agents.commons.TransformContext.recordToTransformContext;
+
 import ai.langstream.ai.agents.commons.TransformContext;
 import ai.langstream.ai.agents.commons.jstl.JstlEvaluator;
 import ai.langstream.api.database.VectorDatabaseWriter;
@@ -115,8 +116,7 @@ public void initialise(Map<String, Object> agentConfiguration) throws Exception
         public CompletableFuture<?> upsert(Record record, Map<String, Object> context) {
             CompletableFuture<?> handle = new CompletableFuture<>();
             try {
-                TransformContext transformContext =
-                        GenAIToolKitAgent.recordToTransformContext(record, true);
+                TransformContext transformContext = recordToTransformContext(record, true);
 
                 JSONObject row = new JSONObject();
                 fields.forEach(
diff --git a/langstream-agents/langstream-vector-agents/src/main/java/ai/langstream/agents/vector/pinecone/PineconeWriter.java b/langstream-agents/langstream-vector-agents/src/main/java/ai/langstream/agents/vector/pinecone/PineconeWriter.java
index 8b07ac749..8b059cda7 100644
--- a/langstream-agents/langstream-vector-agents/src/main/java/ai/langstream/agents/vector/pinecone/PineconeWriter.java
+++ b/langstream-agents/langstream-vector-agents/src/main/java/ai/langstream/agents/vector/pinecone/PineconeWriter.java
@@ -15,7 +15,8 @@
  */
 package ai.langstream.agents.vector.pinecone;
 
-import ai.langstream.ai.agents.GenAIToolKitAgent;
+import static ai.langstream.ai.agents.commons.TransformContext.recordToTransformContext;
+
 import ai.langstream.ai.agents.commons.TransformContext;
 import ai.langstream.ai.agents.commons.jstl.JstlEvaluator;
 import ai.langstream.api.database.VectorDatabaseWriter;
@@ -102,8 +103,7 @@ public void initialise(Map<String, Object> agentConfiguration) {
         public CompletableFuture<?> upsert(Record record, Map<String, Object> context) {
             CompletableFuture<?> handle = new CompletableFuture<>();
             try {
-                TransformContext transformContext =
-                        GenAIToolKitAgent.recordToTransformContext(record, true);
+                TransformContext transformContext = recordToTransformContext(record, true);
                 String id =
                         idFunction != null ? (String) idFunction.evaluate(transformContext) : null;
                 String namespace =
diff --git a/langstream-agents/pom.xml b/langstream-agents/pom.xml
index 7b57aa917..cb13f7814 100644
--- a/langstream-agents/pom.xml
+++ b/langstream-agents/pom.xml
@@ -39,5 +39,6 @@
         <module>langstream-agents-text-processing</module>
         <module>langstream-ai-agents</module>
         <module>langstream-vector-agents</module>
+        <module>langstream-agents-flow-control</module>
     </modules>
 </project>
diff --git a/langstream-api/src/main/java/ai/langstream/api/runner/topics/TopicConnectionProvider.java b/langstream-api/src/main/java/ai/langstream/api/runner/topics/TopicConnectionProvider.java
index ba0ac70ef..5da260edc 100644
--- a/langstream-api/src/main/java/ai/langstream/api/runner/topics/TopicConnectionProvider.java
+++ b/langstream-api/src/main/java/ai/langstream/api/runner/topics/TopicConnectionProvider.java
@@ -20,5 +20,5 @@
 public interface TopicConnectionProvider {
     TopicConsumer createConsumer(String agentId, Map<String, Object> config);
 
-    TopicProducer createProducer(String agentId, Map<String, Object> config);
+    TopicProducer createProducer(String agentId, String topic, Map<String, Object> config);
 }
diff --git a/langstream-api/src/main/java/ai/langstream/api/util/ConfigurationUtils.java b/langstream-api/src/main/java/ai/langstream/api/util/ConfigurationUtils.java
index 595caac1b..b164cdf4f 100644
--- a/langstream-api/src/main/java/ai/langstream/api/util/ConfigurationUtils.java
+++ b/langstream-api/src/main/java/ai/langstream/api/util/ConfigurationUtils.java
@@ -185,6 +185,11 @@ public static void validateEnumField(
             Set<String> values,
             Supplier<String> definition) {
         Object value = configuration.get(name);
+        validateEnumValue(name, values, value, definition);
+    }
+
+    public static void validateEnumValue(
+            String name, Set<String> values, Object value, Supplier<String> definition) {
         if (value == null || value.toString().isEmpty()) {
             return;
         }
diff --git a/langstream-core/src/main/java/ai/langstream/impl/uti/ClassConfigValidator.java b/langstream-core/src/main/java/ai/langstream/impl/uti/ClassConfigValidator.java
index 144d6265e..eb44d4a53 100644
--- a/langstream-core/src/main/java/ai/langstream/impl/uti/ClassConfigValidator.java
+++ b/langstream-core/src/main/java/ai/langstream/impl/uti/ClassConfigValidator.java
@@ -65,6 +65,11 @@ public class ClassConfigValidator {
     static final Map<String, ResourceConfigurationModel> resourceModels = new ConcurrentHashMap<>();
     static final Map<String, AssetConfigurationModel> assetModels = new ConcurrentHashMap<>();
 
+    public static <T> T convertValidatedConfiguration(
+            Map<String, Object> agentConfiguration, Class<T> clazz) {
+        return validatorMapper.convertValue(agentConfiguration, clazz);
+    }
+
     public static AgentConfigurationModel generateAgentModelFromClass(Class clazz) {
         return agentModels.computeIfAbsent(
                 clazz.getName(), k -> generateModelFromClassNoCache(clazz));
@@ -223,7 +228,7 @@ private static void validateModelFromClass(
                 allowUnknownProperties);
 
         try {
-            validatorMapper.convertValue(asMap, modelClazz);
+            convertValidatedConfiguration(asMap, modelClazz);
         } catch (IllegalArgumentException ex) {
             if (ex.getCause() instanceof MismatchedInputException mismatchedInputException) {
                 final String property =
diff --git a/langstream-k8s-runtime/langstream-k8s-runtime-core/src/main/java/ai/langstream/runtime/impl/k8s/agents/FlowControlAgentsProvider.java b/langstream-k8s-runtime/langstream-k8s-runtime-core/src/main/java/ai/langstream/runtime/impl/k8s/agents/FlowControlAgentsProvider.java
new file mode 100644
index 000000000..d5603f13b
--- /dev/null
+++ b/langstream-k8s-runtime/langstream-k8s-runtime-core/src/main/java/ai/langstream/runtime/impl/k8s/agents/FlowControlAgentsProvider.java
@@ -0,0 +1,140 @@
+/*
+ * Copyright DataStax, Inc.
+ *
+ * Licensed 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 ai.langstream.runtime.impl.k8s.agents;
+
+import ai.langstream.api.doc.AgentConfig;
+import ai.langstream.api.doc.ConfigProperty;
+import ai.langstream.api.model.AgentConfiguration;
+import ai.langstream.api.model.Module;
+import ai.langstream.api.model.Pipeline;
+import ai.langstream.api.runtime.ComponentType;
+import ai.langstream.api.runtime.ComputeClusterRuntime;
+import ai.langstream.api.runtime.ExecutionPlan;
+import ai.langstream.api.runtime.PluginsRegistry;
+import ai.langstream.api.util.ConfigurationUtils;
+import ai.langstream.impl.agents.AbstractComposableAgentProvider;
+import ai.langstream.impl.uti.ClassConfigValidator;
+import ai.langstream.runtime.impl.k8s.KubernetesClusterRuntime;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import lombok.Data;
+import lombok.extern.slf4j.Slf4j;
+
+/** Implements support for Flow Control Processing Agents. */
+@Slf4j
+public class FlowControlAgentsProvider extends AbstractComposableAgentProvider {
+
+    protected static final String DISPATCH = "dispatch";
+    private static final Set<String> SUPPORTED_AGENT_TYPES = Set.of(DISPATCH);
+
+    public FlowControlAgentsProvider() {
+        super(SUPPORTED_AGENT_TYPES, List.of(KubernetesClusterRuntime.CLUSTER_TYPE, "none"));
+    }
+
+    @Override
+    protected ComponentType getComponentType(AgentConfiguration agentConfiguration) {
+        return ComponentType.PROCESSOR;
+    }
+
+    @Override
+    protected Map<String, Object> computeAgentConfiguration(
+            AgentConfiguration agentConfiguration,
+            Module module,
+            Pipeline pipeline,
+            ExecutionPlan executionPlan,
+            ComputeClusterRuntime clusterRuntime,
+            PluginsRegistry pluginsRegistry) {
+        DispatchConfig dispatchConfig =
+                ClassConfigValidator.convertValidatedConfiguration(
+                        agentConfiguration.getConfiguration(), DispatchConfig.class);
+        List<RouteConfiguration> routes = dispatchConfig.getRoutes();
+        if (routes != null) {
+            for (RouteConfiguration routeConfiguration : routes) {
+                String action = routeConfiguration.getAction();
+                ConfigurationUtils.validateEnumValue(
+                        "action",
+                        Set.of("dispatch", "drop"),
+                        action,
+                        () -> "route " + routeConfiguration);
+                String destination = routeConfiguration.getDestination();
+                if (destination != null && !destination.isEmpty()) {
+                    if (action.equals("drop")) {
+                        throw new IllegalArgumentException("drop action cannot have a destination");
+                    }
+                    log.info("Validating topic reference {}", destination);
+                    module.resolveTopic(destination);
+                }
+            }
+        }
+        return super.computeAgentConfiguration(
+                agentConfiguration,
+                module,
+                pipeline,
+                executionPlan,
+                clusterRuntime,
+                pluginsRegistry);
+    }
+
+    @Override
+    protected Class getAgentConfigModelClass(String type) {
+        return switch (type) {
+            case DISPATCH -> DispatchConfig.class;
+            default -> throw new IllegalArgumentException("Unsupported agent type: " + type);
+        };
+    }
+
+    @AgentConfig(
+            name = "Dispatch agent",
+            description =
+                    """
+            Dispatches messages to different destinations based on conditions.
+            """)
+    @Data
+    public static class DispatchConfig {
+        @ConfigProperty(
+                description =
+                        """
+                        Routes.
+                                """)
+        List<RouteConfiguration> routes;
+    }
+
+    @Data
+    public static class RouteConfiguration {
+        @ConfigProperty(
+                description =
+                        """
+                        Condition to activate the route. This is a standard EL expression.
+                                """)
+        String when;
+
+        @ConfigProperty(
+                description =
+                        """
+                        Destination of the message.
+                        """)
+        String destination;
+
+        @ConfigProperty(
+                description =
+                        """
+                        Action on the message. Possible values are "dispatch" or "drop".
+                        """,
+                defaultValue = "dispatch")
+        String action = "dispatch";
+    }
+}
diff --git a/langstream-k8s-runtime/langstream-k8s-runtime-core/src/main/resources/META-INF/services/ai.langstream.api.runtime.AgentNodeProvider b/langstream-k8s-runtime/langstream-k8s-runtime-core/src/main/resources/META-INF/services/ai.langstream.api.runtime.AgentNodeProvider
index 214045904..308f54f86 100644
--- a/langstream-k8s-runtime/langstream-k8s-runtime-core/src/main/resources/META-INF/services/ai.langstream.api.runtime.AgentNodeProvider
+++ b/langstream-k8s-runtime/langstream-k8s-runtime-core/src/main/resources/META-INF/services/ai.langstream.api.runtime.AgentNodeProvider
@@ -7,4 +7,5 @@ ai.langstream.runtime.impl.k8s.agents.KafkaConnectAgentsProvider
 ai.langstream.runtime.impl.k8s.agents.KubernetesCompositeAgentProvider
 ai.langstream.runtime.impl.k8s.agents.IdentityAgentProvider
 ai.langstream.runtime.impl.k8s.agents.QueryVectorDBAgentProvider
-ai.langstream.runtime.impl.k8s.agents.ReRankAgentProvider
\ No newline at end of file
+ai.langstream.runtime.impl.k8s.agents.ReRankAgentProvider
+ai.langstream.runtime.impl.k8s.agents.FlowControlAgentsProvider
\ No newline at end of file
diff --git a/langstream-kafka-runtime/src/main/java/ai/langstream/kafka/runner/kafkaconnect/KafkaConnectSourceAgent.java b/langstream-kafka-runtime/src/main/java/ai/langstream/kafka/runner/kafkaconnect/KafkaConnectSourceAgent.java
index 6136409f8..8b16b2dcb 100644
--- a/langstream-kafka-runtime/src/main/java/ai/langstream/kafka/runner/kafkaconnect/KafkaConnectSourceAgent.java
+++ b/langstream-kafka-runtime/src/main/java/ai/langstream/kafka/runner/kafkaconnect/KafkaConnectSourceAgent.java
@@ -299,9 +299,8 @@ public void raiseError(Exception e) {
         topicProducerToOffsetStore =
                 topicConnectionProvider.createProducer(
                         agentId,
+                        offsetTopic,
                         Map.of(
-                                "topic",
-                                offsetTopic,
                                 "key.serializer",
                                 "org.apache.kafka.common.serialization.ByteArraySerializer",
                                 "value.serializer",
diff --git a/langstream-runtime/langstream-runtime-impl/pom.xml b/langstream-runtime/langstream-runtime-impl/pom.xml
index 775dab90b..374acca30 100644
--- a/langstream-runtime/langstream-runtime-impl/pom.xml
+++ b/langstream-runtime/langstream-runtime-impl/pom.xml
@@ -475,6 +475,16 @@
                     <outputDirectory>${project.build.directory}/agents</outputDirectory>
                 </artifactItem>
 
+                <artifactItem>
+                  <groupId>${project.groupId}</groupId>
+                  <artifactId>langstream-agents-flow-control</artifactId>
+                  <version>${project.version}</version>
+                  <type>nar</type>
+                  <classifier>nar</classifier>
+                  <overWrite>false</overWrite>
+                  <outputDirectory>${project.build.directory}/agents</outputDirectory>
+                </artifactItem>
+
                 <artifactItem>
                   <groupId>${project.groupId}</groupId>
                   <artifactId>langstream-agent-grpc</artifactId>
diff --git a/langstream-runtime/langstream-runtime-impl/src/main/java/ai/langstream/runtime/agent/AgentRunner.java b/langstream-runtime/langstream-runtime-impl/src/main/java/ai/langstream/runtime/agent/AgentRunner.java
index cecc95170..f43321561 100644
--- a/langstream-runtime/langstream-runtime-impl/src/main/java/ai/langstream/runtime/agent/AgentRunner.java
+++ b/langstream-runtime/langstream-runtime-impl/src/main/java/ai/langstream/runtime/agent/AgentRunner.java
@@ -55,6 +55,7 @@
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.NoSuchElementException;
@@ -350,7 +351,17 @@ public TopicConsumer createConsumer(
 
                                         @Override
                                         public TopicProducer createProducer(
-                                                String agentId, Map<String, Object> config) {
+                                                String agentId,
+                                                String topic,
+                                                Map<String, Object> config) {
+                                            if (topic != null && !topic.isEmpty()) {
+                                                if (config == null) {
+                                                    config = Map.of("topic", topic);
+                                                } else {
+                                                    config = new HashMap<>(config);
+                                                    config.put("topic", topic);
+                                                }
+                                            }
                                             return topicConnectionsRuntime.createProducer(
                                                     agentId,
                                                     configuration.streamingCluster(),
@@ -932,9 +943,17 @@ public long getTotalOut() {
     }
 
     private static class NoopTopicProducer implements TopicProducer {
+        private final AtomicLong totalIn = new AtomicLong();
+
+        @Override
+        public CompletableFuture<?> write(Record record) {
+            totalIn.incrementAndGet();
+            return CompletableFuture.completedFuture(null);
+        }
+
         @Override
         public long getTotalIn() {
-            return 0;
+            return totalIn.get();
         }
     }
 


Test patch:
diff --git a/langstream-agents/langstream-agents-flow-control/src/test/resources/logback-test.xml b/langstream-agents/langstream-agents-flow-control/src/test/resources/logback-test.xml
new file mode 100644
index 000000000..fdfe741f8
--- /dev/null
+++ b/langstream-agents/langstream-agents-flow-control/src/test/resources/logback-test.xml
@@ -0,0 +1,34 @@
+<?xml version="1.0" encoding="UTF-8" ?>
+<!--
+
+    Copyright DataStax, Inc.
+
+    Licensed 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.
+
+-->
+
+<!DOCTYPE configuration>
+<configuration>
+    <import class="ch.qos.logback.classic.encoder.PatternLayoutEncoder"/>
+    <import class="ch.qos.logback.core.ConsoleAppender"/>
+
+    <appender name="STDOUT" class="ConsoleAppender">
+        <encoder class="PatternLayoutEncoder">
+            <pattern>%d{HH:mm:ss.SSS} [%thread] %-5level %logger{36} -%kvp- %msg%n</pattern>
+        </encoder>
+    </appender>
+
+    <root level="info">
+        <appender-ref ref="STDOUT"/>
+    </root>
+</configuration>
diff --git a/langstream-agents/langstream-ai-agents/src/test/java/com/datastax/oss/pulsar/functions/transforms/TransformFunction.java b/langstream-agents/langstream-ai-agents/src/test/java/com/datastax/oss/pulsar/functions/transforms/TransformFunction.java
index 30fccf983..ff3f42d03 100644
--- a/langstream-agents/langstream-ai-agents/src/test/java/com/datastax/oss/pulsar/functions/transforms/TransformFunction.java
+++ b/langstream-agents/langstream-ai-agents/src/test/java/com/datastax/oss/pulsar/functions/transforms/TransformFunction.java
@@ -15,6 +15,7 @@
  */
 package com.datastax.oss.pulsar.functions.transforms;
 
+import static ai.langstream.ai.agents.commons.TransformContext.attemptJsonConversion;
 import static com.datastax.oss.streaming.ai.util.TransformFunctionUtil.buildStep;
 
 import ai.langstream.ai.agents.commons.TransformContext;
@@ -315,10 +316,9 @@ public static TransformContext newTransformContext(
             transformContext.setValueObject(value);
         }
         if (attemptJsonConversion) {
-            transformContext.setKeyObject(
-                    TransformFunctionUtil.attemptJsonConversion(transformContext.getKeyObject()));
+            transformContext.setKeyObject(attemptJsonConversion(transformContext.getKeyObject()));
             transformContext.setValueObject(
-                    TransformFunctionUtil.attemptJsonConversion(transformContext.getValueObject()));
+                    attemptJsonConversion(transformContext.getValueObject()));
         }
         return transformContext;
     }
diff --git a/langstream-agents/langstream-ai-agents/src/test/java/com/datastax/oss/streaming/ai/Utils.java b/langstream-agents/langstream-ai-agents/src/test/java/com/datastax/oss/streaming/ai/Utils.java
index c263c1257..1abd01afb 100644
--- a/langstream-agents/langstream-ai-agents/src/test/java/com/datastax/oss/streaming/ai/Utils.java
+++ b/langstream-agents/langstream-ai-agents/src/test/java/com/datastax/oss/streaming/ai/Utils.java
@@ -15,6 +15,7 @@
  */
 package com.datastax.oss.streaming.ai;
 
+import static ai.langstream.ai.agents.commons.TransformContext.attemptJsonConversion;
 import static com.datastax.oss.streaming.ai.FlattenStep.AVRO_READ_OFFSET_PROP;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
@@ -23,7 +24,6 @@
 
 import ai.langstream.ai.agents.commons.TransformContext;
 import ai.langstream.ai.agents.commons.TransformSchemaType;
-import com.datastax.oss.streaming.ai.util.TransformFunctionUtil;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.nio.ByteBuffer;
@@ -428,10 +428,9 @@ public static TransformContext newTransformContext(
             transformContext.setValueObject(value);
         }
         if (attemptJsonConversion) {
-            transformContext.setKeyObject(
-                    TransformFunctionUtil.attemptJsonConversion(transformContext.getKeyObject()));
+            transformContext.setKeyObject(attemptJsonConversion(transformContext.getKeyObject()));
             transformContext.setValueObject(
-                    TransformFunctionUtil.attemptJsonConversion(transformContext.getValueObject()));
+                    attemptJsonConversion(transformContext.getValueObject()));
         }
         return transformContext;
     }
diff --git a/langstream-runtime/langstream-runtime-impl/src/test/java/ai/langstream/kafka/FlowControlRunnerIT.java b/langstream-runtime/langstream-runtime-impl/src/test/java/ai/langstream/kafka/FlowControlRunnerIT.java
new file mode 100644
index 000000000..620adae67
--- /dev/null
+++ b/langstream-runtime/langstream-runtime-impl/src/test/java/ai/langstream/kafka/FlowControlRunnerIT.java
@@ -0,0 +1,238 @@
+/*
+ * Copyright DataStax, Inc.
+ *
+ * Licensed 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 ai.langstream.kafka;
+
+import ai.langstream.AbstractApplicationRunner;
+import java.util.List;
+import java.util.Map;
+import lombok.extern.slf4j.Slf4j;
+import org.apache.kafka.clients.consumer.KafkaConsumer;
+import org.apache.kafka.clients.producer.KafkaProducer;
+import org.apache.kafka.common.header.internals.RecordHeader;
+import org.junit.jupiter.api.Test;
+import org.testcontainers.junit.jupiter.Testcontainers;
+
+@Slf4j
+@Testcontainers
+class FlowControlRunnerIT extends AbstractApplicationRunner {
+
+    @Test
+    public void testSimpleFlowControl() throws Exception {
+        String tenant = "tenant";
+        String[] expectedAgents = {"app-step1"};
+
+        Map<String, String> application =
+                Map.of(
+                        "module.yaml",
+                        """
+                                topics:
+                                  - name: "input-topic"
+                                    creation-mode: create-if-not-exists
+                                  - name: "topic1"
+                                    creation-mode: create-if-not-exists
+                                  - name: "topic2"
+                                    creation-mode: create-if-not-exists
+                                  - name: "default-topic"
+                                    creation-mode: create-if-not-exists
+                                pipeline:
+                                  - name: "Dispatch"
+                                    type: "dispatch"
+                                    input: input-topic
+                                    output: default-topic
+                                    id: step1
+                                    configuration:
+                                      routes:
+                                         - when: properties.language == "en"
+                                           destination: topic1
+                                         - when: properties.language == "fr"
+                                           destination: topic2
+                                         - when: properties.language == "none"
+                                           action: drop
+                                """);
+
+        // query the database with re-rank
+        try (ApplicationRuntime applicationRuntime =
+                deployApplication(
+                        tenant, "app", application, buildInstanceYaml(), expectedAgents)) {
+            try (KafkaProducer<String, String> producer = createProducer();
+                    KafkaConsumer<String, String> consumer = createConsumer("default-topic");
+                    KafkaConsumer<String, String> consumer1 = createConsumer("topic1");
+                    KafkaConsumer<String, String> consumer2 = createConsumer("topic2")) {
+
+                sendMessage(
+                        "input-topic",
+                        "for-default",
+                        List.of(new RecordHeader("language", "it".getBytes())),
+                        producer);
+                sendMessage(
+                        "input-topic",
+                        "for-topic1",
+                        List.of(new RecordHeader("language", "en".getBytes())),
+                        producer);
+                sendMessage(
+                        "input-topic",
+                        "for-topic2",
+                        List.of(new RecordHeader("language", "fr".getBytes())),
+                        producer);
+                executeAgentRunners(applicationRuntime);
+                waitForMessages(consumer, List.of("for-default"));
+                waitForMessages(consumer1, List.of("for-topic1"));
+                waitForMessages(consumer2, List.of("for-topic2"));
+            }
+        }
+    }
+
+    /**
+     * This test validates that the default destination is not mandatory. if you don't set an
+     * "output" the processor behaves like a sink (no output)
+     */
+    @Test
+    public void testSimpleFlowControlNoDefaultDestination() throws Exception {
+        String tenant = "tenant";
+        String[] expectedAgents = {"app-step1"};
+
+        Map<String, String> application =
+                Map.of(
+                        "module.yaml",
+                        """
+                                topics:
+                                  - name: "input-topic-no-default"
+                                    creation-mode: create-if-not-exists
+                                  - name: "topic1-no-default"
+                                    creation-mode: create-if-not-exists
+                                  - name: "topic2-no-default"
+                                    creation-mode: create-if-not-exists
+                                pipeline:
+                                  - name: "Dispatch"
+                                    type: "dispatch"
+                                    input: input-topic-no-default
+                                    id: step1
+                                    configuration:
+                                      routes:
+                                         - when: properties.language == "en"
+                                           destination: topic1-no-default
+                                         - when: properties.language == "fr"
+                                           destination: topic2-no-default
+                                         - when: properties.language == "none"
+                                           action: drop
+                                """);
+
+        // query the database with re-rank
+        try (ApplicationRuntime applicationRuntime =
+                deployApplication(
+                        tenant, "app", application, buildInstanceYaml(), expectedAgents)) {
+            try (KafkaProducer<String, String> producer = createProducer();
+                    KafkaConsumer<String, String> consumer1 = createConsumer("topic1-no-default");
+                    KafkaConsumer<String, String> consumer2 = createConsumer("topic2-no-default")) {
+
+                sendMessage(
+                        "input-topic-no-default",
+                        "for-default",
+                        List.of(new RecordHeader("language", "it".getBytes())),
+                        producer);
+                sendMessage(
+                        "input-topic-no-default",
+                        "for-topic1",
+                        List.of(new RecordHeader("language", "en".getBytes())),
+                        producer);
+                sendMessage(
+                        "input-topic-no-default",
+                        "for-topic2",
+                        List.of(new RecordHeader("language", "fr".getBytes())),
+                        producer);
+                // executeAgentRunners validates that the source has been fully committed
+                executeAgentRunners(applicationRuntime);
+                waitForMessages(consumer1, List.of("for-topic1"));
+                waitForMessages(consumer2, List.of("for-topic2"));
+            }
+        }
+    }
+
+    @Test
+    public void testSimpleFlowControlDefaultToAnotherAgent() throws Exception {
+        String tenant = "tenant";
+        String[] expectedAgents = {"app-step1"};
+
+        Map<String, String> application =
+                Map.of(
+                        "module.yaml",
+                        """
+                                topics:
+                                  - name: "input-topic-to-agent"
+                                    creation-mode: create-if-not-exists
+                                  - name: "topic1-to-agent"
+                                    creation-mode: create-if-not-exists
+                                  - name: "topic2-to-agent"
+                                    creation-mode: create-if-not-exists
+                                  - name: "default-topic-to-agent"
+                                    creation-mode: create-if-not-exists
+                                pipeline:
+                                  - name: "Dispatch"
+                                    type: "dispatch"
+                                    input: input-topic-to-agent
+                                    id: step1
+                                    configuration:
+                                      routes:
+                                         - when: properties.language == "en"
+                                           destination: topic1-to-agent
+                                           action: dispatch
+                                         - when: properties.language == "fr"
+                                           destination: topic2-to-agent
+                                         - when: properties.language == "none"
+                                           action: drop
+                                  - name: "Compute"
+                                    type: "compute"
+                                    output: default-topic-to-agent
+                                    id: step1
+                                    configuration:
+                                      fields:
+                                         - name: "value"
+                                           expression: "'modified'"
+                                """);
+
+        // query the database with re-rank
+        try (ApplicationRuntime applicationRuntime =
+                deployApplication(
+                        tenant, "app", application, buildInstanceYaml(), expectedAgents)) {
+            try (KafkaProducer<String, String> producer = createProducer();
+                    KafkaConsumer<String, String> consumer =
+                            createConsumer("default-topic-to-agent");
+                    KafkaConsumer<String, String> consumer1 = createConsumer("topic1-to-agent");
+                    KafkaConsumer<String, String> consumer2 = createConsumer("topic2-to-agent")) {
+
+                sendMessage(
+                        "input-topic-to-agent",
+                        "for-default",
+                        List.of(new RecordHeader("language", "it".getBytes())),
+                        producer);
+                sendMessage(
+                        "input-topic-to-agent",
+                        "for-topic1",
+                        List.of(new RecordHeader("language", "en".getBytes())),
+                        producer);
+                sendMessage(
+                        "input-topic-to-agent",
+                        "for-topic2",
+                        List.of(new RecordHeader("language", "fr".getBytes())),
+                        producer);
+                executeAgentRunners(applicationRuntime);
+                waitForMessages(consumer, List.of("modified"));
+                waitForMessages(consumer1, List.of("for-topic1"));
+                waitForMessages(consumer2, List.of("for-topic2"));
+            }
+        }
+    }
+}