syntax = "proto3";

package nominal.ai.v1;

import 'google/protobuf/timestamp.proto';

// AIAgentService provides AI-powered assistance for general operations
service AIAgentService {
  // StreamChat handles bidirectional streaming chat for AI agent
  rpc StreamChat(StreamChatRequest) returns (stream StreamChatResponse);
  // GetConversation [DEPRACATED] handles getting a complete conversation list, with an optional limit on number of messages returned
  rpc GetConversation(GetConversationRequest) returns (GetConversationResponse) {
    option deprecated = true;
  }
  // GetConversationMetadata handles getting the conversation metadata like title, current token count etc
  rpc GetConversationMetadata(GetConversationMetadataRequest) returns (GetConversationMetadataResponse);
  // GetConversationMessages handles retrieving the set of conversation messages. Supports pagination
  rpc GetConversationMessages(GetConversationMessagesRequest) returns (GetConversationMessagesResponse);
  // ListConversations handles getting the list of conversation ids ordered by most recently updated
  rpc ListConversations(ListConversationsRequest) returns (ListConversationsResponse);
  // CreateConversation handles creating a conversation and assigning it a conversation rid
  rpc CreateConversation(CreateConversationRequest) returns (CreateConversationResponse);
  // UpdateConversationMetadata handles updating any metadata associated with a conversation
  rpc UpdateConversationMetadata(UpdateConversationMetadataRequest) returns (UpdateConversationMetadataResponse);
  // DeleteConversation handles deleting a specific conversation by conversation rid
  rpc DeleteConversation(DeleteConversationRequest) returns (DeleteConversationResponse);
  // GetSnapshotRidByUserMessageId handles resolving the snapshot rid of the workbook at the time the message is sent
  rpc GetSnapshotRidByUserMessageId(GetSnapshotRidByUserMessageIdRequest) returns (GetSnapshotRidByUserMessageIdResponse);
  // CompactConversation handles compacting the conversation context into approximately half its original size
  rpc CompactConversation(CompactConversationRequest) returns (CompactConversationResponse);
}

message GetSnapshotRidByUserMessageIdRequest {
  string conversation_rid = 1;
  string message_id = 2;
}

// Will return an empty response body in the case where the message id exists, but there is no associated snapshot
// This occurs in the instance where a message was sent in a non-workbook context
message GetSnapshotRidByUserMessageIdResponse {
  optional string snapshot_rid = 1;
}

// ReadOnlyMode configures read-only mode where edit tools are shadowed/disabled
message ReadOnlyMode {}

// EditMode configures edit mode where all tools are available
message EditMode {
  // when set to true, we auto accept edits for any tools typically requiring approval
  optional bool auto_accept = 1;
}

// ConversationMode specifies the mode of the conversation
message ConversationMode {
  oneof mode {
    ReadOnlyMode read_only = 1;
    EditMode edit = 2;
  }
}

// When the agent makes a request to use a tool, the user responds
// with one of these for every request - mapping a tool id to its approval/denial
message ToolApprovalResult {
  string tool_call_id = 1; // identifies the tool call
  oneof response {
    ToolApprovedResponse approved = 2;
    ToolDeniedResponse denied = 3;
  }
}

message ToolApprovedResponse {
  // json string representation of the override argument if the user
  // needs it to be changed in some way
  optional string override_args = 1;
}

message ToolDeniedResponse {
  string denial_reason = 2;
}

// RetryRequest retries the last request (e.g., if it was interrupted/failed part-way through)
message RetryRequest {}

// UserPromptRequest contains a new user message
message UserPromptRequest {
  UserModelMessage message = 1;
  // Optional: image files to provide to the agent
  repeated ImagePart images = 2;
}

// ToolApprovalRequest contains tool approval results
message ToolApprovalRequest {
  repeated ToolApprovalResult tool_approvals = 1;
}

// StreamChatRequest is a request to stream chat messages for AI agent.
message StreamChatRequest {
  // The conversation ID
  string conversation_rid = 1;

  // DEPRECATED: use request_type.user_prompt.message instead
  UserModelMessage message = 2 [deprecated = true];
  // DEPRECATED: use request_type.user_prompt.images instead
  repeated ImagePart images = 3 [deprecated = true];
  // DEPRECATED: use request_type.tool_approval instead
  repeated ToolApprovalResult tool_approvals = 6 [deprecated = true];

  // The type of request - exactly one should be set
  oneof request_type {
    RetryRequest retry = 7;
    UserPromptRequest user_prompt = 8;
    ToolApprovalRequest tool_approval = 9;
  }

  // Context-specific fields based on the oneofKind.
  oneof context {
    WorkbookContext workbook = 4;
    GlobalContext global = 5;
    ChecklistContext checklist = 10;
    TemplateContext template = 11;
  }
}

