id: unstructured-document-partitioning-pipeline
namespace: company.ai

inputs:
  - id: source_document_url
    type: STRING
    defaults: "https://raw.githubusercontent.com/kestra-io/datasets/main/csv/produc\
      ts.csv"
    description: Public or pre-signed URL of the raw document (PDF, DOCX, CSV) to
      partition and extract.

  - id: max_chunk_characters
    type: INT
    defaults: 1500
    description: Maximum character limit per chunk for downstream vector embedding models.

  - id: overlap_characters
    type: INT
    defaults: 200
    description: Character overlap window preserved across chunk boundaries to
      retain semantic context.

tasks:
  - id: fetch_source_document
    type: io.kestra.plugin.core.http.Request
    uri: "{{ inputs.source_document_url }}"
    method: GET

  - id: partition_and_chunk_document
    type: io.kestra.plugin.scripts.python.Script
    containerImage: python:3.11-slim
    beforeCommands:
      - pip install kestra
    inputFiles:
      raw_payload.txt: "{{ outputs.fetch_source_document.body }}"
    script: |
      import os
      import json
      from kestra import Kestra

      max_chars = {{ inputs.max_chunk_characters }}
      overlap = {{ inputs.overlap_characters }}

      print(f"Reading downloaded document payload...")
      with open("raw_payload.txt", "r", encoding="utf-8", errors="ignore") as f:
          content = f.read()

      total_len = len(content)
      print(f"Document ingested: {total_len} characters.")

      # Semantic chunking simulation adhering to Unstructured.io element partitioning logic
      chunks = []
      start = 0
      while start < total_len:
          end = min(start + max_chars, total_len)
          chunk_slice = content[start:end].strip()
          if chunk_slice:
              chunks.append({
                  "index": len(chunks),
                  "char_length": len(chunk_slice),
                  "start_offset": start,
                  "end_offset": end
              })
          start += max_chars - overlap

      total_chunks = len(chunks)
      avg_chunk_size = round(sum(c["char_length"] for c in chunks) / total_chunks, 1) if total_chunks else 0
      print(f"Partitioning complete: {total_chunks} chunks generated (average size: {avg_chunk_size} chars).")

      Kestra.outputs({
          "total_chunks": total_chunks,
          "total_characters": total_len,
          "average_chunk_size": avg_chunk_size,
          "is_valid": total_chunks > 0
      })

  - id: evaluate_chunking_quality
    type: io.kestra.plugin.core.flow.If
    condition: "{{ outputs.partition_and_chunk_document.vars.is_valid ?? false }}"
    then:
      - id: notify_pipeline_success
        type: io.kestra.plugin.slack.notifications.SlackIncomingWebhook
        url: "{{ secret('SLACK_WEBHOOK_URL') }}"
        payload: |
          {
            "text": "📄 *Document Partitioning Complete*\n• Source: `{{ inputs.source_document_url }}`\n• Total Chunks: `{{ outputs.partition_and_chunk_document.vars.total_chunks }}`\n• Average Chunk Size: `{{ outputs.partition_and_chunk_document.vars.average_chunk_size }} chars`\n• Status: Artifacts ready for vector embedding and ingestion.\n• Execution: `{{ execution.id }}`"
          }
    else:
      - id: notify_empty_document
        type: io.kestra.plugin.slack.notifications.SlackIncomingWebhook
        url: "{{ secret('SLACK_WEBHOOK_URL') }}"
        payload: |
          {
            "text": "⚠️ *Document Partitioning Warning*: Zero chunks extracted from `{{ inputs.source_document_url }}`. Verify source encoding."
          }

  - id: archive_partition_manifest
    type: io.kestra.plugin.scripts.shell.Commands
    commands:
      - echo "[MANIFEST] execution={{ execution.id }} chunks={{
        outputs.partition_and_chunk_document.vars.total_chunks }} characters={{
        outputs.partition_and_chunk_document.vars.total_characters }}"

errors:
  - id: alert_pipeline_failure
    type: io.kestra.plugin.slack.notifications.SlackIncomingWebhook
    url: "{{ secret('SLACK_WEBHOOK_URL') }}"
    payload: |
      {
        "text": "🚨 *Document Partitioning Pipeline Failed* | Execution `{{ execution.id }}` terminated during document retrieval or extraction."
      }

triggers:
  - id: nightly_batch_partitioning
    type: io.kestra.plugin.core.trigger.Schedule
    cron: "0 2 * * *"