epub-translator 0.1.1__py3-none-any.whl → 0.1.4__py3-none-any.whl

epub_translator/xml_translator/translator.py

@@ -1,82 +1,144 @@
-from collections.abc import Generator, Iterable
-from typing import TypeVar
+from collections.abc import Callable, Generator, Iterable
+from dataclasses import dataclass
+from typing import Generic, TypeVar
 from xml.etree.ElementTree import Element
 
-from ..iter_sync import IterSync
 from ..llm import LLM, Message, MessageRole
-from ..xml import encode_friendly
-from .fill import XMLFill
-from .format import ValidationError, _extract_xml_element
-from .group import XMLGroupContext
-from .progressive_locking import ProgressiveLockingValidator
-from .text_segment import TextSegment
+from ..segment import BlockSegment, InlineSegment, TextSegment
+from ..xml import decode_friendly, encode_friendly
+from .callbacks import Callbacks, FillFailedEvent, warp_callbacks
+from .hill_climbing import HillClimbing
+from .stream_mapper import InlineSegmentMapping, XMLStreamMapper
+from .submitter import SubmitKind, submit
 
 T = TypeVar("T")
 
 
+@dataclass
+class TranslationTask(Generic[T]):
+    element: Element
+    action: SubmitKind
+    payload: T
+
+
 class XMLTranslator:
     def __init__(
         self,
-        llm: LLM,
-        group_context: XMLGroupContext,
+        translation_llm: LLM,
+        fill_llm: LLM,
         target_language: str,
         user_prompt: str | None,
         ignore_translated_error: bool,
         max_retries: int,
         max_fill_displaying_errors: int,
+        max_group_tokens: int,
+        cache_seed_content: str | None = None,
     ) -> None:
-        self._llm: LLM = llm
-        self._group_context: XMLGroupContext = group_context
+        self._translation_llm: LLM = translation_llm
+        self._fill_llm: LLM = fill_llm
         self._target_language: str = target_language
         self._user_prompt: str | None = user_prompt
         self._ignore_translated_error: bool = ignore_translated_error
         self._max_retries: int = max_retries
         self._max_fill_displaying_errors: int = max_fill_displaying_errors
+        self._cache_seed_content: str | None = cache_seed_content
+        self._stream_mapper: XMLStreamMapper = XMLStreamMapper(
+            encoding=translation_llm.encoding,
+            max_group_tokens=max_group_tokens,
+        )
 
-    def translate_to_element(self, element: Element) -> Element:
-        for translated, _, _ in self.translate_to_text_segments(((element, None),)):
+    def translate_element(
+        self,
+        task: TranslationTask[T],
+        interrupt_source_text_segments: Callable[[Iterable[TextSegment]], Iterable[TextSegment]] | None = None,
+        interrupt_translated_text_segments: Callable[[Iterable[TextSegment]], Iterable[TextSegment]] | None = None,
+        interrupt_block_element: Callable[[Element], Element] | None = None,
+        on_fill_failed: Callable[[FillFailedEvent], None] | None = None,
+    ) -> tuple[Element, T]:
+        for translated in self.translate_elements(
+            tasks=((task),),
+            interrupt_source_text_segments=interrupt_source_text_segments,
+            interrupt_translated_text_segments=interrupt_translated_text_segments,
+            interrupt_block_element=interrupt_block_element,
+            on_fill_failed=on_fill_failed,
+        ):
             return translated
+
         raise RuntimeError("Translation failed unexpectedly")
 
-    def translate_to_text_segments(
-        self, items: Iterable[tuple[Element, T]]
-    ) -> Generator[tuple[Element, list[TextSegment], T], None, None]:
-        sync: IterSync[tuple[Element, T]] = IterSync()
-        text_segments: list[TextSegment] = []
+    def translate_elements(
+        self,
+        tasks: Iterable[TranslationTask[T]],
+        interrupt_source_text_segments: Callable[[Iterable[TextSegment]], Iterable[TextSegment]] | None = None,
+        interrupt_translated_text_segments: Callable[[Iterable[TextSegment]], Iterable[TextSegment]] | None = None,
+        interrupt_block_element: Callable[[Element], Element] | None = None,
+        on_fill_failed: Callable[[FillFailedEvent], None] | None = None,
+    ) -> Generator[tuple[Element, T], None, None]:
+        element2task: dict[int, TranslationTask[T]] = {}
+        callbacks = warp_callbacks(
+            interrupt_source_text_segments=interrupt_source_text_segments,
+            interrupt_translated_text_segments=interrupt_translated_text_segments,
+            interrupt_block_element=interrupt_block_element,
+            on_fill_failed=on_fill_failed,
+        )
 
