Merge pull request #122 from iautolab/refactor/clean-retry-logic

refactor: clean up _translate_chunk retry logic and add retry streak …

Author: MaleicAcid

openlrc/translate.py

@@ -32,6 +32,7 @@ class LLMTranslator(Translator):
     """
 
     CHUNK_SIZE = 30
+    RETRY_STREAK = 10  # Number of consecutive chunks to use retry model after a failure
 
     def __init__(
         self,
@@ -86,6 +87,11 @@ def make_chunks(texts: list[str], chunk_size: int = 30) -> list[list[tuple[int,
 
         return chunks
 
+    @staticmethod
+    def _is_valid_translation(translated: list[str] | None, expected_len: int) -> bool:
+        """Check whether a translation result is usable (non-empty and correct line count)."""
+        return translated is not None and len(translated) == expected_len
+
     def _translate_chunk(
         self,
         translator_agent: ChunkedTranslatorAgent,
@@ -97,56 +103,65 @@ def _translate_chunk(
         """
         Translate a single chunk of text, with retry mechanism.
 
-        This method attempts to translate the chunk using the primary translator agent.
-        If the translation fails or is inconsistent, it may use a retry agent or remove the glossary.
+        Tries the primary agent first, then optionally falls back to a retry agent.
+        Each agent attempt includes a glossary-removal retry when the line count
+        is inconsistent.
 
