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;
}

grouping morpheus_pipelines {
    list morpheus_pipelines {
        key name;
        uses morpheus_pipeline;
    }
}

grouping morpheus_pipeline {
    leaf 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 name;
        uses morpheus_pipeline_stage;
    }
}

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

typedef stage {  
    type enumeration {
        enum "FileSourceStage";
        enum "DeserializeStage";
        enum "AbpPcapPreprocessingStage";
        enum "MonitorStage";
        enum "TritonInferenceStage";
        enum "AddClassificationsStage";
        enum "SerializeStage";
        enum "WriteToFileStage";
    }
}

grouping prob {
    leaf value {
        type string;
        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 name { type string; }
    leaf stage_type {type stage; }
    list configuration {key "mode"; uses conf;} 
    leaf filename {type string;}
    leaf iterative {type boolean;}
    leaf file_type {type files;}
    leaf filter_null {type boolean;}
    leaf descriptions {type string;}
    leaf model_name {type string;}
    leaf server_url {type string;}
    leaf force_convert_inputs {type boolean;}
    leaf unit {type string;}
    list labels {key "value"; uses prob;}
    leaf kwargs {type empty;}
    leaf overwrite {type boolean;}
}
}