-        for text_segment in self._translate_text_segments(
-            elements=(e for e, _ in sync.iter(items)),
+        def generate_elements():
+            for task in tasks:
+                element2task[id(task.element)] = task
+                yield task.element
+
+        for element, mappings in self._stream_mapper.map_stream(
+            elements=generate_elements(),
+            callbacks=callbacks,
+            map=lambda inline_segments: self._translate_inline_segments(
+                inline_segments=inline_segments,
+                callbacks=callbacks,
+            ),
         ):
-            while True:
-                if sync.tail is None:
-                    break
-                tail_element, _ = sync.tail
-                if id(tail_element) == id(text_segment.root):
-                    break
-                tail_element, payload = sync.take()
-                yield tail_element, text_segments, payload
-                text_segments = []
-            text_segments.append(text_segment)
-
-        while sync.tail is not None:
-            tail_element, payload = sync.take()
-            yield tail_element, text_segments, payload
-            text_segments = []
-
-    def _translate_text_segments(self, elements: Iterable[Element]):
-        for group in self._group_context.split_groups(elements):
-            text_segments = list(group)
-            fill = XMLFill(text_segments)
-            source_text = "".join(self._render_text_segments(text_segments))
-            translated_text = self._translate_text(source_text)
-            self._fill_into_xml(
-                fill=fill,
-                source_text=source_text,
-                translated_text=translated_text,
-            )
-            yield from group.body
+            task = element2task.get(id(element), None)
+            if task:
+                translated_element = submit(
+                    element=element,
+                    action=task.action,
+                    mappings=mappings,
+                )
+                yield translated_element, task.payload
+
+    def _translate_inline_segments(
+        self,
+        inline_segments: list[InlineSegment],
+        callbacks: Callbacks,
+    ) -> list[InlineSegmentMapping | None]:
+        hill_climbing = HillClimbing(
+            encoding=self._fill_llm.encoding,
+            max_fill_displaying_errors=self._max_fill_displaying_errors,
+            block_segment=BlockSegment(
+                root_tag="xml",
+                inline_segments=inline_segments,
+            ),
+        )
+        text_segments = (text for inline in inline_segments for text in inline)
+        source_text = "".join(self._render_text_segments(text_segments))
+        translated_text = self._translate_text(source_text)
+
+        self._request_and_submit(
+            hill_climbing=hill_climbing,
+            source_text=source_text,
+            translated_text=translated_text,
+            callbacks=callbacks,
+        )
+        mappings: list[InlineSegmentMapping | None] = []
+        for mapping in hill_climbing.gen_mappings():
+            if mapping:
+                _, text_segments = mapping
+                if not text_segments:
+                    mapping = None
+            mappings.append(mapping)
+
+        return mappings
 
     def _render_text_segments(self, segments: Iterable[TextSegment]):
+        # TODO: 没必要，直接按照新的 inline segment 组织就行了
         iterator = iter(segments)
         segment = next(iterator, None)
         if segment is None:
@@ -92,88 +154,94 @@ class XMLTranslator:
         yield segment.text
 
     def _translate_text(self, text: str) -> str:
-        return self._llm.request(
-            input=[
-                Message(
-                    role=MessageRole.SYSTEM,
-                    message=self._llm.template("translate").render(
-                        target_language=self._target_language,
-                        user_prompt=self._user_prompt,
+        with self._translation_llm.context(cache_seed_content=self._cache_seed_content) as ctx:
+            return ctx.request(
+                input=[
+                    Message(
+                        role=MessageRole.SYSTEM,
+                        message=self._translation_llm.template("translate").render(
+                            target_language=self._target_language,
+                            user_prompt=self._user_prompt,
+                        ),
                     ),
-                ),
-                Message(role=MessageRole.USER, message=text),
-            ]
-        )
+                    Message(role=MessageRole.USER, message=text),
+                ]
+            )
 
-    def _fill_into_xml(self, fill: XMLFill, source_text: str, translated_text: str) -> Element:
+    def _request_and_submit(
+        self,
+        hill_climbing: HillClimbing,
+        source_text: str,
+        translated_text: str,
+        callbacks: Callbacks,
+    ) -> None:
         user_message = (
             f"Source text:\n{source_text}\n\n"
-            f"XML template:\n```XML\n{encode_friendly(fill.request_element)}\n```\n\n"
+            f"XML template:\n```XML\n{encode_friendly(hill_climbing.request_element())}\n```\n\n"
             f"Translated text:\n{translated_text}"
         )
         fixed_messages: list[Message] = [
             Message(
                 role=MessageRole.SYSTEM,
-                message=self._llm.template("fill").render(),
+                message=self._fill_llm.template("fill").render(),
             ),
             Message(
                 role=MessageRole.USER,
                 message=user_message,
             ),
         ]
-
-        validator = ProgressiveLockingValidator()
         conversation_history: list[Message] = []
-        latest_error: ValidationError | None = None
 
-        with self._llm.context() as llm_context:
-            for _ in range(self._max_retries):
-                # Request LLM response
-                response = llm_context.request(
-                    input=fixed_messages + conversation_history,
-                )
+        with self._fill_llm.context(cache_seed_content=self._cache_seed_content) as llm_context:
+            error_message: str | None = None
 
-                try:
-                    # Extract XML from response
-                    validated_element = _extract_xml_element(response)
+            for retry_count in range(self._max_retries):
+                response = llm_context.request(fixed_messages + conversation_history)
+                validated_element = self._extract_xml_element(response)
+                error_message = None
+                if isinstance(validated_element, str):
+                    error_message = validated_element
+                elif isinstance(validated_element, Element):
+                    error_message = hill_climbing.submit(validated_element)
 
-                    # Validate with progressive locking
-                    is_complete, error_message, newly_locked = validator.validate_with_locking(
-                        template_ele=fill.request_element,
-                        validated_ele=validated_element,
-                        errors_limit=self._max_fill_displaying_errors,
+                if error_message is None:
+                    break
+
+                callbacks.on_fill_failed(
+                    FillFailedEvent(
+                        error_message=error_message,
+                        retried_count=retry_count + 1,
+                        over_maximum_retries=False,
+                    )
+                )
+                conversation_history = [
+                    Message(role=MessageRole.ASSISTANT, message=response),
+                    Message(role=MessageRole.USER, message=error_message),
+                ]
+            if error_message is not None:
+                callbacks.on_fill_failed(
+                    FillFailedEvent(
+                        error_message=error_message,
+                        retried_count=self._max_retries,
+                        over_maximum_retries=True,
                     )
+                )
 
-                    if is_complete:
-                        # All nodes locked, fill successful
-                        fill._fill_submitted_texts(  # pylint: disable=protected-access
-                            generated_ids_stack=[],
-                            element=validated_element,
-                        )
-                        return validated_element
-
-                    # Not complete yet, construct error message with progress info
-                    progress_msg = f"Progress: {len(validator.locked_ids)} nodes locked"
-                    if newly_locked:
-                        progress_msg += f", {len(newly_locked)} newly locked this round"
-
-                    full_error_message = f"{progress_msg}\n\n{error_message}"
-
-                    conversation_history = [
-                        Message(role=MessageRole.ASSISTANT, message=response),
-                        Message(role=MessageRole.USER, message=full_error_message),
-                    ]
-
-                except ValidationError as error:
-                    # XML extraction or basic validation failed
-                    latest_error = error
-                    conversation_history = [
-                        Message(role=MessageRole.ASSISTANT, message=response),
-                        Message(role=MessageRole.USER, message=str(error)),
-                    ]
-
-            message = f"Failed to get valid XML structure after {self._max_retries} attempts"
-            if latest_error is None:
-                raise ValueError(message)
-            else:
-                raise ValueError(message) from latest_error
+    def _extract_xml_element(self, text: str) -> Element | str:
+        first_xml_element: Element | None = None
+        all_xml_elements: int = 0
+
+        for xml_element in decode_friendly(text, tags="xml"):
+            if first_xml_element is None:
+                first_xml_element = xml_element
+            all_xml_elements += 1
+
+        if first_xml_element is None:
+            return "No complete <xml>...</xml> block found. Please ensure you have properly closed the XML with </xml> tag."  # noqa: E501
+
+        if all_xml_elements > 1:
+            return (
+                f"Found {all_xml_elements} <xml>...</xml> blocks. "
+                "Please return only one XML block without any examples or explanations."
+            )
+        return first_xml_element