refactoring

Author: GreenWizard2015

core/CAIAssistant.py

@@ -44,39 +44,49 @@ def _extractParts(self, text):
     tmp = [x for x in tmp if (len(x) == 2) and (len(x[0]) > 0) and (len(x[1]) > 0)]
     return {k: v for k, v in tmp}
 
-  def translate(self, text, fastTranslation, language):
-    # run shallow translation
-    res = self._translateShallow.run({
-      'UserInput': text,
-      'FastTranslation': fastTranslation,
-      'Language': language
-    })
+  def _executePrompt(self, prompt, variables):
+    res = prompt.run(variables)
     logging.info(res)
-    # extract translation
-    translation = self._extractParts(res)
-    logging.info(json.dumps(translation, indent=2))
+    res = self._extractParts(res)
+    logging.info(json.dumps(res, indent=2))
     flags = {
       k: v.lower() == 'yes'
-      for k, v in translation.items()
+      for k, v in res.items()
       if v.lower() in ['yes', 'no']
     }
+    # remove flags from result
+    res = {k: v for k, v in res.items() if k not in flags}
+    # add flags as separate variable
+    res['Flags'] = flags
+    return res
+  
+  def translate(self, text, fastTranslation, language):
+    # run shallow translation
+    res = self._executePrompt(
+      self._translateShallow,
+      {
+        'UserInput': text,
+        'FastTranslation': fastTranslation,
+        'Language': language
+      }
+    )
+    flags = res['Flags']
     totalIssues = sum([int(v) for k, v in flags.items()])
     if totalIssues < 2:
-      yield translation['Translation']
+      yield res['Translation']
       return # all ok, no need to run deep translation
-    yield translation['Translation'] + '\n\n\n' + translation.get('Notification', '')
+    yield res['Translation'] + '\n\n\n' + res.get('Notification', '')
 
     # run deep translation
-    res = self._translateDeep.run({
-      'UserInput': text,
-      'FastTranslation': translation['Translation'], # use shallow translation as reference
-      'Language': language,
-      'InputLanguage': translation.get('Input language', 'unknown'),
-      'Flags': ', '.join([k for k, v in flags.items() if v])
-    })
-    logging.info(res)
-    # extract final translation
-    translation = self._extractParts(res)
-    logging.info(json.dumps(translation, indent=2))
-    yield translation['Translation']
+    res = self._executePrompt(
+      self._translateDeep,
+      {
+        'UserInput': text,
+        'FastTranslation': res['Translation'], # use shallow translation as reference
+        'Language': language,
+        'InputLanguage': res.get('Input language', 'unknown'),
+        'Flags': ', '.join([k for k, v in flags.items() if v])
+      }
+    )
+    yield res['Translation']
     return