indico_toolkit 6.14.0__py3-none-any.whl

indico_toolkit/pipelines/file_processing.py

@@ -0,0 +1,157 @@
+import json
+import os
+import shutil
+from os.path import isfile
+from pathlib import Path
+from typing import Iterable, List, Tuple, Union
+
+
+class FileProcessing:
+    """
+    Class to support common file processing operations
+    """
+
+    def __init__(self, file_paths: List[str] = None):
+        if file_paths is None:
+            file_paths = []
+        self.file_paths: List[str] = file_paths
+        self.invalid_suffix_paths = set()
+
+    def get_file_paths_from_dir(
+        self,
+        path_to_dir: str,
+        accepted_types: Tuple[str] = ("pdf", "tiff", "tif", "doc", "docx"),
+        recursive_search: bool = False,
+    ):
+        """
+        Recursively find all file in specified types within a target directory
+
+        Args:
+            path_to_dir (str): root directory containing files
+            accepted_types (Tuple[str], optional): Valid extensions types to process.
+                Defaults to ("pdf", "tiff", "tif", "doc", "docx").
+            recursive_search (bool): search sub directories as well. Defaults to False.
+
+        Raises:
+            Exception: Must have at least 1 valid file
+        """
+        if recursive_search:
+            self._recursive_file_search(path_to_dir, accepted_types)
+        else:
+            self._non_recursive_file_search(path_to_dir, accepted_types)
+
+        if len(self.file_paths) == 0:
+            raise RuntimeError(
+                f"There are no files ending with {accepted_types} in {path_to_dir}"
+            )
+
+        print(
+            f"Found {len(self.file_paths)} valid files and "
+            f"{len(self.invalid_suffix_paths)} paths with invalid suffixes."
+        )
+
+    def move_all_file_paths(
+        self,
+        origin_dir: str,
+        destination_dir: str,
+        accepted_types: Tuple[str],
+        copy_files: bool = False,
+    ):
+        self.get_file_paths_from_dir(origin_dir, accepted_types, True)
+        destination_dir = Path(destination_dir)
+        if destination_dir.is_dir:
+            for file in self.file_paths:
+                initial_filepath = Path(file)
+                file_to_be_moved = initial_filepath.name
+                new_path_name = destination_dir / file_to_be_moved
+                if not copy_files:
+                    initial_filepath.rename(new_path_name)
+                else:
+                    shutil.copyfile(initial_filepath, new_path_name)
+        else:
+            raise RuntimeError(f"{destination_dir} is not a valid directory")
+
+    def batch_files(self, batch_size: int = 20) -> List[str]:
+        for i in range(0, len(self.file_paths), batch_size):
+            yield self.file_paths[i : i + batch_size]
+
+    def remove_files_if_processed(self, processed_files: Iterable[str]):
+        """
+        Removes files from self.file_paths if they are part of provided file iterable
+        Args:
+            processed_files (Iterable[str]): iterable of file names, NOT full paths,
+                e.g. ["invoice.pdf",]
+        """
+        unprocessed_filepaths = []
+        for filepath in self.file_paths:
+            if self.file_name_from_path(filepath) not in processed_files:
+                unprocessed_filepaths.append(filepath)
+        print(
+            f"Removing {len(self.file_paths) - len(unprocessed_filepaths)} "
+            "files from file_paths"
+        )
+        self.file_paths = unprocessed_filepaths
+
+    @staticmethod
+    def read_json(path_to_json: str) -> Union[dict, list]:
+        with open(path_to_json, "r") as f:
+            return json.load(f)
+
+    @staticmethod
+    def file_exists(path_to_file: str) -> bool:
+        return os.path.isfile(path_to_file)
+
+    @property
+    def parent_directory_of_filepaths(self) -> List[str]:
+        return [Path(i).parent.name for i in self.file_paths]
+
+    @staticmethod
+    def join_paths(start_path: str, end_path: str) -> str:
+        return os.path.join(start_path, end_path)
+
+    @staticmethod
+    def get_file_name_without_suffix(filepath: str) -> str:
+        return Path(filepath).stem
+
+    @staticmethod
+    def get_file_path_suffix(filepath: str) -> str:
+        return Path(filepath).suffix
+
+    @staticmethod
+    def file_name_from_path(filepath: str) -> str:
+        return Path(filepath).name
+
+    @staticmethod
+    def get_parent_path(filepath: str) -> str:
+        return str(Path(filepath).parent)
+
+    def __iter__(self):
+        """
+        Object is iterable on file_paths
+        """
+        current = 0
+        while current < len(self.file_paths):
+            yield self.file_paths[current]
+            current += 1
+
+    def _recursive_file_search(self, path_to_dir: str, accepted_types: Tuple[str]):
+        for root, _, files in os.walk(path_to_dir):
+            for name in files:
+                if self._check_acceptable_suffix(name, accepted_types):
+                    self.file_paths.append(os.path.join(root, name))
+                else:
+                    self.invalid_suffix_paths.add(os.path.join(root, name))
+
+    def _non_recursive_file_search(self, path_to_dir: str, accepted_types: Tuple[str]):
+        files = [os.path.join(path_to_dir, f) for f in os.listdir(path_to_dir)]
+        for fpath in files:
+            if self._check_acceptable_suffix(fpath, accepted_types):
+                self.file_paths.append(fpath)
+            elif isfile(fpath):
+                self.invalid_suffix_paths.add(fpath)
+
+    @staticmethod
+    def _check_acceptable_suffix(string: str, accepted_suffixes: Tuple[str]) -> bool:
+        if string.lower().endswith(accepted_suffixes):
+            return True
+        return False

