# Generated by the protocol buffer compiler.  DO NOT EDIT!
# sources: workflow.proto
# plugin: python-betterproto
from dataclasses import dataclass
from typing import Dict, List

import betterproto

from . import pql


@dataclass
class ExecutionPlan(betterproto.Message):
    """
    A logical description of the query. This is the top-level object passed
    down to the engine for execution on the data.
    """

    context: pql.ParseContext = betterproto.message_field(1)
    statements: List[pql.ParseStatement] = betterproto.message_field(2)
    metadata: "QueryMetadata" = betterproto.message_field(4)


@dataclass
class TrainModelRequest(betterproto.Message):
    """Create a specific proto for the train request NEXT ID: 11"""

    connection_info: pql.ConnectionInfo = betterproto.message_field(1)
    dataset_mapping: pql.DatasetMapping = betterproto.message_field(2)
    model_config: str = betterproto.string_field(3)
    model_id: int = betterproto.uint64_field(4)
    model_url: str = betterproto.string_field(5)
    external_connection_info: Dict[str, pql.ConnectionInfo] = betterproto.map_field(
        6, betterproto.TYPE_STRING, betterproto.TYPE_MESSAGE
    )
    dataset_size_bytes: int = betterproto.uint64_field(7)
    model_uuid: str = betterproto.string_field(8)
    reset_policy: str = betterproto.string_field(9)
    attempt: int = betterproto.int32_field(10)


@dataclass
class EvaluateModelRequest(betterproto.Message):
    """
    Request used to run evaluation activities following successful model
    training
    """

    train_request: "TrainModelRequest" = betterproto.message_field(1)
    run_id: str = betterproto.string_field(2)
    ground_truth_url: str = betterproto.string_field(3)


@dataclass
class QueryMetadata(betterproto.Message):
    """Query metadata."""

    query_id: int = betterproto.uint64_field(1)
    results_url: str = betterproto.string_field(2)
    params_url: str = betterproto.string_field(3)


@dataclass
class MinioConfig(betterproto.Message):
    endpoint_url: str = betterproto.string_field(1)
    aws_access_key_id: str = betterproto.string_field(2)
    aws_secret_access_key: str = betterproto.string_field(3)


@dataclass
class ConnectionConfig(betterproto.Message):
    protocol: str = betterproto.string_field(1)
    url: str = betterproto.string_field(2)
    data_format: str = betterproto.string_field(3)
    key: str = betterproto.string_field(4)
    secret: str = betterproto.string_field(5)
    token: str = betterproto.string_field(6)
    client_kwargs: "MinioConfig" = betterproto.message_field(7)


@dataclass
class ResourceConfig(betterproto.Message):
    """This is a cut"""

    cpu: str = betterproto.string_field(1)
    memory: str = betterproto.string_field(2)
    gpu: int = betterproto.uint32_field(3)


@dataclass
class ExportModelRequest(betterproto.Message):
    connection_config: "ConnectionConfig" = betterproto.message_field(1)
    model_uuid: str = betterproto.string_field(2)
    run_id: str = betterproto.string_field(3)
    path: str = betterproto.string_field(4)
    resources: "ResourceConfig" = betterproto.message_field(5)
    properties: Dict[str, str] = betterproto.map_field(
        6, betterproto.TYPE_STRING, betterproto.TYPE_STRING
    )
