-
Notifications
You must be signed in to change notification settings - Fork 145
Commit
Signed-off-by: Yaliang Wu <[email protected]>
- Loading branch information
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,18 @@ | ||
/* | ||
* Copyright OpenSearch Contributors | ||
* SPDX-License-Identifier: Apache-2.0 | ||
*/ | ||
|
||
package org.opensearch.ml.common.transport.agent; | ||
|
||
import org.opensearch.action.ActionType; | ||
|
||
public class MLRegisterAgentAction extends ActionType<MLRegisterAgentResponse> { | ||
public static MLRegisterAgentAction INSTANCE = new MLRegisterAgentAction(); | ||
public static final String NAME = "cluster:admin/opensearch/ml/agents/register"; | ||
|
||
private MLRegisterAgentAction() { | ||
super(NAME, MLRegisterAgentResponse::new); | ||
} | ||
|
||
} |
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,77 @@ | ||
/* | ||
* Copyright OpenSearch Contributors | ||
* SPDX-License-Identifier: Apache-2.0 | ||
*/ | ||
|
||
package org.opensearch.ml.common.transport.agent; | ||
|
||
import lombok.AccessLevel; | ||
import lombok.Builder; | ||
import lombok.Getter; | ||
import lombok.ToString; | ||
import lombok.experimental.FieldDefaults; | ||
import org.opensearch.action.ActionRequest; | ||
import org.opensearch.action.ActionRequestValidationException; | ||
import org.opensearch.core.common.io.stream.InputStreamStreamInput; | ||
import org.opensearch.core.common.io.stream.OutputStreamStreamOutput; | ||
import org.opensearch.core.common.io.stream.StreamInput; | ||
import org.opensearch.core.common.io.stream.StreamOutput; | ||
import org.opensearch.ml.common.agent.MLAgent; | ||
|
||
import java.io.ByteArrayInputStream; | ||
import java.io.ByteArrayOutputStream; | ||
import java.io.IOException; | ||
import java.io.UncheckedIOException; | ||
|
||
import static org.opensearch.action.ValidateActions.addValidationError; | ||
|
||
@Getter | ||
@FieldDefaults(makeFinal = true, level = AccessLevel.PRIVATE) | ||
@ToString | ||
public class MLRegisterAgentRequest extends ActionRequest { | ||
|
||
MLAgent mlAgent; | ||
|
||
@Builder | ||
public MLRegisterAgentRequest(MLAgent mlAgent) { | ||
this.mlAgent = mlAgent; | ||
} | ||
|
||
public MLRegisterAgentRequest(StreamInput in) throws IOException { | ||
super(in); | ||
this.mlAgent = new MLAgent(in); | ||
} | ||
|
||
@Override | ||
public ActionRequestValidationException validate() { | ||
ActionRequestValidationException exception = null; | ||
if (mlAgent == null) { | ||
exception = addValidationError("ML agent can't be null", exception); | ||
} | ||
|
||
return exception; | ||
} | ||
|
||
@Override | ||
public void writeTo(StreamOutput out) throws IOException { | ||
super.writeTo(out); | ||
this.mlAgent.writeTo(out); | ||
} | ||
|
||
public static MLRegisterAgentRequest fromActionRequest(ActionRequest actionRequest) { | ||
if (actionRequest instanceof MLRegisterAgentRequest) { | ||
return (MLRegisterAgentRequest) actionRequest; | ||
Check warning on line 63 in common/src/main/java/org/opensearch/ml/common/transport/agent/MLRegisterAgentRequest.java Codecov / codecov/patchcommon/src/main/java/org/opensearch/ml/common/transport/agent/MLRegisterAgentRequest.java#L63
|
||
} | ||
|
||
try (ByteArrayOutputStream baos = new ByteArrayOutputStream(); | ||
OutputStreamStreamOutput osso = new OutputStreamStreamOutput(baos)) { | ||
actionRequest.writeTo(osso); | ||
try (StreamInput input = new InputStreamStreamInput(new ByteArrayInputStream(baos.toByteArray()))) { | ||
return new MLRegisterAgentRequest(input); | ||
} | ||
} catch (IOException e) { | ||
throw new UncheckedIOException("Failed to parse ActionRequest into MLRegisterModelRequest", e); | ||
Check warning on line 73 in common/src/main/java/org/opensearch/ml/common/transport/agent/MLRegisterAgentRequest.java Codecov / codecov/patchcommon/src/main/java/org/opensearch/ml/common/transport/agent/MLRegisterAgentRequest.java#L72-L73
|
||
} | ||
|
||
} | ||
} |
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,65 @@ | ||
/* | ||
* Copyright OpenSearch Contributors | ||
* SPDX-License-Identifier: Apache-2.0 | ||
*/ | ||
|
||
package org.opensearch.ml.common.transport.agent; | ||
|
||
import lombok.Getter; | ||
import org.opensearch.core.action.ActionResponse; | ||
import org.opensearch.core.common.io.stream.InputStreamStreamInput; | ||
import org.opensearch.core.common.io.stream.OutputStreamStreamOutput; | ||
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 java.io.ByteArrayInputStream; | ||
import java.io.ByteArrayOutputStream; | ||
import java.io.IOException; | ||
import java.io.UncheckedIOException; | ||
|
||
@Getter | ||
public class MLRegisterAgentResponse extends ActionResponse implements ToXContentObject { | ||
public static final String AGENT_ID_FIELD = "agent_id"; | ||
|
||
private String agentId; | ||
|
||
public MLRegisterAgentResponse(StreamInput in) throws IOException { | ||
super(in); | ||
this.agentId = in.readString(); | ||
} | ||
|
||
public MLRegisterAgentResponse(String agentId) { | ||
this.agentId= agentId; | ||
} | ||
|
||
@Override | ||
public void writeTo(StreamOutput out) throws IOException { | ||
out.writeString(agentId); | ||
} | ||
|
||
@Override | ||
public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { | ||
builder.startObject(); | ||
builder.field(AGENT_ID_FIELD, agentId); | ||
builder.endObject(); | ||
return builder; | ||
} | ||
|
||
public static MLRegisterAgentResponse fromActionResponse(ActionResponse actionResponse) { | ||
if (actionResponse instanceof MLRegisterAgentResponse) { | ||
return (MLRegisterAgentResponse) actionResponse; | ||
Check warning on line 52 in common/src/main/java/org/opensearch/ml/common/transport/agent/MLRegisterAgentResponse.java Codecov / codecov/patchcommon/src/main/java/org/opensearch/ml/common/transport/agent/MLRegisterAgentResponse.java#L52
|
||
} | ||
|
||
try (ByteArrayOutputStream baos = new ByteArrayOutputStream(); | ||
OutputStreamStreamOutput osso = new OutputStreamStreamOutput(baos)) { | ||
actionResponse.writeTo(osso); | ||
try (StreamInput input = new InputStreamStreamInput(new ByteArrayInputStream(baos.toByteArray()))) { | ||
return new MLRegisterAgentResponse(input); | ||
} | ||
} catch (IOException e) { | ||
throw new UncheckedIOException("failed to parse ActionResponse into MLRegisterAgentResponse", e); | ||
Check warning on line 62 in common/src/main/java/org/opensearch/ml/common/transport/agent/MLRegisterAgentResponse.java Codecov / codecov/patchcommon/src/main/java/org/opensearch/ml/common/transport/agent/MLRegisterAgentResponse.java#L61-L62
|
||
} | ||
} | ||
} |