diff --git a/.env.example b/.env.example index 3a2c720..9a33129 100644 --- a/.env.example +++ b/.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. diff --git a/abogen/kokoro_text_normalization.py b/abogen/kokoro_text_normalization.py index 264445b..f99dc18 100644 --- a/abogen/kokoro_text_normalization.py +++ b/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, *, diff --git a/abogen/llm_client.py b/abogen/llm_client.py index c5f92bb..b6c075b 100644 --- a/abogen/llm_client.py +++ b/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") diff --git a/abogen/normalization_settings.py b/abogen/normalization_settings.py index a73316f..5fd8323 100644 --- a/abogen/normalization_settings.py +++ b/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] = { diff --git a/abogen/web/routes.py b/abogen/web/routes.py index 9acb254..195e32e 100644 --- a/abogen/web/routes.py +++ b/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 = {