Skip to content

Commit 60d996e

Browse files
tshepomareditshepomaredi
committed
feat(aws): add Bedrock Knowledge Base connector
Adds a new outbound connector for AWS Bedrock Knowledge Base Retrieve API. Enables semantic search over documents indexed in a Bedrock Knowledge Base. - BedrockAgentRuntimeClient-based implementation - Retrieve operation with query and numberOfResults parameters - Results returned as Camunda Document (JSON) - Error handling: throttling retry, KB errors, serialization errors - Element template with authentication, configuration, and operation groups - Unit tests: executor (2), input validation (3) — all passing
1 parent b360ed6 commit 60d996e

17 files changed

+1196
-0
lines changed

‎connectors/aws/aws-bedrock-knowledgebase/element-templates/aws-bedrock-knowledgebase-outbound-connector.json‎

Lines changed: 285 additions & 0 deletions
Large diffs are not rendered by default.

‎connectors/aws/aws-bedrock-knowledgebase/element-templates/hybrid/aws-bedrock-knowledgebase-outbound-connector-hybrid.json‎

Lines changed: 290 additions & 0 deletions
Large diffs are not rendered by default.

‎connectors/aws/aws-bedrock-knowledgebase/pom.xml‎

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,100 @@
1+
<?xml version="1.0" encoding="UTF-8"?>
2+
<project xmlns="http://maven.apache.org/POM/4.0.0"
3+
xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
4+
xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
5+
<modelVersion>4.0.0</modelVersion>
6+
7+
<parent>
8+
<groupId>io.camunda.connector</groupId>
9+
<artifactId>connector-aws-parent</artifactId>
10+
<version>8.9.0-SNAPSHOT</version>
11+
</parent>
12+
13+
<artifactId>connector-aws-bedrock-knowledgebase</artifactId>
14+
<name>connector-aws-bedrock-knowledgebase</name>
15+
<description>Camunda Connector for AWS Bedrock Knowledge Base retrieval</description>
16+
<packaging>jar</packaging>
17+
18+
<licenses>
19+
<license>
20+
<name>Camunda Self-Managed Free Edition license</name>
21+
<url>
22+
https://camunda.com/legal/terms/cloud-terms-and-conditions/camunda-cloud-self-managed-free-edition-terms/
23+
</url>
24+
</license>
25+
<license>
26+
<name>Camunda Self-Managed Enterprise Edition license</name>
27+
</license>
28+
</licenses>
29+
30+
<dependencies>
31+
<dependency>
32+
<groupId>io.camunda.connector</groupId>
33+
<artifactId>connector-aws-base</artifactId>
34+
<version>${project.version}</version>
35+
</dependency>
36+
37+
<dependency>
38+
<groupId>software.amazon.awssdk</groupId>
39+
<artifactId>bedrockagentruntime</artifactId>
40+
<version>${version.aws-sdk2}</version>
41+
</dependency>
42+
<dependency>
43+
<groupId>software.amazon.awssdk</groupId>
44+
<artifactId>auth</artifactId>
45+
</dependency>
46+
<dependency>
47+
<groupId>software.amazon.awssdk</groupId>
48+
<artifactId>regions</artifactId>
49+
</dependency>
50+
<dependency>
51+
<groupId>software.amazon.awssdk</groupId>
52+
<artifactId>sdk-core</artifactId>
53+
</dependency>
54+
<dependency>
55+
<groupId>software.amazon.awssdk</groupId>
56+
<artifactId>aws-core</artifactId>
57+
</dependency>
58+
59+
<dependency>
60+
<groupId>io.camunda.connector</groupId>
61+
<artifactId>connector-object-mapper</artifactId>
62+
</dependency>
63+
64+
<dependency>
65+
<groupId>io.camunda.connector</groupId>
66+
<artifactId>element-template-generator-annotations</artifactId>
67+
</dependency>
68+
</dependencies>
69+
70+
<build>
71+
<plugins>
72+
<plugin>
73+
<groupId>io.camunda.connector</groupId>
74+
<artifactId>element-template-generator-maven-plugin</artifactId>
75+
<version>${project.version}</version>
76+
<configuration>
77+
<connectors>
78+
<connector>
79+
<connectorClass>
80+
io.camunda.connector.aws.bedrock.knowledgebase.BedrockKnowledgeBaseConnectorFunction
81+
</connectorClass>
82+
<files>
83+
<file>
84+
<templateId>io.camunda.connectors.aws.bedrock.knowledgebase.v1</templateId>
85+
<templateFileName>aws-bedrock-knowledgebase-outbound-connector.json</templateFileName>
86+
</file>
87+
</files>
88+
<generateHybridTemplates>true</generateHybridTemplates>
89+
</connector>
90+
</connectors>
91+
<versionHistoryEnabled>true</versionHistoryEnabled>
92+
<includeDependencies>
93+
<includeDependency>io.camunda.connector:connector-aws-base</includeDependency>
94+
</includeDependencies>
95+
</configuration>
96+
</plugin>
97+
</plugins>
98+
</build>
99+
100+
</project>

