cleverswarm-python-client/src/cleverswarm_python_client/cswarm_text_to_kg_client.py

"""
CleverSwarm - unstructured text to knowledge graph triplets extractor client.

 Copyright (c) 2016 - present Syncleus, Inc.
 Copyright (c) 2016 - present CleverThis, Inc.

 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at

     http://www.apache.org/licenses/LICENSE-2.0

 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
"""

import argparse
import logging
import os
import sys
from pathlib import Path
from typing import Any, Dict

from cleverswarm_python_client.libs.cleverswarm_client import CleverSwarmClient
from cleverswarm_python_client.libs.exceptions import ClientException
from cleverswarm_python_client.libs.file_type_enum import FileTypeAPI
from cleverswarm_python_client.libs.job_enums import JobStatus, JobType

logger = logging.getLogger("Text-to-KG-client main")


class TextoToKGCLI(object):
    def __init__(
        self,
        base_url: str,
        detailed: bool = False,
        input_prefix: str = None,
        output_prefix: str = None,
        username: str = None,
        token: str = None,
    ):

        self._input_prefix: Path = Path(input_prefix) if input_prefix else Path("../..")
        self._output_prefix: Path = (
            Path(output_prefix) if output_prefix else Path("../..")
        )
        self._username: str = username

        self._input_prefix = self._input_prefix.resolve()
        self._output_prefix = self._output_prefix.resolve()

        if not self._input_prefix.exists() or not self._input_prefix.is_dir():
            raise ClientException(
                f"Input prefix does not exist or is not a folder: {input_prefix}"
            )

        self._output_prefix.mkdir(parents=True, exist_ok=True)

        if not self._output_prefix.exists() or not self._output_prefix.is_dir():
            raise ClientException(
                f"Input prefix could not be created or is not a folder: {output_prefix}"
            )

        self._client: CleverSwarmClient = CleverSwarmClient(base_url, token, username)

        self._detailed = detailed

    def create_job(
        self,
        input_text: str,
        ontology_json: str,
        ontology_owl: str,
        wildcards: str = None,
        force_filetype: FileTypeAPI = FileTypeAPI.AutoDetect,
    ):
        input_text_path: Path = self._input_prefix / input_text
        # Ontology JSON file is optional - only create path if it was provided
        ontology_json_path: Path | None = (
            self._input_prefix / ontology_json if ontology_json else None
        )
        ontology_owl_path: Path = self._input_prefix / ontology_owl
        wildcards_path: Path = Path(wildcards) if wildcards else None

        if not input_text_path.exists() or not input_text_path.is_file():
            raise ClientException(
                f"Unstructured input text does not exist or is not a file: {str(input_text_path)}"
            )

        # If ontology_json was provided, validate that the file exists
        if (
            ontology_json
            and ontology_json_path
            and (not ontology_json_path.exists() or not ontology_json_path.is_file())
        ):
            raise ClientException(
                f"Specified optional Ontology JSON file does not exist or is not a file: {str(ontology_json_path)}"
            )

        # Ontology JSON file is optional - use the path if it was provided and exists
        ontology_json_path_final: Path | None = ontology_json_path

        if not ontology_owl_path.exists() or not ontology_owl_path.is_file():
            raise ClientException(
                f"Input ontology OWL file does not exist or is not a file: {str(ontology_owl_path)}"
            )

        wildcards_job: bool = False
        if wildcards_path is not None:
            wildcards_job = True
            if not wildcards_path.exists() or not wildcards_path.is_file():
                raise ClientException(
                    f"Input wildcards file does not exist or is not a file: {str(wildcards_path)}"
                )

        if wildcards_job:
            job_id: str = self._client.create_unstructured_to_kg_wildcards_job(
                input_text_path,
                ontology_json_path_final,
                ontology_owl_path,
                wildcards_path,
                force_filetype=force_filetype,
            )
        else:
            job_id: str = self._client.create_unstructured_to_kg_job(
                input_text_path,
                ontology_json_path_final,
                ontology_owl_path,
                force_filetype=force_filetype,
            )

        return job_id

    def poll_job_to_completion(self, job_id: str) -> None:
        """
        Polls a job by its job id and waits for the job completion in the server.
        It throws an exception if the job failed or is in a state that cannot be executed.
        :param job_id: the job id of the job to be executed.
        """
        status: JobStatus = self._client.get_job_status(job_id)
        if status == JobStatus.Failed:
            raise ClientException(f"Job with ID: {job_id} finished without completing.")
        if status == JobStatus.Completed:
            return
        if status != JobStatus.ReadyForProcessing and status != JobStatus.Processing:
            raise ClientException(
                f"Job with ID: {job_id} is not processing, nor is it ready for processing."
            )
        final_status: JobStatus = self._client.poll_job_ready_or_failed(job_id)
        if final_status != JobStatus.Completed:
            raise ClientException(f"Job with ID: {job_id} finished without completing.")

    def download_job_result(self, job_id: str) -> None:
        details: Dict[str, Any] = self._client.get_job_details(job_id)

        if len(details["unstructured_sources"]) == 0:
            raise ClientException(
                f"Unstructured to KG job: {job_id} is not valid. It has no input files."
            )
        if len(details["unstructured_sources"]) > 1:
            raise ClientException(
                f"Unstructured to KG job: {job_id} is not valid. It has more than one input files."
            )
        input_filename: str = details["unstructured_sources"][0]
        output_filename: str = input_filename + ".json"
        output_filepath: Path = self._output_prefix / output_filename
        self._client.retrieve_unstructured_to_kg_files(job_id, output_filepath)
        if self._detailed:
            print("Job result was downloaded to: {}".format(str(output_filepath)))

    def delete_server_job(self, job_id: str) -> None:
        """
        Deletes a given job by its ID. A job can only be deleted if it is in a state that allows it.
        :param job_id: the server job id
        """
        del_status = self._client.delete_job(job_id)
        if not del_status:
            print("Failed to delete Job")
        else:
            print("Job was successfully deleted")

    def list_server_jobs(self) -> Dict:
        """
        List all text to KG extraction jobs in the server that are owned by the current logged-in user.
        :return: JSON dictionary with job listing and respective states.
        """
        json_data = self._client.get_jobs_list(is_benchmark=False)
        return json_data

    def print_server_jobs(self, is_detailed: bool = False) -> None:
        """
        Prints a summarized or detailed status of the text to KG extraction jobs in the server.
        """
        jobs = self.list_server_jobs()
        for job in jobs:
            print("Job ID: {}".format(job["id"]))
            print("Created at: {}".format(job["created"]))
            print("Type: {}".format(job["type"]))
            print("Status: {}".format(job["status"]))
            print("Retries: {}".format(job["retries_count"]))
            print("Machine ID: {}".format(job["machine_id"]))
            print("Worker ID: {}".format(job["worker_id"]))
            print("Input file type: {}".format((job["file_type"])))
            if is_detailed:
                print("Unstructured text input files:")
                for source in job["unstructured_sources"]:
                    print(f"   {source}")
                print("Ontology JSON input files:")
                for source in job["ontology_sources"]:
                    print(f"   {source}")
                print("Ontology OWL/XML input files:")
                for source in job["ontology_spec_sources"]:
                    print(f"   {source}")
                if JobType(job["type"]) == JobType.UnstructuredWithWildcards:
                    print("Wildcards input files:")
                    for source in job["wildcards_sources"]:
                        print(f"   {source}")
            print("----------------------------")


