Skip to content

Commit 97890bc

Browse files
pmbrullclaude
pmbrull
and
claude
committed
feat(llm): Bedrock completion uses Converse API for model-agnostic support
Replaces InvokeModel+hardcoded Anthropic JSON body with the Converse API, enabling all Bedrock model families (Llama, Mistral, Titan, Cohere, Anthropic). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 parent 89d3e18 commit 97890bc

2 files changed

Lines changed: 91 additions & 63 deletions

File tree

‎openmetadata-service/src/main/java/org/openmetadata/service/llm/BedrockCompletionClient.java‎

Lines changed: 67 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -1,29 +1,40 @@
1+
/*
2+
* Copyright 2021 Collate
3+
* Licensed under the Apache License, Version 2.0 (the "License");
4+
* you may not use this file except in compliance with the License.
5+
* You may obtain a copy of the License at
6+
* http://www.apache.org/licenses/LICENSE-2.0
7+
* Unless required by applicable law or agreed to in writing, software
8+
* distributed under the License is distributed on an "AS IS" BASIS,
9+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
10+
* See the License for the specific language governing permissions and
11+
* limitations under the License.
12+
*/
13+
114
package org.openmetadata.service.llm;
215

3-
import com.fasterxml.jackson.databind.JsonNode;
4-
import com.fasterxml.jackson.databind.ObjectMapper;
5-
import com.fasterxml.jackson.databind.node.ArrayNode;
6-
import com.fasterxml.jackson.databind.node.ObjectNode;
7-
import java.io.IOException;
816
import java.time.Duration;
917
import lombok.extern.slf4j.Slf4j;
1018
import org.openmetadata.schema.configuration.LLMBedrockConfig;
1119
import org.openmetadata.schema.configuration.LLMConfiguration;
1220
import org.openmetadata.schema.security.credentials.AWSBaseConfig;
1321
import org.openmetadata.service.util.AwsCredentialsUtil;
1422
import software.amazon.awssdk.awscore.exception.AwsServiceException;
15-
import software.amazon.awssdk.core.SdkBytes;
1623
import software.amazon.awssdk.core.exception.SdkClientException;
1724
import software.amazon.awssdk.regions.Region;
1825
import software.amazon.awssdk.services.bedrockruntime.BedrockRuntimeClient;
19-
import software.amazon.awssdk.services.bedrockruntime.model.InvokeModelRequest;
20-
import software.amazon.awssdk.services.bedrockruntime.model.InvokeModelResponse;
26+
import software.amazon.awssdk.services.bedrockruntime.model.ContentBlock;
27+
import software.amazon.awssdk.services.bedrockruntime.model.ConversationRole;
28+
import software.amazon.awssdk.services.bedrockruntime.model.ConverseRequest;
29+
import software.amazon.awssdk.services.bedrockruntime.model.ConverseResponse;
30+
import software.amazon.awssdk.services.bedrockruntime.model.InferenceConfiguration;
31+
import software.amazon.awssdk.services.bedrockruntime.model.Message;
32+
import software.amazon.awssdk.services.bedrockruntime.model.SystemContentBlock;
33+
import software.amazon.awssdk.services.bedrockruntime.model.TokenUsage;
2134

