Skip to content

Commit b048c6f

Browse files
andreadimaioandreadimaio
committed
Polish watsonx code
1 parent a2f33cb commit b048c6f

File tree

4 files changed

+50
-40
lines changed

4 files changed

+50
-40
lines changed

‎model-providers/watsonx/runtime/src/main/java/io/quarkiverse/langchain4j/watsonx/bean/TextChatMessage.java‎

Lines changed: 6 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package io.quarkiverse.langchain4j.watsonx.bean;
22

33
import static io.quarkiverse.langchain4j.watsonx.WatsonxUtils.base64Image;
4+
import static java.util.Objects.isNull;
45

56
import java.util.ArrayList;
67
import java.util.List;
@@ -131,8 +132,9 @@ public static TextChatMessageUser of(UserMessage userMessage) {
131132
}
132133
case IMAGE -> {
133134
var imageContent = ImageContent.class.cast(content);
134-
var base64 = "data:image/%s;base64,%s".formatted(
135-
imageContent.image().mimeType(),
135+
var mimeType = imageContent.image().mimeType();
136+
var base64 = "data:%s;base64,%s".formatted(
137+
isNull(mimeType) ? "image" : mimeType,
136138
base64Image(imageContent.image()));
137139
values.add(Map.of(
138140
"type", "image_url",
@@ -225,11 +227,8 @@ public record TextChatParameterFunction(String name, String description, Map<Str
225227
*/
226228
public static TextChatParameterTool of(ToolSpecification toolSpecification) {
227229
var toolParams = JsonSchemaElementHelper.toMap(toolSpecification.parameters());
228-
229-
var parameters = new TextChatParameterFunction(toolSpecification.name(), toolSpecification.description(), Map.of(
230-
"type", toolParams.get("type"),
231-
"properties", toolParams.get("properties"),
232-
"required", toolParams.get("required")));
230+
var parameters = new TextChatParameterFunction(toolSpecification.name(), toolSpecification.description(),
231+
toolParams);
233232
return new TextChatParameterTool("function", parameters);
234233
}
235234
}

‎model-providers/watsonx/runtime/src/main/java/io/quarkiverse/langchain4j/watsonx/client/WatsonxRestApi.java‎

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -142,6 +142,8 @@ class WatsonClientLogger implements ClientLogger {
142142

143143
private static final Logger log = Logger.getLogger(WatsonClientLogger.class);
144144
private static final Pattern BEARER_PATTERN = Pattern.compile("(Bearer\\s*)(\\w{4})(\\w+)(\\w{4})");
145+
private static final Pattern BASE64_IMAGE_PATTERN = Pattern.compile("(data:.+;base64,)(.{15})(.+)(.{15})([\\s\\S]*)");
146+
145147
private final boolean logRequests;
146148
private final boolean logResponses;
147149

@@ -205,7 +207,7 @@ private String bodyToString(Buffer body) {
205207
if (body == null) {
206208
return "";
207209
}
208-
return body.toString();
210+
return formatBase64ImageForLogging(body.toString());
209211
}
210212

211213
private String inOneLine(MultiMap headers) {
@@ -240,6 +242,26 @@ private static String maskAuthorizationHeaderValue(String authorizationHeaderVal
240242
}
241243
}
242244

245+
private static String formatBase64ImageForLogging(String body) {
246+
try {
247+
248+
if (body == null || body.isBlank())
249+
return body;
250+
251+
Matcher matcher = BASE64_IMAGE_PATTERN.matcher(body);
252+
253+
StringBuilder sb = new StringBuilder();
254+
while (matcher.find()) {
255+
matcher.appendReplacement(sb,
256+
matcher.group(1) + matcher.group(2) + "..." + matcher.group(4) + matcher.group(5));
257+
}
258+
259+
return sb.isEmpty() ? body : sb.toString();
260+
} catch (Exception e) {
261+
return "Failed to format the base64 image value.";
262+
}
263+
}
264+
243265
private static String maskApiKeyHeaderValue(String apiKeyHeaderValue) {
244266
try {
245267
if (apiKeyHeaderValue.length() <= 4) {

‎model-providers/watsonx/runtime/src/main/java/io/quarkiverse/langchain4j/watsonx/client/filter/BearerTokenHeaderFactory.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@
1313
*/
1414
public class BearerTokenHeaderFactory extends ReactiveClientHeadersFactory {
1515

16-
private WatsonxTokenGenerator tokenGenerator;
16+
private final WatsonxTokenGenerator tokenGenerator;
1717

1818
public BearerTokenHeaderFactory(WatsonxTokenGenerator tokenGenerator) {
1919
this.tokenGenerator = tokenGenerator;

‎model-providers/watsonx/runtime/src/main/java/io/quarkiverse/langchain4j/watsonx/runtime/WatsonxRecorder.java‎

Lines changed: 20 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -156,10 +156,7 @@ public Supplier<EmbeddingModel> embeddingModel(LangChain4jWatsonxConfig runtimeC
156156
throw new ConfigValidationException(configProblems.toArray(EMPTY_PROBLEMS));
157157
}
158158

159-
String iamUrl = watsonConfig.iam().baseUrl().toExternalForm();
160-
WatsonxTokenGenerator tokenGenerator = tokenGeneratorCache.computeIfAbsent(iamUrl,
161-
createTokenGenerator(watsonConfig.iam(),
162-
firstOrDefault(null, watsonConfig.apiKey(), runtimeConfig.defaultConfig().apiKey())));
159+
String apiKey = firstOrDefault(null, watsonConfig.apiKey(), runtimeConfig.defaultConfig().apiKey());
163160

164161
URL url;
165162
try {
@@ -170,7 +167,7 @@ public Supplier<EmbeddingModel> embeddingModel(LangChain4jWatsonxConfig runtimeC
170167

171168
EmbeddingModelConfig embeddingModelConfig = watsonConfig.embeddingModel();
172169
var builder = WatsonxEmbeddingModel.builder()
173-
.tokenGenerator(tokenGenerator)
170+
.tokenGenerator(createTokenGenerator(watsonConfig.iam(), apiKey))
174171
.url(url)
175172
.timeout(watsonConfig.timeout().orElse(Duration.ofSeconds(10)))
176173
.logRequests(firstOrDefault(false, embeddingModelConfig.logRequests(), watsonConfig.logRequests()))
@@ -209,10 +206,7 @@ public Supplier<ScoringModel> scoringModel(LangChain4jWatsonxConfig runtimeConfi
209206
throw new ConfigValidationException(configProblems.toArray(EMPTY_PROBLEMS));
210207
}
211208

212-
String iamUrl = watsonConfig.iam().baseUrl().toExternalForm();
213-
WatsonxTokenGenerator tokenGenerator = tokenGeneratorCache.computeIfAbsent(iamUrl,
214-
createTokenGenerator(watsonConfig.iam(),
215-
firstOrDefault(null, watsonConfig.apiKey(), runtimeConfig.defaultConfig().apiKey())));
209+
String apiKey = firstOrDefault(null, watsonConfig.apiKey(), runtimeConfig.defaultConfig().apiKey());
216210

217211
URL url;
218212
try {
@@ -223,7 +217,7 @@ public Supplier<ScoringModel> scoringModel(LangChain4jWatsonxConfig runtimeConfi
223217

224218
ScoringModelConfig rerankModelConfig = watsonConfig.scoringModel();
225219
var builder = WatsonxScoringModel.builder()
226-
.tokenGenerator(tokenGenerator)
220+
.tokenGenerator(createTokenGenerator(watsonConfig.iam(), apiKey))
227221
.url(url)
228222
.timeout(watsonConfig.timeout().orElse(Duration.ofSeconds(10)))
229223
.logRequests(firstOrDefault(false, rerankModelConfig.logRequests(), watsonConfig.logRequests()))
@@ -242,17 +236,6 @@ public WatsonxScoringModel get() {
242236
};
243237
}
244238

245-
private Function<? super String, ? extends WatsonxTokenGenerator> createTokenGenerator(IAMConfig iamConfig, String apiKey) {
246-
return new Function<String, WatsonxTokenGenerator>() {
247-
248-
@Override
249-
public WatsonxTokenGenerator apply(String iamUrl) {
250-
return new WatsonxTokenGenerator(iamConfig.baseUrl(), iamConfig.timeout().orElse(Duration.ofSeconds(10)),
251-
iamConfig.grantType(), apiKey);
252-
}
253-
};
254-
}
255-
256239
private WatsonxChatModel.Builder chatBuilder(LangChain4jWatsonxConfig runtimeConfig, String configName) {
257240
LangChain4jWatsonxConfig.WatsonConfig watsonConfig = correspondingWatsonRuntimeConfig(runtimeConfig, configName);
258241

@@ -263,10 +246,7 @@ private WatsonxChatModel.Builder chatBuilder(LangChain4jWatsonxConfig runtimeCon
263246
}
264247

265248
ChatModelConfig chatModelConfig = watsonConfig.chatModel();
266-
String iamUrl = watsonConfig.iam().baseUrl().toExternalForm();
267-
WatsonxTokenGenerator tokenGenerator = tokenGeneratorCache.computeIfAbsent(iamUrl,
268-
createTokenGenerator(watsonConfig.iam(),
269-
firstOrDefault(null, watsonConfig.apiKey(), runtimeConfig.defaultConfig().apiKey())));
249+
String apiKey = firstOrDefault(null, watsonConfig.apiKey(), runtimeConfig.defaultConfig().apiKey());
270250

271251
URL url;
272252
try {
@@ -276,7 +256,7 @@ private WatsonxChatModel.Builder chatBuilder(LangChain4jWatsonxConfig runtimeCon
276256
}
277257

278258
return WatsonxChatModel.builder()
279-
.tokenGenerator(tokenGenerator)
259+
.tokenGenerator(createTokenGenerator(watsonConfig.iam(), apiKey))
280260
.url(url)
281261
.timeout(watsonConfig.timeout().orElse(Duration.ofSeconds(10)))
282262
.logRequests(firstOrDefault(false, chatModelConfig.logRequests(), watsonConfig.logRequests()))
@@ -306,10 +286,7 @@ private WatsonxGenerationModel.Builder generationBuilder(LangChain4jWatsonxConfi
306286
}
307287

308288
GenerationModelConfig generationModelConfig = watsonConfig.generationModel();
309-
String iamUrl = watsonConfig.iam().baseUrl().toExternalForm();
310-
WatsonxTokenGenerator tokenGenerator = tokenGeneratorCache.computeIfAbsent(iamUrl,
311-
createTokenGenerator(watsonConfig.iam(),
312-
firstOrDefault(null, watsonConfig.apiKey(), runtimeConfig.defaultConfig().apiKey())));
289+
String apiKey = firstOrDefault(null, watsonConfig.apiKey(), runtimeConfig.defaultConfig().apiKey());
313290

314291
URL url;
315292
try {
@@ -323,7 +300,7 @@ private WatsonxGenerationModel.Builder generationBuilder(LangChain4jWatsonxConfi
323300
String promptJoiner = generationModelConfig.promptJoiner();
324301

325302
return WatsonxGenerationModel.builder()
326-
.tokenGenerator(tokenGenerator)
303+
.tokenGenerator(createTokenGenerator(watsonConfig.iam(), apiKey))
327304
.url(url)
328305
.timeout(watsonConfig.timeout().orElse(Duration.ofSeconds(10)))
329306
.logRequests(firstOrDefault(false, generationModelConfig.logRequests(), watsonConfig.logRequests()))
@@ -348,6 +325,18 @@ private WatsonxGenerationModel.Builder generationBuilder(LangChain4jWatsonxConfi
348325
.promptJoiner(promptJoiner);
349326
}
350327

328+
private WatsonxTokenGenerator createTokenGenerator(IAMConfig iamConfig, String apiKey) {
329+
return tokenGeneratorCache.computeIfAbsent(apiKey,
330+
new Function<String, WatsonxTokenGenerator>() {
331+
@Override
332+
public WatsonxTokenGenerator apply(String iamUrl) {
333+
return new WatsonxTokenGenerator(iamConfig.baseUrl(),
334+
iamConfig.timeout().orElse(Duration.ofSeconds(10)),
335+
iamConfig.grantType(), apiKey);
336+
}
337+
});
338+
}
339+
351340
private LangChain4jWatsonxConfig.WatsonConfig correspondingWatsonRuntimeConfig(LangChain4jWatsonxConfig runtimeConfig,
352341
String configName) {
353342
LangChain4jWatsonxConfig.WatsonConfig watsonConfig;

0 commit comments

Comments
 (0)