‎connectors/aws/aws-bedrock-knowledgebase/src/main/java/io/camunda/connector/aws/bedrock/knowledgebase/BedrockKnowledgeBaseConnectorFunction.java‎

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
/*
2+
* Copyright Camunda Services GmbH and/or licensed to Camunda Services GmbH
3+
* under one or more contributor license agreements. Licensed under a proprietary license.
4+
* See the License.txt file for more information. You may not use this file
5+
* except in compliance with the proprietary license.
6+
*/
7+
package io.camunda.connector.aws.bedrock.knowledgebase;
8+
9+
import io.camunda.connector.api.annotation.OutboundConnector;
10+
import io.camunda.connector.api.outbound.OutboundConnectorContext;
11+
import io.camunda.connector.api.outbound.OutboundConnectorFunction;
12+
import io.camunda.connector.aws.CredentialsProviderSupportV2;
13+
import io.camunda.connector.aws.bedrock.knowledgebase.model.request.BedrockKnowledgeBaseRequest;
14+
import io.camunda.connector.generator.java.annotation.ElementTemplate;
15+
import io.camunda.connector.generator.java.annotation.ElementTemplate.PropertyGroup;
16+
import io.camunda.connector.jackson.ConnectorsObjectMapperSupplier;
17+
import java.net.URI;
18+
import software.amazon.awssdk.regions.Region;
19+
import software.amazon.awssdk.services.bedrockagentruntime.BedrockAgentRuntimeClient;
20+
21+
@OutboundConnector(
22+
name = "AWS Bedrock Knowledge Base",
23+
inputVariables = {
24+
"authentication",
25+
"configuration",
26+
"knowledgeBaseId",
27+
"operation",
28+
"operationDiscriminator"
29+
},
30+
type = "io.camunda:aws-bedrock-knowledgebase:1")
31+
@ElementTemplate(
32+
engineVersion = "^8.7",
33+
id = "io.camunda.connectors.aws.bedrock.knowledgebase.v1",
34+
name = "AWS Bedrock Knowledge Base Outbound Connector",
35+
description = "Retrieve relevant documents from an AWS Bedrock Knowledge Base",
36+
inputDataClass = BedrockKnowledgeBaseRequest.class,
37+
version = 1,
38+
propertyGroups = {
39+
@PropertyGroup(id = "authentication", label = "Authentication"),
40+
@PropertyGroup(id = "configuration", label = "Configuration"),
41+
@PropertyGroup(id = "operation", label = "Operation"),
42+
@PropertyGroup(id = "retrieve", label = "Retrieve from Knowledge Base"),
43+
},
44+
icon = "icon.png")
45+
public class BedrockKnowledgeBaseConnectorFunction implements OutboundConnectorFunction {
46+
47+
@Override
48+
public Object execute(OutboundConnectorContext context) {
49+
var request = context.bindVariables(BedrockKnowledgeBaseRequest.class);
50+
try (var client = buildClient(request)) {
51+
return new BedrockKnowledgeBaseExecutor(client, ConnectorsObjectMapperSupplier.getCopy())
52+
.execute(request, context::create);
53+
}
54+
}
55+
56+
private BedrockAgentRuntimeClient buildClient(BedrockKnowledgeBaseRequest request) {
57+
var builder =
58+
BedrockAgentRuntimeClient.builder()
59+
.credentialsProvider(CredentialsProviderSupportV2.credentialsProvider(request));
60+
var config = request.getConfiguration();
61+
if (config != null && config.region() != null) {
62+
builder.region(Region.of(config.region()));
63+
}
64+
if (config != null && config.endpoint() != null && !config.endpoint().isBlank()) {
65+
builder.endpointOverride(URI.create(config.endpoint()));
66+
}
67+
return builder.build();
68+
}
69+
}

