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

// This file contains a subset of the messages defined in go/zmdhe.
// Add additional messages as needed.

edition = "2023";

option optimize_for = LITE_RUNTIME;

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

package private_ai.proto;

message PrivateAiRequest {
  FeatureName feature_name = 1;

  oneof request {
    GenerateContentRequest generate_content_request = 4;
    AnonymousTokenRequest anonymous_token_request = 8;
    PaicMessage paic_request = 14;
  }

  // Unique id used to represent the request which can be used to correlate the
  // request and response. Useful when multiple concurrent requests and
  // potentially out-of-order responses (if they were sequential, and server
  // guaranteed in-order delivery of responses, this wouldn't be necessary) in a
  // single session.
  int32 request_id = 5;
}

message PrivateAiResponse {
  oneof response {
    GenerateContentResponse generate_content_response = 3;
    PaicMessage paic_response = 11;
  }

  // Unique id used to represent the response which can be used to correlate the
  // request and response. - Useful when multiple concurrent requests and
  // potentially out-of-order responses (if they were sequential, and server
  // guaranteed in-order delivery of responses, this wouldn't be necessary) in a
  // single session.
  int32 request_id = 4;
}

// LINT.IfChange(FeatureName)
enum FeatureName {
  reserved 702;

  FEATURE_NAME_UNSPECIFIED = 0;
  FEATURE_NAME_DEMO_GEMINI_GENERATE_CONTENT = 3;
  FEATURE_NAME_CHROME_ZERO_STATE_SUGGESTION = 700;
  FEATURE_NAME_CHROME_FORMS_AI = 701;

  FEATURE_NAME_CHROME_CONTEXTUAL_CUEING = 703;
  FEATURE_NAME_CHROME_AUTOMATED_PASSWORD_CHANGE = 704;
}
// LINT.ThenChange(
//     //tools/metrics/histograms/enums.xml:PrivateAiFeatureName,
//     //tools/metrics/histograms/metadata/private_ai/histograms.xml:PrivateAiFeatures
//)

message PaicMessage {
  FeatureName feature_name = 1;

  // extensions 530000000 to max;
  // The original PaicMessage defines extensions. However extensions create
  // static initializers which are not allowed in Chromium, so we use optional
  // fields instead which have the same serialized format.
  optimization_guide.proto.ExecuteRequest execute_request_ext = 530000002;
  optimization_guide.proto.ExecuteResponse execute_response_ext = 530000003;
  optimization_guide.proto.PrivateExecuteRequest private_execute_request_ext =
      530000016;
}

message GenerateContentRequest {
  // This field is meant for fast prototyping and will only be respected when
  // the feature name is FEATURE_NAME_DEMO_GEMINI_GENERATE_CONTENT.
  string model = 1;

  repeated Content contents = 2;
}

message Content {
  // The producer of the content. Must be either 'user' or 'model'.
  //
  // Useful to set for multi-turn conversations, otherwise can be left blank
  // or unset.
  string role = 1;

  // Ordered `Parts` that constitute a single message. Parts may have different
  // IANA MIME types.
  repeated Part parts = 2;
}

message Part {
  oneof data {
    // Text part (can be code).
    string text = 1;
    // Inlined bytes data.
    Blob inline_data = 2;
  }
}

// Content blob.
message Blob {
  // The IANA standard MIME type of the source data.
  string mime_type = 1;

  // Raw bytes.
  bytes data = 2;

  // When set indicates that the next message in a streaming connection
  // is to be interpreted as a continuation of this Blob.
  // The next message must contain only the Content of the same role as the
  // the one containing this Part, and must contain a single Part with
  // inline_data and the same mime_type as this Blob.
  bool will_continue = 3;
}

message GenerateContentResponse {
  repeated Candidate candidates = 1;
}

message Candidate {
  // Index of the candidate.
  int32 index = 1;

  // Content parts of the candidate.
  Content content = 2;

  enum FinishReason {
    FINISH_REASON_UNSPECIFIED = 0;
    STOP = 1;
    SAFETY = 2;
    BLOCKLIST = 3;
    PROHIBITED_CONTENT = 4;
  }

  FinishReason finish_reason = 3;
}

// Chrome-specific metadata for traffic segmentation.
// See go/privateinference-traffic-segmentation for the design.
message ChromeClientMetadata {
  // The browser product name (e.g., "Google Chrome", "Chromium").
  string product_name = 1;

  // The OS platform the client is running on.
  enum Platform {
    PLATFORM_UNKNOWN = 0;
    PLATFORM_WINDOWS = 1;
    PLATFORM_MAC = 2;
    PLATFORM_LINUX = 3;
    PLATFORM_CHROMEOS = 4;
    PLATFORM_ANDROID = 5;
    PLATFORM_IOS = 6;
  }
  Platform platform = 2;

  // The release channel of the browser.
  enum Channel {
    CHANNEL_UNKNOWN = 0;
    CHANNEL_STABLE = 1;
    CHANNEL_BETA = 2;
    CHANNEL_DEV = 3;
    CHANNEL_CANARY = 4;
  }
  Channel channel = 3;
}

// General client metadata wrapper for traffic segmentation.
// See go/privateinference-traffic-segmentation for the design.
message ClientMetadata {
  // Ensure only one platform's metadata is set at a time.
  oneof platform_metadata {
    ChromeClientMetadata chrome_client_metadata = 2;
  }
}

message AnonymousTokenRequest {
  bytes anonymous_token = 1;

  // Session-level metadata for traffic segmentation.
  // Sent inside the encrypted session, visible only within the TCB.
  ClientMetadata client_metadata = 15;
}