def main():
    """Main entry point for the text-to-KG client console script."""

    logging.basicConfig(
        level=os.environ.get("PYTHON_LOGGING_LEVEL", logging.INFO),
        stream=sys.stderr,
        datefmt="%Y-%m-%dT%H:%M:%S",
        format=("%(asctime)s - %(name)s - %(levelname)s - %(message)s"),
    )

    parser = argparse.ArgumentParser(
        description="CleverSwarm unstructured text to Knowledge Graph triplets client for REST API."
    )
    parser.add_argument(
        "--input_prefix",
        type=str,
        default=".",
        help="Path prefix for the common input folder structure.",
    )
    parser.add_argument(
        "--output_prefix", type=str, default=".", help="Path to the output folder."
    )
    parser.add_argument(
        "--username", type=str, help="Username for logging in into server"
    )
    parser.add_argument(
        "--token",
        type=str,
        help="Valid Bearer Token value to avoid the need for authentication",
    )
    parser.add_argument(
        "--server_url",
        type=str,
        default="http://localhost:8000/api/v0/",
        help="REST server URL",
    )
    parser.add_argument(
        "--action",
        choices=[
            "create-and-exit",
            "delete-and-exit",
            "download-kg-and-exit",
            "list-kg-jobs-and-exit",
            "create-download-kg-delete",
            "create-download-kg",
        ],
        type=str,
        help="Define action to be performed",
        default="create-download-kg-delete",
    )
    parser.add_argument(
        "--job_id",
        type=str,
        help="Identify a specific job. To be used with create-and-exit, delete-and-exit, download-metrics-and-exit",
    )
    parser.add_argument(
        "--unstructured_text",
        type=str,
        help="Unstructured text input file for the Knowledge Graph triplet extraction",
    )
    parser.add_argument(
        "--ontology_json",
        type=str,
        help="Optional enriched ontology file with descriptions in JSON format",
    )
    parser.add_argument(
        "--ontology_owl",
        type=str,
        help="OWL ontology file in XML format",
    )
    parser.add_argument(
        "--wildcards",
        type=str,
        help="Wildcards ontology file. If specified, a wildcards filtering job will be created instead "
        "of a regular KG extraction.",
    )
    parser.add_argument(
        "--detailed",
        default=False,
        action="store_true",
    )

    args = parser.parse_args()

    create = False
    poll = False
    download = False
    delete = False
    list_jobs = False

    if args.action == "create-and-exit":
        create = True
        poll = False
        if args.unstructured_text is None:
            print("create-and-exit action requires --unstructured_text argument.\n")
            parser.print_help()
            exit(-1)
        if args.ontology_owl is None:
            print("create-and-exit action requires --ontology_owl argument.\n")
            parser.print_help()
            exit(-1)
    elif args.action == "delete-and-exit":
        delete = True
        if args.job_id is None:
            parser.print_help()
            exit(-1)
    elif args.action == "download-kg-and-exit":
        poll = True
        download = True
        if args.job_id is None:
            print("download-kg-and-exit action requires --job_id argument.\n")
            parser.print_help()
            exit(-1)
    elif args.action == "list-kg-jobs-and-exit":
        list_jobs = True
    elif args.action == "create-download-kg-delete":
        create = True
        poll = True
        download = True
        delete = True
        if args.unstructured_text is None:
            print(
                "create-download-kg-delete action requires --unstructured_text argument.\n"
            )
            parser.print_help()
            exit(-1)
        if args.ontology_owl is None:
            print(
                "create-download-kg-delete action requires --ontology_owl argument.\n"
            )
            parser.print_help()
            exit(-1)
    elif args.action == "create-download-kg":
        create = True
        poll = True
        download = True
        if args.unstructured_text is None:
            print(
                "create-download-kg-delete action requires --unstructured_text argument.\n"
            )
            parser.print_help()
            exit(-1)
        if args.ontology_owl is None:
            print(
                "create-download-kg-delete action requires --ontology_owl argument.\n"
            )
            parser.print_help()
            exit(-1)

    url: str = args.server_url
    if not url.endswith("/"):
        url += "/"

    try:
        job_id = args.job_id
        cli = TextoToKGCLI(
            base_url=url,
            detailed=args.detailed,
            input_prefix=args.input_prefix,
            output_prefix=args.output_prefix,
            username=args.username,
            token=args.token,
        )
        if create:
            job_id = cli.create_job(
                args.unstructured_text,
                args.ontology_json,
                args.ontology_owl,
                args.wildcards,
            )

        if poll:
            cli.poll_job_to_completion(job_id)

        if download:
            cli.download_job_result(job_id)

        if delete:
            cli.delete_server_job(job_id)

        if list_jobs:
            cli.print_server_jobs(args.detailed)

    except ClientException as msg:
        print(str(msg))
        exit(-1)
    except Exception as error:
        logger.error(f"Unexpected error occurred: {error}", exc_info=error)
        print("Unexpected error occurred, check the logs for more details")
        exit(-1)


if __name__ == "__main__":
    main()