id: hugging_face
namespace: blueprint

inputs:
  - id: message_id
    type: STRING

tasks:
  - id: retrieve_data
    type: io.kestra.plugin.jdbc.postgresql.Query
    sql: SELECT user_id, message FROM customer.message WHERE message_id = '{{
      inputs.message_id }}'
    fetchType: FETCH_ONE

  - id: classification
    type: io.kestra.plugin.huggingface.Inference
    model: facebook/bart-large-mnli
    apiKey: "{{ secret('HUGGINGFACE_API_KEY') }}"
    inputs: "{{ json(outputs.retrieve_data.row).message }}"
    parameters:
      candidate_labels:
        - "support"
        - "warranty"
        - "upsell"
        - "help"

  - id: insert_category
    type: io.kestra.plugin.jdbc.postgresql.Query
    sql: UPDATE customer.message SET category = '{{
      json(outputs.classification.output).labels[0] }}' WHERE message_id = '{{
      inputs.message_id }}'

pluginDefaults:
  - type: io.kestra.plugin.jdbc.postgresql
    values:
      url: jdbc:postgresql://"{{secret('POSTGRES_HOST')}}"
      username: "{{secret('POSTGRES_USERNAME')}}"
      password: "{{secret('POSTGRES_PASSWORD')}}"