22-
/** AWS Bedrock chat-completion client (Anthropic messages API). Mirrors BedrockEmbeddingClient. */
35+
/** AWS Bedrock chat-completion client (Converse API — works with all Bedrock model families). */
2336
@Slf4j
2437
public final class BedrockCompletionClient extends LLMCompletionClient implements AutoCloseable {
25-
private static final ObjectMapper MAPPER = new ObjectMapper();
26-
private static final String ANTHROPIC_VERSION = "bedrock-2023-05-31";
2738

2839
private final BedrockRuntimeClient bedrockClient;
2940
private final String modelId;
@@ -61,59 +72,62 @@ protected CompletionResult doComplete(
6172
int timeout = options.timeoutSecondsOr(this.timeoutSeconds);
6273
CompletionResult result;
6374
try {
64-
InvokeModelRequest request =
65-
InvokeModelRequest.builder()
66-
.modelId(model)
67-
.contentType("application/json")
68-
.accept("application/json")
69-
.body(
70-
SdkBytes.fromUtf8String(buildRequestBody(systemPrompt, userPrompt, tokens, temp)))
71-
.overrideConfiguration(c -> c.apiCallTimeout(Duration.ofSeconds(timeout)))
72-
.build();
73-
InvokeModelResponse response = bedrockClient.invokeModel(request);
74-
result = parseResult(response.body().asUtf8String());
75+
ConverseResponse response =
76+
bedrockClient.converse(
77+
buildRequest(model, systemPrompt, userPrompt, tokens, temp, timeout));
78+
result = toResult(response);
7579
} catch (AwsServiceException | SdkClientException e) {
7680
throw new LLMCompletionException("Bedrock completion failed", e);
7781
}
7882
return result;
7983
}
8084

81-
static String buildRequestBody(
82-
String systemPrompt, String userPrompt, int maxTokens, double temperature) {
83-
String result;
84-
try {
85-
ObjectNode payload = MAPPER.createObjectNode();
86-
payload.put("anthropic_version", ANTHROPIC_VERSION);
87-
payload.put("max_tokens", maxTokens);
88-
payload.put("temperature", temperature);
89-
payload.put("system", systemPrompt);
90-
ArrayNode messages = payload.putArray("messages");
91-
messages.addObject().put("role", "user").put("content", userPrompt);
92-
result = MAPPER.writeValueAsString(payload);
93-
} catch (IOException e) {
94-
throw new LLMCompletionException("Failed to build Bedrock request body", e);
85+
static ConverseRequest buildRequest(
86+
String model,
87+
String systemPrompt,
88+
String userPrompt,
89+
int maxTokens,
90+
double temperature,
91+
int timeoutSeconds) {
92+
return ConverseRequest.builder()
93+
.modelId(model)
94+
.system(SystemContentBlock.fromText(systemPrompt))
95+
.messages(
96+
Message.builder()
97+
.role(ConversationRole.USER)
98+
.content(ContentBlock.fromText(userPrompt))
99+
.build())
100+
.inferenceConfig(
101+
InferenceConfiguration.builder()
102+
.maxTokens(maxTokens)
103+
.temperature((float) temperature)
104+
.build())
105+
.overrideConfiguration(c -> c.apiCallTimeout(Duration.ofSeconds(timeoutSeconds)))
106+
.build();
107+
}
108+
109+
static CompletionResult toResult(ConverseResponse response) {
110+
String text = extractText(response);
111+
if (text.isEmpty()) {
112+
throw new LLMCompletionException("Invalid Bedrock response: no text content returned");
95113
}
96-
return result;
114+
TokenUsage usage = response.usage();
115+
int inputTokens = usage != null && usage.inputTokens() != null ? usage.inputTokens() : 0;
116+
int outputTokens = usage != null && usage.outputTokens() != null ? usage.outputTokens() : 0;
117+
return new CompletionResult(text, inputTokens, outputTokens);
97118
}
98119

99-
static CompletionResult parseResult(String responseBody) {
100-
CompletionResult result;
101-
try {
102-
JsonNode root = MAPPER.readTree(responseBody);
103-
JsonNode content = root.get("content");
104-
if (content == null || !content.isArray() || content.isEmpty()) {
105-
throw new LLMCompletionException("Invalid Bedrock response: no content returned");
120+
private static String extractText(ConverseResponse response) {
121+
String text = "";
122+
if (response.output() != null && response.output().message() != null) {
123+
for (ContentBlock block : response.output().message().content()) {
124+
if (block.text() != null) {
125+
text = block.text();
126+
break;
127+
}
106128
}
107-
JsonNode usage = root.path("usage");
108-
result =
109-
new CompletionResult(
110-
content.get(0).get("text").asText(),
111-
usage.path("input_tokens").asInt(0),
112-
usage.path("output_tokens").asInt(0));
113-
} catch (IOException e) {
114-
throw new LLMCompletionException("Failed to parse Bedrock response", e);
115129
}
116-
return result;
130+
return text;
117131
}
118132

119133
@Override

‎openmetadata-service/src/test/java/org/openmetadata/service/llm/BedrockCompletionClientTest.java‎

Lines changed: 24 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -14,28 +14,42 @@
1414
package org.openmetadata.service.llm;
1515

1616
import static org.junit.jupiter.api.Assertions.assertEquals;
17-
import static org.junit.jupiter.api.Assertions.assertTrue;
1817

1918
import org.junit.jupiter.api.Test;
19+
import software.amazon.awssdk.services.bedrockruntime.model.ContentBlock;
20+
import software.amazon.awssdk.services.bedrockruntime.model.ConversationRole;
21+
import software.amazon.awssdk.services.bedrockruntime.model.ConverseOutput;
22+
import software.amazon.awssdk.services.bedrockruntime.model.ConverseRequest;
23+
import software.amazon.awssdk.services.bedrockruntime.model.ConverseResponse;
24+
import software.amazon.awssdk.services.bedrockruntime.model.Message;
25+
import software.amazon.awssdk.services.bedrockruntime.model.TokenUsage;
2026

2127
class BedrockCompletionClientTest {
2228

23-
private static final String RESPONSE =
24-
"{\"content\":[{\"type\":\"text\",\"text\":\"ok\"}],"
25-
+ "\"usage\":{\"input_tokens\":30,\"output_tokens\":9}}";
26-
2729
@Test
2830
void parsesTextAndUsage() {
29-
CompletionResult result = BedrockCompletionClient.parseResult(RESPONSE);
31+
ConverseResponse response =
32+
ConverseResponse.builder()
33+
.output(
34+
ConverseOutput.fromMessage(
35+
Message.builder()
36+
.role(ConversationRole.ASSISTANT)
37+
.content(ContentBlock.fromText("ok"))
38+
.build()))
39+
.usage(TokenUsage.builder().inputTokens(30).outputTokens(9).build())
40+
.build();
41+
CompletionResult result = BedrockCompletionClient.toResult(response);
3042
assertEquals("ok", result.text());
3143
assertEquals(30, result.inputTokens());
3244
assertEquals(9, result.outputTokens());
3345
}
3446

3547
@Test
36-
void requestBodyHonorsOverrides() {
37-
String body = BedrockCompletionClient.buildRequestBody("sys", "user", 256, 0.0);
38-
assertTrue(body.contains("\"max_tokens\":256"));
39-
assertTrue(body.contains("\"anthropic_version\""));
48+
void requestHonorsOverrides() {
49+
ConverseRequest request =
50+
BedrockCompletionClient.buildRequest("override-model", "sys", "user", 256, 0.0, 15);
51+
assertEquals("override-model", request.modelId());
52+
assertEquals(256, request.inferenceConfig().maxTokens());
53+
assertEquals(0.0f, request.inferenceConfig().temperature(), 0.0001f);
4054
}
4155
}

0 commit comments

Comments
 (0)