trustgraph/trustgraph-cli/scripts/tg-load-pdf

#!/usr/bin/env python3

"""
Loads a PDF document into TrustGraph processing.
"""

import pulsar
from pulsar.schema import JsonSchema
from trustgraph.schema import Document, document_ingest_queue
from trustgraph.schema import Value, Metadata, Triple
import base64
import hashlib
import argparse
import os
import time
import uuid

from trustgraph.log_level import LogLevel

default_user = 'trustgraph'
default_collection = 'default'

IS_A = 'http://www.w3.org/1999/02/22-rdf-syntax-ns#type'
LABEL = 'http://www.w3.org/2000/01/rdf-schema#label'

DIGITAL_DOCUMENT = 'https://schema.org/DigitalDocument'
PUBLICATION_EVENT = 'https://schema.org/PublicationEvent'
ORGANIZATION = 'https://schema.org/Organization'

NAME = 'https://schema.org/name'
DESCRIPTION = 'https://schema.org/description'
COPYRIGHT_NOTICE = 'https://schema.org/copyrightNotice'
COPYRIGHT_HOLDER = 'https://schema.org/copyrightHolder'
COPYRIGHT_YEAR = 'https://schema.org/copyrightYear'
LICENSE = 'https://schema.org/license'
PUBLICATION = 'https://schema.org/publication'
START_DATE = 'https://schema.org/startDate'
END_DATE = 'https://schema.org/endDate'
PUBLISHED_BY = 'https://schema.org/publishedBy'
DATE_PUBLISHED = 'https://schema.org/datePublished'
PUBLICATION = 'https://schema.org/publication'
DATE_PUBLISHED = 'https://schema.org/datePublished'
URL = 'https://schema.org/url'
IDENTIFIER = 'https://schema.org/identifier'
KEYWORD = 'https://schema.org/keywords'

def hash(data):

    if isinstance(data, str):
        data = data.encode("utf-8")

    # Create a SHA256 hash from the data
    id = hashlib.sha256(data).hexdigest()

    # Convert into a UUID, 64-byte hash becomes 32-byte UUID
    id = str(uuid.UUID(id[::2]))

    return id

def curi(pref, id):
    return f"https://trustgraph.ai/{pref}/{id}"

class DigitalDocument:
    def __init__(
            self, id, name=None, description=None, copyright_notice=None,
            copyright_holder=None, copyright_year=None, license=None,
            identifier=None,
            publication=None, url=None, keywords=[]
    ):
        self.id = id
        self.name = name
        self.description = description
        self.copyright_notice = copyright_notice
        self.copyright_holder = copyright_holder
        self.copyright_year = copyright_year
        self.license = license
        self.publication = publication
        self.url = url
        self.identifier = identifier
        self.keywords = keywords

    def emit(self, emit):

        emit(Triple(Value(self.id), Value(IS_A), Value(DIGITAL_DOCUMENT)))

        if self.name:
            emit(Triple(Value(self.id), Value(LABEL), Value(self.name)))
            emit(Triple(Value(self.id), Value(NAME), Value(self.name)))

        if self.identifier:
            emit(Triple(Value(id), Value(IDENTIFIER), Value(self.identifier)))

        if self.description:
            emit(Triple(
                Value(self.id), Value(DESCRIPTION), Value(self.description)
            ))

        if self.copyright_notice:
            emit(Triple(
                Value(self.id), Value(COPYRIGHT_NOTICE),
                Value(self.copyright_notice)
            ))

        if self.copyright_holder:
            emit(Triple(
                Value(self.id), Value(COPYRIGHT_HOLDER),
                Value(self.copyright_holder)
            ))

        if self.copyright_year:
            emit(Triple(
                Value(self.id), Value(COPYRIGHT_YEAR),
                Value(self.copyright_year)
            ))

        if self.license:
            emit(Triple(
                Value(self.id), Value(LICENSE), Value(self.license)
            ))

        if self.keywords:
            for k in self.keywords:
                emit(Triple(Value(self.id), Value(KEYWORD), Value(k)))

        if self.publication:
            emit(Triple(
                Value(self.id), Value(PUBLICATION, Value(self.publication.id))
            ))
            self.publication.emit(emit)

        if self.url:
            emit(Triple(Value(self.id), Value(URL), Value(self.url)))

class PublicationEvent:
    def __init__(
            self, id, organization=None, name=None, description=None,
            start_date=None, end_date=None,
    ):
        self.id = id
        self.organization = organization
        self.name = name
        self.description = description
        self.start_date = start_date
        self.end_date = end_date

    def emit(self, emit):

        emit(Triple(Value(self.id), Value(IS_A), Value(PUBLICATION_EVENT)))

        if self.name:
            emit(Triple(Value(self.id), Value(LABEL), Value(self.name)))
            emit(Triple(Value(self.id), Value(NAME), Value(self.name)))

        if self.description:
            emit(Triple(
                Value(self.id), Value(DESCRIPTION), Value(self.description)
            ))

        if self.organization:
            emit(Triple(
                Value(self.id), Value(PUBLISHED_BY),
                Value(self.organization.id)
            ))
            self.organization.emit(emit)

        if self.start_date:
            emit(Triple(
                Value(self.id), Value(START_DATE), Value(self.start_date)
            ))

        if self.end_date:
            emit(Triple(Value(self.id), Value(END_DATE), Value(self.end_date)))

