docling/.actor/actor.sh

#!/bin/bash

# Function to upload content to the key-value store
upload_to_kvs() {
    local content_file="$1"
    local key_name="$2"
    local content_type="$3"
    local description="$4"

    # Find the Apify CLI command
    find_apify_cmd
    local apify_cmd="$FOUND_APIFY_CMD"

    if [ -n "$apify_cmd" ]; then
        echo "Uploading $description to key-value store (key: $key_name)..."

        # Create a temporary home directory with write permissions
        setup_temp_environment

        # Use the --no-update-notifier flag if available
        if $apify_cmd --help | grep -q "\--no-update-notifier"; then
            if $apify_cmd --no-update-notifier actor:set-value "$key_name" --contentType "$content_type" < "$content_file"; then
                echo "Successfully uploaded $description to key-value store"
                local url="https://api.apify.com/v2/key-value-stores/${APIFY_DEFAULT_KEY_VALUE_STORE_ID}/records/$key_name"
                echo "$description available at: $url"
                cleanup_temp_environment
                return 0
            fi
        else
            # Fall back to regular command if flag isn't available
            if $apify_cmd actor:set-value "$key_name" --contentType "$content_type" < "$content_file"; then
                echo "Successfully uploaded $description to key-value store"
                local url="https://api.apify.com/v2/key-value-stores/${APIFY_DEFAULT_KEY_VALUE_STORE_ID}/records/$key_name"
                echo "$description available at: $url"
                cleanup_temp_environment
                return 0
            fi
        fi

        echo "ERROR: Failed to upload $description to key-value store"
        cleanup_temp_environment
        return 1
    else
        echo "ERROR: Apify CLI not found for $description upload"
        return 1
    fi
}

# Function to find Apify CLI command
find_apify_cmd() {
    FOUND_APIFY_CMD=""
    for cmd in "apify" "actor" "/usr/local/bin/apify" "/usr/bin/apify" "/opt/apify/cli/bin/apify"; do
        if command -v "$cmd" &> /dev/null; then
            FOUND_APIFY_CMD="$cmd"
            break
        fi
    done
}

# Function to set up temporary environment for Apify CLI
setup_temp_environment() {
    export TMPDIR="/tmp/apify-home-${RANDOM}"
    mkdir -p "$TMPDIR"
    export APIFY_DISABLE_VERSION_CHECK=1
    export NODE_OPTIONS="--no-warnings"
    export HOME="$TMPDIR"  # Override home directory to writable location
}

# Function to clean up temporary environment
cleanup_temp_environment() {
    rm -rf "$TMPDIR" 2>/dev/null || true
}

# Function to push data to Apify dataset
push_to_dataset() {
    local document_url="$1"
    local result_url="$2"

    # Find Apify CLI command
    find_apify_cmd
    local apify_cmd="$FOUND_APIFY_CMD"

    if [ -n "$apify_cmd" ]; then
        echo "Adding record to dataset..."
        setup_temp_environment

        # Use the --no-update-notifier flag if available
        if $apify_cmd --help | grep -q "\--no-update-notifier"; then
            if $apify_cmd --no-update-notifier actor:push-data "{\"url\": \"${document_url}\", \"output_file\": \"${result_url}\", \"status\": \"success\"}"; then
                echo "Successfully added record to dataset"
            else
                echo "Warning: Failed to add record to dataset"
            fi
        else
            # Fall back to regular command
            if $apify_cmd actor:push-data "{\"url\": \"${document_url}\", \"output_file\": \"${result_url}\", \"status\": \"success\"}"; then
                echo "Successfully added record to dataset"
            else
                echo "Warning: Failed to add record to dataset"
            fi
        fi

        cleanup_temp_environment
    fi
}


# --- Setup logging and error handling ---

LOG_FILE="/tmp/docling.log"
touch "$LOG_FILE" || {
    echo "Fatal: Cannot create log file at $LOG_FILE"
    exit 1
}

# Log to both console and file
exec 1> >(tee -a "$LOG_FILE")
exec 2> >(tee -a "$LOG_FILE" >&2)

# Exit codes
readonly ERR_API_UNAVAILABLE=15
readonly ERR_INVALID_INPUT=16


# --- Debug environment ---

echo "Date: $(date)"
echo "Python version: $(python --version 2>&1)"
echo "Docling-serve path: $(which docling-serve 2>/dev/null || echo 'Not found')"
echo "Working directory: $(pwd)"


# --- Setup tools ---

echo "Setting up tools..."
TOOLS_DIR="/tmp/docling-tools"
mkdir -p "$TOOLS_DIR"