‎connectors/aws/aws-bedrock-knowledgebase/src/main/java/io/camunda/connector/aws/bedrock/knowledgebase/BedrockKnowledgeBaseExecutor.java‎

Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,122 @@
1+
/*
2+
* Copyright Camunda Services GmbH and/or licensed to Camunda Services GmbH
3+
* under one or more contributor license agreements. Licensed under a proprietary license.
4+
* See the License.txt file for more information. You may not use this file
5+
* except in compliance with the proprietary license.
6+
*/
7+
package io.camunda.connector.aws.bedrock.knowledgebase;
8+
9+
import com.fasterxml.jackson.core.JsonProcessingException;
10+
import com.fasterxml.jackson.databind.ObjectMapper;
11+
import io.camunda.connector.api.document.Document;
12+
import io.camunda.connector.api.document.DocumentCreationRequest;
13+
import io.camunda.connector.api.error.ConnectorException;
14+
import io.camunda.connector.api.error.ConnectorRetryException;
15+
import io.camunda.connector.aws.bedrock.knowledgebase.model.request.BedrockKnowledgeBaseRequest;
16+
import io.camunda.connector.aws.bedrock.knowledgebase.model.request.RetrieveOperation;
17+
import io.camunda.connector.aws.bedrock.knowledgebase.model.response.BedrockKnowledgeBaseResponse;
18+
import io.camunda.connector.aws.bedrock.knowledgebase.model.response.KnowledgeBaseRetrievalResult;
19+
import io.camunda.connector.aws.bedrock.knowledgebase.model.response.RetrievalResultEntry;
20+
import java.time.Instant;
21+
import java.util.Map;
22+
import java.util.function.Function;
23+
import java.util.stream.Collectors;
24+
import software.amazon.awssdk.services.bedrockagentruntime.BedrockAgentRuntimeClient;
25+
import software.amazon.awssdk.services.bedrockagentruntime.model.BedrockAgentRuntimeException;
26+
import software.amazon.awssdk.services.bedrockagentruntime.model.KnowledgeBaseRetrievalConfiguration;
27+
import software.amazon.awssdk.services.bedrockagentruntime.model.KnowledgeBaseVectorSearchConfiguration;
28+
import software.amazon.awssdk.services.bedrockagentruntime.model.RetrieveRequest;
29+
import software.amazon.awssdk.services.bedrockagentruntime.model.RetrieveResponse;
30+
import software.amazon.awssdk.services.bedrockagentruntime.model.ThrottlingException;
31+
32+
public class BedrockKnowledgeBaseExecutor {
33+
34+
private final BedrockAgentRuntimeClient client;
35+
private final ObjectMapper objectMapper;
36+
37+
public BedrockKnowledgeBaseExecutor(BedrockAgentRuntimeClient client, ObjectMapper objectMapper) {
38+
this.client = client;
39+
this.objectMapper = objectMapper;
40+
}
41+
42+
public KnowledgeBaseRetrievalResult execute(
43+
BedrockKnowledgeBaseRequest request,
44+
Function<DocumentCreationRequest, Document> documentFactory) {
45+
return switch (request.getOperation()) {
46+
case RetrieveOperation op -> retrieve(op, request, documentFactory);
47+
};
48+
}
49+
50+
private KnowledgeBaseRetrievalResult retrieve(
51+
RetrieveOperation op,
52+
BedrockKnowledgeBaseRequest request,
53+
Function<DocumentCreationRequest, Document> documentFactory) {
54+
try {
55+
var sdkRequestBuilder =
56+
RetrieveRequest.builder()
57+
.knowledgeBaseId(request.getKnowledgeBaseId())
58+
.retrievalQuery(q -> q.text(op.query()));
59+
60+
if (op.numberOfResults() != null) {
61+
sdkRequestBuilder.retrievalConfiguration(
62+
KnowledgeBaseRetrievalConfiguration.builder()
63+
.vectorSearchConfiguration(
64+
KnowledgeBaseVectorSearchConfiguration.builder()
65+
.numberOfResults(op.numberOfResults())
66+
.build())
67+
.build());
68+
}
69+
70+
RetrieveResponse sdkResponse = client.retrieve(sdkRequestBuilder.build());
71+
72+
var response = mapSdkResponse(sdkResponse);
73+
byte[] json = objectMapper.writeValueAsBytes(response);
74+
75+
Document doc =
76+
documentFactory.apply(
77+
DocumentCreationRequest.from(json)
78+
.contentType("application/json")
79+
.fileName("kb-retrieval-" + Instant.now().toEpochMilli() + ".json")
80+
.build());
81+
82+
return new KnowledgeBaseRetrievalResult(doc, response.results().size(), response.nextToken());
83+
84+
} catch (ThrottlingException e) {
85+
throw ConnectorRetryException.builder()
86+
.errorCode("THROTTLED")
87+
.message("Bedrock Knowledge Base request was throttled: " + e.getMessage())
88+
.cause(e)
89+
.build();
90+
} catch (BedrockAgentRuntimeException e) {
91+
var errorMsg =
92+
e.awsErrorDetails() != null ? e.awsErrorDetails().errorMessage() : e.getMessage();
93+
throw new ConnectorException(
94+
"KB_RETRIEVAL_FAILED", "Bedrock Knowledge Base error: " + errorMsg, e);
95+
} catch (JsonProcessingException e) {
96+
throw new ConnectorException(
97+
"SERIALIZATION_ERROR", "Failed to serialize retrieval results to JSON", e);
98+
}
99+
}
100+
101+
private BedrockKnowledgeBaseResponse mapSdkResponse(RetrieveResponse sdkResponse) {
102+
var results =
103+
sdkResponse.retrievalResults().stream()
104+
.map(
105+
r ->
106+
new RetrievalResultEntry(
107+
r.content() != null ? r.content().text() : null,
108+
r.score(),
109+
r.location() != null && r.location().s3Location() != null
110+
? r.location().s3Location().uri()
111+
: null,
112+
r.hasMetadata()
113+
? r.metadata().entrySet().stream()
114+
.collect(
115+
Collectors.toMap(
116+
Map.Entry::getKey,
117+
e -> e.getValue() != null ? e.getValue().toString() : ""))
118+
: Map.of()))
119+
.toList();
120+
return new BedrockKnowledgeBaseResponse(results, sdkResponse.nextToken());
121+
}
122+
}