indico_toolkit/polling/__init__.py

@@ -0,0 +1,8 @@
+from .autoreview import AutoReviewed, AutoReviewPoller
+from .downstream import DownstreamPoller
+
+__all__ = (
+    "AutoReviewed",
+    "AutoReviewPoller",
+    "DownstreamPoller",
+)

indico_toolkit/polling/autoreview.py

@@ -0,0 +1,204 @@
+import asyncio
+import logging
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from indico import AsyncIndicoClient, IndicoConfig  # type: ignore[import-untyped]
+from indico.queries import (  # type: ignore[import-untyped]
+    GetSubmission,
+    JobStatus,
+    RetrieveStorageObject,
+    SubmitReview,
+)
+
+from .. import etloutput, results
+from ..etloutput import EtlOutput
+from ..results import Document, Result
+from ..retry import retry
+from .queries import SubmissionIdsPendingAutoReview
+
+if TYPE_CHECKING:
+    from collections.abc import Awaitable, Callable
+    from typing import Any, NoReturn, TypeAlias
+
+    SubmissionId: TypeAlias = int
+    Worker: TypeAlias = asyncio.Task[None]
+    WorkerQueue: TypeAlias = asyncio.Queue[tuple[SubmissionId, Worker]]
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class AutoReviewed:
+    changes: "dict[str, Any] | list[dict[str, Any]]"
+    reject: bool = False
+    stp: bool = False
+
+
+class AutoReviewPoller:
+    """
+    Polls for submissions requiring auto review, processes them,
+    and submits the review results concurrently.
+    """
+
+    def __init__(
+        self,
+        config: IndicoConfig,
+        workflow_id: int,
+        auto_review: "Callable[[Result, dict[Document, EtlOutput]], Awaitable[AutoReviewed]]",  # noqa: E501
+        *,
+        worker_count: int = 8,
+        spawn_rate: float = 1,
+        poll_delay: float = 30,
+        load_etl_output: bool = True,
+        load_text: bool = True,
+        load_tokens: bool = True,
+        load_tables: bool = False,
+        retry_count: int = 4,
+        retry_wait: float = 1,
+        retry_backoff: float = 4,
+        retry_jitter: float = 0.5,
+    ):
+        self._config = config
+        self._workflow_id = workflow_id
+        self._auto_review = auto_review
+        self._worker_count = worker_count
+        self._spawn_rate = spawn_rate
+        self._poll_delay = poll_delay
+        self._load_etl_output = load_etl_output
+        self._load_text = load_text
+        self._load_tokens = load_tokens
+        self._load_tables = load_tables
+
+        self._retry = retry(
+            Exception,
+            count=retry_count,
+            wait=retry_wait,
+            backoff=retry_backoff,
+            jitter=retry_jitter,
+        )
+        self._worker_slots = asyncio.Semaphore(worker_count)
+        self._worker_queue: "WorkerQueue" = asyncio.Queue(1)
+        self._processing_submission_ids: "set[SubmissionId]" = set()
+
+    async def poll_forever(self) -> "NoReturn":  # type: ignore[misc]
+        logger.info(
+            "Starting auto review poller for: "
+            f"host={self._config.host} "
+            f"workflow_id={self._workflow_id} "
+            f"worker_count={self._worker_count}"
+        )
+
+        async with AsyncIndicoClient(self._config) as client:
+            self._client_call = self._retry(client.call)
+            await asyncio.gather(
+                self._spawn_workers(),
+                *(self._reap_workers() for _ in range(self._worker_count)),
+            )
+
+    async def _retrieve_storage_object(self, url: str) -> object:
+        return await self._client_call(RetrieveStorageObject(url))
+
+    async def _spawn_workers(self) -> None:
+        """
+        Poll for submissions pending auto review and spawn workers to process them.
+        `self._worker_slots` limits the number of workers that can run concurrently.
+        Submission IDs in progress are tracked with `self._processing_submission_ids`.
+        """
+        logger.info(
+            f"Polling submissions pending auto review every {self._poll_delay} seconds"
+        )
+
+        while True:
+            try:
+                submission_ids: "set[SubmissionId]" = await self._client_call(
+                    SubmissionIdsPendingAutoReview(self._workflow_id)
+                )
+            except Exception:
+                logger.exception("Error occurred while polling submissions")
+                await asyncio.sleep(self._poll_delay)
+                continue
+
+            submission_ids -= self._processing_submission_ids
+
+            if not submission_ids:
+                await asyncio.sleep(self._poll_delay)
+                continue
+
+            for submission_id in submission_ids:
+                await self._worker_slots.acquire()
+                logger.info(f"Spawning worker for {submission_id=}")
+                self._processing_submission_ids.add(submission_id)
+                worker = asyncio.create_task(self._worker(submission_id))
+                await self._worker_queue.put((submission_id, worker))
+                await asyncio.sleep(1 / self._spawn_rate)
+
+    async def _worker(self, submission_id: "SubmissionId") -> None:
+        """
+        Process a single submission by retrieving submission metadata, the result file,
+        etl output, calling `self._auto_review`, and submitting changes.
+        """
+        logger.info(f"Retrieving metadata for {submission_id=}")
+        submission = await self._client_call(GetSubmission(submission_id))
+
+        logger.info(f"Retrieving results for {submission_id=}")
+        result = await results.load_async(
+            submission.result_file,
+            reader=self._retrieve_storage_object,
+        )
+
+        if self._load_etl_output:
+            logger.info(f"Retrieving etl output for {submission_id=}")
+            etl_outputs = {
+                document: await etloutput.load_async(
+                    document.etl_output_uri,
+                    reader=self._retrieve_storage_object,
+                    text=self._load_text,
+                    tokens=self._load_tokens,
+                    tables=self._load_tables,
+                )
+                for document in result.documents
+                if not document.failed
+            }
+        else:
+            logger.info(f"Skipping etl output for {submission_id=}")
+            etl_outputs = {}
+
+        logger.info(f"Applying auto review for {submission_id=}")
+        auto_reviewed = await self._auto_review(result, etl_outputs)
+
+        logger.info(f"Submitting auto review for {submission_id=}")
+        job = await self._client_call(
+            SubmitReview(
+                submission_id,
+                changes=auto_reviewed.changes,
+                rejected=auto_reviewed.reject,
+                force_complete=auto_reviewed.stp,
+            )
+        )
+        job = await self._client_call(JobStatus(job.id))
+
+        if job.status == "SUCCESS":
+            logger.info(f"Completed auto review of {submission_id=}")
+        else:
+            logger.error(
+                f"Submit failed for {submission_id=}: "
+                f"{job.status=!r} {job.result=!r}"
+            )
+
+    async def _reap_workers(self) -> None:
+        """
+        Reap completed workers, releasing their slots for new tasks. Log errors for
+        submissions that failed to process. Remove their submission IDs from
+        `self._processing_submission_ids` to be retried.
+        """
+        while True:
+            submission_id, worker = await self._worker_queue.get()
+
+            try:
+                await worker
+            except Exception:
+                logger.exception(f"Error occurred while processing {submission_id=}")
+
+            self._processing_submission_ids.remove(submission_id)
+            self._worker_slots.release()

