syntax = "proto3";

package exa.opensearch_clients_pb;

import "buf/validate/validate.proto";
import "exa/chat_pb/chat.proto";
import "exa/codeium_common_pb/codeium_common.proto";
import "exa/index_pb/index.proto";
import "google/protobuf/timestamp.proto";

option go_package = "github.com/Exafunction/Exafunction/exa/opensearch_clients_pb";

service KnowledgeBaseService {
    rpc KnowledgeBaseSearch(KnowledgeBaseSearchRequest) returns (KnowledgeBaseSearchResponse);
    rpc GetKnowledgeBaseScopeItems(GetKnowledgeBaseScopeItemsRequest) returns (GetKnowledgeBaseScopeItemsResponse);
    rpc GetKnowledgeBaseItemsFromScopeItems(GetKnowledgeBaseItemsFromScopeItemsRequest) returns (GetKnowledgeBaseItemsFromScopeItemsResponse);
    rpc IngestSlackData(IngestSlackDataRequest) returns (IngestSlackDataResponse);
    rpc IngestGithubData(IngestGithubDataRequest) returns (IngestGithubDataResponse);
    rpc IngestGoogleDriveData(IngestGoogleDriveDataRequest) returns (IngestGoogleDriveDataResponse);
    rpc IngestJiraData(IngestJiraDataRequest) returns (IngestJiraDataResponse);
    rpc IngestJiraPayload(IngestJiraPayloadRequest) returns (IngestJiraPayloadResponse);
    rpc ForwardSlackPayload(ForwardSlackPayloadRequest) returns (ForwardSlackPayloadResponse);
    rpc IngestSlackPayload(IngestSlackPayloadRequest) returns (IngestSlackPayloadResponse);
    rpc ConnectKnowledgeBaseAccount(ConnectKnowledgeBaseAccountRequest) returns (ConnectKnowledgeBaseAccountResponse);
    rpc DeleteKnowledgeBaseConnection(DeleteKnowledgeBaseConnectionRequest) returns (DeleteKnowledgeBaseConnectionResponse);
    rpc UpdateConnectorConfig(UpdateConnectorConfigRequest) returns (UpdateConnectorConfigResponse);
    rpc CancelKnowledgeBaseJobs(CancelKnowledgeBaseJobsRequest) returns (CancelKnowledgeBaseJobsResponse);
    rpc GetKnowledgeBaseConnectorState(GetKnowledgeBaseConnectorStateRequest) returns (GetKnowledgeBaseConnectorStateResponse);
    rpc GetKnowledgeBaseJobStates(GetKnowledgeBaseJobStatesRequest) returns (GetKnowledgeBaseJobStatesResponse);
    rpc AddUsers(AddUsersRequest) returns (AddUsersResponse);
    rpc AddGithubUsers(AddGithubUsersRequest) returns (AddGithubUsersResponse);
    rpc GetKnowledgeBaseWebhookUrl(GetKnowledgeBaseWebhookUrlRequest) returns (GetKnowledgeBaseWebhookUrlResponse);
    rpc GetConnectorInternalConfig(GetConnectorInternalConfigRequest) returns (GetConnectorInternalConfigResponse);
}

service CodeIndexService {
    rpc OpenSearchAddRepository(OpenSearchAddRepositoryRequest) returns (OpenSearchAddRepositoryResponse);
    rpc OpenSearchGetIndex(OpenSearchGetIndexRequest) returns (OpenSearchGetIndexResponse);
    rpc HybridSearch(HybridSearchRequest) returns (HybridSearchResponse);
    rpc GraphSearch(GraphSearchRequest) returns (GraphSearchResponse);
}

message TimeRange {
    google.protobuf.Timestamp start = 1;
    google.protobuf.Timestamp end = 2;
}

message GithubUser {
    string auth_uid = 1;
    string username = 2;
}

message AddGithubUsersRequest {
    repeated GithubUser users = 1;
}