‎connectors/aws/aws-bedrock-knowledgebase/src/main/java/io/camunda/connector/aws/bedrock/knowledgebase/model/request/BedrockKnowledgeBaseRequest.java‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
/*
2+
* Copyright Camunda Services GmbH and/or licensed to Camunda Services GmbH
3+
* under one or more contributor license agreements. Licensed under a proprietary license.
4+
* See the License.txt file for more information. You may not use this file
5+
* except in compliance with the proprietary license.
6+
*/
7+
package io.camunda.connector.aws.bedrock.knowledgebase.model.request;
8+
9+
import io.camunda.connector.api.annotation.FEEL;
10+
import io.camunda.connector.aws.model.impl.AwsBaseRequest;
11+
import io.camunda.connector.generator.java.annotation.TemplateProperty;
12+
import io.camunda.connector.generator.java.annotation.TemplateProperty.PropertyConstraints;
13+
import jakarta.validation.Valid;
14+
import jakarta.validation.constraints.NotBlank;
15+
import jakarta.validation.constraints.NotNull;
16+
17+
public class BedrockKnowledgeBaseRequest extends AwsBaseRequest {
18+
19+
@FEEL
20+
@NotBlank
21+
@TemplateProperty(
22+
group = "configuration",
23+
label = "Knowledge Base ID",
24+
description = "The ID of the Bedrock Knowledge Base to query.",
25+
constraints = @PropertyConstraints(notEmpty = true))
26+
private String knowledgeBaseId;
27+
28+
@Valid @NotNull private KnowledgeBaseOperation operation;
29+
30+
public String getKnowledgeBaseId() {
31+
return knowledgeBaseId;
32+
}
33+
34+
public void setKnowledgeBaseId(String knowledgeBaseId) {
35+
this.knowledgeBaseId = knowledgeBaseId;
36+
}
37+
38+
public KnowledgeBaseOperation getOperation() {
39+
return operation;
40+
}
41+
42+
public void setOperation(KnowledgeBaseOperation operation) {
43+
this.operation = operation;
44+
}
45+
}

‎connectors/aws/aws-bedrock-knowledgebase/src/main/java/io/camunda/connector/aws/bedrock/knowledgebase/model/request/KnowledgeBaseOperation.java‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
/*
2+
* Copyright Camunda Services GmbH and/or licensed to Camunda Services GmbH
3+
* under one or more contributor license agreements. Licensed under a proprietary license.
4+
* See the License.txt file for more information. You may not use this file
5+
* except in compliance with the proprietary license.
6+
*/
7+
package io.camunda.connector.aws.bedrock.knowledgebase.model.request;
8+
9+
import com.fasterxml.jackson.annotation.JsonSubTypes;
10+
import com.fasterxml.jackson.annotation.JsonTypeInfo;
11+
import io.camunda.connector.generator.java.annotation.TemplateDiscriminatorProperty;
12+
13+
@JsonTypeInfo(use = JsonTypeInfo.Id.NAME, property = "operationDiscriminator")
14+
@JsonSubTypes({
15+
@JsonSubTypes.Type(value = RetrieveOperation.class, name = "retrieve"),
16+
})
17+
@TemplateDiscriminatorProperty(
18+
label = "Operation",
19+
group = "operation",
20+
name = "operationDiscriminator",
21+
defaultValue = "retrieve")
22+
public sealed interface KnowledgeBaseOperation permits RetrieveOperation {}

0 commit comments

Comments
 (0)