mseep-txtai 9.1.1__py3-none-any.whl

txtai/workflow/task/image.py

@@ -0,0 +1,36 @@
+"""
+ImageTask module
+"""
+
+import re
+
+# Conditional import
+try:
+    from PIL import Image
+
+    PIL = True
+except ImportError:
+    PIL = False
+
+from .file import FileTask
+
+
+class ImageTask(FileTask):
+    """
+    Task that processes image file urls
+    """
+
+    def register(self):
+        """
+        Checks if required dependencies are installed.
+        """
+
+        if not PIL:
+            raise ImportError('ImageTask is not available - install "workflow" extra to enable')
+
+    def accept(self, element):
+        # Only accept image files
+        return super().accept(element) and re.search(r"\.(gif|bmp|jpg|jpeg|png|webp)$", element.lower())
+
+    def prepare(self, element):
+        return Image.open(super().prepare(element))

txtai/workflow/task/retrieve.py

@@ -0,0 +1,61 @@
+"""
+RetrieveTask module
+"""
+
+import os
+import tempfile
+
+from urllib.request import urlretrieve
+from urllib.parse import urlparse
+
+from .url import UrlTask
+
+
+class RetrieveTask(UrlTask):
+    """
+    Task that retrieves urls (local or remote) to a local directory.
+    """
+
+    def register(self, directory=None, flatten=True):
+        """
+        Adds retrieve parameters to task.
+
+        Args:
+            directory: local directory used to store retrieved files
+            flatten: flatten input directory structure, defaults to True
+        """
+
+        # pylint: disable=W0201
+        # Create default temporary directory if not specified
+        if not directory:
+            # Save tempdir to prevent content from being deleted until this task is out of scope
+            # pylint: disable=R1732
+            self.tempdir = tempfile.TemporaryDirectory()
+            directory = self.tempdir.name
+
+        # Create output directory if necessary
+        os.makedirs(directory, exist_ok=True)
+
+        self.directory = directory
+        self.flatten = flatten
+
+    def prepare(self, element):
+        # Extract file path from URL
+        path = urlparse(element).path
+
+        if self.flatten:
+            # Flatten directory structure (default)
+            path = os.path.join(self.directory, os.path.basename(path))
+        else:
+            # Derive output path
+            path = os.path.join(self.directory, os.path.normpath(path.lstrip("/")))
+            directory = os.path.dirname(path)
+
+            # Create local directory, if necessary
+            os.makedirs(directory, exist_ok=True)
+
+        # Retrieve URL
+        urlretrieve(element, path)
+
+        # Return new file path
+        return path

txtai/workflow/task/service.py

@@ -0,0 +1,102 @@
+"""
+ServiceTask module
+"""
+
+# Conditional import
+try:
+    import requests
+    import xmltodict
+
+    XML_TO_DICT = True
+except ImportError:
+    XML_TO_DICT = False
+
+from .base import Task
+
+
+class ServiceTask(Task):
+    """
+    Task to runs requests against remote service urls.
+    """
+
+    def register(self, url=None, method=None, params=None, batch=True, extract=None):
+        """
+        Adds service parameters to task. Checks if required dependencies are installed.
+
+        Args:
+            url: url to connect to
+            method: http method, GET or POST
+            params: default query parameters
+            batch: if True, all elements are passed in a single batch request, otherwise a service call is executed per element
+            extract: list of sections to extract from response
+        """
+
+        if not XML_TO_DICT:
+            raise ImportError('ServiceTask is not available - install "workflow" extra to enable')
+
+        # pylint: disable=W0201
+        # Save URL, method and parameter defaults
+        self.url = url
+        self.method = method
+        self.params = params
+
+        # If True, all elements are passed in a single batch request, otherwise a service call is executed per element
+        self.batch = batch
+
+        # Save sections to extract. Supports both a single string and a hierarchical list of sections.
+        self.extract = extract
+        if self.extract:
+            self.extract = [self.extract] if isinstance(self.extract, str) else self.extract
+
+    def execute(self, elements, executor=None):
+        if self.batch:
+            elements = self.request(elements)
+        else:
+            elements = [self.request(element) for element in elements]
+
+        return super().execute(elements, executor)
+
+    def request(self, data):
+        """
+        Execute service request.
+
+        Args:
+            url: service url
+            method: method (get or post)
+            params: dict of constant parameters to pass to request
+            data: dynamic data for this specific request
+
+        Returns:
+            response as JSON
+        """
+
+        if not self.params:
+            params = data
+        else:
+            # Create copy of parameters
+            params = self.params.copy()
+
+            # Add data to parameters
+            for key in params:
+                if not params[key]:
+                    params[key] = data
+
+        # Run request
+        if self.method and self.method.lower() == "get":
+            response = requests.get(self.url, params=params)
+        else:
+            response = requests.post(self.url, json=params)
+
+        # Parse data based on content-type
+        mimetype = response.headers["Content-Type"].split(";")[0]
+        if mimetype.lower().endswith("xml"):
+            data = xmltodict.parse(response.text)
+        else:
+            data = response.json()
+
+        # Extract content from response, if necessary
+        if self.extract:
+            for tag in self.extract:
+                data = data[tag]
+
+        return data