class Organization:
    def __init__(self, id, name=None, description=None):
        self.id = id
        self.name = name
        self.description = description

    def emit(self, emit):

        emit(Triple(Value(self.id), Value(IS_A), Value(ORGANIZATION)))

        if self.name:
            emit(Triple(Value(self.id), Value(LABEL), Value(self.name)))
            emit(Triple(Value(self.id), Value(NAME), Value(self.name)))

        if self.description:
            emit(Triple(
                Value(self.id), Value(DESCRIPTION), Value(self.description)
            ))

class Loader:

    def __init__(
            self,
            pulsar_host,
            output_queue,
            user,
            collection,
            log_level,
            metadata,
    ):

        self.client = pulsar.Client(
            pulsar_host,
            logger=pulsar.ConsoleLogger(log_level.to_pulsar())
        )

        self.producer = self.client.create_producer(
            topic=output_queue,
            schema=JsonSchema(Document),
            chunking_enabled=True,
        )

        self.user = user
        self.collection = collection
        self.metadata = metadata

    def load(self, files):

        for file in files:
            self.load_file(file)

    def load_file(self, file):

        try:

            path = file
            data = open(path, "rb").read()

            # Create a SHA256 hash from the data
            id = hash(data)

            id = curi("doc", id)

            triples = []

            def emit(t):
                triples.append(t)

            self.metadata.id = id
            self.metadata.emit(emit)

            r = Document(
                metadata=Metadata(
                    id=id,
                    metadata = triples,
                    user=self.user,
                    collection=self.collection,
                ),
                data=base64.b64encode(data),
            )

            self.producer.send(r)

            print(f"{file}: Loaded successfully.")

        except Exception as e:
            print(f"{file}: Failed: {str(e)}", flush=True)
            
    def __del__(self):
        self.client.close()

def main():

    parser = argparse.ArgumentParser(
        prog='loader',
        description=__doc__,
    )

    default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://localhost:6650')
    default_output_queue = document_ingest_queue

    parser.add_argument(
        '-p', '--pulsar-host',
        default=default_pulsar_host,
        help=f'Pulsar host (default: {default_pulsar_host})',
    )

    parser.add_argument(
        '-o', '--output-queue',
        default=default_output_queue,
        help=f'Output queue (default: {default_output_queue})'
    )

    parser.add_argument(
        '-u', '--user',
        default=default_user,
        help=f'User ID (default: {default_user})'
    )

    parser.add_argument(
        '-c', '--collection',
        default=default_collection,
        help=f'Collection ID (default: {default_collection})'
    )

    parser.add_argument(
        '--name', help=f'Document name'
    )

    parser.add_argument(
        '--description', help=f'Document description'
    )

    parser.add_argument(
        '--copyright-notice', help=f'Copyright notice'
    )

    parser.add_argument(
        '--copyright-holder', help=f'Copyright holder'
    )

    parser.add_argument(
        '--copyright-year', help=f'Copyright year'
    )

    parser.add_argument(
        '--license', help=f'Copyright license'
    )

    parser.add_argument(
        '--publication-organization', help=f'Publication organization'
    )

    parser.add_argument(
        '--publication-description', help=f'Publication description'
    )

    parser.add_argument(
        '--publication-date', help=f'Publication date'
    )

    parser.add_argument(
        '--url', help=f'Document URL'
    )

    parser.add_argument(
        '--keyword', nargs='+', help=f'Keyword'
    )

    parser.add_argument(
        '--identifier', '--id', help=f'Document ID'
    )

    parser.add_argument(
        '-l', '--log-level',
        type=LogLevel,
        default=LogLevel.ERROR,
        choices=list(LogLevel),
        help=f'Output queue (default: info)'
    )

    parser.add_argument(
        'files', nargs='+',
        help=f'File to load'
    )

    args = parser.parse_args()

    while True:

        try:

            document = DigitalDocument(
                id,
                name=args.name,
                description=args.description,
                copyright_notice=args.copyright_notice,
                copyright_holder=args.copyright_holder,
                copyright_year=args.copyright_year,
                license=args.license,
                url=args.url,
                keywords=args.keyword,
            )

            if args.publication_organization:
                org = Organization(
                    id=curi("org", hash(args.publication_organization)),
                    name=args.publication_organization,
                )
                document.publication = PublicationEvent(
                    id = curi("pubev", str(uuid.uuid4())),
                    organization=org,
                    description=args.publication_description,
                    start_date=args.publication_date,
                    end_date=args.publication_date,
                )

            p = Loader(
                pulsar_host=args.pulsar_host,
                output_queue=args.output_queue,
                user=args.user,
                collection=args.collection,
                log_level=args.log_level,
                metadata=document,
            )

            p.load(args.files)

            print("All done.")
            break

        except Exception as e:

            print("Exception:", e, flush=True)
            print("Will retry...", flush=True)

        time.sleep(10)

main()