AI Document Data Extraction: Parse PDFs and Web Pages into Structured JSON with OpenAI

id: ai-document-data-extraction
namespace: company.team

inputs:
  - id: document_url
    type: STRING
    displayName: Document URL
    description: Public URL to a PDF, HTML page, or text document to extract data from
  - id: extraction_schema
    type: STRING
    displayName: Extraction schema (JSON)
    description: |
      JSON Schema defining the fields to extract. Example:
      {"invoice_number": "string", "vendor_name": "string", "total_amount": "number", "line_items": "array of objects with description and amount"}
    defaults: '{"invoice_number": "string", "vendor_name": "string", "invoice_date":
      "string", "total_amount": "number", "currency": "string", "line_items":
      "array of objects with description, quantity, unit_price, and total"}'
  - id: document_type
    type: SELECT
    displayName: Document type
    description: Type of document, used to calibrate the extraction prompt
    values:
      - invoice
      - contract
      - receipt
      - report
      - form
      - other
    defaults: "invoice"
  - id: output_destination
    type: SELECT
    displayName: Output destination
    description: Where to send the extracted structured data
    values:
      - postgres
      - slack
      - both
    defaults: "both"

tasks:
  - id: download_document
    type: io.kestra.plugin.core.http.Download
    uri: "{{ inputs.document_url }}"
    method: GET

  - id: extract_text
    type: io.kestra.plugin.scripts.python.Script
    inputFiles:
      document: "{{ outputs.download_document.uri }}"
    dependencies:
      - pypdf2
      - beautifulsoup4
      - requests
    script: |
      import os
      import sys
      import json
      from kestra import Kestra

      doc_path = "document"
      content = ""

      # Try to detect file type and extract text
      with open(doc_path, "rb") as f:
          header = f.read(8)

      if header[:4] == b'%PDF':
          # PDF extraction
          try:
              import PyPDF2
              with open(doc_path, "rb") as f:
                  reader = PyPDF2.PdfReader(f)
                  pages = []
                  for page in reader.pages:
                      pages.append(page.extract_text() or "")
              content = "\n".join(pages)
          except Exception as e:
              content = f"[PDF extraction error: {e}]"
      else:
          # Try HTML/text
          try:
              from bs4 import BeautifulSoup
              with open(doc_path, "r", encoding="utf-8", errors="replace") as f:
                  raw = f.read()
              if raw.strip().startswith("<"):
                  soup = BeautifulSoup(raw, "html.parser")
                  # Remove scripts and styles
                  for tag in soup(["script", "style", "nav", "footer"]):
                      tag.decompose()
                  content = soup.get_text(separator="\n", strip=True)
              else:
                  content = raw
          except Exception as e:
              content = f"[Text extraction error: {e}]"

      # Truncate to avoid token limits (keep first 8000 chars)
      content_truncated = content[:8000]
      char_count = len(content)

      Kestra.outputs({
          "text_content": content_truncated,
          "char_count": char_count,
          "truncated": char_count > 8000
      })

  - id: extract_structured_data
    type: io.kestra.plugin.openai.ChatCompletion
    apiKey: "{{ secret('OPENAI_API_KEY') }}"
    model: gpt-4o
    maxTokens: 1500
    clientTimeout: 60
    jsonResponseSchema: |
      {
        "type": "object",
        "properties": {
          "extracted_fields": {
            "type": "object",
            "description": "The extracted fields as defined in the extraction schema"
          },
          "confidence": {
            "type": "number",
            "minimum": 0,
            "maximum": 1,
            "description": "Overall extraction confidence score"
          },
          "missing_fields": {
            "type": "array",
            "items": {"type": "string"},
            "description": "List of schema fields that could not be found in the document"
          },
          "extraction_notes": {
            "type": "string",
            "description": "Any caveats, ambiguities, or important observations about the extraction"
          }
        },
        "required": ["extracted_fields", "confidence", "missing_fields"]
      }
    messages:
      - role: system
        content: |
          You are a document data extraction specialist. Extract structured data from the provided
          document text according to the given schema. Be precise, only extract values that are
          explicitly present in the document. For missing fields, list them in `missing_fields`.
          Return null for fields you cannot find rather than guessing.

          Document type: {{ inputs.document_type }}
          Extraction schema requested: {{ inputs.extraction_schema }}
      - role: user
        content: |
          Extract the following data from this {{ inputs.document_type }} document:

          Schema: {{ inputs.extraction_schema }}

          Document text:
          {{ outputs.extract_text.vars.text_content }}

  - id: save_to_postgres
    type: io.kestra.plugin.jdbc.postgresql.Queries
    fetchType: NONE
    runIf: "{{ inputs.output_destination == 'postgres' or inputs.output_destination
      == 'both' }}"
    sql: |
      CREATE TABLE IF NOT EXISTS document_extractions (
          id SERIAL PRIMARY KEY,
          document_url TEXT,
          document_type TEXT,
          extracted_data JSONB,
          confidence NUMERIC,
          missing_fields JSONB,
          extraction_notes TEXT,
          extracted_at TIMESTAMP DEFAULT NOW()
      );

      INSERT INTO document_extractions
          (document_url, document_type, extracted_data, confidence, missing_fields, extraction_notes)
      VALUES (
          '{{ inputs.document_url }}',
          '{{ inputs.document_type }}',
          '{{ outputs.extract_structured_data.choices[0].message.content | json | jq(".extracted_fields") | first | json }}'::jsonb,
          {{ outputs.extract_structured_data.choices[0].message.content | json | jq(".confidence") | first }},
          '{{ outputs.extract_structured_data.choices[0].message.content | json | jq(".missing_fields") | first | json }}'::jsonb,
          '{{ outputs.extract_structured_data.choices[0].message.content | json | jq(".extraction_notes // \"\"") | first | replace("'", "''") }}'
      );

  - id: notify_slack
    type: io.kestra.plugin.slack.notifications.SlackIncomingWebhook
    url: "{{ secret('SLACK_WEBHOOK_URL') }}"
    runIf: "{{ inputs.output_destination == 'slack' or inputs.output_destination ==
      'both' }}"
    messageText: |
      :page_facing_up: *Document Extraction Complete*

      *Type:* {{ inputs.document_type | title }}
      *Source:* {{ inputs.document_url }}
      *Confidence:* {{ outputs.extract_structured_data.choices[0].message.content | json | jq("(.confidence * 100 | floor | tostring) + \"%\"") | first }}
      *Characters processed:* {{ outputs.extract_text.vars.char_count }}{{ outputs.extract_text.vars.truncated | ternary(' (truncated)', '') }}

      *Extracted Data:*
      ```{{ outputs.extract_structured_data.choices[0].message.content | json | jq(".extracted_fields") | first | json }}```

      *Missing Fields:* {{ outputs.extract_structured_data.choices[0].message.content | json | jq(".missing_fields | if length == 0 then \"none\" else join(\", \") end") | first }}

pluginDefaults:
  - type: io.kestra.plugin.jdbc.postgresql.Queries
    values:
      url: "{{ secret('POSTGRES_URL') }}"
      username: "{{ secret('POSTGRES_USER') }}"
      password: "{{ secret('POSTGRES_PASSWORD') }}"