// Copyright (c) Meta Platforms, Inc. and affiliates.
// All rights reserved.
//
// This source code is licensed under the BSD-style license found in the
// LICENSE file in the root directory of this source tree.

syntax = "proto3";
import "google/protobuf/struct.proto";

package matrix.app_server.llm;

option java_multiple_files  = true;
option java_outer_classname = "global";


message ChatCompletionLogProb {
  string token = 1;
  float logprob = 2;
  repeated int32 bytes = 3;
}

message ChatCompletionLogProbsContent {
  string token = 1;
  float logprob = 2;
  repeated int32 bytes = 3;
  repeated ChatCompletionLogProb top_logprobs = 4;
}

message ChatCompletionLogProbs {
  repeated ChatCompletionLogProbsContent content = 1;
}

message CompletionLogProbs {
  repeated int32 text_offset = 1;
  repeated float token_logprobs = 2;
  repeated string tokens = 3;
}

// Usage Represents the total token usage per request to OpenAI.
message Usage {
  int32 prompt_tokens = 1;
  int32 completion_tokens = 2;
  int32 total_tokens = 3;
}

message FunctionCall {
  string name = 1;
  string arguments = 2;
}
message ToolCall {
  string id = 1;
  string type = 2; // Only "function" is valid
  FunctionCall function = 3;
}

// CompletionMessage is the message that is sent to the server.
message CompletionMessage {
  string role = 1;
  string content = 2;
  string refusal = 3;
  repeated FunctionCall function_call = 4;
  repeated ToolCall tool_calls = 5;

  // vLLM-specific fields that are not in OpenAI spec
  string reasoning_content = 6;
  string tool_call_id = 7;
}

message ChatCompletionResponseChoice {
  int64 index = 1;
  CompletionMessage message = 2;
  ChatCompletionLogProbs logprobs = 3;
  string finish_reason = 4;
  string stop_reason = 5;
}

message CompletionResponseChoice {
  int64 index = 1;
  string text = 2;
  CompletionLogProbs logprobs = 3;
  string finish_reason = 4;
  string stop_reason = 5;
  repeated TokenLogProbs prompt_logprobs = 6;
}

message TokenInfo {
  float logprob = 1;
  int32 rank = 2;
  string decoded_token = 3;
}

message TokenLogProbs {
  map<int32, TokenInfo> token_map = 1;
}

/**************** Tool Related ****************/

message ChatCompletionToolsParam {
  string type = 1; // vllm only support "function" https://github.com/vllm-project/vllm/blob/main/vllm/entrypoints/openai/protocol.py#L193
  FunctionDefinition function = 2;  
}

message FunctionDefinition {
  string name = 1;
  string description = 2;
  google.protobuf.Struct parameters = 3; // simple example; adjust to your needs
}

/**************** Request ****************/

// ChatCompletionRequest represents a request structure for chat completion API.
message ChatCompletionRequest {
  string model = 1;
  repeated CompletionMessage messages = 2;
  float temperature = 3;
  float top_p = 4;
  int32 n = 5;
  bool logprobs = 6;
  string stop = 7;
  int32 max_tokens = 8;
  float presence_penalty = 9;
  float frequency_penalty = 10;
  map<string, int64> logit_bias = 11;
  string user = 12;
  int32 seed = 13;
  repeated ChatCompletionToolsParam tools = 14;
  string tool_choice = 15;
}

// CompletionRequest represents a request structure for chat completion API.
message CompletionRequest {
  string model = 1;
  string prompt = 2;
  float temperature = 3;
  float top_p = 4;
  int32 n = 5;
  bool logprobs = 6;
  string stop = 7;
  int32 max_tokens = 8;
  float presence_penalty = 9;
  float frequency_penalty = 10;
  map<string, int64> logit_bias = 11;
  string user = 12;
  int32 seed = 13;
  int32 prompt_logprobs = 14;
}


/**************** Response ****************/

// ChatCompletionResponse represents a response structure for chat completion API.
message ChatCompletionResponse {
  string id = 1;
  string object = 2;
  int64 created = 3;
  string model = 4;
  repeated ChatCompletionResponseChoice choices = 5;
  Usage usage = 6;
}

// CompletionResponse represents a response structure for completion API.
message CompletionResponse {
  string id = 1;
  string object = 2;
  int64 created = 3;
  string model = 4;
  repeated CompletionResponseChoice choices = 5;
  Usage usage = 6;
}

// OpenaiService The GPT service definition.
service OpenaiService {
  rpc CreateChatCompletion(ChatCompletionRequest) returns (ChatCompletionResponse) {}
  rpc CreateCompletion(CompletionRequest) returns (CompletionResponse) {}
}