"Refine" button

Author: GreenWizard2015

core/CAIAssistant.py

@@ -5,6 +5,13 @@
 from langchain.prompts import PromptTemplate
 from langchain.chains import LLMChain
 
+from collections import namedtuple
+
+CAITranslationResult = namedtuple(
+  'CAITranslationResult',
+  ['translation', 'pending', 'text', 'language', 'InputLanguage', 'Flags']
+)
+
 class CAIAssistant:
   def __init__(self, promptsFolder=None, openai_api_key=None):
     if promptsFolder is None:
@@ -56,6 +63,8 @@ def _extractParts(self, text):
     return {k: v for k, v in tmp}
 
   def _executePrompt(self, prompt, variables):
+    if not self._connected: raise Exception('Not connected to API')
+
     rawPrompt = prompt.prompt.format_prompt(**variables).to_string()
     logging.info('Raw prompt: ' + rawPrompt)
     res = prompt.run(variables)
@@ -106,18 +115,35 @@ def _translateDeep(self, text, translation, language, inputLanguage, flags):
     return res['Translation']
 
   def translate(self, text, fastTranslation, language):
-    if not self._connected:
-      raise Exception('Not connected to API')
     # run shallow translation
     raw, translation, done = self._translateShallow(
       text=text, translation=fastTranslation, language=language
     )
-    yield translation, not done
-    if done: return
-    # run deep translation
-    yield self._translateDeep(
-      text=text, translation=translation, language=language,
-      inputLanguage=raw.get('Input language', 'unknown'),
-      flags=raw['Flags']
-    ), False # no more pending translations
-    return
+    translationResult = CAITranslationResult(
+      translation=translation, pending=not done,
+      text=text, language=language,
+      InputLanguage=raw.get('Input language', 'unknown'),
+      Flags=raw['Flags'],
+    )
+    yield translationResult
+    if not done: # run deep translation
+      yield self.refine(translationResult)
+    return
+  
+  def refine(self, previousTranslation: CAITranslationResult):
+    res = self._translateDeep(
+      text=previousTranslation.text,
+      translation=previousTranslation.translation,
+      language=previousTranslation.language,
+      inputLanguage=previousTranslation.InputLanguage,
+      flags=previousTranslation.Flags,
+    )
+
+    translationResult = CAITranslationResult(
+      translation=res,
+      text=previousTranslation.text, language=previousTranslation.language,
+      InputLanguage=previousTranslation.InputLanguage,
+      Flags=previousTranslation.Flags,
+      pending=False, # no more steps
+    )
+    return translationResult

core/worker.py

@@ -11,6 +11,9 @@ def __init__(self, events):
     self._forceTranslateEvent = threading.Event()
     self._translatorFast = Translator(service_urls=['translate.google.com'])
     self._assistant = CAIAssistant()
+
+    # TODO: Fix this hack. It's not thread safe. Also check related issues with UI during refinement.
+    self._refinement = None
     return
 
   def run(self):
@@ -22,6 +25,10 @@ def run(self):
       isForceTranslate = self._forceTranslateEvent.wait(5)
       self._forceTranslateEvent.clear()
 
+      if self._refinement is not None:
+        self._doRefine()
+        continue
+
       userInput = self._events.userInput()
       text = userInput['text']
       language = userInput['language']
@@ -63,22 +70,23 @@ def _performTranslate(self, text, force, language):
     text = text.strip()
     try:
       self._events.startTranslate(force)
-      fastText = self._fastTranslate(text, language=language)
+      fastText = self._fastTranslate(text, languageCode=language['code'])
       self._events.fastTranslated(fastText)
       if not force: return
 
       translationProcess = self._fullTranslate(text, fastTranslation=fastText, language=language)
-      for fullText, hasMore in translationProcess:
-        self._events.fullTranslated(fullText, pending=hasMore)
+      for translationResult in translationProcess:
+        self._events.fullTranslated(translationResult)
         if self._forceTranslateEvent.is_set(): break # stop if force another translate
         continue
     finally:
       self._events.endTranslate()
     return
 