// WorkbookContext contains workbook-specific context fields
message WorkbookContext {
  // RID of the workbook to use for context
  string workbook_rid = 1;
  // The user's presence in the workbook
  optional WorkbookUserPresence user_presence = 2 [deprecated = true];
}

// ChecklistContext for use when the agent is being messaged in the context of a
// checklist currently being edited
message ChecklistContext {
  // RID of the checklist being edited
  string checklist_rid = 1;
  // Name of the branch that this edit is being done on
  // This is equivalent to the 'draft name' for a check being edited in the UI
  string branch_name = 2;
  // A checklist opened in edit mode will always have some resource by which it
  // is referencing for viewing. It can either be an asset or a run
  oneof reference_rid {
    string asset = 3;
    string run = 4;
  }
}

// TemplateContext for use when the agent is being messaged in the context of a
// template currently being edited.
// This provides the agent with the same toolset as the workbook tools
message TemplateContext {
  // RID of template being edited
  string template_rid = 1;
  // Name of the branch this edit is being done on
  string branch_name = 2;
  // Defines the datascope in view in the editor
  oneof reference_rid {
    string asset = 3;
    string run = 4;
  }
}

// DefaultContext (no context)
message GlobalContext {}

// WorkbookUserPresence contains the user's presence in the workbook
// which is used to describe what the user is viewing at the time of the message.
message WorkbookUserPresence {
  int32 tab_index = 1;
  TimeRange range = 2;
}

// CreateConversation request will create a new conversation thread
// if old conversation id is not set, a brand new, clear chat is created
// If old conversation id is set without a previous message id, the full conversation thread will be copied
// if previous message id is set with a previous message id, the conversation thread up until that message will be copied
// the above case is useful for branching a conversation into a new thread
message CreateConversationRequest {
  string title = 1;
  string workspace_rid = 2;
  optional string old_conversation_rid = 3;
  optional string previous_message_id = 4;
  ConversationMode conversation_mode = 5;
}

// CreateConversationResponse will return the conversation id for the new conversation
message CreateConversationResponse {
  string new_conversation_rid = 1;
}

// Updates the fields if specified (optional means no change for that field)
message UpdateConversationMetadataRequest {
  optional string title = 1;
  string conversation_rid = 2;
  optional ConversationMode conversation_mode = 3;
}

message UpdateConversationMetadataResponse {}

message DeleteConversationRequest {
  string conversation_rid = 1;
}

message DeleteConversationResponse {}

// a GetConversationRequest allows you to retrieve a subset of messages from a conversation thread represented
// by provided rid. To start from a particular message - you can also provide a message id.
message GetConversationRequest {
  string conversation_rid = 1;
  optional string page_start_message_id = 2;
  optional int32 max_message_count = 3;
}

// a CompactConversationRequest allows you to reduce the token count in your conversation by up to half
// will be a no op if current conversation has not yet reached at least half of the max token count
message CompactConversationRequest {
  string conversation_rid = 1;
}

// returns the new token count of the now compacted conversation
message CompactConversationResponse {
  ContextStatus context = 1;
}

// Model message with id allows you to identify the message ID of a given message
message ModelMessageWithId {
  oneof content {
    ModelMessage message = 1;
    ToolAction tool_action = 2;
    ToolActionConfirmation tool_action_confirmation = 6;
  }
  string message_id = 3;
  // WB agent user messages can have snapshot rids associated with them
  optional string snapshot_rid = 4;
  repeated ToolCallDescription tool_approval_requests = 5;
}

message GetConversationResponse {
  repeated ModelMessageWithId ordered_messages = 1;
  ConversationMetadata conversation_metadata = 2;
}

message GetConversationMetadataRequest {
  string conversation_rid = 1;
}

message GetConversationMetadataResponse {
  ConversationMetadata conversation_metadata = 1;
}

message GetConversationMessagesRequest {
  string conversation_rid = 1;
  optional string page_start_message_id = 2;
  optional int32 max_message_count = 3;
}

message GetConversationMessagesResponse {
  repeated ModelMessageWithId ordered_messages = 1;
  optional string next_page_token = 2;
}

// Will generate all conversation threads that this user has in this workspace
message ListConversationsRequest {
  string workspace_rid = 1;
  optional string next_page_token = 2;
  optional int32 page_size = 3;
}

message ConversationMetadata {
  string conversation_rid = 1;
  string title = 2;
  google.protobuf.Timestamp created_at = 3;
  google.protobuf.Timestamp last_updated_at = 4;
  ConversationMode mode = 5;
  ContextStatus current_context = 6;
}

// ListConversationsResponse is a list of conversations that can be used in a call to GetConversationRequest
// to get a full conversation from storage. These are ordered by creation time.
message ListConversationsResponse {
  repeated ConversationMetadata conversations = 1;
  optional string next_page_token = 2;
}

message TimeRange {
  Timestamp range_start = 1;
  Timestamp range_end = 2;
}

message Timestamp {
  int32 seconds = 1;
  int32 nanoseconds = 2;
}

