Examples with MLModelGetResponse org.opensearch.ml.common.transport.model.MLModelGetResponse used on opensource projects

Example 1 with MLModelGetResponse

use of org.opensearch.ml.common.transport.model.MLModelGetResponse in project ml-commons by opensearch-project.

the class GetModelTransportAction method doExecute.

@Override
protected void doExecute(Task task, ActionRequest request, ActionListener<MLModelGetResponse> actionListener) {
    MLModelGetRequest mlModelGetRequest = MLModelGetRequest.fromActionRequest(request);
    String modelId = mlModelGetRequest.getModelId();
    GetRequest getRequest = new GetRequest(ML_MODEL_INDEX).id(modelId);
    try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
        client.get(getRequest, ActionListener.wrap(r -> {
            log.info("Completed Get Model Request, id:{}", modelId);
            if (r != null && r.isExists()) {
                try (XContentParser parser = createXContentParserFromRegistry(xContentRegistry, r.getSourceAsBytesRef())) {
                    ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.nextToken(), parser);
                    MLModel mlModel = MLModel.parse(parser);
                    actionListener.onResponse(MLModelGetResponse.builder().mlModel(mlModel).build());
                } catch (Exception e) {
                    log.error("Failed to parse ml model" + r.getId(), e);
                    actionListener.onFailure(e);
                }
            } else {
                actionListener.onFailure(new MLResourceNotFoundException("Fail to find model"));
            }
        }, e -> {
            if (e instanceof IndexNotFoundException) {
                actionListener.onFailure(new MLResourceNotFoundException("Fail to find model"));
            } else {
                log.error("Failed to get ML model " + modelId, e);
                actionListener.onFailure(e);
            }
        }));
    } catch (Exception e) {
        log.error("Failed to get ML model " + modelId, e);
        actionListener.onFailure(e);
    }
}