Skip to content

Commit

Permalink
agent meta classes in common (#1757)
Browse files Browse the repository at this point in the history
* agent meta classes

Signed-off-by: Jing Zhang <[email protected]>

* address comments

Signed-off-by: Jing Zhang <[email protected]>

* add one more UT

Signed-off-by: Jing Zhang <[email protected]>

* add more UT

Signed-off-by: Jing Zhang <[email protected]>

* more UT

Signed-off-by: Jing Zhang <[email protected]>

---------

Signed-off-by: Jing Zhang <[email protected]>
  • Loading branch information
jngz-es authored Dec 13, 2023
1 parent 01cf00c commit 3949ef9
Show file tree
Hide file tree
Showing 9 changed files with 1,111 additions and 0 deletions.
116 changes: 116 additions & 0 deletions common/src/main/java/org/opensearch/ml/common/CommonValue.java
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,25 @@

package org.opensearch.ml.common;

import org.opensearch.ml.common.agent.MLAgent;
import org.opensearch.ml.common.connector.AbstractConnector;

import static org.opensearch.ml.common.conversation.ConversationalIndexConstants.APPLICATION_TYPE_FIELD;
import static org.opensearch.ml.common.conversation.ConversationalIndexConstants.INTERACTIONS_ADDITIONAL_INFO_FIELD;
import static org.opensearch.ml.common.conversation.ConversationalIndexConstants.INTERACTIONS_CONVERSATION_ID_FIELD;
import static org.opensearch.ml.common.conversation.ConversationalIndexConstants.INTERACTIONS_CREATE_TIME_FIELD;
import static org.opensearch.ml.common.conversation.ConversationalIndexConstants.INTERACTIONS_INDEX_SCHEMA_VERSION;
import static org.opensearch.ml.common.conversation.ConversationalIndexConstants.INTERACTIONS_INPUT_FIELD;
import static org.opensearch.ml.common.conversation.ConversationalIndexConstants.INTERACTIONS_ORIGIN_FIELD;
import static org.opensearch.ml.common.conversation.ConversationalIndexConstants.INTERACTIONS_PROMPT_TEMPLATE_FIELD;
import static org.opensearch.ml.common.conversation.ConversationalIndexConstants.INTERACTIONS_RESPONSE_FIELD;
import static org.opensearch.ml.common.conversation.ConversationalIndexConstants.INTERACTIONS_TRACE_NUMBER_FIELD;
import static org.opensearch.ml.common.conversation.ConversationalIndexConstants.META_CREATED_TIME_FIELD;
import static org.opensearch.ml.common.conversation.ConversationalIndexConstants.META_INDEX_SCHEMA_VERSION;
import static org.opensearch.ml.common.conversation.ConversationalIndexConstants.META_NAME_FIELD;
import static org.opensearch.ml.common.conversation.ConversationalIndexConstants.META_UPDATED_TIME_FIELD;
import static org.opensearch.ml.common.conversation.ConversationalIndexConstants.PARENT_INTERACTIONS_ID_FIELD;
import static org.opensearch.ml.common.conversation.ConversationalIndexConstants.USER_FIELD;
import static org.opensearch.ml.common.model.MLModelConfig.ALL_CONFIG_FIELD;
import static org.opensearch.ml.common.model.MLModelConfig.MODEL_TYPE_FIELD;
import static org.opensearch.ml.common.model.TextEmbeddingModelConfig.EMBEDDING_DIMENSION_FIELD;
Expand Down Expand Up @@ -44,6 +61,12 @@ public class CommonValue {
public static final String ML_CONFIG_INDEX = ".plugins-ml-config";
public static final Integer ML_CONFIG_INDEX_SCHEMA_VERSION = 2;
public static final String ML_MAP_RESPONSE_KEY = "response";
public static final String ML_AGENT_INDEX = ".plugins-ml-agent";
public static final Integer ML_AGENT_INDEX_SCHEMA_VERSION = 1;
public static final String ML_MEMORY_META_INDEX = ".plugins-ml-memory-meta";
public static final Integer ML_MEMORY_META_INDEX_SCHEMA_VERSION = 1;
public static final String ML_MEMORY_MESSAGE_INDEX = ".plugins-ml-memory-message";
public static final Integer ML_MEMORY_MESSAGE_INDEX_SCHEMA_VERSION = 1;
public static final String USER_FIELD_MAPPING = " \""
+ CommonValue.USER
+ "\": {\n"
Expand Down Expand Up @@ -326,4 +349,97 @@ public class CommonValue {
+ "\": {\"type\": \"date\", \"format\": \"strict_date_time||epoch_millis\"}\n"
+ " }\n"
+ "}";

public static final String ML_AGENT_INDEX_MAPPING = "{\n"
+ " \"_meta\": {\"schema_version\": "
+ ML_AGENT_INDEX_SCHEMA_VERSION
+ "},\n"
+ " \"properties\": {\n"
+ " \""
+ MLAgent.AGENT_NAME_FIELD
+ "\" : {\"type\":\"text\",\"fields\":{\"keyword\":{\"type\":\"keyword\",\"ignore_above\":256}}},\n"
+ " \""
+ MLAgent.AGENT_TYPE_FIELD
+ "\" : {\"type\":\"keyword\"},\n"
+ " \""
+ MLAgent.DESCRIPTION_FIELD
+ "\" : {\"type\": \"text\"},\n"
+ " \""
+ MLAgent.LLM_FIELD
+ "\" : {\"type\": \"flat_object\"},\n"
+ " \""
+ MLAgent.TOOLS_FIELD
+ "\" : {\"type\": \"flat_object\"},\n"
+ " \""
+ MLAgent.PARAMETERS_FIELD
+ "\" : {\"type\": \"flat_object\"},\n"
+ " \""
+ MLAgent.MEMORY_FIELD
+ "\" : {\"type\": \"flat_object\"},\n"
+ " \""
+ MLAgent.CREATED_TIME_FIELD
+ "\": {\"type\": \"date\", \"format\": \"strict_date_time||epoch_millis\"},\n"
+ " \""
+ MLAgent.LAST_UPDATED_TIME_FIELD
+ "\": {\"type\": \"date\", \"format\": \"strict_date_time||epoch_millis\"}\n"
+ " }\n"
+ "}";

public static final String ML_MEMORY_META_INDEX_MAPPING = "{\n"
+ " \"_meta\": {\n"
+ " \"schema_version\": " + META_INDEX_SCHEMA_VERSION + "\n"
+ " },\n"
+ " \"properties\": {\n"
+ " \""
+ META_NAME_FIELD
+ "\": {\"type\": \"text\"},\n"
+ " \""
+ META_CREATED_TIME_FIELD
+ "\": {\"type\": \"date\", \"format\": \"strict_date_time||epoch_millis\"},\n"
+ " \""
+ META_UPDATED_TIME_FIELD
+ "\": {\"type\": \"date\", \"format\": \"strict_date_time||epoch_millis\"},\n"
+ " \""
+ USER_FIELD
+ "\": {\"type\": \"keyword\"},\n"
+ " \""
+ APPLICATION_TYPE_FIELD
+ "\": {\"type\": \"keyword\"}\n"
+ " }\n"
+ "}";

public static final String ML_MEMORY_MESSAGE_INDEX_MAPPING = "{\n"
+ " \"_meta\": {\n"
+ " \"schema_version\": " + INTERACTIONS_INDEX_SCHEMA_VERSION + "\n"
+ " },\n"
+ " \"properties\": {\n"
+ " \""
+ INTERACTIONS_CONVERSATION_ID_FIELD
+ "\": {\"type\": \"keyword\"},\n"
+ " \""
+ INTERACTIONS_CREATE_TIME_FIELD
+ "\": {\"type\": \"date\", \"format\": \"strict_date_time||epoch_millis\"},\n"
+ " \""
+ INTERACTIONS_INPUT_FIELD
+ "\": {\"type\": \"text\"},\n"
+ " \""
+ INTERACTIONS_PROMPT_TEMPLATE_FIELD
+ "\": {\"type\": \"text\"},\n"
+ " \""
+ INTERACTIONS_RESPONSE_FIELD
+ "\": {\"type\": \"text\"},\n"
+ " \""
+ INTERACTIONS_ORIGIN_FIELD
+ "\": {\"type\": \"keyword\"},\n"
+ " \""
+ INTERACTIONS_ADDITIONAL_INFO_FIELD
+ "\": {\"type\": \"flat_object\"},\n"
+ " \""
+ PARENT_INTERACTIONS_ID_FIELD
+ "\": {\"type\": \"keyword\"},\n"
+ " \""
+ INTERACTIONS_TRACE_NUMBER_FIELD
+ "\": {\"type\": \"long\"}\n"
+ " }\n"
+ "}";
}
102 changes: 102 additions & 0 deletions common/src/main/java/org/opensearch/ml/common/agent/LLMSpec.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,102 @@
/*
* Copyright OpenSearch Contributors
* SPDX-License-Identifier: Apache-2.0
*/

package org.opensearch.ml.common.agent;

import lombok.Builder;
import lombok.Getter;
import org.opensearch.core.common.io.stream.StreamInput;
import org.opensearch.core.common.io.stream.StreamOutput;
import org.opensearch.core.xcontent.ToXContentObject;
import org.opensearch.core.xcontent.XContentBuilder;
import org.opensearch.core.xcontent.XContentParser;

import java.io.IOException;
import java.util.Map;

import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
import static org.opensearch.ml.common.utils.StringUtils.getParameterMap;


@Getter
public class LLMSpec implements ToXContentObject {
public static final String MODEL_ID_FIELD = "model_id";
public static final String PARAMETERS_FIELD = "parameters";

private String modelId;
private Map<String, String> parameters;


@Builder(toBuilder = true)
public LLMSpec(String modelId, Map<String, String> parameters) {
if (modelId == null) {
throw new IllegalArgumentException("model id is null");
}
this.modelId = modelId;
this.parameters = parameters;
}

public LLMSpec(StreamInput input) throws IOException{
modelId = input.readString();
if (input.readBoolean()) {
parameters = input.readMap(StreamInput::readString, StreamInput::readOptionalString);
}
}

public void writeTo(StreamOutput out) throws IOException {
out.writeString(modelId);
if (parameters != null && parameters.size() > 0) {
out.writeBoolean(true);
out.writeMap(parameters, StreamOutput::writeString, StreamOutput::writeOptionalString);
} else {
out.writeBoolean(false);
}
}

@Override
public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
builder.startObject();
if (modelId != null) {
builder.field(MODEL_ID_FIELD, modelId);
}
if (parameters != null && parameters.size() > 0) {
builder.field(PARAMETERS_FIELD, parameters);
}
builder.endObject();
return builder;
}

public static LLMSpec parse(XContentParser parser) throws IOException {
String modelId = null;
Map<String, String> parameters = null;

ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.currentToken(), parser);
while (parser.nextToken() != XContentParser.Token.END_OBJECT) {
String fieldName = parser.currentName();
parser.nextToken();

switch (fieldName) {
case MODEL_ID_FIELD:
modelId = parser.text();
break;
case PARAMETERS_FIELD:
parameters = getParameterMap(parser.map());
break;
default:
parser.skipChildren();
break;
}
}
return LLMSpec.builder()
.modelId(modelId)
.parameters(parameters)
.build();
}

public static LLMSpec fromStream(StreamInput in) throws IOException {
LLMSpec toolSpec = new LLMSpec(in);
return toolSpec;
}
}
Loading

0 comments on commit 3949ef9

Please sign in to comment.