// ModelMessage is a discriminated union of system, user, assistant, and tool messages.
// Each message type has its own structure and content.
message ModelMessage {
  oneof kind {
    UserModelMessage user = 1;
    AssistantModelMessage assistant = 2;
  }
}

// A user message containing text
message UserModelMessage {
  repeated UserContentPart text = 1;
}

// An assistant message containing text
message AssistantModelMessage {
  repeated AssistantContentPart content_parts = 1;
}

message UserContentPart {
  oneof part {
    TextPart text = 1;
    ImagePart image = 2;
  }
}

// Content part for assistant messages: can be text, reasoning, or mutation.
message AssistantContentPart {
  oneof part {
    TextPart text = 1;
    ReasoningPart reasoning = 2;
  }
}

// Text part for user or assistant messages.
message TextPart {
  string text = 1;
}

// User-supplied image part.
message ImagePart {
  // The media type of the image (e.g. "image/png", "image/jpeg")
  string media_type = 2;
  // Attachment RID for image content stored in S3
  string attachment_rid = 4;
}

// Reasoning part for assistant messages.
message ReasoningPart {
  string reasoning = 1;
}

// StreamChatResponse is a discriminated union response to a StreamChatRequest
message StreamChatResponse {
  oneof response {
    Finish finish = 1;
    Error error = 2;
    TextStart text_start = 3;
    TextDelta text_delta = 4;
    TextEnd text_end = 5;
    ReasoningStart reasoning_start = 6;
    ReasoningDelta reasoning_delta = 7;
    ReasoningEnd reasoning_end = 8;
    ToolAction tool_action = 10;
    ToolActionConfirmation tool_action_confirmation = 11;
  }
}

enum ToolCallStatus {
  TOOL_CALL_STATUS_UNSPECIFIED = 0;
  TOOL_CALL_STATUS_APPROVED = 1;
  TOOL_CALL_STATUS_DENIED = 2;
  TOOL_CALL_STATUS_AWAITING_APPROVAL = 3;
}

message ToolCallDescription {
  string tool_call_id = 1;
  string tool_name = 2;
  // string representation of the proposed tool args for display
  string tool_args_json_string = 3;

  // used to conditionally render an approval button based the outcome of the tool call
  ToolCallStatus status = 4;
}

// Indicates the end of a chat session
message Finish {
  // The message ids in order of all generated messages for this agent run
  // These ids can be used to branch a message from that specific message
  repeated string ordered_message_ids = 1;

  // In the case that this is the first agent run in a conversation thread, we also
  // return the new conversation title generated
  optional string new_title = 2;

  repeated ToolCallDescription tool_approval_requests = 3;

  optional ContextStatus updated_context = 4;
}

// An error that occurred during the chat session
message Error {
  string message = 1;
}

// Indicates the start of a text message from the agent
message TextStart {
  // uniquely identifies the text message (e.g. uuid) so that the client can
  // merge parallel message streams (if it happens).
  string id = 1;
}

// A delta (continuation) of a text message from the agent
message TextDelta {
  string id = 1;
  // The next chunk of text
  string delta = 2;
}

// Indicates the end of a text message from the agent
message TextEnd {
  string id = 1;
}

// Indicates the start of a reasoning message from the agent
message ReasoningStart {
  string id = 1;
}

// A delta (continuation) of a reasoning message from the agent
message ReasoningDelta {
  string id = 1;
  // The next chunk of reasoning
  string delta = 2;
}

// Indicates the end of a reasoning message from the agent
message ReasoningEnd {
  string id = 1;
}

// ToolAction is concerned with the display verb and target used to identify a tool
// call without using its exact internal tool name.
//
// For example:
//   verb=Add target=Variable may represent tool_name = add_channel_variable
message ToolAction {
  string tool_action_verb = 2;
  optional string tool_target = 3;
  ToolCallDescription tool_call_description = 5;

  // Deprecated: use tool_call_description.tool_call_id.
  string id = 1 [deprecated = true];
  // Deprecated: use tool_call_description.tool_name.
  string tool_name = 4 [deprecated = true];
}

// When we see a ToolAction in the stream, we wait for the corresponding ToolActionConfirmation
// to indicate whether or not the tool call has successfully executed
message ToolActionConfirmation {
  string id = 1;
  oneof outcome {
    ToolActionSuccess success = 2;
    ToolActionFailure failure = 3;
  }
}

message ToolActionSuccess {
  string tool_success_message = 1;
}

message ToolActionFailure {
  string tool_error_message = 1;
}

// ContextStatus represents the current token usage of a conversation relative to the selected model's context limit.
// NOTE: A curr_token_count of zero means the token count is not yet known (e.g., no prompt has been sent yet,
// or the conversation was just branched). It does not necessarily mean the conversation is empty.
message ContextStatus {
  int32 curr_token_count = 1;
  int32 model_context_limit = 2;
}