message AddGithubUsersResponse {
    
}

message UserInfo {
    string auth_uid = 1;
    string email = 2;
    string name = 3;
    string photo_url = 4;
}

message AddUsersRequest {
    repeated UserInfo users = 1;
}

message AddUsersResponse {
    
}

message KnowledgeBaseSearchRequest {
    int64 max_results = 2;
    repeated string queries = 3;
    codeium_common_pb.Metadata metadata = 4;
    repeated string urls = 12;
    repeated string document_ids = 13;
    repeated string aggregate_ids = 5;
    repeated chat_pb.ChatMessagePrompt chat_message_prompts = 6;
    TimeRange time_range = 7;
    repeated codeium_common_pb.DocumentType document_types = 14;
    SearchMode search_mode = 9;
    bool disable_reranking = 10;
    bool disable_contextual_lookup = 11;
    repeated codeium_common_pb.IndexChoice index_choices = 8;
    string query = 1 [deprecated = true];
}

message KnowledgeBaseSearchResponse {
    repeated codeium_common_pb.KnowledgeBaseGroup knowledge_base_groups = 1;
}

message GetKnowledgeBaseScopeItemsRequest {
    string query = 1;
    codeium_common_pb.Metadata metadata = 3;
    repeated codeium_common_pb.DocumentType document_types = 5;
    repeated codeium_common_pb.IndexChoice index_choices = 4 [deprecated = true];
    repeated string index_names = 2 [deprecated = true];
}

message GetKnowledgeBaseScopeItemsResponse {
    repeated codeium_common_pb.KnowledgeBaseScopeItem scope_items = 1;
}

message GetKnowledgeBaseItemsFromScopeItemsRequest {
    codeium_common_pb.Metadata metadata = 2;
    repeated codeium_common_pb.KnowledgeBaseScopeItem scope_items = 3;
    
    reserved 1;
}

message GetKnowledgeBaseItemsFromScopeItemsResponse {
    repeated codeium_common_pb.KnowledgeBaseItemWithMetadata knowledge_base_items_with_metadata = 1;
}

message IngestSlackDataRequest {
    index_pb.ManagementMetadata metadata = 1;
    repeated string channel_ids = 2;
}

message IngestSlackDataResponse {
    
}

message IngestGithubDataRequest {
    index_pb.ManagementMetadata metadata = 3;
    string organization = 1;
    string repository = 2;
}

message IngestGithubDataResponse {
    
}

message IngestGoogleDriveDataRequest {
    index_pb.ManagementMetadata metadata = 2;
    repeated string folder_ids = 3;
    
    reserved 1;
}

message IngestGoogleDriveDataResponse {
    
}

message IngestJiraDataRequest {
    index_pb.ManagementMetadata metadata = 4;
    
    reserved 1, 2, 3;
}

message IngestJiraDataResponse {
    
}

message IngestJiraPayloadRequest {
    string body = 3;
}

message IngestJiraPayloadResponse {
    
}

message ForwardResult {
    ForwardStatus status = 1;
    
    oneof _error {
        string error = 2;
    }
}

message ForwardSlackPayloadRequest {
    repeated string bodies = 1;
}

message ForwardSlackPayloadResponse {
    repeated ForwardResult results = 1;
}

message IngestSlackPayloadRequest {
    repeated SlackPayload payload = 1;
}

message IngestSlackPayloadResponse {
    
}

message CommonDocument {
    string document_id = 1;
    string text = 2;
}

message CommonDocumentWithScore {
    CommonDocument document = 1;
    float score = 2;
}

message SearchResult {
    string text = 1;
    string url = 2;
}

message QuerySearchResponse {
    repeated CommonDocumentWithScore document_with_scores = 1;
}

message OpenSearchAddRepositoryRequest {
    index_pb.ManagementMetadata metadata = 1;
    index_pb.RepositoryConfig config = 2;
    index_pb.RequestIndexVersion initial_index = 3;
}

