1  """
  2  Template module
  3  """
  4  
  5  from string import Formatter
  6  
  7  from ...util import TemplateFormatter
  8  from .file import Task
  9  
 10  
 11  class TemplateTask(Task):
 12      """
 13      Task that generates text from a template and task inputs. Templates can be used to prepare data for a number of tasks
 14      including generating large language model (LLM) prompts.
 15      """
 16  
 17      def register(self, template=None, rules=None, strict=True):
 18          """
 19          Read template parameters.
 20  
 21          Args:
 22              template: prompt template
 23              rules: parameter rules
 24              strict: requires all task inputs to be consumed by template, defaults to True
 25          """
 26  
 27          # pylint: disable=W0201
 28          # Template text
 29          self.template = template if template else self.defaulttemplate()
 30  
 31          # Template processing rules
 32          self.rules = rules if rules else self.defaultrules()
 33  
 34          # Create formatter
 35          self.formatter = TemplateFormatter() if strict else Formatter()
 36  
 37      def prepare(self, element):
 38          # Check if element matches any processing rules
 39          match = self.match(element)
 40          if match:
 41              return match
 42  
 43          # Apply template processing, if applicable
 44          if self.template:
 45              # Pass dictionary as named prompt template parameters
 46              if isinstance(element, dict):
 47                  return self.formatter.format(self.template, **element)
 48  
 49              # Pass tuple as prompt template parameters (arg0 - argN)
 50              if isinstance(element, tuple):
 51                  return self.formatter.format(self.template, **{f"arg{i}": x for i, x in enumerate(element)})
 52  
 53              # Default behavior is to use input as {text} parameter in prompt template
 54              return self.formatter.format(self.template, text=element)
 55  
 56          # Return original inputs when no prompt provided
 57          return element
 58  
 59      def defaulttemplate(self):
 60          """
 61          Generates a default template for this task. Base method returns None.
 62  
 63          Returns:
 64              default template
 65          """
 66  
 67          return None
 68  
 69      def defaultrules(self):
 70          """
 71          Generates a default rules for this task. Base method returns an empty dictionary.
 72  
 73          Returns:
 74              default rules
 75          """
 76  
 77          return {}
 78  
 79      def match(self, element):
 80          """
 81          Check if element matches any processing rules.
 82  
 83          Args:
 84              element: input element
 85  
 86          Returns:
 87              matching value if found, None otherwise
 88          """
 89  
 90          if self.rules and isinstance(element, dict):
 91              # Check if any rules are matched
 92              for key, value in self.rules.items():
 93                  if element[key] == value:
 94                      return element[key]
 95  
 96          return None
 97  
 98  
 99  class RagTask(TemplateTask):
100      """
101      Template task that prepares input for a rag pipeline.
102      """
103  
104      def prepare(self, element):
105          # Apply prompt template using all variables except "query" and use output as question
106          if isinstance(element, dict):
107              # Make a copy without query and run through template
108              params = dict(element)
109              params.pop("query", None)
110              params["text"] = params.pop("question")
111  
112              element["question"] = super().prepare(params)
113              return element
114  
115          # Default mode is to use element text for both query and question
116          return {"query": element, "question": super().prepare(element)}