epub-translator 0.1.0__py3-none-any.whl → 0.1.3__py3-none-any.whl

epub_translator/llm/core.py

@@ -1,7 +1,5 @@
 import datetime
-import hashlib
-import json
-from collections.abc import Callable, Generator
+from collections.abc import Generator
 from importlib.resources import files
 from logging import DEBUG, FileHandler, Formatter, Logger, getLogger
 from os import PathLike
@@ -11,9 +9,10 @@ from jinja2 import Environment, Template
 from tiktoken import Encoding, get_encoding
 
 from ..template import create_env
+from .context import LLMContext
 from .executor import LLMExecutor
 from .increasable import Increasable
-from .types import Message, MessageRole, R
+from .types import Message
 
 
 class LLM:
@@ -23,42 +22,28 @@ class LLM:
         url: str,
         model: str,
         token_encoding: str,
-        cache_path: PathLike | None = None,
         timeout: float | None = None,
         top_p: float | tuple[float, float] | None = None,
         temperature: float | tuple[float, float] | None = None,
         retry_times: int = 5,
         retry_interval_seconds: float = 6.0,
-        log_dir_path: PathLike | None = None,
-    ):
+        cache_path: PathLike | str | None = None,
+        log_dir_path: PathLike | str | None = None,
+    ) -> None:
         prompts_path = Path(str(files("epub_translator"))) / "data"
         self._templates: dict[str, Template] = {}
         self._encoding: Encoding = get_encoding(token_encoding)
         self._env: Environment = create_env(prompts_path)
-        self._logger_save_path: Path | None = None
-        self._cache_path: Path | None = None
-
-        if cache_path is not None:
-            self._cache_path = Path(cache_path)
-            if not self._cache_path.exists():
-                self._cache_path.mkdir(parents=True, exist_ok=True)
-            elif not self._cache_path.is_dir():
-                self._cache_path = None
-
-        if log_dir_path is not None:
-            self._logger_save_path = Path(log_dir_path)
-            if not self._logger_save_path.exists():
-                self._logger_save_path.mkdir(parents=True, exist_ok=True)
-            elif not self._logger_save_path.is_dir():
-                self._logger_save_path = None
+        self._top_p: Increasable = Increasable(top_p)
+        self._temperature: Increasable = Increasable(temperature)
+        self._cache_path: Path | None = self._ensure_dir_path(cache_path)
+        self._logger_save_path: Path | None = self._ensure_dir_path(log_dir_path)
 
         self._executor = LLMExecutor(
             url=url,
             model=model,
             api_key=key,
             timeout=timeout,
-            top_p=Increasable(top_p),
-            temperature=Increasable(temperature),
             retry_times=retry_times,
             retry_interval_seconds=retry_interval_seconds,
             create_logger=self._create_logger,
@@ -68,41 +53,29 @@ class LLM:
     def encoding(self) -> Encoding:
         return self._encoding
 
+    def context(self, cache_seed_content: str | None = None) -> LLMContext:
+        return LLMContext(
+            executor=self._executor,
+            cache_path=self._cache_path,
+            cache_seed_content=cache_seed_content,
+            top_p=self._top_p,
+            temperature=self._temperature,
+        )
+
     def request(
         self,
         input: str | list[Message],
-        parser: Callable[[str], R] = lambda x: x,
         max_tokens: int | None = None,
-    ) -> R:
-        messages: list[Message]
-        if isinstance(input, str):
-            messages = [Message(role=MessageRole.USER, message=input)]
-        else:
-            messages = input
-
-        # Check cache if cache_path is set
-        if self._cache_path is not None:
-            cache_key = self._compute_messages_hash(messages)
-            cache_file = self._cache_path / f"{cache_key}.txt"
-
-            if cache_file.exists():
-                cached_content = cache_file.read_text(encoding="utf-8")
-                return parser(cached_content)
-
-        # Make the actual request
-        response = self._executor.request(
-            messages=messages,
-            parser=lambda x: x,
-            max_tokens=max_tokens,
-        )
-
-        # Save to cache if cache_path is set
-        if self._cache_path is not None:
-            cache_key = self._compute_messages_hash(messages)
-            cache_file = self._cache_path / f"{cache_key}.txt"
-            cache_file.write_text(response, encoding="utf-8")
-
-        return parser(response)
+        temperature: float | None = None,
+        top_p: float | None = None,
+    ) -> str:
+        with self.context() as ctx:
+            return ctx.request(
+                input=input,
+                max_tokens=max_tokens,
+                temperature=temperature,
+                top_p=top_p,
+            )
 
     def template(self, template_name: str) -> Template:
         template = self._templates.get(template_name, None)
