syntax = "proto3";

package node_service;

service NodeService {
  rpc SendPrompt (PromptRequest) returns (Tensor) {}
  rpc SendTensor (TensorRequest) returns (Tensor) {}
  rpc ResetShard (ResetShardRequest) returns (Empty) {}
  rpc GetInferenceResult (GetInferenceResultRequest) returns (InferenceResult) {}
  rpc CollectTopology (CollectTopologyRequest) returns (Topology) {}
}

message Shard {
  string model_id = 1;
  int32 start_layer = 2;
  int32 end_layer = 3;
  int32 n_layers = 4;
}

message PromptRequest {
  Shard shard = 1;
  string prompt = 2;
  optional string request_id = 3;
}

message TensorRequest {
  Shard shard = 1;
  Tensor tensor = 2;
  optional string request_id = 3;
}

message GetInferenceResultRequest {
  string request_id = 1;
}

message InferenceResult {
  optional Tensor tensor = 1;
  bool is_finished = 2;
}

message Tensor {
  bytes tensor_data = 1;
  repeated int32 shape = 2;
  string dtype = 3;
}

message ResetShardRequest {
  Shard shard = 1;
}

message CollectTopologyRequest {
  int32 max_depth = 1;
}

message Topology {
  map<string, DeviceCapabilities> nodes = 1;
  map<string, Peers> peer_graph = 2;
}

message Peers {
    repeated string peer_ids = 1;
}

message DeviceCapabilities {
  string model = 1;
  string chip = 2;
  int32 memory = 3;
}

message Empty {}