fnMap = new LinkedHashMap<>();
+ fnMap.put("name", toolUse.name());
+ fnMap.put("arguments", documentToMap(toolUse.input()));
+ callMap.put("function", fnMap);
+ callMap.put("original_id", toolUse.toolUseId());
+ toolCalls.add(callMap);
+ }
+ }
+
+ ChatMessage result = ChatMessage.assistant(textContent.toString());
+ if (!toolCalls.isEmpty()) {
+ result.setToolCalls(toolCalls);
+ } else {
+ // Only strip markdown fences for non-tool-call responses.
+ result = ChatMessage.assistant(stripMarkdownFences(textContent.toString()));
+ }
+ return result;
+ }
+
+ /**
+ * Strip markdown code fences from text responses. Some Bedrock models wrap JSON output in
+ * markdown fences like {@code ```json ... ```}.
+ *
+ * Only strips code fences; does not extract JSON from arbitrary text, as that could corrupt
+ * normal prose responses containing braces.
+ */
+ static String stripMarkdownFences(String text) {
+ if (text == null) return null;
+ String trimmed = text.trim();
+ if (trimmed.startsWith("```")) {
+ int firstNewline = trimmed.indexOf('\n');
+ if (firstNewline >= 0) {
+ trimmed = trimmed.substring(firstNewline + 1);
+ }
+ if (trimmed.endsWith("```")) {
+ trimmed = trimmed.substring(0, trimmed.length() - 3).trim();
+ }
+ return trimmed;
+ }
+ return trimmed;
+ }
+
+ @SuppressWarnings("unchecked")
+ private Document toDocument(Object obj) {
+ if (obj == null) {
+ return Document.fromNull();
+ }
+ if (obj instanceof Map) {
+ Map docMap = new LinkedHashMap<>();
+ ((Map) obj).forEach((k, v) -> docMap.put(k, toDocument(v)));
+ return Document.fromMap(docMap);
+ }
+ if (obj instanceof List) {
+ return Document.fromList(
+ ((List) obj)
+ .stream().map(this::toDocument).collect(Collectors.toList()));
+ }
+ if (obj instanceof String) {
+ return Document.fromString((String) obj);
+ }
+ if (obj instanceof Number) {
+ return Document.fromNumber(SdkNumber.fromBigDecimal(new BigDecimal(obj.toString())));
+ }
+ if (obj instanceof Boolean) {
+ return Document.fromBoolean((Boolean) obj);
+ }
+ return Document.fromString(obj.toString());
+ }
+
+ private Map documentToMap(Document doc) {
+ if (doc == null || !doc.isMap()) {
+ return Collections.emptyMap();
+ }
+ Map result = new LinkedHashMap<>();
+ doc.asMap().forEach((k, v) -> result.put(k, documentToObject(v)));
+ return result;
+ }
+
+ private Object documentToObject(Document doc) {
+ if (doc == null || doc.isNull()) return null;
+ if (doc.isString()) return doc.asString();
+ if (doc.isNumber()) return doc.asNumber().bigDecimalValue();
+ if (doc.isBoolean()) return doc.asBoolean();
+ if (doc.isList()) {
+ return doc.asList().stream().map(this::documentToObject).collect(Collectors.toList());
+ }
+ if (doc.isMap()) return documentToMap(doc);
+ return doc.toString();
+ }
+}
diff --git a/integrations/chat-models/bedrock/src/main/java/org/apache/flink/agents/integrations/chatmodels/bedrock/BedrockChatModelSetup.java b/integrations/chat-models/bedrock/src/main/java/org/apache/flink/agents/integrations/chatmodels/bedrock/BedrockChatModelSetup.java
new file mode 100644
index 00000000..cbcd380b
--- /dev/null
+++ b/integrations/chat-models/bedrock/src/main/java/org/apache/flink/agents/integrations/chatmodels/bedrock/BedrockChatModelSetup.java
@@ -0,0 +1,89 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.agents.integrations.chatmodels.bedrock;
+
+import org.apache.flink.agents.api.chat.model.BaseChatModelSetup;
+import org.apache.flink.agents.api.resource.Resource;
+import org.apache.flink.agents.api.resource.ResourceDescriptor;
+import org.apache.flink.agents.api.resource.ResourceType;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Optional;
+import java.util.function.BiFunction;
+
+/**
+ * Chat model setup for AWS Bedrock Converse API.
+ *
+ * Supported parameters:
+ *
+ *
+ * connection (required): name of the BedrockChatModelConnection resource
+ * model (required): Bedrock model ID (e.g. us.anthropic.claude-sonnet-4-20250514-v1:0)
+ * temperature (optional): sampling temperature (default 0.1)
+ * max_tokens (optional): maximum tokens in the response
+ * prompt (optional): prompt resource name
+ * tools (optional): list of tool resource names
+ *
+ *
+ * Example usage:
+ *
+ *
{@code
+ * @ChatModelSetup
+ * public static ResourceDescriptor bedrockModel() {
+ * return ResourceDescriptor.Builder.newBuilder(BedrockChatModelSetup.class.getName())
+ * .addInitialArgument("connection", "bedrockConnection")
+ * .addInitialArgument("model", "us.anthropic.claude-sonnet-4-20250514-v1:0")
+ * .addInitialArgument("temperature", 0.1)
+ * .addInitialArgument("max_tokens", 4096)
+ * .build();
+ * }
+ * }
+ */
+public class BedrockChatModelSetup extends BaseChatModelSetup {
+
+ private final Double temperature;
+ private final Integer maxTokens;
+
+ public BedrockChatModelSetup(
+ ResourceDescriptor descriptor, BiFunction getResource) {
+ super(descriptor, getResource);
+ this.temperature =
+ Optional.ofNullable(descriptor.getArgument("temperature"))
+ .map(Number::doubleValue)
+ .orElse(0.1);
+ this.maxTokens =
+ Optional.ofNullable(descriptor.getArgument("max_tokens"))
+ .map(Number::intValue)
+ .orElse(null);
+ }
+
+ @Override
+ public Map getParameters() {
+ Map params = new HashMap<>();
+ if (model != null) {
+ params.put("model", model);
+ }
+ params.put("temperature", temperature);
+ if (maxTokens != null) {
+ params.put("max_tokens", maxTokens);
+ }
+ return params;
+ }
+}
diff --git a/integrations/chat-models/bedrock/src/test/java/org/apache/flink/agents/integrations/chatmodels/bedrock/BedrockChatModelConnectionTest.java b/integrations/chat-models/bedrock/src/test/java/org/apache/flink/agents/integrations/chatmodels/bedrock/BedrockChatModelConnectionTest.java
new file mode 100644
index 00000000..84c29481
--- /dev/null
+++ b/integrations/chat-models/bedrock/src/test/java/org/apache/flink/agents/integrations/chatmodels/bedrock/BedrockChatModelConnectionTest.java
@@ -0,0 +1,124 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.agents.integrations.chatmodels.bedrock;
+
+import org.apache.flink.agents.api.chat.messages.ChatMessage;
+import org.apache.flink.agents.api.chat.messages.MessageRole;
+import org.apache.flink.agents.api.chat.model.BaseChatModelConnection;
+import org.apache.flink.agents.api.resource.Resource;
+import org.apache.flink.agents.api.resource.ResourceDescriptor;
+import org.apache.flink.agents.api.resource.ResourceType;
+import org.junit.jupiter.api.DisplayName;
+import org.junit.jupiter.api.Test;
+
+import java.util.*;
+import java.util.function.BiFunction;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+
+/** Tests for {@link BedrockChatModelConnection}. */
+class BedrockChatModelConnectionTest {
+
+ private static final BiFunction NOOP = (a, b) -> null;
+
+ private static ResourceDescriptor descriptor(String region, String model) {
+ ResourceDescriptor.Builder b =
+ ResourceDescriptor.Builder.newBuilder(BedrockChatModelConnection.class.getName());
+ if (region != null) b.addInitialArgument("region", region);
+ if (model != null) b.addInitialArgument("model", model);
+ return b.build();
+ }
+
+ @Test
+ @DisplayName("Constructor creates client with default region")
+ void testConstructorDefaultRegion() {
+ BedrockChatModelConnection conn =
+ new BedrockChatModelConnection(
+ descriptor(null, "us.anthropic.claude-sonnet-4-20250514-v1:0"), NOOP);
+ assertNotNull(conn);
+ }
+
+ @Test
+ @DisplayName("Constructor creates client with explicit region")
+ void testConstructorExplicitRegion() {
+ BedrockChatModelConnection conn =
+ new BedrockChatModelConnection(
+ descriptor("us-west-2", "us.anthropic.claude-sonnet-4-20250514-v1:0"),
+ NOOP);
+ assertNotNull(conn);
+ }
+
+ @Test
+ @DisplayName("Extends BaseChatModelConnection")
+ void testInheritance() {
+ BedrockChatModelConnection conn =
+ new BedrockChatModelConnection(descriptor("us-east-1", "test-model"), NOOP);
+ assertThat(conn).isInstanceOf(BaseChatModelConnection.class);
+ }
+
+ @Test
+ @DisplayName("Chat throws when no model specified")
+ void testChatThrowsWithoutModel() {
+ BedrockChatModelConnection conn =
+ new BedrockChatModelConnection(descriptor("us-east-1", null), NOOP);
+ List msgs = List.of(new ChatMessage(MessageRole.USER, "hello"));
+ assertThatThrownBy(() -> conn.chat(msgs, null, Collections.emptyMap()))
+ .isInstanceOf(RuntimeException.class);
+ }
+
+ @Test
+ @DisplayName("stripMarkdownFences: normal text with braces is not modified")
+ void testStripMarkdownFencesPreservesTextWithBraces() {
+ assertThat(
+ BedrockChatModelConnection.stripMarkdownFences(
+ "Use the format {key: value} for config"))
+ .isEqualTo("Use the format {key: value} for config");
+ }
+
+ @Test
+ @DisplayName("stripMarkdownFences: clean JSON passes through")
+ void testStripMarkdownFencesCleanJson() {
+ assertThat(
+ BedrockChatModelConnection.stripMarkdownFences(
+ "{\"score\": 5, \"reasons\": []}"))
+ .isEqualTo("{\"score\": 5, \"reasons\": []}");
+ }
+
+ @Test
+ @DisplayName("stripMarkdownFences: strips ```json fences")
+ void testStripMarkdownFencesJsonBlock() {
+ assertThat(BedrockChatModelConnection.stripMarkdownFences("```json\n{\"score\": 5}\n```"))
+ .isEqualTo("{\"score\": 5}");
+ }
+
+ @Test
+ @DisplayName("stripMarkdownFences: strips plain ``` fences")
+ void testStripMarkdownFencesPlainBlock() {
+ assertThat(BedrockChatModelConnection.stripMarkdownFences("```\n{\"id\": \"P001\"}\n```"))
+ .isEqualTo("{\"id\": \"P001\"}");
+ }
+
+ @Test
+ @DisplayName("stripMarkdownFences: null returns null")
+ void testStripMarkdownFencesNull() {
+ assertThat(BedrockChatModelConnection.stripMarkdownFences(null)).isNull();
+ }
+}
diff --git a/integrations/chat-models/bedrock/src/test/java/org/apache/flink/agents/integrations/chatmodels/bedrock/BedrockChatModelSetupTest.java b/integrations/chat-models/bedrock/src/test/java/org/apache/flink/agents/integrations/chatmodels/bedrock/BedrockChatModelSetupTest.java
new file mode 100644
index 00000000..05094f02
--- /dev/null
+++ b/integrations/chat-models/bedrock/src/test/java/org/apache/flink/agents/integrations/chatmodels/bedrock/BedrockChatModelSetupTest.java
@@ -0,0 +1,77 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.agents.integrations.chatmodels.bedrock;
+
+import org.apache.flink.agents.api.chat.model.BaseChatModelSetup;
+import org.apache.flink.agents.api.resource.Resource;
+import org.apache.flink.agents.api.resource.ResourceDescriptor;
+import org.apache.flink.agents.api.resource.ResourceType;
+import org.junit.jupiter.api.DisplayName;
+import org.junit.jupiter.api.Test;
+
+import java.util.Map;
+import java.util.function.BiFunction;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/** Tests for {@link BedrockChatModelSetup}. */
+class BedrockChatModelSetupTest {
+
+ private static final BiFunction NOOP = (a, b) -> null;
+
+ @Test
+ @DisplayName("getParameters includes model and default temperature")
+ void testGetParametersDefaults() {
+ ResourceDescriptor desc =
+ ResourceDescriptor.Builder.newBuilder(BedrockChatModelSetup.class.getName())
+ .addInitialArgument("connection", "conn")
+ .addInitialArgument("model", "us.anthropic.claude-sonnet-4-20250514-v1:0")
+ .build();
+ BedrockChatModelSetup setup = new BedrockChatModelSetup(desc, NOOP);
+
+ Map params = setup.getParameters();
+ assertThat(params).containsEntry("model", "us.anthropic.claude-sonnet-4-20250514-v1:0");
+ assertThat(params).containsEntry("temperature", 0.1);
+ }
+
+ @Test
+ @DisplayName("getParameters uses custom temperature")
+ void testGetParametersCustomTemperature() {
+ ResourceDescriptor desc =
+ ResourceDescriptor.Builder.newBuilder(BedrockChatModelSetup.class.getName())
+ .addInitialArgument("connection", "conn")
+ .addInitialArgument("model", "test-model")
+ .addInitialArgument("temperature", 0.7)
+ .build();
+ BedrockChatModelSetup setup = new BedrockChatModelSetup(desc, NOOP);
+
+ assertThat(setup.getParameters()).containsEntry("temperature", 0.7);
+ }
+
+ @Test
+ @DisplayName("Extends BaseChatModelSetup")
+ void testInheritance() {
+ ResourceDescriptor desc =
+ ResourceDescriptor.Builder.newBuilder(BedrockChatModelSetup.class.getName())
+ .addInitialArgument("connection", "conn")
+ .addInitialArgument("model", "m")
+ .build();
+ assertThat(new BedrockChatModelSetup(desc, NOOP)).isInstanceOf(BaseChatModelSetup.class);
+ }
+}
diff --git a/integrations/chat-models/pom.xml b/integrations/chat-models/pom.xml
index 20b1b425..e5f4b9d4 100644
--- a/integrations/chat-models/pom.xml
+++ b/integrations/chat-models/pom.xml
@@ -33,6 +33,7 @@ under the License.
anthropic
azureai
+ bedrock
ollama
openai
diff --git a/integrations/embedding-models/bedrock/pom.xml b/integrations/embedding-models/bedrock/pom.xml
new file mode 100644
index 00000000..353c32c8
--- /dev/null
+++ b/integrations/embedding-models/bedrock/pom.xml
@@ -0,0 +1,48 @@
+
+
+
+ 4.0.0
+
+
+ org.apache.flink
+ flink-agents-integrations-embedding-models
+ 0.3-SNAPSHOT
+ ../pom.xml
+
+
+ flink-agents-integrations-embedding-models-bedrock
+ Flink Agents : Integrations: Embedding Models: Bedrock
+ jar
+
+
+
+ org.apache.flink
+ flink-agents-api
+ ${project.version}
+
+
+
+ software.amazon.awssdk
+ bedrockruntime
+ ${aws.sdk.version}
+
+
+
+
diff --git a/integrations/embedding-models/bedrock/src/main/java/org/apache/flink/agents/integrations/embeddingmodels/bedrock/BedrockEmbeddingModelConnection.java b/integrations/embedding-models/bedrock/src/main/java/org/apache/flink/agents/integrations/embeddingmodels/bedrock/BedrockEmbeddingModelConnection.java
new file mode 100644
index 00000000..2366d123
--- /dev/null
+++ b/integrations/embedding-models/bedrock/src/main/java/org/apache/flink/agents/integrations/embeddingmodels/bedrock/BedrockEmbeddingModelConnection.java
@@ -0,0 +1,188 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.agents.integrations.embeddingmodels.bedrock;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.fasterxml.jackson.databind.node.ObjectNode;
+import org.apache.flink.agents.api.embedding.model.BaseEmbeddingModelConnection;
+import org.apache.flink.agents.api.resource.Resource;
+import org.apache.flink.agents.api.resource.ResourceDescriptor;
+import org.apache.flink.agents.api.resource.ResourceType;
+import software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider;
+import software.amazon.awssdk.core.SdkBytes;
+import software.amazon.awssdk.regions.Region;
+import software.amazon.awssdk.services.bedrockruntime.BedrockRuntimeClient;
+import software.amazon.awssdk.services.bedrockruntime.model.InvokeModelRequest;
+import software.amazon.awssdk.services.bedrockruntime.model.InvokeModelResponse;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.function.BiFunction;
+
+/**
+ * Bedrock embedding model connection using Amazon Titan Text Embeddings V2.
+ *
+ * Uses the InvokeModel API to generate embeddings. Supports configurable dimensions (256, 512,
+ * or 1024) and normalization. Since Titan V2 processes one text per API call, batch embedding is
+ * parallelized via a configurable thread pool.
+ *
+ *
Supported connection parameters:
+ *
+ *
+ * region (optional): AWS region, defaults to us-east-1
+ * model (optional): default model ID, defaults to amazon.titan-embed-text-v2:0
+ * embed_concurrency (optional): thread pool size for parallel embedding (default: 4)
+ *
+ *
+ * Example usage:
+ *
+ *
{@code
+ * @EmbeddingModelConnection
+ * public static ResourceDescriptor bedrockEmbedding() {
+ * return ResourceDescriptor.Builder.newBuilder(BedrockEmbeddingModelConnection.class.getName())
+ * .addInitialArgument("region", "us-east-1")
+ * .addInitialArgument("model", "amazon.titan-embed-text-v2:0")
+ * .addInitialArgument("embed_concurrency", 8)
+ * .build();
+ * }
+ * }
+ */
+public class BedrockEmbeddingModelConnection extends BaseEmbeddingModelConnection {
+
+ private static final ObjectMapper MAPPER = new ObjectMapper();
+ private static final String DEFAULT_MODEL = "amazon.titan-embed-text-v2:0";
+ private static final int MAX_RETRIES = 5;
+
+ private final BedrockRuntimeClient client;
+ private final String defaultModel;
+ private final ExecutorService embedPool;
+
+ public BedrockEmbeddingModelConnection(
+ ResourceDescriptor descriptor, BiFunction getResource) {
+ super(descriptor, getResource);
+
+ String region = descriptor.getArgument("region");
+ if (region == null || region.isBlank()) {
+ region = "us-east-1";
+ }
+
+ this.client =
+ BedrockRuntimeClient.builder()
+ .region(Region.of(region))
+ .credentialsProvider(DefaultCredentialsProvider.create())
+ .build();
+
+ String model = descriptor.getArgument("model");
+ this.defaultModel = (model != null && !model.isBlank()) ? model : DEFAULT_MODEL;
+
+ Integer concurrency = descriptor.getArgument("embed_concurrency");
+ int threads = concurrency != null ? concurrency : 4;
+ this.embedPool = Executors.newFixedThreadPool(threads);
+ }
+
+ @Override
+ public float[] embed(String text, Map parameters) {
+ String model = (String) parameters.getOrDefault("model", defaultModel);
+ Integer dimensions = (Integer) parameters.get("dimensions");
+
+ ObjectNode body = MAPPER.createObjectNode();
+ body.put("inputText", text);
+ if (dimensions != null) {
+ body.put("dimensions", dimensions);
+ }
+ body.put("normalize", true);
+
+ for (int attempt = 0; ; attempt++) {
+ try {
+ InvokeModelResponse response =
+ client.invokeModel(
+ InvokeModelRequest.builder()
+ .modelId(model)
+ .contentType("application/json")
+ .body(SdkBytes.fromUtf8String(body.toString()))
+ .build());
+
+ JsonNode result = MAPPER.readTree(response.body().asUtf8String());
+ JsonNode embeddingNode = result.get("embedding");
+ float[] embedding = new float[embeddingNode.size()];
+ for (int i = 0; i < embeddingNode.size(); i++) {
+ embedding[i] = (float) embeddingNode.get(i).asDouble();
+ }
+ return embedding;
+ } catch (Exception e) {
+ if (attempt < MAX_RETRIES && isRetryable(e)) {
+ try {
+ long delay =
+ (long) (Math.pow(2, attempt) * 200 * (0.5 + Math.random() * 0.5));
+ Thread.sleep(delay);
+ } catch (InterruptedException ie) {
+ Thread.currentThread().interrupt();
+ }
+ } else {
+ throw new RuntimeException("Failed to generate Bedrock embedding.", e);
+ }
+ }
+ }
+ }
+
+ private static boolean isRetryable(Exception e) {
+ String msg = e.toString();
+ return msg.contains("ThrottlingException")
+ || msg.contains("ModelErrorException")
+ || msg.contains("429")
+ || msg.contains("424")
+ || msg.contains("503");
+ }
+
+ @Override
+ public List embed(List texts, Map parameters) {
+ if (texts.size() <= 1) {
+ List results = new ArrayList<>(texts.size());
+ for (String text : texts) {
+ results.add(embed(text, parameters));
+ }
+ return results;
+ }
+ @SuppressWarnings("unchecked")
+ CompletableFuture[] futures =
+ texts.stream()
+ .map(
+ text ->
+ CompletableFuture.supplyAsync(
+ () -> embed(text, parameters), embedPool))
+ .toArray(CompletableFuture[]::new);
+ CompletableFuture.allOf(futures).join();
+ List results = new ArrayList<>(texts.size());
+ for (CompletableFuture f : futures) {
+ results.add(f.join());
+ }
+ return results;
+ }
+
+ @Override
+ public void close() throws Exception {
+ this.embedPool.shutdown();
+ this.client.close();
+ }
+}
diff --git a/integrations/embedding-models/bedrock/src/main/java/org/apache/flink/agents/integrations/embeddingmodels/bedrock/BedrockEmbeddingModelSetup.java b/integrations/embedding-models/bedrock/src/main/java/org/apache/flink/agents/integrations/embeddingmodels/bedrock/BedrockEmbeddingModelSetup.java
new file mode 100644
index 00000000..90dc1934
--- /dev/null
+++ b/integrations/embedding-models/bedrock/src/main/java/org/apache/flink/agents/integrations/embeddingmodels/bedrock/BedrockEmbeddingModelSetup.java
@@ -0,0 +1,80 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.agents.integrations.embeddingmodels.bedrock;
+
+import org.apache.flink.agents.api.embedding.model.BaseEmbeddingModelSetup;
+import org.apache.flink.agents.api.resource.Resource;
+import org.apache.flink.agents.api.resource.ResourceDescriptor;
+import org.apache.flink.agents.api.resource.ResourceType;
+
+import java.util.HashMap;
+import java.util.Map;
+import java.util.function.BiFunction;
+
+/**
+ * Embedding model setup for Bedrock Titan Text Embeddings.
+ *
+ * Supported parameters:
+ *
+ *
+ * connection (required): name of the BedrockEmbeddingModelConnection resource
+ * model (optional): model ID (default: amazon.titan-embed-text-v2:0)
+ * dimensions (optional): embedding dimensions (256, 512, or 1024)
+ *
+ *
+ * Example usage:
+ *
+ *
{@code
+ * @EmbeddingModelSetup
+ * public static ResourceDescriptor bedrockEmbeddingSetup() {
+ * return ResourceDescriptor.Builder.newBuilder(BedrockEmbeddingModelSetup.class.getName())
+ * .addInitialArgument("connection", "bedrockEmbedding")
+ * .addInitialArgument("model", "amazon.titan-embed-text-v2:0")
+ * .addInitialArgument("dimensions", 1024)
+ * .build();
+ * }
+ * }
+ */
+public class BedrockEmbeddingModelSetup extends BaseEmbeddingModelSetup {
+
+ private final Integer dimensions;
+
+ public BedrockEmbeddingModelSetup(
+ ResourceDescriptor descriptor, BiFunction getResource) {
+ super(descriptor, getResource);
+ this.dimensions = descriptor.getArgument("dimensions");
+ }
+
+ @Override
+ public Map getParameters() {
+ Map params = new HashMap<>();
+ if (model != null) {
+ params.put("model", model);
+ }
+ if (dimensions != null) {
+ params.put("dimensions", dimensions);
+ }
+ return params;
+ }
+
+ @Override
+ public BedrockEmbeddingModelConnection getConnection() {
+ return (BedrockEmbeddingModelConnection) super.getConnection();
+ }
+}
diff --git a/integrations/embedding-models/bedrock/src/test/java/org/apache/flink/agents/integrations/embeddingmodels/bedrock/BedrockEmbeddingModelTest.java b/integrations/embedding-models/bedrock/src/test/java/org/apache/flink/agents/integrations/embeddingmodels/bedrock/BedrockEmbeddingModelTest.java
new file mode 100644
index 00000000..3d2d3d07
--- /dev/null
+++ b/integrations/embedding-models/bedrock/src/test/java/org/apache/flink/agents/integrations/embeddingmodels/bedrock/BedrockEmbeddingModelTest.java
@@ -0,0 +1,99 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.agents.integrations.embeddingmodels.bedrock;
+
+import org.apache.flink.agents.api.embedding.model.BaseEmbeddingModelConnection;
+import org.apache.flink.agents.api.embedding.model.BaseEmbeddingModelSetup;
+import org.apache.flink.agents.api.resource.Resource;
+import org.apache.flink.agents.api.resource.ResourceDescriptor;
+import org.apache.flink.agents.api.resource.ResourceType;
+import org.junit.jupiter.api.DisplayName;
+import org.junit.jupiter.api.Test;
+
+import java.util.Map;
+import java.util.function.BiFunction;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+
+/** Tests for {@link BedrockEmbeddingModelConnection} and {@link BedrockEmbeddingModelSetup}. */
+class BedrockEmbeddingModelTest {
+
+ private static final BiFunction NOOP = (a, b) -> null;
+
+ private static ResourceDescriptor connDescriptor(String region) {
+ ResourceDescriptor.Builder b =
+ ResourceDescriptor.Builder.newBuilder(
+ BedrockEmbeddingModelConnection.class.getName());
+ if (region != null) b.addInitialArgument("region", region);
+ return b.build();
+ }
+
+ @Test
+ @DisplayName("Connection constructor creates client with defaults")
+ void testConnectionDefaults() {
+ BedrockEmbeddingModelConnection conn =
+ new BedrockEmbeddingModelConnection(connDescriptor(null), NOOP);
+ assertNotNull(conn);
+ assertThat(conn).isInstanceOf(BaseEmbeddingModelConnection.class);
+ }
+
+ @Test
+ @DisplayName("Connection constructor with explicit region and concurrency")
+ void testConnectionExplicitParams() {
+ ResourceDescriptor desc =
+ ResourceDescriptor.Builder.newBuilder(
+ BedrockEmbeddingModelConnection.class.getName())
+ .addInitialArgument("region", "eu-west-1")
+ .addInitialArgument("embed_concurrency", 8)
+ .build();
+ BedrockEmbeddingModelConnection conn = new BedrockEmbeddingModelConnection(desc, NOOP);
+ assertNotNull(conn);
+ }
+
+ @Test
+ @DisplayName("Setup getParameters includes model and dimensions")
+ void testSetupParameters() {
+ ResourceDescriptor desc =
+ ResourceDescriptor.Builder.newBuilder(BedrockEmbeddingModelSetup.class.getName())
+ .addInitialArgument("connection", "conn")
+ .addInitialArgument("model", "amazon.titan-embed-text-v2:0")
+ .addInitialArgument("dimensions", 1024)
+ .build();
+ BedrockEmbeddingModelSetup setup = new BedrockEmbeddingModelSetup(desc, NOOP);
+
+ Map params = setup.getParameters();
+ assertThat(params).containsEntry("model", "amazon.titan-embed-text-v2:0");
+ assertThat(params).containsEntry("dimensions", 1024);
+ assertThat(setup).isInstanceOf(BaseEmbeddingModelSetup.class);
+ }
+
+ @Test
+ @DisplayName("Setup getParameters omits null dimensions")
+ void testSetupParametersNoDimensions() {
+ ResourceDescriptor desc =
+ ResourceDescriptor.Builder.newBuilder(BedrockEmbeddingModelSetup.class.getName())
+ .addInitialArgument("connection", "conn")
+ .addInitialArgument("model", "amazon.titan-embed-text-v2:0")
+ .build();
+ BedrockEmbeddingModelSetup setup = new BedrockEmbeddingModelSetup(desc, NOOP);
+
+ assertThat(setup.getParameters()).doesNotContainKey("dimensions");
+ }
+}
diff --git a/integrations/embedding-models/pom.xml b/integrations/embedding-models/pom.xml
index f1bc6c08..1845a480 100644
--- a/integrations/embedding-models/pom.xml
+++ b/integrations/embedding-models/pom.xml
@@ -31,6 +31,7 @@ under the License.
pom
+ bedrock
ollama
diff --git a/integrations/pom.xml b/integrations/pom.xml
index 0e5df222..9989a5f0 100644
--- a/integrations/pom.xml
+++ b/integrations/pom.xml
@@ -35,6 +35,7 @@ under the License.
8.19.0
4.8.0
2.11.1
+ 2.32.16