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

epub_translator/llm/core.py

@@ -0,0 +1,175 @@
+import datetime
+import hashlib
+import json
+from collections.abc import Callable, Generator
+from importlib.resources import files
+from logging import DEBUG, FileHandler, Formatter, Logger, getLogger
+from os import PathLike
+from pathlib import Path
+
+from jinja2 import Environment, Template
+from tiktoken import Encoding, get_encoding
+
+from ..template import create_env
+from .executor import LLMExecutor
+from .increasable import Increasable
+from .types import Message, MessageRole, R
+
+
+class LLM:
+    def __init__(
+        self,
+        key: str,
+        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,
+    ):
+        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._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,
+        )
+
+    @property
+    def encoding(self) -> Encoding:
+        return self._encoding
+
+    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)
+
+    def template(self, template_name: str) -> Template:
+        template = self._templates.get(template_name, None)
+        if template is None:
+            template = self._env.get_template(template_name)
+            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 _create_logger(self) -> Logger | None:
+        if self._logger_save_path is None:
+            return None
+
+        now = datetime.datetime.now(datetime.timezone.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}")
+        logger.setLevel(DEBUG)
+        handler = FileHandler(file_path, encoding="utf-8")
+        handler.setLevel(DEBUG)
+        handler.setFormatter(Formatter("%(asctime)s    %(message)s", "%H:%M:%S"))
+        logger.addHandler(handler)
+
+        return logger
+
+    def _search_quotes(self, kind: str, response: str) -> Generator[str, None, None]:
+        start_marker = f"```{kind}"
+        end_marker = "```"
+        start_index = 0
+
+        while True:
+            start_index = self._find_ignore_case(
+                raw=response,
+                sub=start_marker,
+                start=start_index,
+            )
+            if start_index == -1:
+                break
+
+            end_index = self._find_ignore_case(
+                raw=response,
+                sub=end_marker,
+                start=start_index + len(start_marker),
+            )
+            if end_index == -1:
+                break
+
+            extracted_text = response[start_index + len(start_marker) : end_index].strip()
+            yield extracted_text
+            start_index = end_index + len(end_marker)
+
+    def _find_ignore_case(self, raw: str, sub: str, start: int = 0):
+        if not sub:
+            return 0 if 0 >= start else -1
+
+        raw_len, sub_len = len(raw), len(sub)
+        for i in range(start, raw_len - sub_len + 1):
+            match = True
+            for j in range(sub_len):
+                if raw[i + j].lower() != sub[j].lower():
+                    match = False
+                    break
+            if match:
+                return i
+        return -1

epub_translator/llm/error.py

@@ -1,49 +1,52 @@
-import openai
 import httpx
+import openai
 import requests
 
 
 def is_retry_error(err: Exception) -> bool:
-  if _is_openai_retry_error(err):
-    return True
-  if _is_httpx_retry_error(err):
-    return True
-  if _is_request_retry_error(err):
-    return True
-  return False
+    if _is_openai_retry_error(err):
+        return True
+    if _is_httpx_retry_error(err):
+        return True
+    if _is_request_retry_error(err):
+        return True
+    return False
+
 
 # https://help.openai.com/en/articles/6897213-openai-library-error-types-guidance
 def _is_openai_retry_error(err: Exception) -> bool:
-  if isinstance(err, openai.Timeout):
-    return True
-  if isinstance(err, openai.APIConnectionError):
-    return True
-  if isinstance(err, openai.InternalServerError):
-    return err.status_code in (502, 503, 504)
-  return False
+    if isinstance(err, openai.Timeout):
+        return True
+    if isinstance(err, openai.APIConnectionError):
+        return True
+    if isinstance(err, openai.InternalServerError):
+        return err.status_code in (502, 503, 504)
+    return False
+
 
 # https://www.python-httpx.org/exceptions/
 def _is_httpx_retry_error(err: Exception) -> bool:
-  if isinstance(err, httpx.RemoteProtocolError):
-    return True
-  if isinstance(err, httpx.StreamError):
-    return True
-  if isinstance(err, httpx.TimeoutException):
-    return True
-  if isinstance(err, httpx.NetworkError):
-    return True
-  if isinstance(err, httpx.ProtocolError):
-    return True
-  return False
+    if isinstance(err, httpx.RemoteProtocolError):
+        return True
+    if isinstance(err, httpx.StreamError):
+        return True
+    if isinstance(err, httpx.TimeoutException):
+        return True
+    if isinstance(err, httpx.NetworkError):
+        return True
+    if isinstance(err, httpx.ProtocolError):
+        return True
+    return False
+
 
 # https://requests.readthedocs.io/en/latest/api/#exceptions
 def _is_request_retry_error(err: Exception) -> bool:
-  if isinstance(err, requests.ConnectionError):
-    return True
-  if isinstance(err, requests.ConnectTimeout):
-    return True
-  if isinstance(err, requests.ReadTimeout):
-    return True
-  if isinstance(err, requests.Timeout):
-    return True
-  return False
+    if isinstance(err, requests.ConnectionError):
+        return True
+    if isinstance(err, requests.ConnectTimeout):
+        return True
+    if isinstance(err, requests.ReadTimeout):
+        return True
+    if isinstance(err, requests.Timeout):
+        return True
+    return False

epub_translator/llm/executor.py

@@ -1,150 +1,173 @@
-from typing import cast, Any, Callable
+from collections.abc import Callable
 from io import StringIO
-from time import sleep
-from pydantic import SecretStr
 from logging import Logger
-from langchain_core.messages import SystemMessage, HumanMessage, AIMessage
-from langchain_core.language_models import LanguageModelInput
-from langchain_openai import ChatOpenAI
+from time import sleep
+from typing import cast
+
+from openai import OpenAI
+from openai.types.chat import ChatCompletionMessageParam
 
-from .increasable import Increasable, Increaser
 from .error import is_retry_error
+from .increasable import Increasable, Increaser
+from .types import Message, MessageRole, R
 
 
 class LLMExecutor:
-  def __init__(
-    self,
-    api_key: SecretStr,
-    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._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
-    self._model = ChatOpenAI(
-      api_key=cast(SecretStr, api_key),
-      base_url=url,
-      model=model,
-      timeout=timeout,
-    )
-
-  def request(self, input: LanguageModelInput, parser: Callable[[str], Any], max_tokens: int | None) -> Any:
-    result: Any | None = None
-    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:
-      logger.debug(f"[[Request]]:\n{self._input2str(input)}\n")
-
-    try:
-      for i in range(self._retry_times + 1):
-        try:
-          response = self._invoke_model(
-            input=input,
-            top_p=top_p.current,
-            temperature=temperature.current,
-            max_tokens=max_tokens,
-          )
-          if logger is not None:
-            logger.debug(f"[[Response]]:\n{response}\n")
+    def __init__(
+        self,
+        api_key: str,
+        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
+        self._client = OpenAI(
+            api_key=api_key,
+            base_url=url,
+            timeout=timeout,
+        )
+
+    def request(self, messages: list[Message], parser: Callable[[str], R], max_tokens: int | None) -> R:
+        result: R | None = None
+        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:
+            logger.debug(f"[[Request]]:\n{self._input2str(messages)}\n")
 
-        except Exception as err:
-          last_error = err
-          if not is_retry_error(err):
+        try:
+            for i in range(self._retry_times + 1):
+                try:
+                    response = self._invoke_model(
+                        input_messages=messages,
+                        top_p=top_p.current,
+                        temperature=temperature.current,
+                        max_tokens=max_tokens,
+                    )
+                    if logger is not None:
+                        logger.debug(f"[[Response]]:\n{response}\n")
+
+                except Exception as err:
+                    last_error = err
+                    if not is_retry_error(err):
+                        raise err
+                    if logger is not None:
+                        logger.warning(f"request failed with connection error, retrying... ({i + 1} times)")
+                    if self._retry_interval_seconds > 0.0 and i < self._retry_times:
+                        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
+
+        except KeyboardInterrupt as err:
+            if last_error is not None and logger is not None:
+                logger.debug(f"[[Error]]:\n{last_error}\n")
             raise err
-          if logger is not None:
-            logger.warning(f"request failed with connection error, retrying... ({i + 1} times)")
-          if self._retry_interval_seconds > 0.0 and \
-            i < self._retry_times:
-            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
-
-    except KeyboardInterrupt as err:
-      if last_error is not None and logger is not None:
-        logger.debug(f"[[Error]]:\n{last_error}\n")
-      raise err
-
-    if not did_success:
-      if last_error is None:
-        raise RuntimeError("Request failed with unknown error")
-      else:
-        raise last_error
-
-    return result
-
-  def _input2str(self, input: LanguageModelInput) -> str:
-    if isinstance(input, str):
-      return input
-    if not isinstance(input, list):
-      raise ValueError(f"Unsupported input type: {type(input)}")
-
-    buffer = StringIO()
-    is_first = True
-    for message in input:
-      if not is_first:
-        buffer.write("\n\n")
-      if isinstance(message, SystemMessage):
-        buffer.write("System:\n")
-        buffer.write(message.content)
-      elif isinstance(message, HumanMessage):
-        buffer.write("User:\n")
-        buffer.write(message.content)
-      elif isinstance(message, AIMessage):
-        buffer.write("Assistant:\n")
-        buffer.write(message.content)
-      else:
-        buffer.write(str(message))
-      is_first = False
-
-    return buffer.getvalue()
-
-  def _invoke_model(
+        if not did_success:
+            if last_error is None:
+                raise RuntimeError("Request failed with unknown error")
+            else:
+                raise last_error
+
+        return cast(R, result)
+
+    def _input2str(self, input: str | list[Message]) -> str:
+        if isinstance(input, str):
+            return input
+        if not isinstance(input, list):
+            raise ValueError(f"Unsupported input type: {type(input)}")
+
+        buffer = StringIO()
+        is_first = True
+        for message in input:
+            if not is_first:
+                buffer.write("\n\n")
+            if message.role == MessageRole.SYSTEM:
+                buffer.write("System:\n")
+                buffer.write(message.message)
+            elif message.role == MessageRole.USER:
+                buffer.write("User:\n")
+                buffer.write(message.message)
+            elif message.role == MessageRole.ASSISTANT:
+                buffer.write("Assistant:\n")
+                buffer.write(message.message)
+            else:
+                buffer.write(str(message))
+            is_first = False
+
+        return buffer.getvalue()
+
+    def _invoke_model(
         self,
-        input: LanguageModelInput,
+        input_messages: list[Message],
         top_p: float | None,
         temperature: float | None,
         max_tokens: int | None,
-      ):
-    stream = self._model.stream(
-      input=input,
-      timeout=self._timeout,
-      top_p=top_p,
-      temperature=temperature,
-      max_tokens=max_tokens,
-    )
-    buffer = StringIO()
-    for chunk in stream:
-      data = str(chunk.content)
-      buffer.write(data)
-    return buffer.getvalue()
+    ):
+        messages: list[ChatCompletionMessageParam] = []
+        for item in input_messages:
+            if item.role == MessageRole.SYSTEM:
+                messages.append(
+                    {
+                        "role": "system",
+                        "content": item.message,
+                    }
+                )
+            elif item.role == MessageRole.USER:
+                messages.append(
+                    {
+                        "role": "user",
+                        "content": item.message,
+                    }
+                )
+            elif item.role == MessageRole.ASSISTANT:
+                messages.append(
+                    {
+                        "role": "assistant",
+                        "content": item.message,
+                    }
+                )
+
+        stream = self._client.chat.completions.create(
+            model=self._model_name,
+            messages=messages,
+            stream=True,
+            top_p=top_p,
+            temperature=temperature,
+            max_tokens=max_tokens,
+        )
+        buffer = StringIO()
+        for chunk in stream:
+            if chunk.choices and chunk.choices[0].delta.content:
+                buffer.write(chunk.choices[0].delta.content)
+        return buffer.getvalue()

epub_translator/llm/increasable.py

@@ -1,35 +1,35 @@
 class Increaser:
-  def __init__(self, value_range: tuple[float, float] | None):
-    self._value_range: tuple[float, float] | None = value_range
-    self._current: float | None = value_range[0] if value_range is not None else None
+    def __init__(self, value_range: tuple[float, float] | None):
+        self._value_range: tuple[float, float] | None = value_range
+        self._current: float | None = value_range[0] if value_range is not None else None
 
-  @property
-  def current(self) -> float | None:
-    return self._current
+    @property
+    def current(self) -> float | None:
+        return self._current
+
+    def increase(self):
+        if self._value_range is not None and self._current is not None:
+            _, end_value = self._value_range
+            self._current = self._current + 0.5 * (end_value - self._current)
 
-  def increase(self):
-    if self._value_range is None:
-      return
-    _, end_value = self._value_range
-    self._current = self._current + 0.5 * (end_value - self._current)
 
 class Increasable:
-  def __init__(self, param: float | tuple[float, float] | None):
-    self._value_range: tuple[float, float] | None = None
+    def __init__(self, param: float | tuple[float, float] | None):
+        self._value_range: tuple[float, float] | None = None
 
-    if isinstance(param, int):
-      param = float(param)
-    if isinstance(param, float):
-      param = (param, param)
-    if isinstance(param, tuple):
-      if len(param) != 2:
-        raise ValueError(f"Expected a tuple of length 2, got {len(param)}")
-      begin, end = param
-      if isinstance(begin, int):
-        begin = float(begin)
-      if isinstance(end, int):
-        end = float(end)
-      self._value_range = (begin, end)
+        if isinstance(param, int):
+            param = float(param)
+        if isinstance(param, float):
+            param = (param, param)
+        if isinstance(param, tuple):
+            if len(param) != 2:
+                raise ValueError(f"Expected a tuple of length 2, got {len(param)}")
+            begin, end = param
+            if isinstance(begin, int):
+                begin = float(begin)
+            if isinstance(end, int):
+                end = float(end)
+            self._value_range = (begin, end)
 
-  def context(self) -> Increaser:
-    return Increaser(self._value_range)
+    def context(self) -> Increaser:
+        return Increaser(self._value_range)

epub_translator/llm/types.py

@@ -0,0 +1,17 @@
+from dataclasses import dataclass
+from enum import Enum, auto
+from typing import TypeVar
+
+R = TypeVar("R")
+
+
+@dataclass
+class Message:
+    role: "MessageRole"
+    message: str
+
+
+class MessageRole(Enum):
+    SYSTEM = auto()
+    USER = auto()
+    ASSISTANT = auto()

epub_translator/serial/__init__.py

@@ -0,0 +1,2 @@
+from .segment import ST, S, Segment, T
+from .splitter import split