# Auto generated by 'inv collect-airflow'
from airfly._vendor.airflow.providers.amazon.aws.operators.base_aws import (
    AwsBaseOperator,
)


class BedrockInvokeModelOperator(AwsBaseOperator):
    model_id: "str"
    input_data: "dict[str, Any]"
    content_type: "str | None"
    accept_type: "str | None"


class BedrockCustomizeModelOperator(AwsBaseOperator):
    job_name: "str"
    custom_model_name: "str"
    role_arn: "str"
    base_model_id: "str"
    training_data_uri: "str"
    output_data_uri: "str"
    hyperparameters: "dict[str, str]"
    ensure_unique_job_name: "bool"
    customization_job_kwargs: "dict[str, Any] | None"
    wait_for_completion: "bool"
    waiter_delay: "int"
    waiter_max_attempts: "int"
    deferrable: "bool"


class BedrockCreateProvisionedModelThroughputOperator(AwsBaseOperator):
    model_units: "int"
    provisioned_model_name: "str"
    model_id: "str"
    create_throughput_kwargs: "dict[str, Any] | None"
    wait_for_completion: "bool"
    waiter_delay: "int"
    waiter_max_attempts: "int"
    deferrable: "bool"


class BedrockCreateKnowledgeBaseOperator(AwsBaseOperator):
    name: "str"
    embedding_model_arn: "str"
    role_arn: "str"
    storage_config: "dict[str, Any]"
    create_knowledge_base_kwargs: "dict[str, Any] | None"
    wait_for_indexing: "bool"
    indexing_error_retry_delay: "int"
    indexing_error_max_attempts: "int"
    wait_for_completion: "bool"
    waiter_delay: "int"
    waiter_max_attempts: "int"
    deferrable: "bool"


class BedrockCreateDataSourceOperator(AwsBaseOperator):
    name: "str"
    knowledge_base_id: "str"
    bucket_name: "str | None"
    create_data_source_kwargs: "dict[str, Any] | None"


class BedrockIngestDataOperator(AwsBaseOperator):
    knowledge_base_id: "str"
    data_source_id: "str"
    ingest_data_kwargs: "dict[str, Any] | None"
    wait_for_completion: "bool"
    waiter_delay: "int"
    waiter_max_attempts: "int"
    deferrable: "bool"


class BedrockRaGOperator(AwsBaseOperator):
    input: "str"
    source_type: "str"
    model_arn: "str"
    prompt_template: "str | None"
    knowledge_base_id: "str | None"
    vector_search_config: "dict[str, Any] | None"
    sources: "list[dict[str, Any]] | None"
    rag_kwargs: "dict[str, Any] | None"


class BedrockRetrieveOperator(AwsBaseOperator):
    retrieval_query: "str"
    knowledge_base_id: "str"
    vector_search_config: "dict[str, Any] | None"
    retrieve_kwargs: "dict[str, Any] | None"
