syntax = "proto3";

package ai.verta.registry;
option go_package = "github.com/VertaAI/modeldb/protos/gen/go/protos/public/registry";

option java_multiple_files = true;

import "google/api/annotations.proto";
import "google/protobuf/field_mask.proto";
import "common/CommonService.proto";
import "modeldb/versioning/Environment.proto";
import "modeldb/versioning/Code.proto";
import "registry/ChecklistService.proto";
import "registry/CustomAttributeService.proto";
import "registry/StageService.proto";
import "uac/Collaborator.proto";

message DataTypeEnum {
  enum DataType {
      UNKNOWN = 0;
      OTHER = 1;
      AUDIO = 2;
      IMAGE = 3;
      TABULAR = 4;
      TEXT = 5;
      VIDEO = 6;
  }
}

message ActionTypeEnum {  
  option deprecated = true;
  enum ActionType {
    option deprecated = true;
    UNKNOWN = 0;
    OTHER = 1;
    CLASSIFICATION = 2;
    CLUSTERING = 3;
    DETECTION = 4;
    REGRESSION = 5;
    TRANSCRIPTION = 6;
    TRANSLATION = 7;
  }
}

message TaskTypeEnum {
  enum TaskType {
    UNKNOWN = 0;
    OTHER = 1;
    CLASSIFICATION = 2;
    CLUSTERING = 3;
    DETECTION = 4;
    REGRESSION = 5;
    TRANSCRIPTION = 6;
    TRANSLATION = 7;
  }
}

message RegisteredModel {
    // internal id
    uint64 id = 1;

    // Unique name for the model within the workspace.
    string name = 2;

    // Timestamp recorded when this ``registered_model`` was created.
    int64 time_created = 3;

    // Timestamp recorded when metadata for this ``registered_model`` was last updated.
    int64 time_updated = 4;

    // Description of this ``registered_model``.
    string description = 5;

    // Labels added to the registered models
    repeated string labels = 6;

    // Visibility of Registered Models
    common.VisibilityEnum.Visibility visibility = 7;

    // id of the workspace this belongs to
    //not storing workspace_type uac should support https://github.com/VertaAI/protos-all/pull/19/files
    string workspace_id = 8;

    //readme
    string readme_text = 9;

    //owner's verta_id
    string owner = 10;
    oneof owner_tracking {
        uint64 owner_id = 21;
        common.GroupId group_owner_id = 22;
    }

    // Other attributes
    repeated common.KeyValue attributes = 11;

    uint64 workspace_service_id = 12;

    // The next field only makes sense when visibility == ORG_CUSTOM
    // They will be empty otherwise
    uac.CollaboratorPermissions custom_permission = 13;
    uac.ResourceVisibility resource_visibility = 14;
    // Other artifacts
    repeated common.Artifact artifacts = 15;

    // Is true if the readme text has been updated at least once
    bool readme_text_updated = 16;

    // Doc for version_number https://github.com/VertaAI/docs/pull/53
    uint64 version_number = 17;

    // Fields for model's task_type. Those fields must be searchable as per our "find" operation
    DataTypeEnum.DataType data_type = 18;
    ActionTypeEnum.ActionType action_type = 19 [deprecated = true];
    TaskTypeEnum.TaskType task_type = 20;

    repeated CustomAttributeValue custom_attribute_values = 23;

    // Is true if the model accepts pii as input
    bool pii = 24;
}

//CRUD for Registred Models
message FindRegisteredModelRequest{
    //User or organization name
    string workspace_name = 1;
    //Filtering supported on
    // 1. name
    // 2. labels
    // 3. description
    repeated common.KeyValueQuery predicates = 2;
    //If not present will default to page 1 page_size 10
    common.Pagination pagination = 3;
    bool ascending = 4;
    string sort_key = 5;
    repeated uint64 ids = 6;

    message Response {
        repeated RegisteredModel registered_models = 1;
        int64 total_records = 2;
    }
}

message GetRegisteredModelRequest {
    RegisteredModelIdentification id = 1;

    message Response {
        RegisteredModel registered_model = 1;
    }
}

message GetRegisteredModelCountRequest {
    string project_id = 1;

    message Response {
        int64 registered_model_count = 1;
    }
}

message RegisteredModelNamedIdentification {
    //Model Name
    string name = 1;
    // WorkspaceName
    string workspace_name = 2;
}