txtai/workflow/task/storage.py

@@ -0,0 +1,110 @@
+"""
+StorageTask module
+"""
+
+import os
+import re
+
+# Conditional import
+try:
+    from libcloud.storage.providers import get_driver
+
+    LIBCLOUD = True
+except ImportError:
+    LIBCLOUD = False
+
+from .base import Task
+
+
+class StorageTask(Task):
+    """
+    Task that processes object storage buckets. Supports local and cloud providers in Apache libcloud.
+    """
+
+    # URL prefix
+    PREFIX = r"(\w+):\/\/.*"
+    PATH = r"\w+:\/\/(.*)"
+
+    def register(self, key=None, secret=None, host=None, port=None, token=None, region=None):
+        """
+        Checks if required dependencies are installed. Reads in cloud storage parameters.
+
+        Args:
+            key: provider-specific access key
+            secret: provider-specific access secret
+            host: server host name
+            port: server port
+            token: temporary session token
+            region: storage region
+        """
+
+        if not LIBCLOUD:
+            raise ImportError('StorageTask is not available - install "workflow" extra to enable')
+
+        # pylint: disable=W0201
+        self.key = key
+        self.secret = secret
+        self.host = host
+        self.port = port
+        self.token = token
+        self.region = region
+
+    def __call__(self, elements, executor=None):
+        # Create aggregated directory listing for all elements
+        outputs = []
+        for element in elements:
+            if self.matches(element):
+                # Get directory listing and run actions
+                outputs.extend(super().__call__(self.list(element), executor))
+            else:
+                outputs.append(element)
+
+        return outputs
+
+    def matches(self, element):
+        """
+        Determines if this element is a storage element.
+
+        Args:
+            element: input storage element
+
+        Returns:
+            True if this is a storage element
+        """
+
+        # Only accept file URLs
+        return re.match(StorageTask.PREFIX, self.upack(element, True).lower())
+
+    def list(self, element):
+        """
+        Gets a list of urls for a object container.
+
+        Args:
+            element: object container
+
+        Returns:
+            list of urls
+        """
+
+        provider = re.sub(StorageTask.PREFIX, r"\1", element.lower())
+        path = re.sub(StorageTask.PATH, r"\1", element)
+
+        # Load key and secret, if applicable
+        key = self.key if self.key is not None else os.environ.get("ACCESS_KEY")
+        secret = self.secret if self.secret is not None else os.environ.get("ACCESS_SECRET")
+
+        # Parse key and container
+        key, container = (os.path.dirname(path), os.path.basename(path)) if key is None else (key, path)
+
+        # Parse optional prefix from container
+        parts = container.split("/", 1)
+        container, prefix = (parts[0], parts[1]) if len(parts) > 1 else (container, None)
+
+        # Get driver for provider
+        driver = get_driver(provider)
+
+        # Get client connection
+        client = driver(key, secret, **{field: getattr(self, field) for field in ["host", "port", "region", "token"] if getattr(self, field)})
+
+        container = client.get_container(container_name=container)
+        return [client.get_object_cdn_url(obj) for obj in client.list_container_objects(container=container, prefix=prefix)]

txtai/workflow/task/stream.py

