// Copyright 2024 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// NOTICE: run `make generate-protos` if making changes to this file

syntax = "proto3";

package jetstream_proto;

// TODO: Merge this with main JetStream core once we settle on an API.

service Orchestrator {
  // Query LLM to generate text or tokens.
  rpc Decode(DecodeRequest) returns (stream DecodeResponse) {}
  // Checks if the model server is live.
  rpc HealthCheck(HealthCheckRequest) returns (HealthCheckResponse) {}
}

message DecodeRequest {
  // The maximum output length of a sequence. It's used in JetStream to control
  // the output/decode length of a sequence. It would not be used in the engine.
  // We should always set max_tokens <= (max_target_length -
  // max_prefill_predict_length). max_target_length is the maximum length of a
  // sequence; max_prefill_predict_length is the maximum length of the
  // input/prefill of a sequence.
  int32 max_tokens = 4;

  message TextContent {
    string text = 1;
  }
  message TokenContent {
    repeated int32 token_ids = 1;
  }

  // The client can pass the inputs either as a string, in which case the server will
  // tokenize it, or as tokens, in which case it's the client's responsibility to
  // ensure they tokenize its input strings with the correct tokenizer.
  oneof content {
    TextContent text_content = 5;
    TokenContent token_content = 6;
  }

  message Metadata {
    float start_time = 1;
  }

  oneof metadata_optional {
    Metadata metadata = 7;
  }

  int32 num_samples = 8;

  reserved 1, 2, 3;
  // Next ID: 9
}

message DecodeResponse {
  // InitialContent supports returning initial one-off response data from the
  // stream. It's a placeholder for future features such as history cache.
  message InitialContent {}
  message StreamContent {
    message Sample {
      // The text string decoded from token id(s).
      string text = 1;
      // List of token ids, one list per sample. When speculative decoding is disabled, the list size should be 1; When speculative decoding is enabled, the list size should be >= 1.
      repeated int32 token_ids = 2;
    }
    // Supports multiple samples in the StreamContent. The Sample list size depends on text generation strategy the engine used.
    repeated Sample samples = 1;
  }

  oneof content {
    InitialContent initial_content = 2;
    StreamContent stream_content = 3;
  }
  reserved 1;
  // Next ID: 4
}

message HealthCheckRequest {}

message HealthCheckResponse {
  // Denotes whether the model server is live
  bool is_live = 1;
}