// changed from oneof because of that bug: https://github.com/grpc-ecosystem/grpc-gateway/issues/753
message RegisteredModelIdentification {
    // Model refered to by workspace name , model nmae combination
    RegisteredModelNamedIdentification named_id = 1;
    // Model refered to by id
    uint64 registered_model_id = 2;
}

message SetRegisteredModel {
    RegisteredModelIdentification id = 1;
    RegisteredModel registered_model = 2;
    google.protobuf.FieldMask update_mask = 3;

    message Response {
        RegisteredModel registered_model = 1;
    }
}

message DeleteRegisteredModelRequest {
    RegisteredModelIdentification id = 1;

    message Response { }
}

message ModelVersionLockLevelEnum {
    enum ModelVersionLockLevel {
        UNKNOWN = 0;
        OPEN = 1;
        CLOSED = 2;
        REDACT = 3;
    }
}

message ModelVersion {

    // internal id
    uint64 id = 1;

    // id of the model this belongs to
    uint64 registered_model_id = 2;

    // Model's tag. Unique for a model nmae
    string version = 3;

    // Timestamp recorded when metadata for this ``model_version`` was last updated.
    int64 time_updated = 4;

    // Timestamp recorded when this ``model_version`` was created.
    int64 time_created = 5;

    // Description of this ``model_version``.
    string description = 6;

    // ModelDB run ID used when creating ``model_version``, if ``model_source`` was generated by an
    // experiment run stored in ModelDB tracking server.
    oneof source {
        string experiment_run_id = 7;
        // string model_url = 15; NOT SUPPORTED YET
    }

    // Labels added to the registered models
    repeated string labels = 8;

    // Model
    common.Artifact model = 9;

    //Requirement
    modeldb.versioning.EnvironmentBlob environment = 10;
    DockerMetadata docker_metadata = 24;

    // Other artifacts
    repeated common.Artifact artifacts = 11;

    //flag to archive a model version currently does nothing other than storing the information in DB
    common.TernaryEnum.Ternary archived = 12;

    //readme
    string readme_text = 13;

    //APIs
    repeated string apis = 15;

    //owner's verta_id
    string owner = 16;
    oneof owner_tracking {
        uint64 owner_id = 31;
        common.GroupId group_owner_id = 32;
    }

    // Other attributes
    repeated common.KeyValue attributes = 17;

    // Stage of the model version. This field must be searchable as per our "find" operation
    StageEnum.Stage stage = 18;

    // Current lock level for this model version
    ModelVersionLockLevelEnum.ModelVersionLockLevel lock_level = 19;

    repeated common.Artifact datasets = 20;
    map<string, modeldb.versioning.CodeBlob> code_blob_map = 21;

    // Is true if the readme text has been updated at least once
    bool readme_text_updated = 22;

    // Doc for version_number https://github.com/VertaAI/docs/pull/53
    uint64 version_number = 23;

    // Reference: https://vertaai.atlassian.net/browse/VRD-329
    string redirect_metadata = 25;

    // Reference: https://vertaai.atlassian.net/browse/CAT-84
    string input_description = 26;

    // Reference: https://vertaai.atlassian.net/browse/CAT-84
    bool hide_input_label = 27;

    // Reference: https://vertaai.atlassian.net/browse/CAT-84
    string output_description = 28;

    // Reference: https://vertaai.atlassian.net/browse/CAT-84
    bool hide_output_label = 29;

    repeated ExternalDeployment external_deployments = 30;

    // Reference: https://vertaai.atlassian.net/browse/CAT-173
    repeated ChecklistItemValue checklist_item_values = 33;

    repeated CustomAttributeValue custom_attribute_values = 34;
}

message DockerMetadata {
    uint32 request_port = 1;
    string request_path = 2;
    string health_path = 3;
}

message SetLockModelVersionRequest {
    //User or organization name
    ModelVersionIdentification id = 1;
    ModelVersionLockLevelEnum.ModelVersionLockLevel lock_level = 2;

    message Response {}
}

message ModelVersionIdentification {
    // ModelVersion referred to by id
    uint64 model_version_id = 1;

    //Registered Model referred to by id
    RegisteredModelIdentification model_id = 2; //Compulsory
}

//  CRUD for Model Version
message GetModelVersionRequest{
    //User or organization name
    ModelVersionIdentification id = 1;

    message Response {
        ModelVersion model_version = 1;
    }
}

