This is an automated email from the ASF dual-hosted git repository. mkataria pushed a commit to branch OAK-11714 in repository https://gitbox.apache.org/repos/asf/jackrabbit-oak.git
commit 07ee6995edd028d3e536ca5506b76d0c8ab1c5fe Author: Mohit Kataria <[email protected]> AuthorDate: Sat May 10 00:11:24 2025 +0530 OAK-11714: Add jmx to expose inferenceConfig --- .../jackrabbit/oak/api/jmx/InferenceMBean.java | 40 +++ .../index/elastic/ElasticIndexProviderService.java | 13 + .../elastic/query/inference/EnricherStatus.java | 26 ++ .../elastic/query/inference/InferenceConfig.java | 57 +++- .../query/inference/InferenceHeaderPayload.java | 8 +- .../query/inference/InferenceIndexConfig.java | 27 +- .../query/inference/InferenceMBeanImpl.java | 49 ++++ .../query/inference/InferenceModelConfig.java | 30 +- .../elastic/query/inference/InferencePayload.java | 21 ++ .../InferenceConfigSerializationTest.java | 301 +++++++++++++++++++++ 10 files changed, 541 insertions(+), 31 deletions(-) diff --git a/oak-api/src/main/java/org/apache/jackrabbit/oak/api/jmx/InferenceMBean.java b/oak-api/src/main/java/org/apache/jackrabbit/oak/api/jmx/InferenceMBean.java new file mode 100644 index 0000000000..2690b6b64d --- /dev/null +++ b/oak-api/src/main/java/org/apache/jackrabbit/oak/api/jmx/InferenceMBean.java @@ -0,0 +1,40 @@ +/* + * 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.jackrabbit.oak.api.jmx; + +import org.osgi.annotation.versioning.ProviderType; + +/** + * An MBean that provides the inference configuration. + */ +@ProviderType +public interface InferenceMBean { + + String TYPE = "Inference"; + + /** + * Get the inference configuration as a Json string. + */ + String getConfigJson(); + + /** + * Get the inference configuration as a Json string. + */ + String getConfigNodeStateJson(); +} diff --git a/oak-search-elastic/src/main/java/org/apache/jackrabbit/oak/plugins/index/elastic/ElasticIndexProviderService.java b/oak-search-elastic/src/main/java/org/apache/jackrabbit/oak/plugins/index/elastic/ElasticIndexProviderService.java index 43565495dd..a53af22e10 100644 --- a/oak-search-elastic/src/main/java/org/apache/jackrabbit/oak/plugins/index/elastic/ElasticIndexProviderService.java +++ b/oak-search-elastic/src/main/java/org/apache/jackrabbit/oak/plugins/index/elastic/ElasticIndexProviderService.java @@ -16,6 +16,7 @@ */ package org.apache.jackrabbit.oak.plugins.index.elastic; +import org.apache.jackrabbit.oak.api.jmx.InferenceMBean; import org.apache.jackrabbit.oak.commons.IOUtils; import org.apache.jackrabbit.oak.osgi.OsgiWhiteboard; import org.apache.jackrabbit.oak.plugins.index.AsyncIndexInfoService; @@ -25,6 +26,7 @@ import org.apache.jackrabbit.oak.plugins.index.elastic.index.ElasticIndexEditorP import org.apache.jackrabbit.oak.plugins.index.elastic.query.ElasticIndexProvider; import org.apache.jackrabbit.oak.plugins.index.elastic.query.inference.InferenceConfig; import org.apache.jackrabbit.oak.plugins.index.elastic.query.inference.InferenceConstants; +import org.apache.jackrabbit.oak.plugins.index.elastic.query.inference.InferenceMBeanImpl; import org.apache.jackrabbit.oak.plugins.index.fulltext.PreExtractedTextProvider; import org.apache.jackrabbit.oak.plugins.index.search.ExtractedTextCache; import org.apache.jackrabbit.oak.query.QueryEngineSettings; @@ -209,6 +211,13 @@ public class ElasticIndexProviderService { ElasticIndexMBean.TYPE, "Elastic Index statistics")); + InferenceMBeanImpl inferenceMBean = new InferenceMBeanImpl(); + oakRegs.add(registerMBean(whiteboard, + InferenceMBean.class, + inferenceMBean, + InferenceMBean.TYPE, + "Inference")); + LOG.info("Registering Index and Editor providers with connection {}", elasticConnection); registerIndexProvider(bundleContext); @@ -284,4 +293,8 @@ public class ElasticIndexProviderService { .withApiKeys(apiKeyId, apiSecretId) .build(); } + + public InferenceConfig getInferenceConfig() { + return InferenceConfig.getInstance(); + } } diff --git a/oak-search-elastic/src/main/java/org/apache/jackrabbit/oak/plugins/index/elastic/query/inference/EnricherStatus.java b/oak-search-elastic/src/main/java/org/apache/jackrabbit/oak/plugins/index/elastic/query/inference/EnricherStatus.java index 678b3daaa7..77121f2419 100644 --- a/oak-search-elastic/src/main/java/org/apache/jackrabbit/oak/plugins/index/elastic/query/inference/EnricherStatus.java +++ b/oak-search-elastic/src/main/java/org/apache/jackrabbit/oak/plugins/index/elastic/query/inference/EnricherStatus.java @@ -18,9 +18,11 @@ */ package org.apache.jackrabbit.oak.plugins.index.elastic.query.inference; +import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.core.type.TypeReference; import com.fasterxml.jackson.databind.ObjectMapper; import org.apache.jackrabbit.oak.commons.PathUtils; +import org.apache.jackrabbit.oak.commons.json.JsopBuilder; import org.apache.jackrabbit.oak.spi.state.NodeState; import org.apache.jackrabbit.oak.spi.state.NodeStore; import org.slf4j.Logger; @@ -85,4 +87,28 @@ public class EnricherStatus { return enricherStatusJsonMapping; } + @Override + public String toString() { + JsopBuilder builder = new JsopBuilder().object(); + // Add the mapping data + builder.key("enricherStatusJsonMapping").value(enricherStatusJsonMapping); + + // Add enricher status data + builder.key("enricherStatusData").object(); + for (Map.Entry<String, Object> entry : enricherStatusData.entrySet()) { + builder.key(entry.getKey()); + if (entry.getValue() instanceof String) { + builder.value((String) entry.getValue()); + } else { + try { + builder.encodedValue(MAPPER.writeValueAsString(entry.getValue())); + } catch (JsonProcessingException e) { + LOG.warn("Failed to serialize value for key {}: {}", entry.getKey(), e.getMessage()); + builder.value(entry.getValue().toString()); + } + } + } + builder.endObject().endObject(); + return JsopBuilder.prettyPrint(builder.toString()); + } } \ No newline at end of file diff --git a/oak-search-elastic/src/main/java/org/apache/jackrabbit/oak/plugins/index/elastic/query/inference/InferenceConfig.java b/oak-search-elastic/src/main/java/org/apache/jackrabbit/oak/plugins/index/elastic/query/inference/InferenceConfig.java index 34730b7904..f5f251690a 100644 --- a/oak-search-elastic/src/main/java/org/apache/jackrabbit/oak/plugins/index/elastic/query/inference/InferenceConfig.java +++ b/oak-search-elastic/src/main/java/org/apache/jackrabbit/oak/plugins/index/elastic/query/inference/InferenceConfig.java @@ -18,7 +18,10 @@ */ package org.apache.jackrabbit.oak.plugins.index.elastic.query.inference; +import com.fasterxml.jackson.core.JsonProcessingException; import org.apache.jackrabbit.oak.commons.PathUtils; +import org.apache.jackrabbit.oak.commons.json.JsopBuilder; +import org.apache.jackrabbit.oak.json.JsonUtils; import org.apache.jackrabbit.oak.plugins.index.IndexName; import org.apache.jackrabbit.oak.spi.state.NodeState; import org.apache.jackrabbit.oak.spi.state.NodeStore; @@ -85,7 +88,7 @@ public class InferenceConfig { reInitialize(nodeStore, inferenceConfigPath, isInferenceEnabled, true); } - public static void reInitialize(){ + public static void reInitialize() { reInitialize(INSTANCE.nodeStore, INSTANCE.inferenceConfigPath, INSTANCE.isInferenceEnabled, true); } @@ -101,7 +104,7 @@ public class InferenceConfig { } } - private static void reInitialize(NodeStore nodeStore, String inferenceConfigPath, boolean isInferenceEnabled, boolean updateActiveInferenceConfig){ + private static void reInitialize(NodeStore nodeStore, String inferenceConfigPath, boolean isInferenceEnabled, boolean updateActiveInferenceConfig) { lock.writeLock().lock(); try { if (updateActiveInferenceConfig) { @@ -156,11 +159,11 @@ public class InferenceConfig { InferenceIndexConfig inferenceIndexConfig; IndexName indexNameObject; Function<String, InferenceIndexConfig> getInferenceIndexConfig = (iName) -> - getIndexConfigs().getOrDefault(iName, InferenceIndexConfig.NOOP); + getIndexConfigs().getOrDefault(iName, InferenceIndexConfig.NOOP); if (!InferenceIndexConfig.NOOP.equals(inferenceIndexConfig = getInferenceIndexConfig.apply(indexName))) { LOG.debug("InferenceIndexConfig for indexName: {} is: {}", indexName, inferenceIndexConfig); } else if ((indexNameObject = IndexName.parse(indexName)) != null && indexNameObject.isLegal() - && indexNameObject.getBaseName() != null + && indexNameObject.getBaseName() != null ) { LOG.debug("InferenceIndexConfig is using baseIndexName {} and is: {}", indexNameObject.getBaseName(), inferenceIndexConfig); inferenceIndexConfig = getInferenceIndexConfig.apply(indexNameObject.getBaseName()); @@ -175,7 +178,7 @@ public class InferenceConfig { public @NotNull InferenceModelConfig getInferenceModelConfig(String inferenceIndexName, String inferenceModelConfigName) { lock.readLock().lock(); try { - if (inferenceModelConfigName == null){ + if (inferenceModelConfigName == null) { return InferenceModelConfig.NOOP; } else if (inferenceModelConfigName.isEmpty()) { return getInferenceIndexConfig(inferenceIndexName).getDefaultEnabledModel(); @@ -188,7 +191,7 @@ public class InferenceConfig { } - public Map<String, Object> getEnricherStatus(){ + public Map<String, Object> getEnricherStatus() { lock.readLock().lock(); try { return INSTANCE.enricherStatus.getEnricherStatus(); @@ -197,7 +200,7 @@ public class InferenceConfig { } } - public String getEnricherStatusMapping(){ + public String getEnricherStatusMapping() { lock.readLock().lock(); try { return INSTANCE.enricherStatus.getEnricherStatusJsonMapping(); @@ -206,11 +209,32 @@ public class InferenceConfig { } } + public String getInferenceConfigNodeState() { + if (nodeStore != null) { + NodeState ns = nodeStore.getRoot(); + for (String elem : PathUtils.elements(inferenceConfigPath)) { + ns = ns.getChildNode(elem); + } + if (!ns.exists()) { + LOG.warn("InferenceConfig: NodeState does not exist for path: " + inferenceConfigPath); + return "{}"; + } + try { + return JsonUtils.nodeStateToJson(ns, 5); + } catch (JsonProcessingException e) { + throw new RuntimeException(e); + } + } else { + LOG.warn("InferenceConfig: NodeStore is null"); + return "{}"; + } + } + private @NotNull Map<String, InferenceIndexConfig> getIndexConfigs() { lock.readLock().lock(); try { return isEnabled() ? - Collections.unmodifiableMap(indexConfigs) : Map.of(); + Collections.unmodifiableMap(indexConfigs) : Map.of(); } finally { lock.readLock().unlock(); } @@ -241,4 +265,21 @@ public class InferenceConfig { return UUID.randomUUID().toString(); } + @Override + public String toString() { + JsopBuilder builder = new JsopBuilder().object(). + key("type").value(TYPE). + key("enabled").value(enabled). + key("inferenceConfigPath").value(inferenceConfigPath). + key("currentInferenceConfig").value(currentInferenceConfig). + key("activeInferenceConfig").value(activeInferenceConfig). + key("isInferenceEnabled").value(isInferenceEnabled). + key("indexConfigs").object(); + // Serialize each index config + for (Map.Entry<String, InferenceIndexConfig> e : indexConfigs.entrySet()) { + builder.key(e.getKey()).encodedValue(e.getValue().toString()); + } + builder.endObject().endObject(); + return JsopBuilder.prettyPrint(builder.toString()); + } } \ No newline at end of file diff --git a/oak-search-elastic/src/main/java/org/apache/jackrabbit/oak/plugins/index/elastic/query/inference/InferenceHeaderPayload.java b/oak-search-elastic/src/main/java/org/apache/jackrabbit/oak/plugins/index/elastic/query/inference/InferenceHeaderPayload.java index 53c387e20d..36e5e8c618 100644 --- a/oak-search-elastic/src/main/java/org/apache/jackrabbit/oak/plugins/index/elastic/query/inference/InferenceHeaderPayload.java +++ b/oak-search-elastic/src/main/java/org/apache/jackrabbit/oak/plugins/index/elastic/query/inference/InferenceHeaderPayload.java @@ -18,6 +18,7 @@ */ package org.apache.jackrabbit.oak.plugins.index.elastic.query.inference; +import org.apache.jackrabbit.oak.commons.json.JsopBuilder; import org.apache.jackrabbit.oak.json.JsonUtils; import org.apache.jackrabbit.oak.plugins.index.elastic.util.EnvironmentVariableProcessorUtil; import org.apache.jackrabbit.oak.spi.state.NodeState; @@ -64,7 +65,12 @@ public class InferenceHeaderPayload { @Override public String toString() { - return inferenceHeaderPayloadMap.toString(); + JsopBuilder builder = new JsopBuilder().object(); + for (Map.Entry<String, String> entry : inferenceHeaderPayloadMap.entrySet()) { + builder.key(entry.getKey()).value(entry.getValue()); + } + builder.endObject(); + return JsopBuilder.prettyPrint(builder.toString()); } } \ No newline at end of file diff --git a/oak-search-elastic/src/main/java/org/apache/jackrabbit/oak/plugins/index/elastic/query/inference/InferenceIndexConfig.java b/oak-search-elastic/src/main/java/org/apache/jackrabbit/oak/plugins/index/elastic/query/inference/InferenceIndexConfig.java index a7243655ed..5a2b78bb9d 100644 --- a/oak-search-elastic/src/main/java/org/apache/jackrabbit/oak/plugins/index/elastic/query/inference/InferenceIndexConfig.java +++ b/oak-search-elastic/src/main/java/org/apache/jackrabbit/oak/plugins/index/elastic/query/inference/InferenceIndexConfig.java @@ -18,6 +18,7 @@ */ package org.apache.jackrabbit.oak.plugins.index.elastic.query.inference; +import org.apache.jackrabbit.oak.commons.json.JsopBuilder; import org.apache.jackrabbit.oak.json.JsonUtils; import org.apache.jackrabbit.oak.spi.state.NodeState; import org.slf4j.Logger; @@ -79,7 +80,7 @@ public class InferenceIndexConfig { this.enricherConfig = getOptionalValue(nodeState, InferenceConstants.ENRICHER_CONFIG, DISABLED_ENRICHER_CONFIG); inferenceModelConfigs = Map.of(); LOG.warn("inference index config for indexName: {} is not valid. Node: {}", - indexName, nodeState); + indexName, nodeState); } } @@ -108,18 +109,26 @@ public class InferenceIndexConfig { */ public InferenceModelConfig getDefaultEnabledModel() { return inferenceModelConfigs.values().stream() - .filter(InferenceModelConfig::isDefault) - .filter(InferenceModelConfig::isEnabled) - .findFirst() - .orElse(InferenceModelConfig.NOOP); + .filter(InferenceModelConfig::isDefault) + .filter(InferenceModelConfig::isEnabled) + .findFirst() + .orElse(InferenceModelConfig.NOOP); } @Override public String toString() { - return TYPE + "{" + - ENRICHER_CONFIG + "='" + enricherConfig + '\'' + - ", " + InferenceModelConfig.TYPE + "=" + inferenceModelConfigs + - '}'; + JsopBuilder builder = new JsopBuilder().object(). + key("type").value(TYPE). + key(ENRICHER_CONFIG).value(enricherConfig). + key(InferenceConstants.ENABLED).value(isEnabled). + key("inferenceModelConfigs").object(); + + // Serialize each model config + for (Map.Entry<String, InferenceModelConfig> e : inferenceModelConfigs.entrySet()) { + builder.key(e.getKey()).encodedValue(e.getValue().toString()); + } + builder.endObject().endObject(); + return JsopBuilder.prettyPrint(builder.toString()); } } \ No newline at end of file diff --git a/oak-search-elastic/src/main/java/org/apache/jackrabbit/oak/plugins/index/elastic/query/inference/InferenceMBeanImpl.java b/oak-search-elastic/src/main/java/org/apache/jackrabbit/oak/plugins/index/elastic/query/inference/InferenceMBeanImpl.java new file mode 100644 index 0000000000..bfd9f5f6fc --- /dev/null +++ b/oak-search-elastic/src/main/java/org/apache/jackrabbit/oak/plugins/index/elastic/query/inference/InferenceMBeanImpl.java @@ -0,0 +1,49 @@ +/* + * 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.jackrabbit.oak.plugins.index.elastic.query.inference; + +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.ObjectMapper; +import org.apache.jackrabbit.oak.api.jmx.InferenceMBean; +import org.apache.jackrabbit.oak.commons.jmx.AnnotatedStandardMBean; +import org.apache.jackrabbit.oak.plugins.index.elastic.ElasticIndexProviderService; +import org.jetbrains.annotations.NotNull; + +import java.util.Objects; + +/** + * An MBean that provides the inference configuration. + */ +public class InferenceMBeanImpl extends AnnotatedStandardMBean implements InferenceMBean { + private static final ObjectMapper MAPPER = new ObjectMapper(); + + public InferenceMBeanImpl() { + super(InferenceMBean.class); + } + + @Override + public String getConfigJson() { + return InferenceConfig.getInstance().toString(); + } + + @Override + public String getConfigNodeStateJson() { + return InferenceConfig.getInstance().getInferenceConfigNodeState(); + } +} diff --git a/oak-search-elastic/src/main/java/org/apache/jackrabbit/oak/plugins/index/elastic/query/inference/InferenceModelConfig.java b/oak-search-elastic/src/main/java/org/apache/jackrabbit/oak/plugins/index/elastic/query/inference/InferenceModelConfig.java index b4f79d0a0f..04a825acd8 100644 --- a/oak-search-elastic/src/main/java/org/apache/jackrabbit/oak/plugins/index/elastic/query/inference/InferenceModelConfig.java +++ b/oak-search-elastic/src/main/java/org/apache/jackrabbit/oak/plugins/index/elastic/query/inference/InferenceModelConfig.java @@ -18,6 +18,7 @@ */ package org.apache.jackrabbit.oak.plugins.index.elastic.query.inference; +import org.apache.jackrabbit.oak.commons.json.JsopBuilder; import org.apache.jackrabbit.oak.plugins.index.elastic.util.EnvironmentVariableProcessorUtil; import org.apache.jackrabbit.oak.spi.query.fulltext.VectorQueryConfig; import org.apache.jackrabbit.oak.spi.state.NodeState; @@ -102,7 +103,7 @@ public class InferenceModelConfig { this.isDefault = getOptionalValue(nodeState, IS_DEFAULT, false); this.model = getOptionalValue(nodeState, MODEL, ""); this.embeddingServiceUrl = EnvironmentVariableProcessorUtil.processEnvironmentVariable( - InferenceConstants.INFERENCE_ENVIRONMENT_VARIABLE_PREFIX, getOptionalValue(nodeState, EMBEDDING_SERVICE_URL, ""), DEFAULT_ENVIRONMENT_VARIABLE_VALUE); + InferenceConstants.INFERENCE_ENVIRONMENT_VARIABLE_PREFIX, getOptionalValue(nodeState, EMBEDDING_SERVICE_URL, ""), DEFAULT_ENVIRONMENT_VARIABLE_VALUE); this.similarityThreshold = getOptionalValue(nodeState, SIMILARITY_THRESHOLD, DEFAULT_SIMILARITY_THRESHOLD); this.minTerms = getOptionalValue(nodeState, MIN_TERMS, DEFAULT_MIN_TERMS); this.timeout = getOptionalValue(nodeState, TIMEOUT, DEFAULT_TIMEOUT_MILLIS); @@ -112,18 +113,21 @@ public class InferenceModelConfig { @Override public String toString() { - return TYPE + "{" + - MODEL + "='" + model + '\'' + - ", " + EMBEDDING_SERVICE_URL + "='" + embeddingServiceUrl + '\'' + - ", " + SIMILARITY_THRESHOLD + similarityThreshold + - ", " + MIN_TERMS + "=" + minTerms + - ", " + IS_DEFAULT + "=" + isDefault + - ", " + ENABLED + "=" + enabled + - ", " + HEADER + "=" + header + - ", " + INFERENCE_PAYLOAD + "=" + payload + - ", " + TIMEOUT + "=" + timeout + - ", " + NUM_CANDIDATES + "=" + numCandidates + - "}"; + JsopBuilder builder = new JsopBuilder().object(). + key("type").value(TYPE). + key(MODEL).value(model). + key(EMBEDDING_SERVICE_URL).value(embeddingServiceUrl). + key(SIMILARITY_THRESHOLD).encodedValue("" + similarityThreshold). + key(MIN_TERMS).value(minTerms). + key(IS_DEFAULT).value(isDefault). + key(ENABLED).value(enabled). + key(HEADER).encodedValue(header.toString()). + key(INFERENCE_PAYLOAD).encodedValue(payload.toString()). + key(TIMEOUT).value(timeout). + key(NUM_CANDIDATES).value(numCandidates). + key(CACHE_SIZE).value(cacheSize); + builder.endObject(); + return JsopBuilder.prettyPrint(builder.toString()); } public String getInferenceModelConfigName() { diff --git a/oak-search-elastic/src/main/java/org/apache/jackrabbit/oak/plugins/index/elastic/query/inference/InferencePayload.java b/oak-search-elastic/src/main/java/org/apache/jackrabbit/oak/plugins/index/elastic/query/inference/InferencePayload.java index ac230014a9..93a6065581 100644 --- a/oak-search-elastic/src/main/java/org/apache/jackrabbit/oak/plugins/index/elastic/query/inference/InferencePayload.java +++ b/oak-search-elastic/src/main/java/org/apache/jackrabbit/oak/plugins/index/elastic/query/inference/InferencePayload.java @@ -20,6 +20,7 @@ package org.apache.jackrabbit.oak.plugins.index.elastic.query.inference; import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.databind.ObjectMapper; +import org.apache.jackrabbit.oak.commons.json.JsopBuilder; import org.apache.jackrabbit.oak.json.JsonUtils; import org.apache.jackrabbit.oak.plugins.index.elastic.util.EnvironmentVariableProcessorUtil; import org.apache.jackrabbit.oak.spi.state.NodeState; @@ -58,6 +59,7 @@ public class InferencePayload { //replace current keys with swapped inferencePayloadMap.putAll(swappedEnvVarsMap); } + /* * Get the inference payload as a json string * @@ -76,4 +78,23 @@ public class InferencePayload { } } + @Override + public String toString() { + JsopBuilder builder = new JsopBuilder().object(); + for (Map.Entry<String, Object> entry : inferencePayloadMap.entrySet()) { + builder.key(entry.getKey()); + if (entry.getValue() instanceof String) { + builder.value((String) entry.getValue()); + } else { + try { + builder.encodedValue(objectMapper.writeValueAsString(entry.getValue())); + } catch (JsonProcessingException e) { + LOG.warn("Failed to serialize value for key {}: {}", entry.getKey(), e.getMessage()); + builder.value(entry.getValue().toString()); + } + } + } + builder.endObject(); + return JsopBuilder.prettyPrint(builder.toString()); + } } \ No newline at end of file diff --git a/oak-search-elastic/src/test/java/org/apache/jackrabbit/oak/plugins/index/elastic/query/inference/InferenceConfigSerializationTest.java b/oak-search-elastic/src/test/java/org/apache/jackrabbit/oak/plugins/index/elastic/query/inference/InferenceConfigSerializationTest.java new file mode 100644 index 0000000000..e68bd5f81c --- /dev/null +++ b/oak-search-elastic/src/test/java/org/apache/jackrabbit/oak/plugins/index/elastic/query/inference/InferenceConfigSerializationTest.java @@ -0,0 +1,301 @@ +/* + * 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.jackrabbit.oak.plugins.index.elastic.query.inference; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import org.apache.jackrabbit.oak.plugins.memory.EmptyNodeState; +import org.apache.jackrabbit.oak.plugins.memory.MemoryNodeBuilder; +import org.apache.jackrabbit.oak.plugins.memory.MemoryNodeStore; +import org.apache.jackrabbit.oak.spi.commit.CommitInfo; +import org.apache.jackrabbit.oak.spi.commit.EmptyHook; +import org.apache.jackrabbit.oak.spi.state.NodeBuilder; +import org.apache.jackrabbit.oak.spi.state.NodeStore; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; + +/** + * Tests for the toString() methods of the inference-related classes which use JsopBuilder + */ +public class InferenceConfigSerializationTest { + + private static final ObjectMapper MAPPER = new ObjectMapper(); + private static final String DEFAULT_CONFIG_PATH = InferenceConstants.DEFAULT_OAK_INDEX_INFERENCE_CONFIG_PATH; + private static final String ENRICHER_CONFIG = "{\"enricher\":{\"config\":{\"vectorSpaces\":{\"semantic\":{\"pipeline\":{\"steps\":[{\"inputFields\":{\"description\":\"STRING\",\"title\":\"STRING\"},\"chunkingConfig\":{\"enabled\":true},\"name\":\"sentence-embeddings\",\"model\":\"text-embedding-ada-002\",\"optional\":true,\"type\":\"embeddings\"}]},\"default\":false}},\"version\":\"0.0.1\"}}}"; + private static final String DEFAULT_ENRICHER_STATUS_MAPPING = "{\"properties\":{\"processingTimeMs\":{\"type\":\"date\"},\"latestError\":{\"type\":\"keyword\",\"index\":false},\"errorCount\":{\"type\":\"short\"},\"status\":{\"type\":\"keyword\"}}}"; + private static final String DEFAULT_ENRICHER_STATUS_DATA = "{\"processingTimeMs\":0,\"latestError\":\"\",\"errorCount\":0,\"status\":\"PENDING\"}"; + + private NodeBuilder rootBuilder; + private NodeStore nodeStore; + + @Before + public void setup() { + // Initialize memory node store + rootBuilder = new MemoryNodeBuilder(EmptyNodeState.EMPTY_NODE); + nodeStore = new MemoryNodeStore(rootBuilder.getNodeState()); + } + + @After + public void tearDown() { + rootBuilder = null; + nodeStore = null; + } + + /** + * Test for InferenceConfig.toString() + */ + @Test + public void testInferenceConfigToString() throws Exception { + // Setup: Create a basic inference config + NodeBuilder inferenceConfigBuilder = createNodePath(rootBuilder, DEFAULT_CONFIG_PATH); + inferenceConfigBuilder.setProperty(InferenceConstants.TYPE, InferenceConfig.TYPE); + inferenceConfigBuilder.setProperty(InferenceConstants.ENABLED, true); + + // Add index config + String indexName = "testIndex"; + NodeBuilder indexConfigBuilder = inferenceConfigBuilder.child(indexName); + indexConfigBuilder.setProperty(InferenceConstants.TYPE, InferenceIndexConfig.TYPE); + indexConfigBuilder.setProperty(InferenceConstants.ENABLED, true); + indexConfigBuilder.setProperty(InferenceConstants.ENRICHER_CONFIG, ENRICHER_CONFIG); + + // Commit the changes + nodeStore.merge(rootBuilder, EmptyHook.INSTANCE, CommitInfo.EMPTY); + + // Initialize the inference config + InferenceConfig.reInitialize(nodeStore, DEFAULT_CONFIG_PATH, true); + InferenceConfig inferenceConfig = InferenceConfig.getInstance(); + + // Get the toString representation + String json = inferenceConfig.toString(); + + // Verify it's valid JSON + JsonNode node = MAPPER.readTree(json); + + // Verify the structure + assertTrue("JSON should contain 'type' key", node.has("type")); + assertEquals("Type should be inferenceConfig", InferenceConfig.TYPE, node.get("type").asText()); + assertTrue("JSON should contain 'enabled' key", node.has("enabled")); + assertTrue("enabled should be true", node.get("enabled").asBoolean()); + assertTrue("JSON should contain 'indexConfigs' key", node.has("indexConfigs")); + assertTrue("indexConfigs should be an object", node.get("indexConfigs").isObject()); + assertTrue("indexConfigs should contain testIndex", node.get("indexConfigs").has(indexName)); + } + + /** + * Test for InferenceIndexConfig.toString() + */ + @Test + public void testInferenceIndexConfigToString() throws Exception { + // Create a simple index config + NodeBuilder indexConfigBuilder = rootBuilder.child("testIndex"); + indexConfigBuilder.setProperty(InferenceConstants.TYPE, InferenceIndexConfig.TYPE); + indexConfigBuilder.setProperty(InferenceConstants.ENABLED, true); + indexConfigBuilder.setProperty(InferenceConstants.ENRICHER_CONFIG, ENRICHER_CONFIG); + + // Create the index config object + InferenceIndexConfig indexConfig = new InferenceIndexConfig("testIndex", indexConfigBuilder.getNodeState()); + + // Get the toString representation + String json = indexConfig.toString(); + + // Verify it's valid JSON + JsonNode node = MAPPER.readTree(json); + + // Verify the structure + assertTrue("JSON should contain 'type' key", node.has("type")); + assertEquals("Type should be inferenceIndexConfig", InferenceIndexConfig.TYPE, node.get("type").asText()); + assertTrue("JSON should contain 'enricherConfig' key", node.has(InferenceIndexConfig.ENRICHER_CONFIG)); + assertEquals("Enricher config should match", ENRICHER_CONFIG, node.get(InferenceIndexConfig.ENRICHER_CONFIG).asText()); + assertTrue("JSON should contain 'enabled' key", node.has(InferenceConstants.ENABLED)); + assertTrue("enabled should be true", node.get(InferenceConstants.ENABLED).asBoolean()); + assertTrue("JSON should contain 'inferenceModelConfigs' key", node.has("inferenceModelConfigs")); + assertTrue("inferenceModelConfigs should be an object", node.get("inferenceModelConfigs").isObject()); + } + + /** + * Test for InferenceModelConfig.toString() + */ + @Test + public void testInferenceModelConfigToString() throws Exception { + // Create a model config with header and payload + NodeBuilder modelConfigBuilder = rootBuilder.child("testModel"); + modelConfigBuilder.setProperty(InferenceConstants.TYPE, InferenceModelConfig.TYPE); + modelConfigBuilder.setProperty(InferenceConstants.ENABLED, true); + modelConfigBuilder.setProperty(InferenceModelConfig.IS_DEFAULT, true); + modelConfigBuilder.setProperty(InferenceModelConfig.MODEL, "test-model"); + modelConfigBuilder.setProperty(InferenceModelConfig.EMBEDDING_SERVICE_URL, "http://test-service";); + modelConfigBuilder.setProperty(InferenceModelConfig.SIMILARITY_THRESHOLD, 0.85); + modelConfigBuilder.setProperty(InferenceModelConfig.MIN_TERMS, 3); + modelConfigBuilder.setProperty(InferenceModelConfig.TIMEOUT, 10000); + modelConfigBuilder.setProperty(InferenceModelConfig.NUM_CANDIDATES, 50); + modelConfigBuilder.setProperty(InferenceModelConfig.CACHE_SIZE, 200); + + // Create header node + NodeBuilder headerBuilder = modelConfigBuilder.child(InferenceModelConfig.HEADER); + headerBuilder.setProperty("Authorization", "Bearer test-token"); + headerBuilder.setProperty("Content-Type", "application/json"); + + // Create payload node + NodeBuilder payloadBuilder = modelConfigBuilder.child(InferenceModelConfig.INFERENCE_PAYLOAD); + payloadBuilder.setProperty("model", "text-embedding-ada-002"); + payloadBuilder.setProperty("dimensions", 1536); + + // Create the model config object + InferenceModelConfig modelConfig = new InferenceModelConfig("testModel", modelConfigBuilder.getNodeState()); + + // Get the toString representation + String json = modelConfig.toString(); + + // Verify it's valid JSON + JsonNode node = MAPPER.readTree(json); + + // Verify structure + assertTrue("JSON should contain 'TYPE' key", node.has("type")); + assertEquals("Type should match", InferenceModelConfig.TYPE, node.get("type").asText()); + assertTrue("JSON should contain 'model' key", node.has(InferenceModelConfig.MODEL)); + assertEquals("Model should match", "test-model", node.get(InferenceModelConfig.MODEL).asText()); + assertTrue("JSON should contain 'embeddingServiceUrl' key", node.has(InferenceModelConfig.EMBEDDING_SERVICE_URL)); + assertEquals("Service URL should match", "http://test-service";, node.get(InferenceModelConfig.EMBEDDING_SERVICE_URL).asText()); + assertTrue("JSON should contain 'similarityThreshold' key", node.has(InferenceModelConfig.SIMILARITY_THRESHOLD)); + assertEquals("Similarity threshold should match", 0.85, node.get(InferenceModelConfig.SIMILARITY_THRESHOLD).asDouble(), 0.001); + assertTrue("JSON should contain 'minTerms' key", node.has(InferenceModelConfig.MIN_TERMS)); + assertEquals("Min terms should match", 3, node.get(InferenceModelConfig.MIN_TERMS).asInt()); + assertTrue("JSON should contain 'isDefault' key", node.has(InferenceModelConfig.IS_DEFAULT)); + assertTrue("isDefault should be true", node.get(InferenceModelConfig.IS_DEFAULT).asBoolean()); + assertTrue("JSON should contain 'enabled' key", node.has(InferenceModelConfig.ENABLED)); + assertTrue("enabled should be true", node.get(InferenceModelConfig.ENABLED).asBoolean()); + assertTrue("JSON should contain 'header' key", node.has(InferenceModelConfig.HEADER)); + assertTrue("JSON should contain 'inferencePayload' key", node.has(InferenceModelConfig.INFERENCE_PAYLOAD)); + assertTrue("JSON should contain 'timeout' key", node.has(InferenceModelConfig.TIMEOUT)); + assertEquals("Timeout should match", 10000, node.get(InferenceModelConfig.TIMEOUT).asInt()); + assertTrue("JSON should contain 'numCandidates' key", node.has(InferenceModelConfig.NUM_CANDIDATES)); + assertEquals("Num candidates should match", 50, node.get(InferenceModelConfig.NUM_CANDIDATES).asInt()); + assertTrue("JSON should contain 'cacheSize' key", node.has(InferenceModelConfig.CACHE_SIZE)); + assertEquals("Cache size should match", 200, node.get(InferenceModelConfig.CACHE_SIZE).asInt()); + } + + /** + * Test for InferenceHeaderPayload.toString() + */ + @Test + public void testInferenceHeaderPayloadToString() throws Exception { + // Create a header payload + NodeBuilder headerBuilder = rootBuilder.child("header"); + headerBuilder.setProperty("Authorization", "Bearer test-token"); + headerBuilder.setProperty("Content-Type", "application/json"); + + // Create the header payload object + InferenceHeaderPayload headerPayload = new InferenceHeaderPayload(headerBuilder.getNodeState()); + + // Get the toString representation + String json = headerPayload.toString(); + + // Verify it's valid JSON + JsonNode node = MAPPER.readTree(json); + + // Verify structure + assertTrue("JSON should contain Authorization", node.has("Authorization")); + assertEquals("Authorization should match", "Bearer test-token", node.get("Authorization").asText()); + assertTrue("JSON should contain Content-Type", node.has("Content-Type")); + assertEquals("Content-Type should match", "application/json", node.get("Content-Type").asText()); + } + + /** + * Test for InferencePayload.toString() + */ + @Test + public void testInferencePayloadToString() throws Exception { + // Create a payload + NodeBuilder payloadBuilder = rootBuilder.child("payload"); + payloadBuilder.setProperty("model", "text-embedding-ada-002"); + payloadBuilder.setProperty("dimensions", 1536); + + // Create the payload object + InferencePayload payload = new InferencePayload("testModel", payloadBuilder.getNodeState()); + + // Get the toString representation + String json = payload.toString(); + + // Verify it's valid JSON + JsonNode node = MAPPER.readTree(json); + + // Verify structure + assertTrue("JSON should contain model", node.has("model")); + assertEquals("Model should match", "text-embedding-ada-002", node.get("model").asText()); + assertTrue("JSON should contain dimensions", node.has("dimensions")); + assertEquals("Dimensions should match", 1536, node.get("dimensions").asInt()); + } + + /** + * Test for EnricherStatus.toString() + */ + @Test + public void testEnricherStatusToString() throws Exception { + // Setup: Create a node structure with enricher status data + NodeBuilder inferenceConfigBuilder = createNodePath(rootBuilder, DEFAULT_CONFIG_PATH); + NodeBuilder enrichNode = inferenceConfigBuilder.child(InferenceConstants.ENRICH_NODE); + enrichNode.setProperty(InferenceConstants.ENRICHER_STATUS_MAPPING, DEFAULT_ENRICHER_STATUS_MAPPING); + enrichNode.setProperty(InferenceConstants.ENRICHER_STATUS_DATA, DEFAULT_ENRICHER_STATUS_DATA); + + // Commit the changes + nodeStore.merge(rootBuilder, EmptyHook.INSTANCE, CommitInfo.EMPTY); + + // Create the enricher status object + EnricherStatus status = new EnricherStatus(nodeStore, DEFAULT_CONFIG_PATH); + + // Get the toString representation + String json = status.toString(); + + // Verify it's valid JSON + JsonNode node = MAPPER.readTree(json); + + // Verify structure + assertTrue("JSON should contain enricherStatusJsonMapping", node.has("enricherStatusJsonMapping")); + JsonNode mappingNode = MAPPER.readTree(node.get("enricherStatusJsonMapping").asText()); + assertTrue("Mapping should contain properties", mappingNode.has("properties")); + + assertTrue("JSON should contain enricherStatusData", node.has("enricherStatusData")); + JsonNode statusData = node.get("enricherStatusData"); + assertTrue("Status data should contain processingTimeMs", statusData.has("processingTimeMs")); + assertEquals("Processing time should be 0", 0, statusData.get("processingTimeMs").asInt()); + assertTrue("Status data should contain status", statusData.has("status")); + assertEquals("Status should be PENDING", "PENDING", statusData.get("status").asText()); + assertTrue("Status data should contain errorCount", statusData.has("errorCount")); + assertEquals("Error count should be 0", 0, statusData.get("errorCount").asInt()); + assertTrue("Status data should contain latestError", statusData.has("latestError")); + assertEquals("Latest error should be empty", "", statusData.get("latestError").asText()); + } + + /** + * Helper method to create node paths + */ + private NodeBuilder createNodePath(NodeBuilder rootBuilder, String path) { + NodeBuilder currentBuilder = rootBuilder; + for (String element : path.split("/")) { + if (!element.isEmpty()) { + currentBuilder = currentBuilder.child(element); + } + } + return currentBuilder; + } +} \ No newline at end of file