indico_toolkit/polling/downstream.py

@@ -0,0 +1,145 @@
+import asyncio
+import logging
+from typing import TYPE_CHECKING
+
+from indico import AsyncIndicoClient, IndicoConfig  # type: ignore[import-untyped]
+from indico.queries import (  # type: ignore[import-untyped]
+    GetSubmission,
+    UpdateSubmission,
+)
+from indico.types import Submission  # type: ignore[import-untyped]
+
+from ..retry import retry
+from .queries import SubmissionIdsPendingDownstream
+
+if TYPE_CHECKING:
+    from collections.abc import Awaitable, Callable
+    from typing import NoReturn, TypeAlias
+
+    SubmissionId: TypeAlias = int
+    Worker: TypeAlias = asyncio.Task[None]
+    WorkerQueue: TypeAlias = asyncio.Queue[tuple[SubmissionId, Worker]]
+
+logger = logging.getLogger(__name__)
+
+
+class DownstreamPoller:
+    """
+    Polls for completed and failed submissions pending downstream egestion, processes
+    them concurrently, and marks them as retrieved.
+    """
+
+    def __init__(
+        self,
+        config: IndicoConfig,
+        workflow_id: int,
+        downstream: "Callable[[Submission], Awaitable[None]]",
+        *,
+        worker_count: int = 8,
+        spawn_rate: float = 1,
+        poll_delay: float = 30,
+        retry_count: int = 4,
+        retry_wait: float = 1,
+        retry_backoff: float = 4,
+        retry_jitter: float = 0.5,
+    ):
+        self._config = config
+        self._workflow_id = workflow_id
+        self._downstream = downstream
+        self._worker_count = worker_count
+        self._spawn_rate = spawn_rate
+        self._poll_delay = poll_delay
+
+        self._retry = retry(
+            Exception,
+            count=retry_count,
+            wait=retry_wait,
+            backoff=retry_backoff,
+            jitter=retry_jitter,
+        )
+        self._worker_slots = asyncio.Semaphore(worker_count)
+        self._worker_queue: "WorkerQueue" = asyncio.Queue(1)
+        self._processing_submission_ids: "set[SubmissionId]" = set()
+
+    async def poll_forever(self) -> "NoReturn":  # type: ignore[misc]
+        logger.info(
+            "Starting downstream poller for: "
+            f"host={self._config.host} "
+            f"workflow_id={self._workflow_id} "
+            f"worker_count={self._worker_count}"
+        )
+
+        async with AsyncIndicoClient(self._config) as client:
+            self._client_call = self._retry(client.call)
+            await asyncio.gather(
+                self._spawn_workers(),
+                *(self._reap_workers() for _ in range(self._worker_count)),
+            )
+
+    async def _spawn_workers(self) -> None:
+        """
+        Poll for completed and failed submissions and spawn workers to send them
+        downstream. `self._worker_slots` limits the number of workers that can run
+        concurrently. Submission IDs in progress are tracked with
+        `self._processing_submission_ids`.
+        """
+        logger.info(
+            f"Polling submissions pending downstream every {self._poll_delay} seconds"
+        )
+
+        while True:
+            try:
+                submission_ids: "set[SubmissionId]" = await self._client_call(
+                    SubmissionIdsPendingDownstream(self._workflow_id)
+                )
+            except Exception:
+                logger.exception("Error occurred while polling submissions")
+                await asyncio.sleep(self._poll_delay)
+                continue
+
+            submission_ids -= self._processing_submission_ids
+
+            if not submission_ids:
+                await asyncio.sleep(self._poll_delay)
+                continue
+
+            for submission_id in submission_ids:
+                await self._worker_slots.acquire()
+                logger.info(f"Spawning worker for {submission_id=}")
+                self._processing_submission_ids.add(submission_id)
+                worker = asyncio.create_task(self._worker(submission_id))
+                await self._worker_queue.put((submission_id, worker))
+                await asyncio.sleep(1 / self._spawn_rate)
+
+    async def _worker(self, submission_id: "SubmissionId") -> None:
+        """
+        Process a single submission by retrieving submission metadata and calling
+        `self._downstream`. Once completed, mark the submission retrieved.
+        """
+        logger.info(f"Retrieving metadata for {submission_id=}")
+        submission = await self._client_call(GetSubmission(submission_id))
+
+        logger.info(f"Sending {submission_id=} downstream")
+        await self._downstream(submission)
+
+        logger.info(f"Marking {submission_id=} retrieved")
+        await self._client_call(UpdateSubmission(submission_id, retrieved=True))
+
+        logger.info(f"Completed dowstream of {submission_id=}")
+
+    async def _reap_workers(self) -> None:
+        """
+        Reap completed workers, releasing their slots for new tasks. Log errors for
+        submissions that failed to process. Remove their submission IDs from
+        `self._processing_submission_ids` to be retried.
+        """
+        while True:
+            submission_id, worker = await self._worker_queue.get()
+
+            try:
+                await worker
+            except Exception:
+                logger.exception(f"Error occurred while processing {submission_id=}")
+
+            self._processing_submission_ids.remove(submission_id)
+            self._worker_slots.release()

