module probes-agent {

namespace "urn:probes-agent";
prefix "probes-agent";

import openconfig-probes { prefix oc-probes; }

organization "EVIDEN";
contact "jose.carcel@eviden.com";
description "Basic example of data model for configuring SmartNICs through Morpheus API";
revision "2023-08-03" {
    description "Basic example of data model for configuring SmartNICs through Morpheus API";
}

augment '/oc-probes:probes/oc-probes:probe' {
    leaf probe_type {
        type enumeration {
            enum plain_traffic;
            enum morpheus_pipeline;
        }
    }
}

augment '/oc-probes:probes/oc-probes:probe/oc-probes:tests/oc-probes:test/oc-probes:config' {
    uses morpheus_pipelines;
}

typedef files{
    type enumeration {
        enum "auto";
        enum "csv";
        enum "json";
    }
}

grouping prob{
    leaf value {
        type string; // e.g. "probs"
        description "probs";
    }
}

grouping conf{
    leaf mode {type string; }
    leaf num_threads { type uint16; }
    leaf pipeline_batch_size { type uint64; }
    leaf model_max_batch_size { type uint64; }
    leaf model_fea_length { type uint64; }
    list class_labels {
        key "value";
        uses prob; }
}

grouping morpheus_pipeline_stage {
    leaf stage_name { type string; }
    choice stage_type {
        case FileSourceStage {
            container FileSourceStage {
                list fs_configuration {key "mode"; uses conf;} 
                leaf filename {type string;}
                leaf iterative {type boolean;}
                leaf file_type {type files;}
                leaf filter_null {type boolean;}
            }
        }
        case DeserializeStage {
            container DeserializeStage {
                list ds_configuration {key "mode"; uses conf;} 
            }
        }
        case AbpPcapPreprocessingStage {
            container AbpPcapPreprocessingStage {
                list apps_configuration {key "mode"; uses conf;} 
            }
        }
        case MonitorStage {
            container MonitorStage {
                list ms_configuration {key "mode"; uses conf;} 
                leaf descriptions {type string;}
                leaf unit {type string;}
            }
        }
        case TritonInferenceStage {
            container TritonInferenceStage {
                list tis_configuration {key "mode"; uses conf;} 
                leaf model_name {type string;}
                leaf server_url {type string;}
                leaf force_convert_inputs {type boolean;}
            }
        }
        case AddClassificationsStage{
            container AddClassificationsStage {
                list acs_configuration {key "mode"; uses conf;} 
                list labels {key "value"; uses prob;}
            }
        }
        case SerializeStage {
            container SerializeStage {
                list ss_configuration {key "mode"; uses conf;} 
                leaf kwargs {type empty;}
            }
        }
        case WriteToFileStage{
            container WriteToFileStage {
                list wfs_configuration {key "mode"; uses conf;} 
                leaf wfs_filename {type string;}
                leaf overwrite {type boolean;}
            }
        }
        case CustomStage {
            list custom {
                key "field_name";
                leaf field_name {type string;}
                leaf field_value {type string;}
            } 
        }
    }
}

grouping morpheus_pipeline {
    leaf pipeline_name {type string;}
    leaf num_threads {type uint16;}
    leaf pipeline_batch_size {type uint64; }
    leaf model_max_batch_size {type uint64; }
    leaf input_file {type string;}
    leaf output_file {type string;}
    leaf model_fea_length {type uint16;}
    leaf model_name {type string;}
    leaf iterative {type boolean;}
    leaf server_url {type string;}
    leaf file_type {type files;}
    list stages {
        key "stage_name";
        uses morpheus_pipeline_stage;
    }
}

grouping morpheus_pipelines {
    list morpheus_pipeline {
        key "pipeline_name";
        uses morpheus_pipeline;
    }
}

container morpheus_pipelines {
    uses morpheus_pipelines;
}

}