@@ -0,0 +1,33 @@
+"""
+StreamTask module
+"""
+
+from .base import Task
+
+
+class StreamTask(Task):
+    """
+    Task that calls a task action and yields results.
+    """
+
+    def register(self, batch=False):
+        """
+        Adds stream parameters to task.
+
+        Args:
+            batch: all elements are passed to a single action call if True, otherwise an action call is executed per element, defaults to False
+        """
+
+        # pylint: disable=W0201
+        # All elements are passed to a single action call if True, otherwise an action call is executed per element, defaults to False
+        self.batch = batch
+
+    def __call__(self, elements, executor=None):
+        for action in self.action:
+            if self.batch:
+                # Single batch call
+                yield from action(elements)
+            else:
+                # Call action for each element
+                for x in elements:
+                    yield from action(x)

txtai/workflow/task/template.py

@@ -0,0 +1,116 @@
+"""
+Template module
+"""
+
+from string import Formatter
+
+from ...util import TemplateFormatter
+from .file import Task
+
+
+class TemplateTask(Task):
+    """
+    Task that generates text from a template and task inputs. Templates can be used to prepare data for a number of tasks
+    including generating large language model (LLM) prompts.
+    """
+
+    def register(self, template=None, rules=None, strict=True):
+        """
+        Read template parameters.
+
+        Args:
+            template: prompt template
+            rules: parameter rules
+            strict: requires all task inputs to be consumed by template, defaults to True
+        """
+
+        # pylint: disable=W0201
+        # Template text
+        self.template = template if template else self.defaulttemplate()
+
+        # Template processing rules
+        self.rules = rules if rules else self.defaultrules()
+
+        # Create formatter
+        self.formatter = TemplateFormatter() if strict else Formatter()
+
+    def prepare(self, element):
+        # Check if element matches any processing rules
+        match = self.match(element)
+        if match:
+            return match
+
+        # Apply template processing, if applicable
+        if self.template:
+            # Pass dictionary as named prompt template parameters
+            if isinstance(element, dict):
+                return self.formatter.format(self.template, **element)
+
+            # Pass tuple as prompt template parameters (arg0 - argN)
+            if isinstance(element, tuple):
+                return self.formatter.format(self.template, **{f"arg{i}": x for i, x in enumerate(element)})
+
+            # Default behavior is to use input as {text} parameter in prompt template
+            return self.formatter.format(self.template, text=element)
+
+        # Return original inputs when no prompt provided
+        return element
+
+    def defaulttemplate(self):
+        """
+        Generates a default template for this task. Base method returns None.
+
+        Returns:
+            default template
+        """
+
+        return None
+
+    def defaultrules(self):
+        """
+        Generates a default rules for this task. Base method returns an empty dictionary.
+
+        Returns:
+            default rules
+        """
+
+        return {}
+
+    def match(self, element):
+        """
+        Check if element matches any processing rules.
+
+        Args:
+            element: input element
+
+        Returns:
+            matching value if found, None otherwise
+        """
+
+        if self.rules and isinstance(element, dict):
+            # Check if any rules are matched
+            for key, value in self.rules.items():
+                if element[key] == value:
+                    return element[key]
+
+        return None
+
+
+class RagTask(TemplateTask):
+    """
+    Template task that prepares input for a rag pipeline.
+    """
+
+    def prepare(self, element):
+        # Apply prompt template using all variables except "query" and use output as question
+        if isinstance(element, dict):
+            # Make a copy without query and run through template
+            params = dict(element)
+            params.pop("query", None)
+            params["text"] = params.pop("question")
+
+            element["question"] = super().prepare(params)
+            return element
+
+        # Default mode is to use element text for both query and question
+        return {"query": element, "question": super().prepare(element)}

txtai/workflow/task/url.py

@@ -0,0 +1,20 @@
+"""
+UrlTask module
+"""
+
+import re
+
+from .base import Task
+
+
+class UrlTask(Task):
+    """
+    Task that processes urls
+    """
+
+    # URL prefix
+    PREFIX = r"\w+:\/\/"
+
+    def accept(self, element):
+        # Only accept elements that start with a url prefix
+        return super().accept(element) and re.match(UrlTask.PREFIX, element.lower())

txtai/workflow/task/workflow.py

@@ -0,0 +1,14 @@
+"""
+WorkflowTask module
+"""
+
+from .base import Task
+
+
+class WorkflowTask(Task):
+    """
+    Task that executes a separate Workflow
+    """
+
+    def process(self, action, inputs):
+        return list(super().process(action, inputs))