@@ -111,17 +84,21 @@ class LLM:
             self._templates[template_name] = template
         return template
 
-    def _compute_messages_hash(self, messages: list[Message]) -> str:
-        """Compute SHA-512 hash of m·essages for cache key."""
-        messages_dict = [{"role": msg.role.value, "message": msg.message} for msg in messages]
-        messages_json = json.dumps(messages_dict, ensure_ascii=False, sort_keys=True)
-        return hashlib.sha512(messages_json.encode("utf-8")).hexdigest()
+    def _ensure_dir_path(self, path: PathLike | str | None) -> Path | None:
+        if path is None:
+            return None
+        dir_path = Path(path)
+        if not dir_path.exists():
+            dir_path.mkdir(parents=True, exist_ok=True)
+        elif not dir_path.is_dir():
+            return None
+        return dir_path.resolve()
 
     def _create_logger(self) -> Logger | None:
         if self._logger_save_path is None:
             return None
 
-        now = datetime.datetime.now(datetime.timezone.utc)
+        now = datetime.datetime.now(datetime.UTC)
         timestamp = now.strftime("%Y-%m-%d %H-%M-%S %f")
         file_path = self._logger_save_path / f"request {timestamp}.log"
         logger = getLogger(f"LLM Request {timestamp}")

epub_translator/llm/executor.py

@@ -2,14 +2,12 @@ from collections.abc import Callable
 from io import StringIO
 from logging import Logger
 from time import sleep
-from typing import cast
 
 from openai import OpenAI
 from openai.types.chat import ChatCompletionMessageParam
 
 from .error import is_retry_error
-from .increasable import Increasable, Increaser
-from .types import Message, MessageRole, R
+from .types import Message, MessageRole
 
 
 class LLMExecutor:
@@ -19,16 +17,12 @@ class LLMExecutor:
         url: str,
         model: str,
         timeout: float | None,
-        top_p: Increasable,
-        temperature: Increasable,
         retry_times: int,
         retry_interval_seconds: float,
         create_logger: Callable[[], Logger | None],
     ) -> None:
         self._model_name: str = model
         self._timeout: float | None = timeout
-        self._top_p: Increasable = top_p
-        self._temperature: Increasable = temperature
         self._retry_times: int = retry_times
         self._retry_interval_seconds: float = retry_interval_seconds
         self._create_logger: Callable[[], Logger | None] = create_logger
@@ -38,15 +32,29 @@ class LLMExecutor:
             timeout=timeout,
         )
 
-    def request(self, messages: list[Message], parser: Callable[[str], R], max_tokens: int | None) -> R:
-        result: R | None = None
+    def request(
+        self,
+        messages: list[Message],
+        max_tokens: int | None,
+        temperature: float | None,
+        top_p: float | None,
+        cache_key: str | None,
+    ) -> str:
+        response: str = ""
         last_error: Exception | None = None
         did_success = False
-        top_p: Increaser = self._top_p.context()
-        temperature: Increaser = self._temperature.context()
         logger = self._create_logger()
 
         if logger is not None:
+            parameters: list[str] = [
+                f"\t\ntemperature={temperature}",
+                f"\t\ntop_p={top_p}",
+                f"\t\nmax_tokens={max_tokens}",
+            ]
+            if cache_key is not None:
+                parameters.append(f"\t\ncache_key={cache_key}")
+
+            logger.debug(f"[[Parameters]]:{''.join(parameters)}\n")
             logger.debug(f"[[Request]]:\n{self._input2str(messages)}\n")
 
         try:
@@ -54,8 +62,8 @@ class LLMExecutor:
                 try:
                     response = self._invoke_model(
                         input_messages=messages,
-                        top_p=top_p.current,
-                        temperature=temperature.current,
+                        temperature=temperature,
+                        top_p=top_p,
                         max_tokens=max_tokens,
                     )
                     if logger is not None:
@@ -71,22 +79,8 @@ class LLMExecutor:
                         sleep(self._retry_interval_seconds)
                     continue
 
-                try:
-                    result = parser(response)
-                    did_success = True
-                    break
-
-                except Exception as err:
-                    last_error = err
-                    warn_message = f"request failed with parsing error, retrying... ({i + 1} times)"
-                    if logger is not None:
-                        logger.warning(warn_message)
-                    print(warn_message)
-                    top_p.increase()
-                    temperature.increase()
-                    if self._retry_interval_seconds > 0.0 and i < self._retry_times:
-                        sleep(self._retry_interval_seconds)
-                    continue
+                did_success = True
+                break
 
         except KeyboardInterrupt as err:
             if last_error is not None and logger is not None:
@@ -99,7 +93,7 @@ class LLMExecutor:
             else:
                 raise last_error
 
-        return cast(R, result)
+        return response
 
     def _input2str(self, input: str | list[Message]) -> str:
         if isinstance(input, str):
@@ -133,7 +127,7 @@ class LLMExecutor:
         top_p: float | None,
         temperature: float | None,
         max_tokens: int | None,
-    ):
+    ) -> str:
         messages: list[ChatCompletionMessageParam] = []
         for item in input_messages:
             if item.role == MessageRole.SYSTEM:

epub_translator/llm/increasable.py

@@ -21,7 +21,7 @@ class Increasable:
             param = float(param)
         if isinstance(param, float):
             param = (param, param)
-        if isinstance(param, tuple):
+        if isinstance(param, (tuple, list)):
             if len(param) != 2:
                 raise ValueError(f"Expected a tuple of length 2, got {len(param)}")
             begin, end = param

epub_translator/llm/types.py

@@ -1,8 +1,5 @@
 from dataclasses import dataclass
 from enum import Enum, auto
-from typing import TypeVar
-
-R = TypeVar("R")
 
 
 @dataclass

epub_translator/segment/__init__.py

@@ -0,0 +1,26 @@
+from .block_segment import (
+    BlockContentError,
+    BlockError,
+    BlockExpectedIDsError,
+    BlockSegment,
+    BlockSubmitter,
+    BlockUnexpectedIDError,
+    BlockWrongTagError,
+)
+from .common import FoundInvalidIDError
+from .inline_segment import (
+    InlineError,
+    InlineExpectedIDsError,
+    InlineLostIDError,
+    InlineSegment,
+    InlineUnexpectedIDError,
+    InlineWrongTagCountError,
+    search_inline_segments,
+)
+from .text_segment import (
+    TextPosition,
+    TextSegment,
+    combine_text_segments,
+    incision_between,
+    search_text_segments,
+)

epub_translator/segment/block_segment.py