indico_toolkit/polling/queries.py

@@ -0,0 +1,68 @@
+from typing import TYPE_CHECKING
+
+from indico.queries import GraphQLRequest  # type: ignore[import-untyped]
+
+if TYPE_CHECKING:
+    from typing import Any
+
+
+class SubmissionIdsPendingAutoReview(GraphQLRequest):  # type: ignore[misc]
+    QUERY = """
+    query SubmissionIdsPendingAutoReview($workflowIds: [Int]) {
+        submissions(
+            desc: false
+            filters: { status: PENDING_AUTO_REVIEW }
+            limit: 1000
+            orderBy: ID
+            workflowIds: $workflowIds
+        ) {
+            submissions {
+                id
+            }
+        }
+    }
+    """
+
+    def __init__(self, workflow_id: int):
+        super().__init__(self.QUERY, {"workflowIds": [workflow_id]})
+
+    def process_response(self, response: "Any") -> set[int]:
+        response = super().process_response(response)
+        return {
+            submission["id"] for submission in response["submissions"]["submissions"]
+        }
+
+
+class SubmissionIdsPendingDownstream(GraphQLRequest):  # type: ignore[misc]
+    QUERY = """
+    query SubmissionIdsPendingDownstream($workflowIds: [Int]) {
+        submissions(
+            desc: false
+            filters: {
+                AND: {
+                    retrieved: false
+                    OR: [
+                        { status: COMPLETE }
+                        { status: FAILED }
+                    ]
+                }
+            }
+            limit: 1000
+            orderBy: ID
+            workflowIds: $workflowIds
+        ) {
+            submissions {
+                id
+            }
+        }
+    }
+    """
+
+    def __init__(self, workflow_id: int):
+        super().__init__(self.QUERY, {"workflowIds": [workflow_id]})
+
+    def process_response(self, response: "Any") -> set[int]:
+        response = super().process_response(response)
+        return {
+            submission["id"] for submission in response["submissions"]["submissions"]
+        }

