|
| 1 | +/* |
| 2 | + * Copyright OpenSearch Contributors |
| 3 | + * SPDX-License-Identifier: Apache-2.0 |
| 4 | + */ |
| 5 | + |
| 6 | +package org.opensearch.ml.common; |
| 7 | + |
| 8 | +import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken; |
| 9 | +import static org.opensearch.ml.common.CommonValue.TENANT_ID_FIELD; |
| 10 | + |
| 11 | +import java.io.IOException; |
| 12 | +import java.time.Instant; |
| 13 | +import java.util.ArrayList; |
| 14 | +import java.util.List; |
| 15 | +import java.util.Map; |
| 16 | + |
| 17 | +import org.opensearch.core.common.io.stream.StreamInput; |
| 18 | +import org.opensearch.core.common.io.stream.StreamOutput; |
| 19 | +import org.opensearch.core.common.io.stream.Writeable; |
| 20 | +import org.opensearch.core.xcontent.ToXContentObject; |
| 21 | +import org.opensearch.core.xcontent.XContentBuilder; |
| 22 | +import org.opensearch.core.xcontent.XContentParser; |
| 23 | + |
| 24 | +import lombok.Builder; |
| 25 | +import lombok.EqualsAndHashCode; |
| 26 | +import lombok.Getter; |
| 27 | + |
| 28 | +/** |
| 29 | + * MLPrompt is the class to store prompt information. |
| 30 | + */ |
| 31 | +@Getter |
| 32 | +@EqualsAndHashCode |
| 33 | +public class MLPrompt implements ToXContentObject, Writeable { |
| 34 | + |
| 35 | + public static final String PROMPT_ID_FIELD = "prompt_id"; |
| 36 | + public static final String NAME_FIELD = "name"; |
| 37 | + public static final String DESCRIPTION_FIELD = "description"; |
| 38 | + public static final String VERSION_FIELD = "version"; |
| 39 | + public static final String PROMPT_FIELD = "prompt"; |
| 40 | + public static final String TAGS_FIELD = "tags"; |
| 41 | + public static final String CREATE_TIME_FIELD = "create_time"; |
| 42 | + public static final String LAST_UPDATE_TIME_FIELD = "last_update_time"; |
| 43 | + |
| 44 | + private String promptId; |
| 45 | + private String name; |
| 46 | + private String description; |
| 47 | + private String version; |
| 48 | + private Map<String, String> prompt; |
| 49 | + private List<String> tags; |
| 50 | + private String tenantId; |
| 51 | + private Instant createTime; |
| 52 | + private Instant lastUpdateTime; |
| 53 | + |
| 54 | + /** |
| 55 | + * Constructor to pass values to the MLPrompt constructor |
| 56 | + * |
| 57 | + * @param promptId The prompt id of the MLPrompt |
| 58 | + * @param name The name of the MLPrompt |
| 59 | + * @param description The description of the MLPrompt |
| 60 | + * @param version The version of the MLPrompt |
| 61 | + * @param prompt The prompt of the MLPrompt -> contains system and user prompts |
| 62 | + * @param tags The tags of the MLPrompt |
| 63 | + * @param tenantId The tenant id of the MLPrompt |
| 64 | + * @param createTime The create time of the MLPrompt |
| 65 | + * @param lastUpdateTime The last update time of the MLPrompt |
| 66 | + */ |
| 67 | + @Builder(toBuilder = true) |
| 68 | + public MLPrompt( |
| 69 | + String promptId, |
| 70 | + String name, |
| 71 | + String description, |
| 72 | + String version, |
| 73 | + Map<String, String> prompt, |
| 74 | + List<String> tags, |
| 75 | + String tenantId, |
| 76 | + Instant createTime, |
| 77 | + Instant lastUpdateTime |
| 78 | + ) { |
| 79 | + this.promptId = promptId; |
| 80 | + this.name = name; |
| 81 | + this.description = description; |
| 82 | + this.version = version; |
| 83 | + this.prompt = prompt; |
| 84 | + this.tags = tags; |
| 85 | + this.tenantId = tenantId; |
| 86 | + this.createTime = createTime; |
| 87 | + this.lastUpdateTime = lastUpdateTime; |
| 88 | + } |
| 89 | + |
| 90 | + /** |
| 91 | + * Deserialize the Stream Input and constructs MLPrompt |
| 92 | + * |
| 93 | + * @param input Abstract class that describes Stream Input |
| 94 | + * @throws IOException if an I/O exception occurred while reading from input stream |
| 95 | + */ |
| 96 | + public MLPrompt(StreamInput input) throws IOException { |
| 97 | + this.promptId = input.readOptionalString(); |
| 98 | + this.name = input.readOptionalString(); |
| 99 | + this.description = input.readOptionalString(); |
| 100 | + this.version = input.readOptionalString(); |
| 101 | + this.prompt = input.readMap(s -> s.readString(), s -> s.readString()); |
| 102 | + this.tags = input.readList(StreamInput::readString); |
| 103 | + this.tenantId = input.readOptionalString(); |
| 104 | + this.createTime = input.readOptionalInstant(); |
| 105 | + this.lastUpdateTime = input.readOptionalInstant(); |
| 106 | + } |
| 107 | + |
| 108 | + /** |
| 109 | + * Serialize and Writes the MLPrompt object to the output stream. |
| 110 | + * |
| 111 | + * @param out Abstract class that describes Stream Output |
| 112 | + * @throws IOException if an I/O exception occurred while writing to the output stream |
| 113 | + */ |
| 114 | + @Override |
| 115 | + public void writeTo(StreamOutput out) throws IOException { |
| 116 | + out.writeOptionalString(promptId); |
| 117 | + out.writeOptionalString(name); |
| 118 | + out.writeOptionalString(description); |
| 119 | + out.writeOptionalString(version); |
| 120 | + out.writeMap(prompt, StreamOutput::writeString, StreamOutput::writeString); |
| 121 | + out.writeCollection(tags, StreamOutput::writeString); |
| 122 | + out.writeOptionalString(tenantId); |
| 123 | + out.writeOptionalInstant(createTime); |
| 124 | + out.writeOptionalInstant(lastUpdateTime); |
| 125 | + } |
| 126 | + |
| 127 | + /** |
| 128 | + * Serialize and Writes the MLPrompt object to the XContentBuilder |
| 129 | + * |
| 130 | + * @param xContentBuilder XContentBuilder |
| 131 | + * @param params Parameters that need to be written to xContentBuilder |
| 132 | + * @return XContentBuilder |
| 133 | + * @throws IOException if an I/O exception occurred while writing field values to xContent |
| 134 | + */ |
| 135 | + @Override |
| 136 | + public XContentBuilder toXContent(XContentBuilder xContentBuilder, Params params) throws IOException { |
| 137 | + XContentBuilder builder = xContentBuilder.startObject(); |
| 138 | + if (promptId != null) { |
| 139 | + builder.field(PROMPT_ID_FIELD, promptId); |
| 140 | + } |
| 141 | + if (name != null) { |
| 142 | + builder.field(NAME_FIELD, name); |
| 143 | + } |
| 144 | + if (description != null) { |
| 145 | + builder.field(DESCRIPTION_FIELD, description); |
| 146 | + } |
| 147 | + if (version != null) { |
| 148 | + builder.field(VERSION_FIELD, version); |
| 149 | + } |
| 150 | + if (prompt != null) { |
| 151 | + builder.field(PROMPT_FIELD, prompt); |
| 152 | + } |
| 153 | + if (tags != null) { |
| 154 | + builder.field(TAGS_FIELD, tags); |
| 155 | + } |
| 156 | + if (tenantId != null) { |
| 157 | + builder.field(TENANT_ID_FIELD, tenantId); |
| 158 | + } |
| 159 | + if (createTime != null) { |
| 160 | + builder.field(CREATE_TIME_FIELD, createTime.toEpochMilli()); |
| 161 | + } |
| 162 | + if (lastUpdateTime != null) { |
| 163 | + builder.field(LAST_UPDATE_TIME_FIELD, lastUpdateTime.toEpochMilli()); |
| 164 | + } |
| 165 | + return builder.endObject(); |
| 166 | + } |
| 167 | + |
| 168 | + /** |
| 169 | + * Creates MLPrompt from stream input |
| 170 | + * |
| 171 | + * @param in Input Stream |
| 172 | + * @return MLPrompt |
| 173 | + * @throws IOException if an I/O exception occurred while reading from input stream |
| 174 | + */ |
| 175 | + public static MLPrompt fromStream(StreamInput in) throws IOException { |
| 176 | + return new MLPrompt(in); |
| 177 | + } |
| 178 | + |
| 179 | + /** |
| 180 | + * Creates MLPrompt from XContentParser |
| 181 | + * |
| 182 | + * @param parser XContentParser |
| 183 | + * @return MLPrompt |
| 184 | + * @throws IOException if an I/O exception occurred while parsing the XContentParser into MLPrompt fields |
| 185 | + */ |
| 186 | + public static MLPrompt parse(XContentParser parser) throws IOException { |
| 187 | + String promptId = null; |
| 188 | + String name = null; |
| 189 | + String description = null; |
| 190 | + String version = null; |
| 191 | + Map<String, String> prompt = null; |
| 192 | + List<String> tags = null; |
| 193 | + String tenantId = null; |
| 194 | + Instant createTime = null; |
| 195 | + Instant lastUpdateTime = null; |
| 196 | + |
| 197 | + ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.currentToken(), parser); |
| 198 | + while (parser.nextToken() != XContentParser.Token.END_OBJECT) { |
| 199 | + String fieldName = parser.currentName(); |
| 200 | + parser.nextToken(); |
| 201 | + switch (fieldName) { |
| 202 | + case PROMPT_ID_FIELD: |
| 203 | + promptId = parser.text(); |
| 204 | + break; |
| 205 | + case NAME_FIELD: |
| 206 | + name = parser.text(); |
| 207 | + break; |
| 208 | + case DESCRIPTION_FIELD: |
| 209 | + description = parser.text(); |
| 210 | + break; |
| 211 | + case VERSION_FIELD: |
| 212 | + version = parser.text(); |
| 213 | + break; |
| 214 | + case PROMPT_FIELD: |
| 215 | + prompt = parser.mapStrings(); |
| 216 | + break; |
| 217 | + case TAGS_FIELD: |
| 218 | + tags = new ArrayList<>(); |
| 219 | + ensureExpectedToken(XContentParser.Token.START_ARRAY, parser.currentToken(), parser); |
| 220 | + while (parser.nextToken() != XContentParser.Token.END_ARRAY) { |
| 221 | + tags.add(parser.text()); |
| 222 | + } |
| 223 | + break; |
| 224 | + case TENANT_ID_FIELD: |
| 225 | + tenantId = parser.text(); |
| 226 | + break; |
| 227 | + case CREATE_TIME_FIELD: |
| 228 | + createTime = Instant.ofEpochMilli(parser.longValue()); |
| 229 | + break; |
| 230 | + case LAST_UPDATE_TIME_FIELD: |
| 231 | + lastUpdateTime = Instant.ofEpochMilli(parser.longValue()); |
| 232 | + break; |
| 233 | + default: |
| 234 | + parser.skipChildren(); |
| 235 | + break; |
| 236 | + } |
| 237 | + } |
| 238 | + return MLPrompt |
| 239 | + .builder() |
| 240 | + .promptId(promptId) |
| 241 | + .name(name) |
| 242 | + .description(description) |
| 243 | + .version(version) |
| 244 | + .prompt(prompt) |
| 245 | + .tags(tags) |
| 246 | + .tenantId(tenantId) |
| 247 | + .createTime(createTime) |
| 248 | + .lastUpdateTime(lastUpdateTime) |
| 249 | + .build(); |
| 250 | + } |
| 251 | +} |
0 commit comments