// Copyright 2024 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/common_types.proto";
import "components/optimization_guide/proto/descriptors.proto";
import "components/optimization_guide/proto/fieldwise_parser_config.proto";
import "components/optimization_guide/proto/redaction.proto";
import "components/optimization_guide/proto/substitution.proto";
import "components/optimization_guide/proto/model_execution.proto";
import "components/optimization_guide/proto/parser_kind.proto";

message OnDeviceModelExecutionConfig {
  // The set of configs for features that leverage the on-device model.
  //
  // It is expected that there is only one feature config per feature.
  repeated OnDeviceModelExecutionFeatureConfig feature_configs = 1;

  // A config used for validating the model and device after download.
  optional OnDeviceModelValidationConfig validation_config = 2;

  // List of capabilities supported by the base model.
  repeated OnDeviceModelCapability capabilities = 3;
}

enum OnDeviceModelCapability {
  ON_DEVICE_MODEL_CAPABILITY_UNSPECIFIED = 0;
  // The model supports taking image input.
  ON_DEVICE_MODEL_CAPABILITY_IMAGE_INPUT = 1;
  // The model supports taking audio input.
  ON_DEVICE_MODEL_CAPABILITY_AUDIO_INPUT = 2;
}

message OnDeviceModelExecutionFeatureConfig {
  // The feature this configuration is for.
  optional ModelExecutionFeature feature = 1;

  // The config used to construct the input for on-device model execution.
  optional OnDeviceModelExecutionInputConfig input_config = 2;

  // The config used to construct the output for on-device model execution.
  optional OnDeviceModelExecutionOutputConfig output_config = 3;

  // The config used to construct the request for server fallback text safety
  // evaluation.
  optional TextSafetyFallbackConfig text_safety_fallback_config = 4;

  // Whether text safety can be skipped for this feature, if not configured.
  optional bool can_skip_text_safety = 5;

  // Sampling parameters to use with the model.
  // These will override global defaults, but not per session configurations.
  optional SamplingParams sampling_params = 6;

  // Feature defined metadata.
  optional Any feature_metadata = 7;

  // The rank of the LORA associated with this feature.
  optional uint32 adaptation_rank = 8;
}

message OnDeviceModelExecutionInputConfig {
  // The base name of the request metadata proto this input config is applicable
  // for.
  optional string request_base_name = 1;

  // An ordered list of substituted strings to apply for input context.
  //
  // These will be concatenated in the order they appear here if the conditions
  // apply based on the input request.
  repeated SubstitutedString input_context_substitutions = 3;

  // An ordered list of substituted strings to apply when the model is executed.
  //
  // These will be concatenated in the order they appear here if the conditions
  // apply based on the input request.
  //
  // It is expected that the resulting string here will be concatenated with the
  // resulting string for the input context if `should_ignore_input_context` is
  // not set on any of the used substitutions.
  repeated SubstitutedString execute_substitutions = 2;

  // The number of tokens generated from `input_context_substitutions` that are
  // guaranteed to be processed before context processing can be cancelled
  // (truncating any unprocessed tokens) to begin execution. A default value is
  // used when this is not set. To disable cancellations, set this to be the
  // same as `max_context_tokens`.
  optional uint32 min_context_tokens = 4;

  // The maximum number of tokens that can be generated from
  // `input_context_substitutions`. Tokens in excess of this number will be
  // truncated. Uses a default value when unset.
  optional uint32 max_context_tokens = 5;

  // The maximum number of tokens that can be generated from
  // `execute_substitutions`. Tokens in excess of this number will be truncated.
  // Uses a default value when unset.
  optional uint32 max_execute_tokens = 6;
}

// Rules that generated output must follow.
message ResponseConstraint {
  oneof format {
    // A JSON Schema describing an llguidance constraint on the output. See
    // https://github.com/guidance-ai/llguidance/blob/main/docs/json_schema.md
    string json_schema = 1;
    // A regex describing the llguidance constraint on the output.
    // See https://docs.rs/regex/latest/regex/#syntax for syntax.
    string regex = 2;
  }
}

message OnDeviceModelExecutionOutputConfig {
  reserved 7;

  // The proto type to use for the response metadata.
  optional string proto_type = 1;

  // The proto field to populate the output string with. Applies for
  // PARSER_KIND_SIMPLE only.
  optional ProtoField proto_field = 2;

  // Rules that result in redacting the output.
  optional RedactRules redact_rules = 3;

  // Which output parsing implementation to use.
  optional ParserKind parser_kind = 4;

  // A config for specifying parsing logic for the fieldwise parser.
  optional FieldwiseParserConfig fieldwise_parser_config = 8;

  // Whether to suppress parsing incomplete output.
  optional bool suppress_parsing_incomplete_output = 5;

  // The maximum number of tokens that can be generated as output.
  optional uint32 max_output_tokens = 6;

  // The response constraint to use when generating output.
  optional ResponseConstraint response_constraint = 9;
}

message TextSafetyFallbackConfig {
  // The proto field in the input request that contains the URL this request was
  // activated on, if any.
  optional ProtoField input_url_proto_field = 1;
}

message OnDeviceModelValidationConfig {
  // Prompts and responses used for basic model validation.
  repeated ValidationPrompt validation_prompts = 1;
}

message ValidationPrompt {
  // The input prompt to send.
  optional string prompt = 1;

  // A string that the output is expected to contain. This is not case
  // sensitive.
  optional string expected_output = 2;
}

message SamplingParams {
  optional uint32 top_k = 1;
  optional float temperature = 2;
}