message OpenSearchAddRepositoryResponse {
    string repo_name = 1;
    string index_id = 2;
}

message OpenSearchGetIndexRequest {
    string index_id = 1;
}

message OpenSearchGetIndexResponse {
    index_pb.IndexingStatus status = 1;
}

message HybridSearchRequest {
    string query = 1;
    codeium_common_pb.Embedding embedding = 2;
    int64 max_results = 3;
}

message HybridSearchResponse {
    repeated CommonDocumentWithScore document_with_scores = 1;
}

message GraphSearchRequest {
    string query = 1;
    codeium_common_pb.Embedding embedding = 2;
    int64 max_results = 3;
}

message GraphSearchResponse {
    repeated CommonDocumentWithScore document_with_scores = 1;
}

message ConnectorConfig {
    oneof config {
        ConnectorConfigSlack slack = 1;
        ConnectorConfigGithub github = 2;
        ConnectorConfigGoogleDrive google_drive = 3;
        ConnectorConfigJira jira = 4;
    }
}

message ConnectorConfigSlack {
    repeated string include_channel_ids = 3;
    repeated string exclude_channel_ids = 4;
    
    reserved 1, 2;
}

message ConnectorConfigGithub {
    reserved 1;
}

message ConnectorConfigGoogleDrive {
    repeated string include_drive_ids = 2;
    repeated string exclude_drive_ids = 3;
    
    reserved 1;
}

message ConnectorConfigJira {
    
}

message ConnectorInternalConfig {
    oneof config {
        ConnectorInternalConfigSlack slack = 1;
        ConnectorInternalConfigGithub github = 2;
        ConnectorInternalConfigGoogleDrive google_drive = 3;
        ConnectorInternalConfigJira jira = 4;
    }
}

message ConnectorInternalConfigSlack {
    string client_id = 2;
    string client_secret = 3;
    string signing_secret = 1;
    
    reserved 4;
}

message GithubRepoConfig {
    string organization = 1;
    string repository = 2;
}

message ConnectorInternalConfigGithub {
    int64 installation_id = 1;
    repeated GithubRepoConfig repo_configs = 2;
}

message ConnectorInternalConfigGoogleDrive {
    reserved 1;
}

message ConnectorInternalConfigJira {
    int64 webhook_id = 1;
}

message ConnectKnowledgeBaseAccountRequest {
    index_pb.ManagementMetadata metadata = 7;
    ConnectorType connector = 2;
    string access_token = 3;
    google.protobuf.Timestamp access_token_expires_at = 4;
    string refresh_token = 5;
    google.protobuf.Timestamp refresh_token_expires_at = 6;
    ConnectorAdditionalParams additional_params = 8;
    
    reserved 1;
}

message DeleteKnowledgeBaseConnectionRequest {
    index_pb.ManagementMetadata metadata = 1;
    ConnectorType connector = 2;
}

message DeleteKnowledgeBaseConnectionResponse {
    
}

message UpdateConnectorConfigRequest {
    index_pb.ManagementMetadata metadata = 1;
    ConnectorType connector = 2;
    ConnectorConfig config = 3;
}

message UpdateConnectorConfigResponse {
    
}

message ConnectorAdditionalParams {
    oneof config {
        ConnectorAdditionalParamsSlack slack = 2;
        ConnectorAdditionalParamsGithub github = 1;
    }
}

message ConnectorAdditionalParamsSlack {
    string client_id = 1;
    string client_secret = 2;
    string signing_secret = 3;
}

message ConnectorAdditionalParamsGithub {
    int64 installation_id = 1;
}

message ConnectKnowledgeBaseAccountResponse {
    
}

message CancelKnowledgeBaseJobsRequest {
    index_pb.ManagementMetadata metadata = 1;
    repeated int64 job_ids = 2;
}

message CancelKnowledgeBaseJobsResponse {
    
}

message DocumentTypeCount {
    codeium_common_pb.DocumentType document_type = 1;
    int64 count = 2;
}

