feat: Update LLM context mode to use sentence-level context and enhance prompt for regex replacements

.env.example

@@ -33,6 +33,6 @@ ABOGEN_LLM_BASE_URL=http://localhost:11434  # Supply the server root; /v1 is add
 ABOGEN_LLM_API_KEY=ollama
 ABOGEN_LLM_MODEL=llama3.1:8b
 ABOGEN_LLM_TIMEOUT=45
-ABOGEN_LLM_CONTEXT_MODE=paragraph
+ABOGEN_LLM_CONTEXT_MODE=sentence
 # For custom prompts, keep the text on a single line or escape newlines.
-#ABOGEN_LLM_PROMPT=You are assisting with audiobook preparation. Rewrite {{sentence}} for narration.
+#ABOGEN_LLM_PROMPT=Provide regex replacements for any apostrophes in {{sentence}} using apply_regex_replacements.

abogen/kokoro_text_normalization.py

@@ -1,13 +1,18 @@
 from __future__ import annotations
+
+import json
 import re
 import unicodedata
 from dataclasses import dataclass
-from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Sequence, Tuple
+from typing import TYPE_CHECKING, Any, Callable, Dict, Iterable, List, Mapping, Optional, Sequence, Tuple
 try:  # pragma: no cover - optional dependency guard
     from num2words import num2words
 except Exception:  # pragma: no cover - graceful degradation
     num2words = None  # type: ignore
 
+if TYPE_CHECKING:  # pragma: no cover - type checking only
+    from abogen.llm_client import LLMCompletion
+
 # ---------- Configuration Dataclass ----------
 
 @dataclass
@@ -600,10 +605,67 @@ DEFAULT_APOSTROPHE_CONFIG = ApostropheConfig()
 
 _MUSTACHE_PATTERN = re.compile(r"{{\s*([a-zA-Z0-9_]+)\s*}}")
 _LLM_SYSTEM_PROMPT = (
-    "You rewrite text for audiobook narration. Expand or clarify contractions and apostrophes "
-    "so the output is explicit and natural to read aloud. Respond with only the rewritten text."
+    "You assist with audiobook preparation. Review the sentence, identify any apostrophes or "
+    "contractions that should be expanded for clarity, and respond by calling the "
+    "apply_regex_replacements tool. Each replacement must target a single token, include a precise "
+    "regex pattern, and provide the exact replacement text. If no changes are required, call the tool "
+    "with an empty replacements list. Do not rewrite the sentence directly."
 )
 
+_LLM_REGEX_TOOL_NAME = "apply_regex_replacements"
+_LLM_REGEX_TOOL = {
+    "type": "function",
+    "function": {
+        "name": _LLM_REGEX_TOOL_NAME,
+        "description": (
+            "Return regex substitutions to normalize apostrophes or contractions in the provided sentence."
+        ),
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "replacements": {
+                    "description": "Ordered substitutions to apply to the sentence.",
+                    "type": "array",
+                    "items": {
+                        "type": "object",
+                        "properties": {
+                            "pattern": {
+                                "type": "string",
+                                "description": "Regular expression that matches the token to replace.",
+                            },
+                            "replacement": {
+                                "type": "string",
+                                "description": "Replacement text for the match.",
+                            },
+                            "flags": {
+                                "type": "array",
+                                "items": {"type": "string"},
+                                "description": "Optional re flags such as IGNORECASE.",
+                            },
+                            "count": {
+                                "type": "integer",
+                                "description": "Optional maximum number of replacements (default all).",
+                            },
+                            "reason": {
+                                "type": "string",
+                                "description": "Short explanation of why the replacement is needed.",
+                            },
+                        },
+                        "required": ["pattern", "replacement"],
+                    },
+                }
+            },
+            "required": ["replacements"],
+        },
+    },
+}
+_LLM_REGEX_TOOL_CHOICE = {"type": "function", "function": {"name": _LLM_REGEX_TOOL_NAME}}
+_LLM_ALLOWED_REGEX_FLAGS = {
+    "IGNORECASE": re.IGNORECASE,
+    "MULTILINE": re.MULTILINE,
+    "DOTALL": re.DOTALL,
+}
+
 
 def _render_mustache(template: str, context: Mapping[str, str]) -> str:
     if not template:
@@ -638,7 +700,6 @@ def _normalize_with_llm(
         raise LLMClientError("LLM configuration is incomplete")
 
     prompt_template = str(settings.get("llm_prompt") or DEFAULT_LLM_PROMPT)
-    context_mode = str(settings.get("llm_context_mode") or "sentence").strip().lower()
     lines = text.splitlines(keepends=True)
     if not lines:
         return text
@@ -661,39 +722,31 @@ def _normalize_with_llm(
         trailing_ws = line_body[len(line_body.rstrip()):]
         core = line_body[len(leading_ws) : len(line_body) - len(trailing_ws)]
 
-        paragraph_context = core
-        if context_mode == "sentence":
-            sentences = _split_sentences_for_llm(core)
-            if not sentences:
-                normalized_lines.append(line_body + newline)
-                continue
-            rewritten_sentences: List[str] = []
-            for sentence in sentences:
-                prompt_context = {
-                    "text": sentence,
-                    "sentence": sentence,
-                    "paragraph": paragraph_context,
-                }
-                prompt = _render_mustache(prompt_template, prompt_context)
-                completion = generate_completion(
-                    llm_config,
-                    system_message=_LLM_SYSTEM_PROMPT,
-                    user_message=prompt,
-                )
-                rewritten_sentences.append(completion.strip())
-            normalized_core = " ".join(filter(None, rewritten_sentences)) or core
-        else:
+        sentences = _split_sentences_for_llm(core)
+        if not sentences:
+            normalized_lines.append(line_body + newline)
+            continue
+
+        rewritten_sentences: List[str] = []
+        for sentence in sentences:
             prompt_context = {
-                "text": core,
-                "sentence": core,
-                "paragraph": paragraph_context,
+                "text": sentence,
+                "sentence": sentence,
+                "paragraph": sentence,
             }
             prompt = _render_mustache(prompt_template, prompt_context)
-            normalized_core = generate_completion(
+            completion = generate_completion(
                 llm_config,
                 system_message=_LLM_SYSTEM_PROMPT,
                 user_message=prompt,
-            ).strip() or core
+                tools=[_LLM_REGEX_TOOL],
+                tool_choice=_LLM_REGEX_TOOL_CHOICE,
+            )
+            rewritten_sentences.append(
+                _apply_llm_regex_replacements(sentence, completion)
+            )
+
+        normalized_core = " ".join(filter(None, rewritten_sentences)) or core
 
         rebuilt = f"{leading_ws}{normalized_core}{trailing_ws}{newline}"
         normalized_lines.append(rebuilt)
@@ -702,6 +755,129 @@ def _normalize_with_llm(
     return result if result else text
 
 
+def _apply_llm_regex_replacements(sentence: str, completion: "LLMCompletion") -> str:
+    replacements = _extract_llm_replacements(completion)
+    if not replacements:
+        return sentence
+
+    updated = sentence
+    for spec in replacements:
+        updated = _apply_single_regex_replacement(updated, spec)
+    return updated
+
+
+def _extract_llm_replacements(completion: "LLMCompletion") -> List[Dict[str, Any]]:
+    if completion is None:
+        return []
+
+    for call in getattr(completion, "tool_calls", ()):  # type: ignore[attr-defined]
+        if getattr(call, "name", None) != _LLM_REGEX_TOOL_NAME:
+            continue
+        payload = _safe_load_json(getattr(call, "arguments", None))
+        replacements = _coerce_replacement_list(payload)
+        if replacements:
+            return replacements
+
+    if getattr(completion, "content", None):
+        payload = _safe_load_json(completion.content)
+        replacements = _coerce_replacement_list(payload)
+        if replacements:
+            return replacements
+
+    return []
+
+
+def _safe_load_json(raw: Optional[str]) -> Any:
+    if not raw:
+        return None
+    try:
+        return json.loads(raw)
+    except json.JSONDecodeError:
+        return None
+
+
+def _coerce_replacement_list(raw: Any) -> List[Dict[str, Any]]:
+    if isinstance(raw, Mapping):
+        candidates = raw.get("replacements")
+    else:
+        candidates = raw
+
+    if not isinstance(candidates, list):
+        return []
+
+    replacements: List[Dict[str, Any]] = []
+    for item in candidates:
+        if not isinstance(item, Mapping):
+            continue
+        pattern = str(item.get("pattern") or "").strip()
+        if not pattern:
+            continue
+        replacement = str(item.get("replacement") or "")
+        entry: Dict[str, Any] = {"pattern": pattern, "replacement": replacement}
+
+        flags = _normalize_flag_field(item.get("flags"))
+        if flags:
+            entry["flags"] = flags
+
+        count = item.get("count")
+        if isinstance(count, int) and count >= 0:
+            entry["count"] = count
+
+        replacements.append(entry)
+
+    return replacements
+
+
+def _normalize_flag_field(raw: Any) -> List[str]:
+    if not raw:
+        return []
+
+    if isinstance(raw, str):
+        raw_iterable: Iterable[Any] = [raw]
+    elif isinstance(raw, Iterable) and not isinstance(raw, (bytes, str, Mapping)):
+        raw_iterable = raw
+    else:
+        return []
+
+    normalized: List[str] = []
+    seen: set[str] = set()
+    for value in raw_iterable:
+        candidate = str(value or "").strip().upper()
+        if not candidate or candidate not in _LLM_ALLOWED_REGEX_FLAGS or candidate in seen:
+            continue
+        seen.add(candidate)
+        normalized.append(candidate)
+    return normalized
+
+
+def _apply_single_regex_replacement(text: str, spec: Mapping[str, Any]) -> str:
+    pattern = str(spec.get("pattern") or "")
+    replacement = str(spec.get("replacement") or "")
+    if not pattern:
+        return text
+
+    flags_value = 0
+    flag_names = spec.get("flags")
+    if isinstance(flag_names, str):
+        flag_iterable: Iterable[Any] = [flag_names]
+    elif isinstance(flag_names, Iterable) and not isinstance(flag_names, (bytes, str, Mapping)):
+        flag_iterable = flag_names
+    else:
+        flag_iterable = []
+
+    for flag_name in flag_iterable:
+        lookup = str(flag_name or "").strip().upper()
+        flags_value |= _LLM_ALLOWED_REGEX_FLAGS.get(lookup, 0)
+
+    count = spec.get("count")
+    count_value = count if isinstance(count, int) and count >= 0 else 0
+
+    try:
+        return re.sub(pattern, replacement, text, count=count_value, flags=flags_value)
+    except re.error:
+        return text
+
+
 def normalize_for_pipeline(
     text: str,
     *,

abogen/llm_client.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import json
 from dataclasses import dataclass
-from typing import Any, Dict, List, Mapping, Optional
+from typing import Any, Dict, List, Mapping, Optional, Sequence, Tuple
 from urllib import error, parse, request
 
 
@@ -21,6 +21,18 @@ class LLMConfiguration:
         return bool(self.base_url.strip() and self.model.strip())
 
 
+@dataclass(frozen=True)
+class LLMToolCall:
+    name: str
+    arguments: str
+
+
+@dataclass(frozen=True)
+class LLMCompletion:
+    content: Optional[str]
+    tool_calls: Tuple[LLMToolCall, ...]
+
+
 _DEFAULT_HEADERS = {
     "Content-Type": "application/json",
     "Accept": "application/json",
@@ -115,7 +127,10 @@ def generate_completion(
     user_message: str,
     temperature: float = 0.2,
     max_tokens: Optional[int] = None,
-) -> str:
+    tools: Optional[Sequence[Mapping[str, Any]]] = None,
+    tool_choice: Optional[Mapping[str, Any]] = None,
+    response_format: Optional[Mapping[str, Any]] = None,
+) -> LLMCompletion:
     if not configuration.is_configured():
         raise LLMClientError("LLM configuration is incomplete")
 
@@ -131,6 +146,12 @@ def generate_completion(
     }
     if max_tokens is not None:
         payload["max_tokens"] = max_tokens
+    if tools:
+        payload["tools"] = list(tools)
+    if tool_choice:
+        payload["tool_choice"] = dict(tool_choice)
+    if response_format:
+        payload["response_format"] = dict(response_format)
 
     response = _perform_request("POST", url, headers=headers, payload=payload, timeout=configuration.timeout)
     if not isinstance(response, Mapping):
@@ -142,11 +163,40 @@ def generate_completion(
     if not isinstance(first, Mapping):
         raise LLMClientError("LLM response choice was invalid")
     message = first.get("message")
+    content: Optional[str] = None
+    tool_calls: List[LLMToolCall] = []
     if isinstance(message, Mapping):
         content = message.get("content")
-        if isinstance(content, str) and content.strip():
-            return content.strip()
+        if isinstance(content, str):
+            stripped = content.strip()
+            if stripped:
+                content = stripped
+            else:
+                content = None
+        tool_call_entries = message.get("tool_calls")
+        if isinstance(tool_call_entries, list):
+            for entry in tool_call_entries:
+                if not isinstance(entry, Mapping):
+                    continue
+                fn = entry.get("function")
+                if not isinstance(fn, Mapping):
+                    continue
+                name = str(fn.get("name") or "").strip()
+                if not name:
+                    continue
+                args = fn.get("arguments", "")
+                if isinstance(args, (dict, list)):
+                    arguments = json.dumps(args)
+                else:
+                    arguments = str(args)
+                tool_calls.append(LLMToolCall(name=name, arguments=arguments))
+    if content:
+        return LLMCompletion(content=content, tool_calls=tuple(tool_calls))
     text = first.get("text")
-    if isinstance(text, str) and text.strip():
-        return text.strip()
+    if isinstance(text, str):
+        stripped = text.strip()
+        if stripped:
+            content = stripped
+    if content or tool_calls:
+        return LLMCompletion(content=content, tool_calls=tuple(tool_calls))
     raise LLMClientError("LLM response did not include text content")

abogen/normalization_settings.py

@@ -10,10 +10,10 @@ from abogen.llm_client import LLMConfiguration
 from abogen.utils import load_config
 
 DEFAULT_LLM_PROMPT = (
-    "You are assisting with audiobook preparation. Rewrite the provided sentence so apostrophes and "
-    "contractions are unambiguous for text-to-speech. Respond with only the rewritten sentence.\n"
-    "Sentence: {{ sentence }}\n"
-    "Context: {{ paragraph }}"
+    "You are assisting with audiobook preparation. Analyze the sentence and identify any apostrophes or "
+    "contractions that should be expanded for clarity. Call the apply_regex_replacements tool with precise "
+    "regex substitutions for only the words that need adjustment. If no changes are required, return an empty list.\n"
+    "Sentence: {{ sentence }}"
 )
 
 _SETTINGS_DEFAULTS: Dict[str, Any] = {

abogen/web/routes.py

@@ -1873,7 +1873,6 @@ _APOSTROPHE_MODE_OPTIONS = [
 
 _LLM_CONTEXT_OPTIONS = [
     {"value": "sentence", "label": "Sentence only"},
-    {"value": "paragraph", "label": "Sentence with paragraph context"},
 ]
 
 BOOLEAN_SETTINGS = {