message FindModelVersionRequest{
    //identity of parent registered model
    RegisteredModelIdentification id = 1;
    //Filtering supported on
    // 1. version
    // 2. labels
    // 3. description
    // 4. model_source
    // 5. experiment_run_id
    // 6. status
    // 7. archived status
    repeated common.KeyValueQuery predicates = 2;
    //If not present will default to page 1 page_size 10
    common.Pagination pagination = 3;
    bool ascending = 4;
    string sort_key = 5;
    repeated uint64 ids = 6;

    message Response {
        repeated ModelVersion model_versions = 1;
        int64 total_records = 2;
    }
}

message SetModelVersion {
    ModelVersionIdentification id = 1;
    ModelVersion model_version = 2;
    google.protobuf.FieldMask update_mask = 3;

    message Response {
        ModelVersion model_version = 1;
    }

}

message DeleteModelVersionRequest {
    ModelVersionIdentification id = 1;

    message Response { }
}

message GetUrlForArtifact {
    // id of model version
    uint64 model_version_id = 1;
    //artifact key
    string key = 2;
    //GET or PUT
    string method = 3;
    //type of artifact
    common.ArtifactTypeEnum.ArtifactType artifact_type = 4;
    // 0 if the url is not required to be multi part
    uint64 part_number = 5;

    message Response {
        // URL based on action
        string url = 1;
        bool multipart_upload_ok = 2;
        map<string, string> fields = 3;
    }
}

message CommitArtifactPart {
    // id of model version
    uint64 model_version_id = 1;
    //artifact key
    string key = 2;
    common.ArtifactPart artifact_part = 3;

    message Response{ }
}

message GetCommittedArtifactParts {
    // id of model version
    uint64 model_version_id = 1;
    //artifact key
    string key = 2;

    message Response {
        repeated common.ArtifactPart artifact_parts = 1;
    }
}

message CommitMultipartArtifact {
    // id of model version
    uint64 model_version_id = 1;
    //artifact key
    string key = 2;

    message Response { }
}

message LogDatasetsInModelVersion {
    uint64 model_version_id = 1;
    repeated common.Artifact datasets = 2;

    message Response { }
}

message LogCodeBlobInModelVersion {
    uint64 model_version_id = 1;
    map<string, modeldb.versioning.CodeBlob> code_blob_map = 2;

    message Response { }
}

message LogAttributesInModelVersion {
    uint64 model_version_id = 1;
    repeated common.KeyValue attributes = 2;

    message Response { }
}

message LogDockerMetadataInModelVersion {
    uint64 model_version_id = 1;
    DockerMetadata docker_metadata = 2;

    message Response { }
}

message LocationTypeEnum {
  enum LocationType {
    UNKNOWN = 0;
    CLOUD_PROVIDER = 1;
    INTERNAL_PRIVATE_NETWORK = 2;
    DEVICE = 3;
  }
}

message ExternalDeployment {    
    // internal id
    uint64 id = 1;

    // id of the model this belongs to
    uint64 model_version_id = 2;

    // Timestamp recorded when this ``external_deployment`` was created.
    int64 time_created = 3;

    // Timestamp recorded when metadata for this ``external_deployment`` was last updated.
    int64 time_updated = 4;

    // Location where container will be deployed
    string location_deployed = 5;

    // Deploy URL, path or title
    string url_path_title = 6;

    // Brief description
    string description = 7;

    // Location Type
    LocationTypeEnum.LocationType location_type = 8;

    // Cloud Provider
    string cloud_provider = 9;
}   

message ExternalDeploymentRequest {    
    uint64 external_deployment_id = 1;

    message Response {
        ExternalDeployment external_deployment = 1;
    }
}

message Empty {
}

message DeleteModelVersionArtifactRequest {
    uint64 model_version_id = 1;
    string artifact_key = 2;
}

