Skip to content

Commit c13b352

Browse files
pmbrullclaude
pmbrull
and
claude
committed
refactor(search): embedding clients read from llmConfiguration.embeddings
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 parent 45c55fe commit c13b352

10 files changed

Lines changed: 441 additions & 364 deletions

File tree

‎openmetadata-service/src/main/java/org/openmetadata/service/jdbi3/SystemRepository.java‎

Lines changed: 87 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -51,17 +51,20 @@
5151
import org.openmetadata.schema.configuration.AssetCertificationSettings;
5252
import org.openmetadata.schema.configuration.ExecutorConfiguration;
5353
import org.openmetadata.schema.configuration.HistoryCleanUpConfiguration;
54+
import org.openmetadata.schema.configuration.LLMConfiguration;
55+
import org.openmetadata.schema.configuration.LLMEmbeddingsConfig;
56+
import org.openmetadata.schema.configuration.LLMGoogleConfig;
57+
import org.openmetadata.schema.configuration.LLMOpenAIConfig;
5458
import org.openmetadata.schema.configuration.SecurityConfiguration;
5559
import org.openmetadata.schema.configuration.WorkflowSettings;
5660
import org.openmetadata.schema.email.SmtpSettings;
5761
import org.openmetadata.schema.entity.app.App;
5862
import org.openmetadata.schema.entity.services.ingestionPipelines.PipelineServiceClientResponse;
5963
import org.openmetadata.schema.security.client.OidcClientConfig;
6064
import org.openmetadata.schema.security.client.OpenMetadataJWTClientConfig;
65+
import org.openmetadata.schema.security.credentials.AWSBaseConfig;
6166
import org.openmetadata.schema.security.scim.ScimConfiguration;
6267
import org.openmetadata.schema.service.configuration.elasticsearch.ElasticSearchConfiguration;
63-
import org.openmetadata.schema.service.configuration.elasticsearch.Google;
64-
import org.openmetadata.schema.service.configuration.elasticsearch.NaturalLanguageSearchConfiguration;
6568
import org.openmetadata.schema.service.configuration.slackApp.SlackAppConfiguration;
6669
import org.openmetadata.schema.services.connections.metadata.AuthProvider;
6770
import org.openmetadata.schema.services.connections.metadata.OpenMetadataConnection;
@@ -83,6 +86,7 @@
8386
import org.openmetadata.service.attachments.AssetService;
8487
import org.openmetadata.service.attachments.AssetServiceFactory;
8588
import org.openmetadata.service.attachments.NoOpAssetService;
89+
import org.openmetadata.service.clients.llm.LlmConfigHolder;
8690
import org.openmetadata.service.config.ObjectStorageConfiguration;
8791
import org.openmetadata.service.events.scheduled.ServicesStatusJobHandler;
8892
import org.openmetadata.service.exception.CustomExceptionMessage;
@@ -882,53 +886,21 @@ private StepValidation validateEmbeddingGeneration(
882886

883887
private String getEmbeddingConfigurationMessage(OpenMetadataApplicationConfig applicationConfig) {
884888
try {
885-
NaturalLanguageSearchConfiguration nlpConfig =
886-
applicationConfig.getElasticSearchConfiguration().getNaturalLanguageSearch();
887-
String provider = nlpConfig.getEmbeddingProvider();
889+
LLMConfiguration llmConfig = LlmConfigHolder.get();
890+
LLMEmbeddingsConfig embeddings = llmConfig != null ? llmConfig.getEmbeddings() : null;
891+
String provider =
892+
embeddings != null && embeddings.getProvider() != null
893+
? embeddings.getProvider().value()
894+
: null;
888895
if (nullOrEmpty(provider)) {
889896
return "Required configuration: embeddingProvider";
890897
}
891898

892899
return switch (provider.toLowerCase()) {
893-
case "djl" -> String.format(
894-
"DJL configuration: embeddingModel: %s", nlpConfig.getDjl().getEmbeddingModel());
895-
case "bedrock" -> String.format(
896-
"Bedrock configuration: region: %s, embeddingModelId: %s, embeddingDimension %s",
897-
nlpConfig.getBedrock().getAwsConfig() != null
898-
? nlpConfig.getBedrock().getAwsConfig().getRegion()
899-
: "not configured",
900-
nlpConfig.getBedrock().getEmbeddingModelId(),
901-
nlpConfig.getBedrock().getEmbeddingDimension());
902-
case "openai" -> {
903-
String openaiEndpoint =
904-
nullOrEmpty(nlpConfig.getOpenai().getEndpoint())
905-
? "api.openai.com"
906-
: nlpConfig.getOpenai().getEndpoint();
907-
String deploymentInfo =
908-
nullOrEmpty(nlpConfig.getOpenai().getDeploymentName())
909-
? ""
910-
: String.format(
911-
", deploymentName: %s", nlpConfig.getOpenai().getDeploymentName());
912-
yield String.format(
913-
"OpenAI configuration: endpoint: %s, embeddingModelId: %s, embeddingDimension: %s%s",
914-
openaiEndpoint,
915-
nlpConfig.getOpenai().getEmbeddingModelId(),
916-
nlpConfig.getOpenai().getEmbeddingDimension(),
917-
deploymentInfo);
918-
}
919-
case "google" -> {
920-
Google googleCfg = nlpConfig.getGoogle();
921-
if (googleCfg == null) {
922-
yield "Google provider selected but google configuration block is missing";
923-
}
924-
String googleEndpoint =
925-
nullOrEmpty(googleCfg.getEndpoint())
926-
? "generativelanguage.googleapis.com"
927-
: googleCfg.getEndpoint();
928-
yield String.format(
929-
"Google configuration: endpoint: %s, embeddingModelId: %s, embeddingDimension: %s",
930-
googleEndpoint, googleCfg.getEmbeddingModelId(), googleCfg.getEmbeddingDimension());
931-
}
900+
case "djl" -> getDjlEmbeddingMessage(embeddings);
901+
case "bedrock" -> getBedrockEmbeddingMessage(llmConfig, embeddings);
902+
case "openai" -> getOpenAiEmbeddingMessage(llmConfig, embeddings);
903+
case "google" -> getGoogleEmbeddingMessage(llmConfig, embeddings);
932904
default -> String.format(
933905
"Unknown provider '%s'. Supported providers: djl, bedrock, openai, google", provider);
934906
};
@@ -938,6 +910,77 @@ private String getEmbeddingConfigurationMessage(OpenMetadataApplicationConfig ap
938910
}
939911
}
940912

913+
private String getDjlEmbeddingMessage(LLMEmbeddingsConfig embeddings) {
914+
String message = "DJL provider selected but djl configuration block is missing";
915+
if (embeddings.getDjl() != null) {
916+
message =
917+
String.format(
918+
"DJL configuration: embeddingModel: %s", embeddings.getDjl().getEmbeddingModel());
919+
}
920+
return message;
921+
}
922+
923+
private String getBedrockEmbeddingMessage(
924+
LLMConfiguration llmConfig, LLMEmbeddingsConfig embeddings) {
925+
String message = "Bedrock provider selected but bedrock configuration block is missing";
926+
if (embeddings.getBedrock() != null) {
927+
AWSBaseConfig awsConfig =
928+
llmConfig.getBedrock() != null ? llmConfig.getBedrock().getAwsConfig() : null;
929+
message =
930+
String.format(
931+
"Bedrock configuration: region: %s, embeddingModelId: %s, embeddingDimension %s",
932+
awsConfig != null && awsConfig.getRegion() != null
933+
? awsConfig.getRegion()
934+
: "not configured",
935+
embeddings.getBedrock().getEmbeddingModelId(),
936+
embeddings.getBedrock().getEmbeddingDimension());
937+
}
938+
return message;
939+
}
940+
941+
private String getOpenAiEmbeddingMessage(
942+
LLMConfiguration llmConfig, LLMEmbeddingsConfig embeddings) {
943+
String message = "OpenAI provider selected but openai configuration block is missing";
944+
if (embeddings.getOpenai() != null) {
945+
LLMOpenAIConfig openaiCreds = llmConfig.getOpenai();
946+
String openaiEndpoint =
947+
openaiCreds == null || nullOrEmpty(openaiCreds.getEndpoint())
948+
? "api.openai.com"
949+
: openaiCreds.getEndpoint();
950+
String deploymentInfo =
951+
openaiCreds == null || nullOrEmpty(openaiCreds.getDeploymentName())
952+
? ""
953+
: String.format(", deploymentName: %s", openaiCreds.getDeploymentName());
954+
message =
955+
String.format(
956+
"OpenAI configuration: endpoint: %s, embeddingModelId: %s, embeddingDimension: %s%s",
957+
openaiEndpoint,
958+
embeddings.getOpenai().getEmbeddingModelId(),
959+
embeddings.getOpenai().getEmbeddingDimension(),
960+
deploymentInfo);
961+
}
962+
return message;
963+
}
964+
965+
private String getGoogleEmbeddingMessage(
966+
LLMConfiguration llmConfig, LLMEmbeddingsConfig embeddings) {
967+
String message = "Google provider selected but google configuration block is missing";
968+
if (embeddings.getGoogle() != null) {
969+
LLMGoogleConfig googleCreds = llmConfig.getGoogle();
970+
String googleEndpoint =
971+
googleCreds == null || nullOrEmpty(googleCreds.getEndpoint())
972+
? "generativelanguage.googleapis.com"
973+
: googleCreds.getEndpoint();
974+
message =
975+
String.format(
976+
"Google configuration: endpoint: %s, embeddingModelId: %s, embeddingDimension: %s",
977+
googleEndpoint,
978+
embeddings.getGoogle().getEmbeddingModelId(),
979+
embeddings.getGoogle().getEmbeddingDimension());
980+
}
981+
return message;
982+
}
983+
941984
private StepValidation getDatabaseValidation(OpenMetadataApplicationConfig applicationConfig) {
942985
try {
943986
dao.testConnection();

‎openmetadata-service/src/main/java/org/openmetadata/service/search/SearchRepository.java‎

Lines changed: 29 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,8 @@
9797
import org.openmetadata.schema.api.lineage.SearchLineageResult;
9898
import org.openmetadata.schema.api.search.SearchSettings;
9999
import org.openmetadata.schema.configuration.AssetCertificationSettings;
100+
import org.openmetadata.schema.configuration.LLMConfiguration;
101+
import org.openmetadata.schema.configuration.LLMEmbeddingsConfig;
100102
import org.openmetadata.schema.dataInsight.DataInsightChartResult;
101103
import org.openmetadata.schema.entity.classification.Tag;
102104
import org.openmetadata.schema.entity.data.Pipeline;
@@ -124,6 +126,7 @@
124126
import org.openmetadata.service.apps.bundles.searchIndex.BulkSink;
125127
import org.openmetadata.service.apps.bundles.searchIndex.ElasticSearchBulkSink;
126128
import org.openmetadata.service.apps.bundles.searchIndex.OpenSearchBulkSink;
129+
import org.openmetadata.service.clients.llm.LlmConfigHolder;
127130
import org.openmetadata.service.events.lifecycle.EntityLifecycleEventDispatcher;
128131
import org.openmetadata.service.events.lifecycle.handlers.SearchIndexHandler;
129132
import org.openmetadata.service.jdbi3.EntityRepository;
@@ -552,8 +555,9 @@ public synchronized void initializeVectorSearchService() {
552555
return;
553556
}
554557

558+
LLMConfiguration llmConfig = LlmConfigHolder.get();
555559
try {
556-
this.embeddingClient = createEmbeddingClient(cfg);
560+
this.embeddingClient = createEmbeddingClient(llmConfig);
557561

558562
if (cfg.getSearchType() == ElasticSearchConfiguration.SearchType.OPENSEARCH) {
559563
os.org.opensearch.client.opensearch.OpenSearchClient osClient =
@@ -575,7 +579,7 @@ public synchronized void initializeVectorSearchService() {
575579

576580
LOG.info(
577581
"Vector search service initialized with provider={}, dimension={}",
578-
cfg.getNaturalLanguageSearch().getEmbeddingProvider(),
582+
resolveEmbeddingProvider(llmConfig),
579583
embeddingClient.getDimension());
580584
} catch (Exception e) {
581585
this.vectorServiceInitError = e.toString();
@@ -3489,43 +3493,55 @@ private String reformatVectorIndexWithDimension(String mapping, int dimension) {
34893493
}
34903494
}
34913495

3492-
protected EmbeddingClient createEmbeddingClient(ElasticSearchConfiguration esConfig) {
3493-
NaturalLanguageSearchConfiguration config = esConfig.getNaturalLanguageSearch();
3496+
protected EmbeddingClient createEmbeddingClient(LLMConfiguration llmConfig) {
3497+
LLMEmbeddingsConfig embeddings = llmConfig != null ? llmConfig.getEmbeddings() : null;
34943498
String provider =
3495-
config.getEmbeddingProvider() != null ? config.getEmbeddingProvider() : "bedrock";
3499+
embeddings != null && embeddings.getProvider() != null
3500+
? embeddings.getProvider().value()
3501+
: "bedrock";
34963502

34973503
return switch (provider.toLowerCase()) {
34983504
case "bedrock" -> {
3499-
if (config.getBedrock() == null) {
3505+
if (embeddings == null || embeddings.getBedrock() == null) {
35003506
throw new IllegalStateException(
35013507
"Bedrock configuration is required when using bedrock provider");
35023508
}
3503-
yield new BedrockEmbeddingClient(esConfig);
3509+
yield new BedrockEmbeddingClient(llmConfig);
35043510
}
35053511
case "openai" -> {
3506-
if (config.getOpenai() == null) {
3512+
if (llmConfig.getOpenai() == null || embeddings.getOpenai() == null) {
35073513
throw new IllegalStateException(
35083514
"OpenAI configuration is required when using openai provider");
35093515
}
3510-
yield new OpenAIEmbeddingClient(esConfig);
3516+
yield new OpenAIEmbeddingClient(llmConfig);
35113517
}
35123518
case "google" -> {
3513-
if (config.getGoogle() == null) {
3519+
if (llmConfig.getGoogle() == null || embeddings.getGoogle() == null) {
35143520
throw new IllegalStateException(
35153521
"Google configuration is required when using google provider");
35163522
}
3517-
yield new GoogleEmbeddingClient(esConfig);
3523+
yield new GoogleEmbeddingClient(llmConfig);
35183524
}
35193525
case "djl" -> {
3520-
if (config.getDjl() == null) {
3526+
if (embeddings.getDjl() == null) {
35213527
throw new IllegalStateException("DJL configuration is required when using djl provider");
35223528
}
3523-
yield new DjlEmbeddingClient(esConfig);
3529+
yield new DjlEmbeddingClient(llmConfig);
35243530
}
35253531
default -> throw new IllegalArgumentException("Unknown embedding provider: " + provider);
35263532
};
35273533
}
35283534

3535+
private static String resolveEmbeddingProvider(LLMConfiguration llmConfig) {
3536+
String provider = "unknown";
3537+
if (llmConfig != null
3538+
&& llmConfig.getEmbeddings() != null
3539+
&& llmConfig.getEmbeddings().getProvider() != null) {
3540+
provider = llmConfig.getEmbeddings().getProvider().value();
3541+
}
3542+
return provider;
3543+
}
3544+
35293545
public String getModelIdentifier() {
35303546
return embeddingClient != null ? embeddingClient.getModelId() : null;
35313547
}

‎openmetadata-service/src/main/java/org/openmetadata/service/search/vector/client/BedrockEmbeddingClient.java‎

Lines changed: 14 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -5,9 +5,9 @@
55
import com.fasterxml.jackson.databind.node.ObjectNode;
66
import java.io.IOException;
77
import lombok.extern.slf4j.Slf4j;
8+
import org.openmetadata.schema.configuration.LLMBedrockEmbeddingConfig;
9+
import org.openmetadata.schema.configuration.LLMConfiguration;
810
import org.openmetadata.schema.security.credentials.AWSBaseConfig;
9-
import org.openmetadata.schema.service.configuration.elasticsearch.ElasticSearchConfiguration;
10-
import org.openmetadata.schema.service.configuration.elasticsearch.NaturalLanguageSearchConfiguration;
1111
import org.openmetadata.service.util.AwsCredentialsUtil;
1212
import software.amazon.awssdk.awscore.exception.AwsServiceException;
1313
import software.amazon.awssdk.core.SdkBytes;
@@ -25,28 +25,30 @@ public final class BedrockEmbeddingClient extends EmbeddingClient implements Aut
2525
private final String modelId;
2626
private final int dimension;
2727

28-
public BedrockEmbeddingClient(ElasticSearchConfiguration config) {
28+
public BedrockEmbeddingClient(LLMConfiguration config) {
2929
super(resolveMaxConcurrent(config));
30-
NaturalLanguageSearchConfiguration nlsCfg = config.getNaturalLanguageSearch();
31-
if (nlsCfg.getBedrock() == null) {
30+
LLMBedrockEmbeddingConfig bedrockEmbedding =
31+
config.getEmbeddings() != null ? config.getEmbeddings().getBedrock() : null;
32+
if (bedrockEmbedding == null) {
3233
throw new IllegalArgumentException("Bedrock configuration is required");
3334
}
34-
if (nlsCfg.getBedrock().getEmbeddingModelId() == null
35-
|| nlsCfg.getBedrock().getEmbeddingModelId().isBlank()) {
35+
if (bedrockEmbedding.getEmbeddingModelId() == null
36+
|| bedrockEmbedding.getEmbeddingModelId().isBlank()) {
3637
throw new IllegalArgumentException("Bedrock embedding model ID is required");
3738
}
38-
if (nlsCfg.getBedrock().getEmbeddingDimension() == null
39-
|| nlsCfg.getBedrock().getEmbeddingDimension() <= 0) {
39+
if (bedrockEmbedding.getEmbeddingDimension() == null
40+
|| bedrockEmbedding.getEmbeddingDimension() <= 0) {
4041
throw new IllegalArgumentException("Bedrock embedding dimension must be positive");
4142
}
4243

43-
AWSBaseConfig awsConfig = nlsCfg.getBedrock().getAwsConfig();
44+
AWSBaseConfig awsConfig =
45+
config.getBedrock() != null ? config.getBedrock().getAwsConfig() : null;
4446
if (awsConfig == null || awsConfig.getRegion() == null || awsConfig.getRegion().isBlank()) {
4547
throw new IllegalArgumentException("AWS region is required for Bedrock");
4648
}
4749

48-
this.modelId = nlsCfg.getBedrock().getEmbeddingModelId();
49-
this.dimension = nlsCfg.getBedrock().getEmbeddingDimension();
50+
this.modelId = bedrockEmbedding.getEmbeddingModelId();
51+
this.dimension = bedrockEmbedding.getEmbeddingDimension();
5052

5153
this.bedrockClient =
5254
BedrockRuntimeClient.builder()

‎openmetadata-service/src/main/java/org/openmetadata/service/search/vector/client/DjlEmbeddingClient.java‎

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -12,8 +12,8 @@
1212
import java.util.ArrayList;
1313
import java.util.List;
1414
import lombok.extern.slf4j.Slf4j;
15-
import org.openmetadata.schema.service.configuration.elasticsearch.ElasticSearchConfiguration;
16-
import org.openmetadata.schema.service.configuration.elasticsearch.NaturalLanguageSearchConfiguration;
15+
import org.openmetadata.schema.configuration.LLMConfiguration;
16+
import org.openmetadata.schema.configuration.LLMDjlEmbeddingConfig;
1717

1818
@Slf4j
1919
public class DjlEmbeddingClient extends EmbeddingClient implements AutoCloseable {
@@ -34,15 +34,16 @@ public EmbeddingGenerationException(String message, Throwable cause) {
3434
private final String modelName;
3535
private final int dimension;
3636

37-
public DjlEmbeddingClient(ElasticSearchConfiguration config) {
37+
public DjlEmbeddingClient(LLMConfiguration config) {
3838
super(resolveMaxConcurrent(config));
39-
NaturalLanguageSearchConfiguration nlsCfg = config.getNaturalLanguageSearch();
40-
if (nlsCfg.getDjl() == null) {
39+
LLMDjlEmbeddingConfig djlCfg =
40+
config.getEmbeddings() != null ? config.getEmbeddings().getDjl() : null;
41+
if (djlCfg == null) {
4142
throw new IllegalArgumentException("DJL configuration is required");
4243
}
4344
this.modelName =
44-
nlsCfg.getDjl().getEmbeddingModel() != null
45-
? nlsCfg.getDjl().getEmbeddingModel()
45+
djlCfg.getEmbeddingModel() != null
46+
? djlCfg.getEmbeddingModel()
4647
: "ai.djl.huggingface.pytorch/sentence-transformers/all-MiniLM-L6-v2";
4748

4849
try {

‎openmetadata-service/src/main/java/org/openmetadata/service/search/vector/client/EmbeddingClient.java‎

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,8 @@
33
import java.util.ArrayList;
44
import java.util.List;
55
import java.util.concurrent.Semaphore;
6-
import org.openmetadata.schema.service.configuration.elasticsearch.ElasticSearchConfiguration;
7-
import org.openmetadata.schema.service.configuration.elasticsearch.NaturalLanguageSearchConfiguration;
6+
import org.openmetadata.schema.configuration.LLMConfiguration;
7+
import org.openmetadata.schema.configuration.LLMEmbeddingsConfig;
88

99
public abstract class EmbeddingClient {
1010
static final int DEFAULT_MAX_CONCURRENT_REQUESTS = 10;
@@ -52,14 +52,15 @@ public List<float[]> embedBatch(List<String> texts) {
5252

5353
public abstract String getModelId();
5454

55-
protected static int resolveMaxConcurrent(ElasticSearchConfiguration config) {
56-
NaturalLanguageSearchConfiguration nlsCfg = config.getNaturalLanguageSearch();
57-
if (nlsCfg != null) {
58-
Integer value = nlsCfg.getMaxConcurrentRequests();
55+
protected static int resolveMaxConcurrent(LLMConfiguration config) {
56+
int result = DEFAULT_MAX_CONCURRENT_REQUESTS;
57+
LLMEmbeddingsConfig embeddings = config != null ? config.getEmbeddings() : null;
58+
if (embeddings != null) {
59+
Integer value = embeddings.getMaxConcurrentRequests();
5960
if (value != null && value > 0) {
60-
return value;
61+
result = value;
6162
}
6263
}
63-
return DEFAULT_MAX_CONCURRENT_REQUESTS;
64+
return result;
6465
}
6566
}

0 commit comments

Comments
 (0)