-  def _fastTranslate(self, text, language):
+  @lru_cache(maxsize=20)
+  def _fastTranslate(self, text, languageCode):
     if 0 == len(text): return ""
-    translated = self._translatorFast.translate(text, dest=language['code'])
+    translated = self._translatorFast.translate(text, dest=languageCode)
     return translated.text
 
   def _fullTranslate(self, text, fastTranslation, language):
@@ -90,9 +98,7 @@ def _fullTranslate(self, text, fastTranslation, language):
       text, language=language['name'],
       fastTranslation=fastTranslation,
     )
-    for translation in translationProcess:
-      yield translation
-      continue
+    for translation in translationProcess: yield translation
     return
 
   @lru_cache(maxsize=None)
@@ -108,4 +114,19 @@ def _updateLocalization(self, languageName, languageCode):
   def bindAPI(self, key):
     self._assistant.bindAPI(key)
     return
-  
+  
+  def refine(self, previousTranslation):
+    self._refinement = previousTranslation
+    self.forceTranslate()
+    return
+  
+  def _doRefine(self):
+    assert self._refinement is not None
+    previousTranslation = self._refinement
+    self._refinement = None
+    try:
+      res = self._assistant.refine(previousTranslation)
+      self._events.fullTranslated(res)
+    except Exception as e:
+      self._events.error(e)
+    return

main.py

@@ -9,12 +9,12 @@
 import dotenv
 
 # main app
-# TODO: figure out how to enforce deep translation via UI
 # TODO: add translation history
 # TODO: add simple translation diff
 class App(tk.Frame):
   def __init__(self, master, languages, configs):
     super().__init__(master)
+    self._lastAIResult = None
     self._configs = configs
     self._localizationMap = {}
     # predefine messages
@@ -91,6 +91,13 @@ def _UI_fullTranslation(self, owner):
       textvariable=self._localization("Slow and improved translation (ChatGPT/AI):")
     )
     label.pack(side="top", fill=tk.X)
+    # Button "Refine" to force deep translation, disabled by default
+    self._refineBtn = btn = tk.Button(
+      owner, state=tk.DISABLED,
+      command=self.onRefine,
+      textvariable=self._localization("Refine"),
+    )
+    btn.pack(side="bottom", padx=5, pady=5, anchor="e")
 
     self._fullOutputText = tkst.ScrolledText(owner, wrap=tk.WORD)
     self._fullOutputText.pack(side="top", fill=tk.BOTH, expand=tk.YES)
@@ -162,11 +169,14 @@ def fastTranslated(self, text):
     self._fastOutputText.insert(tk.END, text)
     return
 
-  def fullTranslated(self, text, pending):
+  def fullTranslated(self, translationResult):
+    self._lastAIResult = translationResult
+    self._refineBtn.config(state=tk.DISABLED if translationResult.pending else tk.NORMAL)
+
     self._fullOutputText.delete("1.0", tk.END)
-    self._fullOutputText.insert(tk.END, text)
+    self._fullOutputText.insert(tk.END, translationResult.translation)
 
-    if pending:
+    if translationResult.pending:
       notification = self._localization(
         "Translation is not accurate and will be updated soon."
       ).get()
@@ -185,7 +195,10 @@ def onSelectLanguage(self, event):
     language = self._language.get()
     code = next((code for code, name in self._languages.items() if name == language), None)
     if code is None: return
- 
+    # discard AI result and disable refine button
+    self._lastAIResult = None
+    self._refineBtn.config(state=tk.DISABLED)
+    # update other stuff
     self._currentLanguage = code
     self._configs['language'] = code
     self._worker.forceTranslate() # hack to force translation
@@ -209,6 +222,14 @@ def onSwitchAPIKey(self):
     return
 
   def configs(self): return self._configs
+
+  def onRefine(self, event=None):
+    # check if refine button is enabled
+    if tk.DISABLED == self._refineBtn['state']: return 'break' # prevent unwanted action
+
+    self._worker.refine(self._lastAIResult)
+    self._refineBtn.config(state=tk.DISABLED)
+    return 'break'
 # End of class
 
 def main():