service RegistryService {
   // CRUD for RegisteredModel
    rpc FindRegisteredModel(FindRegisteredModelRequest) returns (FindRegisteredModelRequest.Response) {
        option (google.api.http) = {
            post: "/v1/registry/workspaces/{workspace_name}/registered_models/find"
            body: "*"
            additional_bindings {
                post: "/v1/registry/registered_models/find"
                body: "*"
            }
        };

    };

    rpc GetRegisteredModel(GetRegisteredModelRequest) returns (GetRegisteredModelRequest.Response) {
        option (google.api.http) = {
            get: "/v1/registry/workspaces/{id.named_id.workspace_name}/registered_models/{id.named_id.name}"
            additional_bindings {
                get: "/v1/registry/registered_models/{id.registered_model_id}"
            }
        };
    };

    rpc GetRegisteredModelCount(GetRegisteredModelCountRequest) returns (GetRegisteredModelCountRequest.Response) {
        option (google.api.http) = {
            get: "/v1/registry/registered_models/count"
        };
    };

    rpc CreateRegisteredModel(SetRegisteredModel) returns (SetRegisteredModel.Response) {
        option (google.api.http) = {
            post: "/v1/registry/workspaces/{id.named_id.workspace_name}/registered_models"
            body: "registered_model"
        };
    }

    rpc UpdateRegisteredModel(SetRegisteredModel) returns (SetRegisteredModel.Response) {
        option (google.api.http) = {
            patch: "/v1/registry/workspaces/{id.named_id.workspace_name}/registered_models/{id.named_id.name}"
            body: "registered_model"
            additional_bindings {
                patch: "/v1/registry/registered_models/{id.registered_model_id}"
                body: "registered_model"
            }
            additional_bindings {
                patch: "/v1/registry/workspaces/{id.named_id.workspace_name}/registered_models/{id.named_id.name}/full_body"
                body: "*"
            }
            additional_bindings {
                patch: "/v1/registry/registered_models/{id.registered_model_id}/full_body"
                body: "*"
            }
            additional_bindings {
                put: "/v1/registry/workspaces/{id.named_id.workspace_name}/registered_models/{id.named_id.name}"
                body: "registered_model"
            }
            additional_bindings {
                put: "/v1/registry/registered_models/{id.registered_model_id}"
                body: "registered_model"
            }
        };
    }

    rpc DeleteRegisteredModel(DeleteRegisteredModelRequest) returns (DeleteRegisteredModelRequest.Response) {
        option (google.api.http) = {
            delete: "/v1/registry/workspaces/{id.named_id.workspace_name}/registered_models/{id.named_id.name}"
            additional_bindings {
                delete: "/v1/registry/registered_models/{id.registered_model_id}"
            }
        };
    }

    //CRUD for Model Version
    rpc FindModelVersion(FindModelVersionRequest) returns (FindModelVersionRequest.Response) {
        option (google.api.http) = {
            post: "/v1/registry/workspaces/{id.named_id.workspace_name}/registered_models/{id.named_id.name}/model_versions/find"
            body: "*"
            additional_bindings {
                post: "/v1/registry/registered_models/{id.registered_model_id}/model_versions/find"
                body: "*"
            }
            additional_bindings {
                post: "/v1/registry/workspaces/{id.named_id.workspace_name}/model_versions/find"
                body: "*"
            }
            additional_bindings {
                post: "/v1/registry/model_versions/find"
                body: "*"
            }
        };
    };

    rpc GetModelVersion(GetModelVersionRequest) returns (GetModelVersionRequest.Response) {
        option (google.api.http) = {
            get: "/v1/registry/workspaces/{id.model_id.named_id.workspace_name}/registered_models/{id.model_id.named_id.name}/model_versions/{id.model_version_id}"
            additional_bindings {
                get: "/v1/registry/model_versions/{id.model_version_id}"
            }
        };
    };


    rpc CreateModelVersion(SetModelVersion) returns (SetModelVersion.Response) {
        option (google.api.http) = {
            post: "/v1/registry/workspaces/{id.model_id.named_id.workspace_name}/registered_models/{id.model_id.named_id.name}/model_versions"
            body: "model_version"
            additional_bindings {
                post: "/v1/registry/registered_models/{id.model_id.registered_model_id}/model_versions"
                body: "model_version"
            }
        };
    }

    rpc UpdateModelVersion(SetModelVersion) returns (SetModelVersion.Response) {
        option (google.api.http) = {
            patch: "/v1/registry/workspaces/{id.model_id.named_id.workspace_name}/registered_models/{id.model_id.named_id.name}/model_versions/{id.model_version_id}"
            body: "model_version"
            additional_bindings {
                patch: "/v1/registry/registered_models/{id.model_id.registered_model_id}/model_versions/{id.model_version_id}"
                body: "model_version"
            }
            additional_bindings {
                patch: "/v1/registry/workspaces/{id.model_id.named_id.workspace_name}/registered_models/{id.model_id.named_id.name}/model_versions/{id.model_version_id}/full_body"
                body: "*"
            }
            additional_bindings {
                patch: "/v1/registry/registered_models/{id.model_id.registered_model_id}/model_versions/{id.model_version_id}/full_body"
                body: "*"
            }
            additional_bindings {
                put: "/v1/registry/workspaces/{id.model_id.named_id.workspace_name}/registered_models/{id.model_id.named_id.name}/model_versions/{id.model_version_id}"
                body: "model_version"
            }
            additional_bindings {
                put: "/v1/registry/registered_models/{id.model_id.registered_model_id}/model_versions/{id.model_version_id}"
                body: "model_version"
            }
        };
    }

    rpc SetLockModelVersion(SetLockModelVersionRequest) returns (SetLockModelVersionRequest.Response) {
        option (google.api.http) = {
            put: "/v1/registry/model_versions/{id.model_version_id}/lock"
            body: "*"
        };
    };

    rpc DeleteModelVersion(DeleteModelVersionRequest) returns (DeleteModelVersionRequest.Response) {
        option (google.api.http) = {
            delete: "/v1/registry/workspaces/{id.model_id.named_id.workspace_name}/registered_models/{id.model_id.named_id.name}/model_versions/{id.model_version_id}"
            additional_bindings {
                delete: "/v1/registry/registered_models/{id.model_id.registered_model_id}/model_versions/{id.model_version_id}"
            }
            additional_bindings {
                delete: "/v1/registry/model_versions/{id.model_version_id}"
            }
        };
    }

    rpc getUrlForArtifact(GetUrlForArtifact) returns (GetUrlForArtifact.Response) {
        option (google.api.http) = {
            post: "/v1/registry/model_versions/{model_version_id}/getUrlForArtifact"
            body: "*"
        };
    };

    rpc commitArtifactPart(CommitArtifactPart) returns (CommitArtifactPart.Response) {
        option (google.api.http) = {
            post: "/v1/registry/model_versions/{model_version_id}/commitArtifactPart"
            body: "*"
        };
    };

    rpc getCommittedArtifactParts(GetCommittedArtifactParts) returns (GetCommittedArtifactParts.Response) {
        option (google.api.http) = {
            get: "/v1/registry/model_versions/{model_version_id}/getCommittedArtifactParts"
        };
    };

    rpc commitMultipartArtifact(CommitMultipartArtifact) returns (CommitMultipartArtifact.Response) {
        option (google.api.http) = {
            post: "/v1/registry/model_versions/{model_version_id}/commitMultipartArtifact"
            body: "*"
        };
    };

    rpc DeleteModelVersionArtifact(DeleteModelVersionArtifactRequest) returns (Empty) {
        option (google.api.http) = {
            delete: "/v1/registry/model_versions/{model_version_id}/artifacts/{artifact_key}"
        };
    }

    rpc logDatasetsInModelVersion(LogDatasetsInModelVersion) returns (LogDatasetsInModelVersion.Response) {
        option (google.api.http) = {
            post: "/v1/registry/model_versions/{model_version_id}/logDatasets"
            body: "*"
        };
    };

    rpc logCodeBlobInModelVersion(LogCodeBlobInModelVersion) returns (LogCodeBlobInModelVersion.Response) {
        option (google.api.http) = {
            post: "/v1/registry/model_versions/{model_version_id}/logCodeBlobInModelVersion"
            body: "*"
        };
    };

    rpc logAttributesInModelVersion(LogAttributesInModelVersion) returns (LogAttributesInModelVersion.Response) {
        option (google.api.http) = {
            post: "/v1/registry/model_versions/{model_version_id}/logAttributes"
            body: "*"
        };
    };

    rpc logDockerMetadataInModelVersion(LogDockerMetadataInModelVersion) returns (LogDockerMetadataInModelVersion.Response) {
        option (google.api.http) = {
            post: "/v1/registry/model_versions/{model_version_id}/logDockerMetadata"
            body: "*"
        };
    };
    
    rpc GetExternalDeployment(ExternalDeploymentRequest) returns (ExternalDeploymentRequest.Response) {
        option (google.api.http) = {
            get: "/v1/registry/registered_models/model_versions/external_deployments/{external_deployment_id}"
        };
    };

   rpc CreateExternalDeployment(ExternalDeployment) returns (ExternalDeploymentRequest.Response) {
        option (google.api.http) = {
            post: "/v1/registry/registered_models/model_versions/external_deployments"
            body: "*"
        };
    }

    rpc UpdateExternalDeployment(ExternalDeployment) returns (ExternalDeploymentRequest.Response) {
        option (google.api.http) = {
            put: "/v1/registry/registered_models/model_versions/external_deployments"
            body: "*"
        };
    }

    rpc DeleteExternalDeployment(ExternalDeploymentRequest) returns (Empty) {
        option (google.api.http) = {
            delete: "/v1/registry/registered_models/model_versions/external_deployments/{external_deployment_id}"
        };
    }
}