message ConnectorState {
    ConnectorType connector = 1;
    bool initialized = 2;
    ConnectorConfig config = 3;
    repeated DocumentTypeCount document_type_counts = 4;
    google.protobuf.Timestamp last_indexed_at = 5;
    google.protobuf.Timestamp unhealthy_since = 6;
    google.protobuf.Timestamp last_configured_at = 7;
}

message GetKnowledgeBaseConnectorStateRequest {
    index_pb.ManagementMetadata metadata = 1;
    
    reserved 2;
}

message GetKnowledgeBaseConnectorStateResponse {
    repeated ConnectorState connector_states = 1;
}

message JobState {
    ConnectorType connector = 1;
    int64 id = 2;
    JobStatus status = 3;
}

message GetKnowledgeBaseJobStatesRequest {
    index_pb.ManagementMetadata metadata = 1;
    repeated ConnectorType connector_types = 2;
}

message GetKnowledgeBaseJobStatesResponse {
    repeated JobState job_states = 1;
}

message SlackMessagePayload {
    string dataset_id = 1;
    string previous_message_dataset_id = 2;
    string type = 3;
    string channel_id = 4;
    string user = 5;
    string text = 6;
    string timestamp = 7;
    string thread_timestamp = 8;
    string channel_name = 9;
    string team_name = 10;
    string team_id = 11;
    bool is_private_channel = 12;
    string team_domain = 13;
    string original_timestamp = 14;
}

message SlackChannelPayload {
    string type = 1;
    string channel_id = 2;
    string channel_name = 4;
    string description = 7;
    string team_id = 8;
    bool is_private_channel = 9;
    
    reserved 3, 5, 6;
}

message SlackPayload {
    oneof payload {
        SlackMessagePayload message = 13;
        SlackChannelPayload channel = 14;
    }
    
    reserved 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12;
}

message GetKnowledgeBaseWebhookUrlRequest {
    index_pb.ManagementMetadata metadata = 1;
}

message GetKnowledgeBaseWebhookUrlResponse {
    string webhook_url = 1;
}

message GetConnectorInternalConfigRequest {
    ConnectorType connector = 1;
}

message GetConnectorInternalConfigResponse {
    ConnectorInternalConfig internal_config = 1;
}

enum SearchMode {
    SEARCH_MODE_UNSPECIFIED = 0;
    SEARCH_MODE_HYBRID = 1;
    SEARCH_MODE_KEYWORD = 2;
    SEARCH_MODE_APPROXIMATE_KNN = 3;
    SEARCH_MODE_BRUTE_FORCE_KNN = 4;
}

enum ForwardStatus {
    FORWARD_STATUS_UNSPECIFIED = 0;
    FORWARD_STATUS_FAILURE = 1;
    FORWARD_STATUS_SAVED = 2;
    FORWARD_STATUS_SUCCESS = 3;
}

enum ConnectorType {
    CONNECTOR_TYPE_UNSPECIFIED = 0;
    CONNECTOR_TYPE_GITHUB = 1;
    CONNECTOR_TYPE_SLACK = 2;
    CONNECTOR_TYPE_GOOGLE_DRIVE = 3;
    CONNECTOR_TYPE_JIRA = 4;
    CONNECTOR_TYPE_CODEIUM = 5;
    CONNECTOR_TYPE_EMAIL = 6;
    CONNECTOR_TYPE_GITHUB_OAUTH = 7;
}

enum JobStatus {
    JOB_STATUS_UNSPECIFIED = 0;
    JOB_STATUS_QUEUED = 1;
    JOB_STATUS_RUNNING = 2;
    JOB_STATUS_COMPLETED = 3;
    JOB_STATUS_CANCELLED = 4;
    JOB_STATUS_CANCELLING = 5;
    JOB_STATUS_ERRORED = 6;
    JOB_STATUS_RETRYABLE = 7;
}