indico_toolkit/results/__init__.py

@@ -0,0 +1,104 @@
+import json
+from typing import TYPE_CHECKING
+
+from .document import Document
+from .errors import ResultError
+from .model import ModelGroup, ModelGroupType
+from .predictionlist import PredictionList
+from .predictions import (
+    NULL_BOX,
+    NULL_CITATION,
+    NULL_SPAN,
+    Box,
+    Classification,
+    DocumentExtraction,
+    Extraction,
+    FormExtraction,
+    FormExtractionType,
+    Group,
+    Prediction,
+    Span,
+    Summarization,
+    Unbundling,
+)
+from .result import Result
+from .review import Review, ReviewType
+from .utilities import get
+
+if TYPE_CHECKING:
+    from collections.abc import Awaitable, Callable
+
+
+__all__ = (
+    "Box",
+    "Classification",
+    "Document",
+    "DocumentExtraction",
+    "Extraction",
+    "FormExtraction",
+    "FormExtractionType",
+    "Group",
+    "load",
+    "load_async",
+    "ModelGroup",
+    "ModelGroupType",
+    "NULL_BOX",
+    "NULL_CITATION",
+    "NULL_SPAN",
+    "Prediction",
+    "PredictionList",
+    "Result",
+    "ResultError",
+    "Review",
+    "ReviewType",
+    "Span",
+    "Summarization",
+    "Unbundling",
+)
+
+
+def load(result: object, *, reader: "Callable[..., object] | None" = None) -> Result:
+    """
+    Load `result` as a Result dataclass.
+
+    `result` can be a dict, a JSON string, or something that can be read with `reader`
+    to produce either.
+
+    ```
+    for result_file in result_folder.glob("*.json"):
+        result = results.load(result_file, reader=Path.read_text)
+    ```
+    """
+    if reader:
+        result = reader(result)
+
+    return _load(result)
+
+
+async def load_async(
+    result: object, *, reader: "Callable[..., Awaitable[object]] | None" = None
+) -> Result:
+    """
+    Load `result` as a Result dataclass.
+
+    `result` can be a dict, a JSON string, or something that can be read with `reader`
+    to produce either.
+    """
+    if reader:
+        result = await reader(result)
+
+    return _load(result)
+
+
+def _load(result: object) -> Result:
+    if isinstance(result, str) and result.strip().startswith("{"):
+        result = json.loads(result)
+
+    file_version = get(result, int, "file_version")
+
+    if file_version == 1:
+        return Result.from_v1_dict(result)
+    elif file_version == 3:
+        return Result.from_v3_dict(result)
+    else:
+        raise ResultError(f"unsupported file version `{file_version!r}`")