-        Args:
-            translator_agent (ChunkedTranslatorAgent): Primary agent for translation.
-            chunk (List[Tuple[int, str]]): Chunk of text to translate.
-            context (TranslationContext): Current translation context.
-            chunk_id (int): ID of the current chunk.
-            retry_agent (Optional[ChunkedTranslatorAgent]): Agent for retry attempts.
-
-        Returns:
-            Tuple[List[str], TranslationContext]: Translated texts and updated context.
+        Returns the best available result. The caller should check whether
+        ``len(translated) == len(chunk)`` to decide if further fallback
+        (e.g. split or atomic translation) is needed.
         """
+        expected = len(chunk)
 
-        def handle_translation(agent: ChunkedTranslatorAgent) -> tuple[list[str] | None, TranslationContext | None]:
+        def _try_agent(agent: ChunkedTranslatorAgent) -> tuple[list[str] | None, TranslationContext | None]:
+            """Single agent attempt: translate, then retry without glossary if line count mismatches."""
             trans: list[str] | None = None
-            updated_context: TranslationContext | None = None
+            ctx: TranslationContext | None = None
             try:
-                trans, updated_context = agent.translate_chunk(chunk_id, chunk, context)
+                trans, ctx = agent.translate_chunk(chunk_id, chunk, context)
             except ChatBotException:
                 logger.error(f"Failed to translate chunk {chunk_id}.")
+                return None, None
 
-            if trans is not None and len(trans) != len(chunk) and agent.info.glossary:
+            if self._is_valid_translation(trans, expected):
+                return trans, ctx
+
+            # Line count mismatch — retry without glossary if applicable.
+            if trans is not None and agent.info.glossary:
                 logger.warning(
                     f"Agent {agent}: Removing glossary for chunk {chunk_id} due to inconsistent translation."
                 )
                 try:
-                    trans, updated_context = agent.translate_chunk(chunk_id, chunk, context, use_glossary=False)
+                    trans, ctx = agent.translate_chunk(chunk_id, chunk, context, use_glossary=False)
                 except ChatBotException:
                     logger.error(f"Failed to translate chunk {chunk_id}.")
 
-            return trans, updated_context
+            return trans, ctx
 
-        translated: list[str] | None
-        updated_ctx: TranslationContext | None
+        translated: list[str] | None = None
+        updated_ctx: TranslationContext | None = None
 
+        # Step 1: Try primary or retry agent based on retry streak.
         if self.use_retry_cnt == 0 or not retry_agent:
-            translated, updated_ctx = handle_translation(translator_agent)
-
-            if retry_agent and (translated is None or len(translated) != len(chunk)):
-                self.use_retry_cnt = 10  # Use retry_agent for the next 10 chunks
-                logger.warning(
-                    f"Using retry agent {retry_agent} for chunk {chunk_id}, and next {self.use_retry_cnt} chunks."
-                )
-                translated, updated_ctx = handle_translation(retry_agent)
+            translated, updated_ctx = _try_agent(translator_agent)
         else:
             logger.info(f"Using retry agent for chunk {chunk_id}, remaining retries: {self.use_retry_cnt}")
-            translated, updated_ctx = handle_translation(retry_agent)
+            translated, updated_ctx = _try_agent(retry_agent)
             self.use_retry_cnt -= 1
 
+        # Step 2: If primary failed and retry agent is available, switch to it.
+        if not self._is_valid_translation(translated, expected) and retry_agent and self.use_retry_cnt == 0:
+            self.use_retry_cnt = self.RETRY_STREAK
+            logger.warning(
+                f"Using retry agent {retry_agent} for chunk {chunk_id}, and next {self.use_retry_cnt} chunks."
+            )
+            translated, updated_ctx = _try_agent(retry_agent)
+
+            # Retry agent also failed — reset streak so next chunk tries primary first.
+            if not self._is_valid_translation(translated, expected):
+                logger.warning(f"Retry agent also failed for chunk {chunk_id}, resetting retry streak.")
+                self.use_retry_cnt = 0
+
         if not translated:
             raise ChatBotException(f"Failed to translate chunk {chunk_id}.")
 
@@ -220,12 +235,8 @@ def translate(
             atomic = False
             translated, context = self._translate_chunk(translator_agent, chunk, context, i, retry_agent=retry_agent)
             chunk_texts = [c[1] for c in chunk]
-            # Proofreader Not fully tested
-            # localized_trans = proofreader.proofread(
-            #     texts=chunk_texts, translations=translated, context=context
-            # )
 
-            if len(translated) != len(chunk):
+            if not self._is_valid_translation(translated, len(chunk)):
                 logger.warning(
                     f"Chunk {i} translation length inconsistent: {len(translated)} vs {len(chunk)},"
                     f" Attempting atomic translation."

tests/test_llm_translator.py

@@ -250,6 +250,50 @@ def test_retry_agent_used_on_primary_failure(self, mock_agent_cls, mock_reviewer
         primary_agent.translate_chunk.assert_called_once()
         retry_agent.translate_chunk.assert_called_once()
 
+    @patch("openlrc.translate.ContextReviewerAgent")
+    @patch("openlrc.translate.ChunkedTranslatorAgent")
+    def test_retry_streak_resets_when_retry_agent_also_fails(self, mock_agent_cls, mock_reviewer_cls):
+        """When retry agent also returns wrong length, use_retry_cnt resets to 0."""
+        texts = ["Hello", "World"]
+
+        primary_agent = MagicMock()
+        primary_agent.cost = 0
+        primary_agent.info.glossary = None
+        primary_agent.translate_chunk.return_value = (
+            ["only_one"],
+            TranslationContext(summary="s", scene="sc", guideline="g"),
+        )
+
+        retry_agent = MagicMock()
+        retry_agent.cost = 0
+        retry_agent.info.glossary = None
+        # Retry also returns wrong length
+        retry_agent.translate_chunk.return_value = (
+            ["still_wrong"],
+            TranslationContext(summary="s", scene="sc", guideline="g"),
+        )
+
+        mock_agent_cls.side_effect = [primary_agent, retry_agent]
+
+        mock_reviewer = mock_reviewer_cls.return_value
+        mock_reviewer.build_context.return_value = "guideline"
+
+        translator = self._make_translator(chunk_size=30, retry_chatbot=_make_mock_chatbot())
+
+        # Mock atomic_translate to provide fallback
+        translator.atomic_translate = MagicMock(return_value=["你好", "世界"])
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            compare_path = Path(tmpdir) / "compare.json"
+            result = translator.translate(texts, "en", "zh", compare_path=compare_path)
+
+        self.assertEqual(result, ["你好", "世界"])
+        # Both agents were tried
+        primary_agent.translate_chunk.assert_called_once()
+        retry_agent.translate_chunk.assert_called_once()
+        # Streak should be reset to 0 after retry agent also failed
+        self.assertEqual(translator.use_retry_cnt, 0)
+
     @patch("openlrc.translate.ContextReviewerAgent")
     @patch("openlrc.translate.ChunkedTranslatorAgent")
     def test_resume_from_compare_file(self, mock_agent_cls, mock_reviewer_cls):