# Copy tools if available
if [ -d "/build-files" ]; then
    echo "Copying tools from /build-files..."
    cp -r /build-files/* "$TOOLS_DIR/"
    export PATH="$TOOLS_DIR/bin:$PATH"
else
    echo "Warning: No build files directory found. Some tools may be unavailable."
fi

# Copy Python processor script to tools directory
PYTHON_SCRIPT_PATH="$(dirname "$0")/docling_processor.py"
if [ -f "$PYTHON_SCRIPT_PATH" ]; then
    echo "Copying Python processor script to tools directory..."
    cp "$PYTHON_SCRIPT_PATH" "$TOOLS_DIR/"
    chmod +x "$TOOLS_DIR/docling_processor.py"
else
    echo "ERROR: Python processor script not found at $PYTHON_SCRIPT_PATH"
    exit 1
fi

# Check OCR directories and ensure they're writable
echo "Checking OCR directory permissions..."
OCR_DIR="/opt/app-root/src/.EasyOCR"
if [ -d "$OCR_DIR" ]; then
    # Test if we can write to the directory
    if touch "$OCR_DIR/test_write" 2>/dev/null; then
        echo "[✓] OCR directory is writable"
        rm "$OCR_DIR/test_write"
    else
        echo "[✗] OCR directory is not writable, setting up alternative in /tmp"

        # Create alternative in /tmp (which is writable)
        mkdir -p "/tmp/.EasyOCR/user_network"
        export EASYOCR_MODULE_PATH="/tmp/.EasyOCR"
    fi
else
    echo "OCR directory not found, creating in /tmp"
    mkdir -p "/tmp/.EasyOCR/user_network"
    export EASYOCR_MODULE_PATH="/tmp/.EasyOCR"
fi


# --- Starting the API ---

echo "Starting docling-serve API..."

# Create a dedicated working directory in /tmp (writable)
API_DIR="/tmp/docling-api"
mkdir -p "$API_DIR"
cd "$API_DIR"
echo "API working directory: $(pwd)"

# Find docling-serve executable
DOCLING_SERVE_PATH=$(which docling-serve)
echo "Docling-serve executable: $DOCLING_SERVE_PATH"

# Start the API with minimal parameters to avoid any issues
echo "Starting docling-serve API..."
"$DOCLING_SERVE_PATH" run --host 0.0.0.0 --port 5001 > "$API_DIR/docling-serve.log" 2>&1 &
API_PID=$!
echo "Started docling-serve API with PID: $API_PID"

# A more reliable wait for API startup
echo "Waiting for API to initialize..."
MAX_TRIES=30
tries=0
started=false

while [ $tries -lt $MAX_TRIES ]; do
    tries=$((tries + 1))

    # Check if process is still running
    if ! ps -p $API_PID > /dev/null; then
        echo "ERROR: docling-serve API process terminated unexpectedly after $tries seconds"
        break
    fi

    # Check log for startup completion or errors
    if grep -q "Application startup complete" "$API_DIR/docling-serve.log" 2>/dev/null; then
        echo "[✓] API startup completed successfully after $tries seconds"
        started=true
        break
    fi

    if grep -q "Permission denied\|PermissionError" "$API_DIR/docling-serve.log" 2>/dev/null; then
        echo "ERROR: Permission errors detected in API startup"
        break
    fi

    # Sleep and check again
    sleep 1

    # Output a progress indicator every 5 seconds
    if [ $((tries % 5)) -eq 0 ]; then
        echo "Still waiting for API startup... ($tries/$MAX_TRIES seconds)"
    fi
done

# Show log content regardless of outcome
echo "docling-serve log output so far:"
tail -n 20 "$API_DIR/docling-serve.log"

# Verify the API is running
if ! ps -p $API_PID > /dev/null; then
    echo "ERROR: docling-serve API failed to start"
    if [ -f "$API_DIR/docling-serve.log" ]; then
        echo "Full log output:"
        cat "$API_DIR/docling-serve.log"
    fi
    exit $ERR_API_UNAVAILABLE
fi

if [ "$started" != "true" ]; then
    echo "WARNING: API process is running but startup completion was not detected"
    echo "Will attempt to continue anyway..."
fi

# Try to verify API is responding at this point
echo "Verifying API responsiveness..."
(python -c "
import sys, time, socket
for i in range(5):
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        s.settimeout(1)
        result = s.connect_ex(('localhost', 5001))
        if result == 0:
            s.close()
            print('Port 5001 is open and accepting connections')
            sys.exit(0)
        s.close()
    except Exception as e:
        pass
    time.sleep(1)
print('Could not connect to API port after 5 attempts')
sys.exit(1)
" && echo "API verification succeeded") || echo "API verification failed, but continuing anyway"

# Define API endpoint
DOCLING_API_ENDPOINT="http://localhost:5001/v1alpha/convert/source"


# --- Processing document ---

echo "Starting document processing..."
echo "Reading input from Apify..."

# Function to handle Actor input detection
get_actor_input() {
    local input=""

    # Create directory if it doesn't exist
    mkdir -p "/tmp/actor-input" || echo "Warning: Could not create /tmp/actor-input directory" >&2

    # If /tmp/actor-input/INPUT exists as a directory, remove it
    if [ -d "/tmp/actor-input/INPUT" ]; then
        echo "Warning: /tmp/actor-input/INPUT exists as a directory. Removing it to create a file." >&2
        rm -rf "/tmp/actor-input/INPUT"
    fi

    # Check multiple potential locations for input file
    if [ -f "/tmp/actor-input/INPUT" ]; then
        echo "Found standard Actor input file at /tmp/actor-input/INPUT" >&2
        input=$(cat "/tmp/actor-input/INPUT")
    elif [ -f "/input/INPUT" ]; then
        echo "Found Actor input file at /input/INPUT" >&2
        input=$(cat "/input/INPUT")

    # Fallback to environment variable
    elif [ -n "$APIFY_INPUT_JSON" ]; then
        echo "Using APIFY_INPUT_JSON environment variable" >&2
        input="$APIFY_INPUT_JSON"

    # Last resort: use test input with md format
    else
        echo "No input found, using test input with md format" >&2
        TEST_INPUT='{"documentUrl":"https://vancura.dev/assets/actor-test/facial-hairstyles-and-filtering-facepiece-respirators.pdf","ocr":true,"outputFormat":"md"}'
        mkdir -p "/tmp/actor-input"
        echo "$TEST_INPUT" > "/tmp/actor-input/INPUT"

        # Read back the test input to ensure we get clean JSON
        input=$(cat "/tmp/actor-input/INPUT")
    fi

    # Return only the JSON content
    echo "$input"
}

# Get actor input
INPUT=$(get_actor_input)
echo "Input content:" >&2
echo "$INPUT" >&2  # Send the raw input to stderr for debugging
echo "$INPUT"      # Send the clean JSON to stdout for processing

# Extract values from INPUT using Python
echo "Using Python to parse input..."
DOCUMENT_URL="$(echo "$INPUT" | python -c "import sys, json; print(json.load(sys.stdin).get('documentUrl', ''))")"
OUTPUT_FORMAT="$(echo "$INPUT" | python -c "import sys, json; print(json.load(sys.stdin).get('outputFormat', 'md'))")"
OCR_ENABLED="$(echo "$INPUT" | python -c "import sys, json; print(str(json.load(sys.stdin).get('ocr', True)).lower())")"

# Validate input schema should already enforce this, but double-check
if [ -z "$DOCUMENT_URL" ]; then
    echo "ERROR: No document URL provided in input"

    # Try to push data to Actor but don't exit if it fails
    find_apify_cmd
    apify_cmd="$FOUND_APIFY_CMD"
    if [ -n "$apify_cmd" ]; then
        echo "Reporting missing document URL to Actor storage..."
        setup_temp_environment
        if $apify_cmd actor:push-data "{\"status\": \"error\", \"error\": \"No document URL provided in input\"}" 2>&1; then
            echo "Successfully pushed error message to Actor storage"
        else
            echo "Warning: Failed to push error message to Actor storage"
        fi
        cleanup_temp_environment
    fi

    # Use default document URL for testing instead of exiting
    echo "Using a default document URL for testing: https://arxiv.org/pdf/2408.09869"
    DOCUMENT_URL="https://arxiv.org/pdf/2408.09869"
fi

if [ -z "$OUTPUT_FORMAT" ]; then
    echo "No output format specified, defaulting to 'md'"
    OUTPUT_FORMAT="md"
fi

# Ensure OCR_ENABLED has a valid boolean value
if [ -z "$OCR_ENABLED" ]; then
    echo "No OCR setting specified, defaulting to true"
    OCR_ENABLED="true"
fi

echo "Input values: documentUrl=$DOCUMENT_URL, outputFormat=$OUTPUT_FORMAT, ocr=$OCR_ENABLED"

# Create the request JSON
REQUEST_JSON="{\"options\":{\"to_formats\":[\"$OUTPUT_FORMAT\"],\"ocr\":$OCR_ENABLED},\"http_sources\":[{\"url\":\"$DOCUMENT_URL\"}]}"
echo "Creating request JSON:" >&2
echo "$REQUEST_JSON" >&2
echo "$REQUEST_JSON" > "$API_DIR/request.json"

# Send the conversion request using our Python script
echo "Sending conversion request to docling-serve API..."
python "$TOOLS_DIR/docling_processor.py" \
    --api-endpoint "$DOCLING_API_ENDPOINT" \
    --request-json "$API_DIR/request.json" \
    --output-dir "$API_DIR" \
    --output-format "$OUTPUT_FORMAT"

PYTHON_EXIT_CODE=$?

# --- Check for various potential output files ---

echo "Checking for output files..."
if [ -f "$API_DIR/output.$OUTPUT_FORMAT" ]; then
    echo "Conversion completed successfully! Output file found."

    # Get content from the converted file
    OUTPUT_SIZE=$(wc -c < "$API_DIR/output.$OUTPUT_FORMAT")
    echo "Output file found with size: $OUTPUT_SIZE bytes"

    # Calculate the access URL for result display
    RESULT_URL="https://api.apify.com/v2/key-value-stores/${APIFY_DEFAULT_KEY_VALUE_STORE_ID}/records/OUTPUT"

    echo "=============================="
    echo "PROCESSING COMPLETE!"
    echo "Document URL: ${DOCUMENT_URL}"
    echo "Output format: ${OUTPUT_FORMAT}"
    echo "Output size: ${OUTPUT_SIZE} bytes"
    echo "=============================="

    # Set the output content type based on format
    CONTENT_TYPE="text/plain"
    case "$OUTPUT_FORMAT" in
        md) CONTENT_TYPE="text/markdown" ;;
        html) CONTENT_TYPE="text/html" ;;
        json) CONTENT_TYPE="application/json" ;;
        text) CONTENT_TYPE="text/plain" ;;
    esac

    # Upload the document content using our function
    upload_to_kvs "$API_DIR/output.$OUTPUT_FORMAT" "OUTPUT" "$CONTENT_TYPE" "Document content"

    # Only proceed with dataset record if document upload succeeded
    if [ $? -eq 0 ]; then
        echo "Your document is available at: ${RESULT_URL}"
        echo "=============================="

        # Push data to dataset
        push_to_dataset "$DOCUMENT_URL" "$RESULT_URL"
    fi
else
    echo "ERROR: No converted output file found at $API_DIR/output.$OUTPUT_FORMAT"

    # Create error metadata
    ERROR_METADATA="{\"status\":\"error\",\"error\":\"No converted output file found\",\"documentUrl\":\"$DOCUMENT_URL\"}"
    echo "$ERROR_METADATA" > "/tmp/actor-output/OUTPUT"
    chmod 644 "/tmp/actor-output/OUTPUT"

    echo "Error information has been saved to /tmp/actor-output/OUTPUT"
fi


# --- Verify output files for debugging ---

echo "=== Final Output Verification ==="
echo "Files in /tmp/actor-output:"
ls -la /tmp/actor-output/ 2>/dev/null || echo "Cannot list /tmp/actor-output/"

echo "All operations completed. The output should be available in the default key-value store."
echo "Content URL: ${RESULT_URL:-No URL available}"


# --- Cleanup function ---

cleanup() {
    echo "Running cleanup..."

    # Stop the API process
    if [ -n "$API_PID" ]; then
        echo "Stopping docling-serve API (PID: $API_PID)..."
        kill $API_PID 2>/dev/null || true
    fi

    # Export log file to KVS if it exists
    # DO THIS BEFORE REMOVING TOOLS DIRECTORY
    if [ -f "$LOG_FILE" ]; then
        if [ -s "$LOG_FILE" ]; then
            echo "Log file is not empty, pushing to key-value store (key: LOG)..."

            # Upload log using our function
            upload_to_kvs "$LOG_FILE" "LOG" "text/plain" "Log file"
        else
            echo "Warning: log file exists but is empty"
        fi
    else
        echo "Warning: No log file found"
    fi

    # Clean up temporary files AFTER log is uploaded
    echo "Cleaning up temporary files..."
    if [ -d "$API_DIR" ]; then
        echo "Removing API working directory: $API_DIR"
        rm -rf "$API_DIR" 2>/dev/null || echo "Warning: Failed to remove $API_DIR"
    fi

    if [ -d "$TOOLS_DIR" ]; then
        echo "Removing tools directory: $TOOLS_DIR"
        rm -rf "$TOOLS_DIR" 2>/dev/null || echo "Warning: Failed to remove $TOOLS_DIR"
    fi

    # Keep log file until the very end
    echo "Script execution completed at $(date)"
    echo "Actor execution completed"
}

# Register cleanup
trap cleanup EXIT