From 790c470f257d210246ec513bd4728ba41934c698 Mon Sep 17 00:00:00 2001 From: miguel-vila Date: Fri, 31 May 2024 08:04:24 -0500 Subject: [PATCH] Add text embedding processor --- .../client/opensearch/ingest/Processor.java | 30 +++ .../opensearch/ingest/ProcessorBuilders.java | 8 + .../ingest/TextEmbeddingProcessor.java | 214 ++++++++++++++++++ 3 files changed, 252 insertions(+) create mode 100644 java-client/src/main/java/org/opensearch/client/opensearch/ingest/TextEmbeddingProcessor.java diff --git a/java-client/src/main/java/org/opensearch/client/opensearch/ingest/Processor.java b/java-client/src/main/java/org/opensearch/client/opensearch/ingest/Processor.java index c315527ee8..3bc8c0afea 100644 --- a/java-client/src/main/java/org/opensearch/client/opensearch/ingest/Processor.java +++ b/java-client/src/main/java/org/opensearch/client/opensearch/ingest/Processor.java @@ -127,6 +127,8 @@ public enum Kind implements JsonEnum { Inference("inference"), + TextEmbedding("text_embedding"), + ; private final String jsonValue; @@ -735,6 +737,23 @@ public InferenceProcessor inference() { return TaggedUnionUtils.get(this, Kind.Inference); } + /** + * Is this variant instance of kind {@code text_embedding}? + */ + public boolean isTextEmbedding() { + return _kind == Kind.TextEmbedding; + } + + /** + * Get the {@code text_embedding} variant value. + * + * @throws IllegalStateException + * if the current variant is not of the {@code text_embedding} kind. + */ + public TextEmbeddingProcessor textEmbedding() { + return TaggedUnionUtils.get(this, Kind.TextEmbedding); + } + @Override @SuppressWarnings("unchecked") public void serialize(JsonGenerator generator, JsonpMapper mapper) { @@ -1086,6 +1105,16 @@ public ObjectBuilder inference(Function textEmbedding(TextEmbeddingProcessor v) { + this._kind = Kind.TextEmbedding; + this._value = v; + return this; + } + + public ObjectBuilder textEmbedding(Function> fn) { + return this.textEmbedding(fn.apply(new TextEmbeddingProcessor.Builder()).build()); + } + public Processor build() { _checkSingleUse(); return new Processor(this); @@ -1128,6 +1157,7 @@ protected static void setupProcessorDeserializer(ObjectDeserializer op) op.add(Builder::drop, DropProcessor._DESERIALIZER, "drop"); op.add(Builder::circle, CircleProcessor._DESERIALIZER, "circle"); op.add(Builder::inference, InferenceProcessor._DESERIALIZER, "inference"); + op.add(Builder::textEmbedding, TextEmbeddingProcessor._DESERIALIZER, "text_embedding"); } diff --git a/java-client/src/main/java/org/opensearch/client/opensearch/ingest/ProcessorBuilders.java b/java-client/src/main/java/org/opensearch/client/opensearch/ingest/ProcessorBuilders.java index 85f95c4772..921ce75f6a 100644 --- a/java-client/src/main/java/org/opensearch/client/opensearch/ingest/ProcessorBuilders.java +++ b/java-client/src/main/java/org/opensearch/client/opensearch/ingest/ProcessorBuilders.java @@ -301,4 +301,12 @@ public static InferenceProcessor.Builder inference() { return new InferenceProcessor.Builder(); } + /** + * Creates a builder for the {@link TextEmbeddingProcessor text_embedding} + * {@code Processor} variant. + */ + public static TextEmbeddingProcessor.Builder textEmbedding() { + return new TextEmbeddingProcessor.Builder(); + } + } diff --git a/java-client/src/main/java/org/opensearch/client/opensearch/ingest/TextEmbeddingProcessor.java b/java-client/src/main/java/org/opensearch/client/opensearch/ingest/TextEmbeddingProcessor.java new file mode 100644 index 0000000000..ce3343f2b2 --- /dev/null +++ b/java-client/src/main/java/org/opensearch/client/opensearch/ingest/TextEmbeddingProcessor.java @@ -0,0 +1,214 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +/* + * Licensed to Elasticsearch B.V. under one or more contributor + * license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright + * ownership. Elasticsearch B.V. 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. + */ + +/* + * Modifications Copyright OpenSearch Contributors. See + * GitHub history for details. + */ + +package org.opensearch.client.opensearch.ingest; + +import jakarta.json.stream.JsonGenerator; +import java.util.Map; +import java.util.function.Function; +import javax.annotation.Nullable; +import org.opensearch.client.json.JsonData; +import org.opensearch.client.json.JsonpDeserializable; +import org.opensearch.client.json.JsonpDeserializer; +import org.opensearch.client.json.JsonpMapper; +import org.opensearch.client.json.ObjectBuilderDeserializer; +import org.opensearch.client.json.ObjectDeserializer; +import org.opensearch.client.util.ApiTypeHelper; +import org.opensearch.client.util.ObjectBuilder; + +// typedef: ingest._types.TextEmbeddingProcessor + +@JsonpDeserializable +public class TextEmbeddingProcessor extends ProcessorBase implements ProcessorVariant { + private final String modelId; + + private final Map fieldMap; + + @Nullable + private final String description; + + // --------------------------------------------------------------------------------------------- + + private TextEmbeddingProcessor(Builder builder) { + super(builder); + + this.modelId = ApiTypeHelper.requireNonNull(builder.modelId, this, "modelId"); + this.fieldMap = ApiTypeHelper.unmodifiableRequired(builder.fieldMap, this, "fieldMap"); + this.description = builder.description; + + } + + public static TextEmbeddingProcessor of(Function> fn) { + return fn.apply(new Builder()).build(); + } + + /** + * Processor variant kind. + */ + @Override + public Processor.Kind _processorKind() { + return Processor.Kind.Inference; + } + + /** + * Required - API name: {@code model_id} + */ + public final String modelId() { + return this.modelId; + } + + /** + * API name: {@code field_map} + */ + public final Map fieldMap() { + return this.fieldMap; + } + + /** + * API name: {@code description} + */ + @Nullable + public final String description() { + return this.description; + } + + protected void serializeInternal(JsonGenerator generator, JsonpMapper mapper) { + + super.serializeInternal(generator, mapper); + generator.writeKey("model_id"); + generator.write(this.modelId); + + if (ApiTypeHelper.isDefined(this.fieldMap)) { + generator.writeKey("field_map"); + generator.writeStartObject(); + for (Map.Entry item0 : this.fieldMap.entrySet()) { + generator.writeKey(item0.getKey()); + item0.getValue().serialize(generator, mapper); + + } + generator.writeEnd(); + + } + if (this.description != null) { + generator.writeKey("description"); + generator.write(this.description); + } + + } + + // --------------------------------------------------------------------------------------------- + + /** + * Builder for {@link TextEmbeddingProcessor}. + */ + + public static class Builder extends ProcessorBase.AbstractBuilder implements ObjectBuilder { + private String modelId; + + @Nullable + private Map fieldMap; + + @Nullable + private String description; + + /** + * Required - API name: {@code model_id} + */ + public final Builder modelId(String value) { + this.modelId = value; + return this; + } + + /** + * API name: {@code field_map} + *

+ * Adds all entries of map to fieldMap. + */ + public final Builder fieldMap(Map map) { + this.fieldMap = _mapPutAll(this.fieldMap, map); + return this; + } + + /** + * API name: {@code field_map} + *

+ * Adds an entry to fieldMap. + */ + public final Builder fieldMap(String key, JsonData value) { + this.fieldMap = _mapPut(this.fieldMap, key, value); + return this; + } + + /** + * API name: {@code description} + */ + public final Builder description(@Nullable String value) { + this.description = value; + return this; + } + + @Override + protected Builder self() { + return this; + } + + /** + * Builds a {@link TextEmbeddingProcessor}. + * + * @throws NullPointerException + * if some of the required fields are null. + */ + public TextEmbeddingProcessor build() { + _checkSingleUse(); + + return new TextEmbeddingProcessor(this); + } + } + + // --------------------------------------------------------------------------------------------- + + /** + * Json deserializer for {@link TextEmbeddingProcessor} + */ + public static final JsonpDeserializer _DESERIALIZER = ObjectBuilderDeserializer.lazy( + Builder::new, + TextEmbeddingProcessor::setupTextEmbeddingProcessorDeserializer + ); + + protected static void setupTextEmbeddingProcessorDeserializer(ObjectDeserializer op) { + ProcessorBase.setupProcessorBaseDeserializer(op); + op.add(Builder::modelId, JsonpDeserializer.stringDeserializer(), "model_id"); + op.add(Builder::fieldMap, JsonpDeserializer.stringMapDeserializer(JsonData._DESERIALIZER), "field_map"); + op.add(Builder::description, JsonpDeserializer.stringDeserializer(), "description"); + } + +}