// Copyright 2023 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

syntax = "proto2";
option optimize_for = LITE_RUNTIME;
option java_package = "org.chromium.components.optimization_guide.proto";
option java_outer_classname = "ModelExecutionProto";

package optimization_guide.proto;

import "components/optimization_guide/proto/model_execution.proto";
import "components/optimization_guide/proto/substitution.proto";

message TextSafetyModelMetadata {
  // The number of categories the model is expected to output.
  optional uint32 num_output_categories = 1;

  // The set of feature configurations used to determine whether the text is
  // safe.
  repeated FeatureTextSafetyConfiguration feature_text_safety_configurations =
      2;
}

enum CheckInputType {
  CHECK_INPUT_TYPE_UNSPECIFIED = 0;
  CHECK_INPUT_TYPE_REQUEST = 1;
  CHECK_INPUT_TYPE_RESPONSE = 2;
}

message CheckInput {
  // Which message to use as input.
  optional CheckInputType input_type = 1;

  // Concatenated template for converting the message to input chunks.
  repeated SubstitutedString templates = 2;
}

message SafetyCategoryThreshold {
  // Label for the category. E.g. 'TOXICITY', 'SEXUAL', 'HEALTH', etc.
  optional string category_label = 1;

  // Index of the category from the output (scores) of the text safety model.
  optional uint32 output_index = 2;

  // Threshold for the category, scores >= to the threshold will be filtered.
  optional float threshold = 3;
}

message LanguageCheck {
  // Requires one of the allowed languages was detected with at least this level
  // of confidence. If this value is 0 or unset, language results will be
  // ignored.
  optional double confidence_threshold = 1;

  // An alternative threshold used when checking incomplete model output.
  // If unspecified, the same threshold is used for all checks.
  optional double partial_threshold = 2;
}

enum RegexFilterType {
  REGEX_FILTER_TYPE_UNSPECIFIED = 0;
  REGEX_FILTER_TYPE_BLOCK_ON_MATCH = 1;
  REGEX_FILTER_TYPE_BLOCK_ON_NO_MATCH = 2;
}

message RegexFilter {
  // The blocking logic for the filter.
  optional RegexFilterType type = 1;

  // The regex string to match against.
  optional string regex = 2;
}

message RequestSafetyCheck {
  // How to generate input for the text safety check from the request message.
  repeated SubstitutedString input_template = 1;

  // Block requests based on regex match result.
  optional RegexFilter regex_filter = 6;

  // The set of thresholds to apply per category, overriding those in the
  // parent.
  repeated SafetyCategoryThreshold safety_category_thresholds = 2;

  // Whether to only perform the allowed language check.
  optional bool check_language_only = 3;

  // Whether to ignore the language result of the safety check.
  // Exclusive with "check_language_only" and "language_check".
  optional bool ignore_language_result = 4;

  // Specify a requirement on the output language.
  // If neither "language_check" nor "ignore_language_result" is specified, the
  // language will be checked with a default threshold.
  optional LanguageCheck language_check = 5;
}

message PartialOutputChecks {
  // The minimum number of tokens required to evaluate partial output safety.
  optional uint32 minimum_tokens = 1;

  // The minimum number of tokens required between partial outputs.
  optional uint32 token_interval = 2;
}

message RawOutputCheck {
  // How to generate input for the text safety check.
  // Should be written against a google.protobuf.StringValue proto with the raw
  // output.
  repeated SubstitutedString input_template = 1;

  // Block raw output based on regex match result.
  optional RegexFilter regex_filter = 3;

  // Specify a requirement on the output language.
  // If unspecified, the language will be checked with a default threshold.
  optional LanguageCheck language_check = 2;
}

message ResponseSafetyCheck {
  // How to generate input for the text safety check based on the request and
  // response.
  repeated CheckInput inputs = 1;

  // Block responses based on regex match result.
  optional RegexFilter regex_filter = 6;

  // The set of thresholds to apply per category, overriding those in the
  // parent.
  repeated SafetyCategoryThreshold safety_category_thresholds = 2;

  // Whether to ignore the language result of the safety check.
  optional bool ignore_language_result = 3;

  // Specify a requirement on the output language.
  // If neither "language_check" nor "ignore_language_result" is specified, the
  // language will be checked with a default threshold.
  optional LanguageCheck language_check = 4;
}

message FeatureTextSafetyConfiguration {
  // The feature this configuration pertains to.
  optional ModelExecutionFeature feature = 1;

  // The default set of thresholds to apply per category.
  repeated SafetyCategoryThreshold safety_category_thresholds = 2;

  // The set of languages allowed for text safety evaluation. If empty, no
  // language constraints are enforced.
  repeated string allowed_languages = 3;

  // Safety checks to run on the feature request message.
  repeated RequestSafetyCheck request_check = 4;

  // Configures when and how partial outputs are checked.
  // If unspecified, all partial outputs are suppressed, and only completed
  // outputs will be checked / returned to the feature code.
  optional PartialOutputChecks partial_output_checks = 8;

  // Configures the checks to run on raw model output.
  optional RawOutputCheck raw_output_check = 5;

  // Configures check that run on parsed model output, with request as context.
  repeated ResponseSafetyCheck response_check = 6;

  // When enabled, the unsafe text does not cancel a pending response, but
  // instead just doesn't send it until it is safe again. If the text is still
  // unsafe on complete, then it cancels the pending response.
  optional bool only_cancel_unsafe_response_on_complete = 7;
}