@@ -0,0 +1,124 @@
+from collections.abc import Generator
+from dataclasses import dataclass
+from typing import cast
+from xml.etree.ElementTree import Element
+
+from .common import FoundInvalidIDError, validate_id_in_element
+from .inline_segment import InlineError, InlineSegment
+from .text_segment import TextSegment
+from .utils import IDGenerator, id_in_element
+
+
+@dataclass
+class BlockSubmitter:
+    id: int
+    origin_text_segments: list[TextSegment]
+    submitted_element: Element
+
+
+@dataclass
+class BlockWrongTagError:
+    block: tuple[int, Element] | None  # (block_id, block_element) | None 表示根元素
+    expected_tag: str
+    instead_tag: str
+
+
+@dataclass
+class BlockUnexpectedIDError:
+    id: int
+    element: Element
+
+
+@dataclass
+class BlockExpectedIDsError:
+    id2element: dict[int, Element]
+
+
+@dataclass
+class BlockContentError:
+    id: int
+    element: Element
+    errors: list[InlineError | FoundInvalidIDError]
+
+
+BlockError = BlockWrongTagError | BlockUnexpectedIDError | BlockExpectedIDsError | BlockContentError
+
+
+class BlockSegment:
+    def __init__(self, root_tag: str, inline_segments: list[InlineSegment]) -> None:
+        id_generator = IDGenerator()
+        for inline_segment in inline_segments:
+            inline_segment.id = id_generator.next_id()
+            inline_segment.recreate_ids(id_generator)
+
+        self._root_tag: str = root_tag
+        self._inline_segments: list[InlineSegment] = inline_segments
+        self._id2inline_segment: dict[int, InlineSegment] = dict((cast(int, s.id), s) for s in self._inline_segments)
+
+    def __iter__(self) -> Generator[InlineSegment, None, None]:
+        yield from self._inline_segments
+
+    def create_element(self) -> Element:
+        root_element = Element(self._root_tag)
+        for inline_segment in self._inline_segments:
+            root_element.append(inline_segment.create_element())
+        return root_element
+
+    def validate(self, validated_element: Element) -> Generator[BlockError | FoundInvalidIDError, None, None]:
+        if validated_element.tag != self._root_tag:
+            yield BlockWrongTagError(
+                block=None,
+                expected_tag=self._root_tag,
+                instead_tag=validated_element.tag,
+            )
+
+        remain_expected_elements: dict[int, Element] = dict(
+            (id, inline_segment.parent) for id, inline_segment in self._id2inline_segment.items()
+        )
+        for child_validated_element in validated_element:
+            element_id = validate_id_in_element(child_validated_element)
+            if isinstance(element_id, FoundInvalidIDError):
+                yield element_id
+            else:
+                inline_segment = self._id2inline_segment.get(element_id, None)
+                if inline_segment is None:
+                    yield BlockUnexpectedIDError(
+                        id=element_id,
+                        element=child_validated_element,
+                    )
+                else:
+                    if inline_segment.parent.tag != child_validated_element.tag:
+                        yield BlockWrongTagError(
+                            block=(cast(int, inline_segment.id), inline_segment.parent),
+                            expected_tag=inline_segment.parent.tag,
+                            instead_tag=child_validated_element.tag,
+                        )
+
+                    remain_expected_elements.pop(element_id, None)
+                    inline_errors = list(inline_segment.validate(child_validated_element))
+
+                    if inline_errors:
+                        yield BlockContentError(
+                            id=element_id,
+                            element=child_validated_element,
+                            errors=inline_errors,
+                        )
+
+        if remain_expected_elements:
+            yield BlockExpectedIDsError(id2element=remain_expected_elements)
+
+    def submit(self, target: Element) -> Generator[BlockSubmitter, None, None]:
+        for child_element in target:
+            element_id = id_in_element(child_element)
+            if element_id is None:
+                continue
+            inline_segment = self._id2inline_segment.get(element_id, None)
+            if inline_segment is None:
+                continue
+            inline_segment_id = inline_segment.id
+            assert inline_segment_id is not None
+            yield BlockSubmitter(
+                id=inline_segment_id,
+                origin_text_segments=list(inline_segment),
+                submitted_element=inline_segment.assign_attributes(child_element),
+            )

epub_translator/segment/common.py

@@ -0,0 +1,29 @@
+from dataclasses import dataclass
+from xml.etree.ElementTree import Element
+
+from ..xml import ID_KEY
+
+
+@dataclass
+class FoundInvalidIDError(Exception):
+    invalid_id: str | None
+    element: Element
+
+
+def validate_id_in_element(element: Element, enable_no_id: bool = False) -> int | FoundInvalidIDError:
+    id_str = element.get(ID_KEY, None)
+    if id_str is None:
+        if enable_no_id:
+            return -1
+        else:
+            return FoundInvalidIDError(
+                invalid_id=None,
+                element=element,
+            )
+    try:
+        return int(id_str)
+    except ValueError:
+        return FoundInvalidIDError(
+            invalid_id=id_str,
+            element=element,
+        )