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

epub_translator/llm/core.py

@@ -0,0 +1,233 @@
+import datetime
+import hashlib
+import json
+import uuid
+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 typing import Self
+
+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 LLMContext:
+    """Context manager for LLM requests with transactional caching."""
+
+    def __init__(
+        self,
+        executor: LLMExecutor,
+        cache_path: Path | None,
+    ) -> None:
+        self._executor = executor
+        self._cache_path = cache_path
+        self._context_id = uuid.uuid4().hex[:12]
+        self._temp_files: list[Path] = []
+
+    def __enter__(self) -> Self:
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        if exc_type is None:
+            # Success: commit all temporary cache files
+            self._commit()
+        else:
+            # Failure: rollback (delete) all temporary cache files
+            self._rollback()
+
+    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
+
+        cache_key: str | None = None
+        if self._cache_path is not None:
+            cache_key = self._compute_messages_hash(messages)
+            permanent_cache_file = self._cache_path / f"{cache_key}.txt"
+            if permanent_cache_file.exists():
+                cached_content = permanent_cache_file.read_text(encoding="utf-8")
+                return parser(cached_content)
+
+            temp_cache_file = self._cache_path / f"{cache_key}.{self._context_id}.txt"
+            if temp_cache_file.exists():
+                cached_content = temp_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 temporary cache if cache_path is set
+        if self._cache_path is not None and cache_key is not None:
+            temp_cache_file = self._cache_path / f"{cache_key}.{self._context_id}.txt"
+            temp_cache_file.write_text(response, encoding="utf-8")
+            self._temp_files.append(temp_cache_file)
+
+        return parser(response)
+
+    def _compute_messages_hash(self, messages: list[Message]) -> str:
+        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 _commit(self) -> None:
+        for temp_file in self._temp_files:
+            if temp_file.exists():
+                # Remove the .[context-id].txt suffix to get permanent name
+                permanent_name = temp_file.name.rsplit(".", 2)[0] + ".txt"
+                permanent_file = temp_file.parent / permanent_name
+                temp_file.rename(permanent_file)
+
+    def _rollback(self) -> None:
+        for temp_file in self._temp_files:
+            if temp_file.exists():
+                temp_file.unlink()
+
+
+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,
+    ) -> 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 context(self) -> LLMContext:
+        return LLMContext(
+            executor=self._executor,
+            cache_path=self._cache_path,
+        )
+
+    def request(
+        self,
+        input: str | list[Message],
+        parser: Callable[[str], R] = lambda x: x,
+        max_tokens: int | None = None,
+    ) -> R:
+        with self.context() as ctx:
+            return ctx.request(input=input, parser=parser, max_tokens=max_tokens)
+
+    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 _create_logger(self) -> Logger | None:
+        if self._logger_save_path is None:
+            return None
+
+        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}")
+        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()