Merge pull request #164 from CyberSecDef/copilot/enforce-target-novel-length

Enforce target novel length via expansion agent and length validation

Author: CyberSecDef

config.py

@@ -10,6 +10,7 @@
 """
 
 from novelforge.config import (  # noqa: F401 – re-exports for backward compatibility
+    CHAPTER_MIN_LENGTH_PCT,
     ConfigurationError,
     IMAGE_API_KEY,
     IMAGE_API_URL,
@@ -26,6 +27,7 @@
     LLM_TIMEOUT,
     LOGS_DIR,
     MAX_CHAPTERS,
+    MAX_EXPANSION_ATTEMPTS,
     MAX_WORD_COUNT,
     NOVELS_DIR,
     EXPORT_DIR,

novelforge/agents/chapter/__init__.py

@@ -30,6 +30,8 @@
     _format_anti_repetition_rules,
     _log_pass_failure,
     _sanitize_for_content_policy,
+    check_chapter_length,
+    expand_chapter,
     format_vocabulary_rules,
     get_forbidden_words,
     get_soft_limited_words,
@@ -50,6 +52,7 @@
 from novelforge.agents.chapter.prompts import (  # noqa: F401
     build_anti_llm_agent_prompt,
     build_chapter_draft_prompt,
+    build_chapter_expansion_prompt,
     build_chapter_revision_prompt,
     build_chapter_rhythm_classifier_prompt,
     build_rhythm_compliance_verifier_prompt,

novelforge/agents/chapter/_helpers.py

@@ -1,4 +1,4 @@
-"""Shared constants, pass-failure helpers, content-policy retry, and vocabulary scanning."""
+"""Shared constants, pass-failure helpers, content-policy retry, vocabulary scanning, and length enforcement."""
 
 import logging
 import re
@@ -395,3 +395,83 @@ def scan_vocabulary_overuse(chapter_text: str, genre: str = "") -> list[str]:
             )
 
     return warnings
+
+
+# ---------------------------------------------------------------------------
+# Chapter length enforcement
+# ---------------------------------------------------------------------------
+
+def check_chapter_length(
+    text: str,
+    target_words: int,
+    min_pct: int = 85,
+) -> tuple[int, int, bool]:
+    """Check whether a chapter meets the minimum length threshold.
+
+    Returns ``(actual_word_count, min_threshold, is_acceptable)`` where
+    *min_threshold* is the minimum word count derived from *target_words*
+    and *min_pct*.
+    """
+    actual = len(text.split())
+    min_threshold = (target_words * min_pct) // 100
+    return actual, min_threshold, actual >= min_threshold
+
+
+def expand_chapter(
+    text: str,
+    *,
+    target_words: int,
+    min_words: int,
+    chapter_num: int,
+    title: str,
+    max_attempts: int = 2,
+) -> tuple[str, int]:
+    """Expand an under-length chapter by calling the expansion agent.
+
+    Tries up to *max_attempts* expansion calls.  Returns
+    ``(expanded_text, final_word_count)``.  If the expansion agent fails or
+    the chapter still doesn't meet the threshold, returns the best result
+    achieved so far rather than raising.
+    """
+    from novelforge.agents.chapter.prompts import build_chapter_expansion_prompt
+
+    current = text
+    current_wc = len(current.split())
+    for attempt in range(1, max_attempts + 1):
+        if current_wc >= min_words:
+            break
+        logger.info(
+            "Chapter %d: expansion attempt %d/%d (%d words, need %d)",
+            chapter_num, attempt, max_attempts, current_wc, min_words,
+        )
+        try:
+            expanded = call_llm(
+                build_chapter_expansion_prompt(
+                    chapter_text=current,
+                    current_words=current_wc,
+                    target_words=target_words,
+                    min_words=min_words,
+                ),
+                action=f"Chapter {chapter_num}: expansion (attempt {attempt})",
+            )
+            new_wc = len(expanded.split())
+            if new_wc > current_wc:
+                current = expanded
+                current_wc = new_wc
+                logger.info(
+                    "Chapter %d: expansion attempt %d produced %d words",
+                    chapter_num, attempt, new_wc,
+                )
+            else:
+                logger.warning(
+                    "Chapter %d: expansion attempt %d did not increase length (%d → %d)",
+                    chapter_num, attempt, current_wc, new_wc,
+                )
+                break
+        except Exception as exc:
+            logger.warning(
+                "Chapter %d: expansion attempt %d failed: %s: %s",
+                chapter_num, attempt, type(exc).__name__, exc,
+            )
+            break
+    return current, current_wc

novelforge/agents/chapter/pipeline.py

@@ -248,6 +248,7 @@ def _run_all_chapter_agents(
     degraded_passes: list[dict] | None = None,
     chapter_rhythm_shape: str = "",
     chapter_rhythm_reason: str = "",
+    target_words: int = 0,
 ) -> tuple[str, str]:
     """
     Run all chapter refinement agents (post-draft) and return:
@@ -468,6 +469,94 @@ def _build_with_vocab_rules(t: str) -> list[dict]:
             text, action=f"Chapter {chapter_num}: vocabulary fix-up",
         )
 
+    # Length enforcement — expand under-length chapters after all refinement
+    # passes have finished so that the expansion doesn't get trimmed.
+    # Expansion uses _safe() so vocabulary rules and content-retry apply,
+    # and a vocabulary scan runs afterwards to catch any reintroduced words.
+    if target_words > 0:
+        from novelforge.agents.chapter._helpers import check_chapter_length
+        from novelforge.agents.chapter.prompts import build_chapter_expansion_prompt
+        import novelforge.config as _cfg
+
+        min_pct = _cfg.CHAPTER_MIN_LENGTH_PCT
+        actual_wc, min_threshold, acceptable = check_chapter_length(text, target_words, min_pct)
+        if not acceptable:
+            for _exp_attempt in range(1, _cfg.MAX_EXPANSION_ATTEMPTS + 1):
+                if actual_wc >= min_threshold:
+                    break
+                _check_deadline()
+                if step_callback:
+                    step_callback(
+                        f"Chapter {chapter_num}: expanding ({actual_wc} words, "
+                        f"need {min_threshold}, attempt {_exp_attempt})"
+                    )
+                logger.info(
+                    "Chapter %d: expansion attempt %d/%d (%d words, need %d)",
+                    chapter_num, _exp_attempt, _cfg.MAX_EXPANSION_ATTEMPTS,
+                    actual_wc, min_threshold,
+                )
+                try:
+                    _cur_wc = actual_wc
+                    _cur_target = target_words
+                    _cur_min = min_threshold
+                    expanded = _safe(
+                        lambda t: build_chapter_expansion_prompt(
+                            chapter_text=t,
+                            current_words=_cur_wc,
+                            target_words=_cur_target,
+                            min_words=_cur_min,
+                        ),
+                        text,
+                        action=f"Chapter {chapter_num}: expansion (attempt {_exp_attempt})",
+                    )
+                    new_wc = len(expanded.split())
+                    if new_wc > actual_wc:
+                        text = expanded
+                        actual_wc = new_wc
+                        logger.info(
+                            "Chapter %d: expansion attempt %d produced %d words",
+                            chapter_num, _exp_attempt, new_wc,
+                        )
+                    else:
+                        logger.warning(
+                            "Chapter %d: expansion attempt %d did not increase length (%d → %d)",
+                            chapter_num, _exp_attempt, actual_wc, new_wc,
+                        )
+                        break
+                except Exception as exc:
+                    logger.warning(
+                        "Chapter %d: expansion attempt %d failed: %s: %s",
+                        chapter_num, _exp_attempt, type(exc).__name__, exc,
+                    )
+                    break
+
+            logger.info(
+                "Chapter %d: post-expansion word count = %d (target=%d, min=%d)",
+                chapter_num, actual_wc, target_words, min_threshold,
+            )
+
+            # Re-run vocabulary scan after expansion since the LLM may have
+            # reintroduced forbidden or overused words.
+            _check_deadline()
+            post_violations = scan_vocabulary_overuse(text, genre=genre)
+            if post_violations:
+                if step_callback:
+                    step_callback(
+                        f"Chapter {chapter_num}: fixing {len(post_violations)} "
+                        f"vocabulary issues (post-expansion)"
+                    )
+                logger.info(
+                    "Chapter %d: post-expansion vocabulary scan found %d violations",
+                    chapter_num, len(post_violations),
+                )
+                text = _safe(
+                    lambda t: build_vocabulary_fix_prompt(
+                        t, chapter_num, title, post_violations,
+                    ),
+                    text,
+                    action=f"Chapter {chapter_num}: vocabulary fix-up (post-expansion)",
+                )
+
     _check_deadline()
     if step_callback:
         step_callback(f"Chapter {chapter_num}: summarising")

novelforge/agents/chapter/prompts.py

@@ -106,6 +106,22 @@ def build_chapter_draft_prompt(
     )
 
 
+def build_chapter_expansion_prompt(
+    chapter_text: str,
+    current_words: int,
+    target_words: int,
+    min_words: int,
+) -> list[dict[str, str]]:
+    """Build the expansion prompt for under-length chapters."""
+    return render_prompt(
+        "chapter_expansion",
+        chapter_text=chapter_text,
+        current_words=f"{current_words:,}",
+        target_words=f"{target_words:,}",
+        min_words=f"{min_words:,}",
+    )
+
+
 # ---------------------------------------------------------------------------
 # Chapter refinement agent prompt builders
 # ---------------------------------------------------------------------------

novelforge/config.py

@@ -52,6 +52,9 @@
     # Input validation limits
     "MAX_CHAPTERS",
     "MAX_WORD_COUNT",
+    # Chapter length enforcement
+    "CHAPTER_MIN_LENGTH_PCT",
+    "MAX_EXPANSION_ATTEMPTS",
     # Flask
     "SECRET_KEY",
     "SESSION_FILE_DIR",
@@ -226,6 +229,33 @@ def _parse_llm_providers() -> list[ProviderConfig]:
 MAX_CHAPTERS = get_env_int("MAX_CHAPTERS", 100, min_value=1)
 MAX_WORD_COUNT = get_env_int("MAX_WORD_COUNT", 500000, min_value=1)
 
+# ---------------------------------------------------------------------------
+# Chapter length enforcement (override via environment variables)
+# ---------------------------------------------------------------------------
+
+
+def _get_percentage_env(
+    name: str, default: int, *, min_value: int = 0, max_value: int = 100
+) -> int:
+    """Return an integer percentage env var constrained to the given range."""
+    value = get_env_int(name, default, min_value=min_value)
+    if value > max_value:
+        _CONFIG_PARSE_ERRORS.append(
+            f"{name} must be <= {max_value} (got {value})."
+        )
+        return default
+    return value
+
+
+# Minimum acceptable chapter word count as a percentage of the per-chapter target.
+# Chapters below this threshold trigger an automatic expansion pass.
+CHAPTER_MIN_LENGTH_PCT = _get_percentage_env(
+    "CHAPTER_MIN_LENGTH_PCT", 85, min_value=50, max_value=100
+)
+
+# Maximum number of expansion attempts per chapter before accepting as-is.
+MAX_EXPANSION_ATTEMPTS = get_env_int("MAX_EXPANSION_ATTEMPTS", 2, min_value=0)
+
 # ---------------------------------------------------------------------------
 # Flask
 # ---------------------------------------------------------------------------

novelforge/progress.py

@@ -58,6 +58,10 @@ class ProgressState(TypedDict, total=False):
     # Token linking a novel entry to its illustration job
     illustration_token: str
 
+    # Length enforcement tracking: list of per-chapter dicts such as
+    # {chapter_num, target, min_threshold, actual, meets_min_threshold, total_words_so_far}.
+    length_enforcement: list[dict[str, Any]]
+
 
 _REQUIRED_CREATION_KEYS: frozenset[str] = frozenset(
     {"status", "current", "total", "step", "chapters_done", "error"}

novelforge/routes/generation/chapters.py

@@ -203,6 +203,8 @@ def _run_chapter_generation_internal(
     consequence_log: list[dict] = []          # tracks irreversible consequences per chapter
     exposition_log: list[dict] = []           # tracks dramatization ratio per chapter
     thematic_phrase_log: list[str] = []       # tracks thematic phrases to avoid repeating
+    length_enforcement_log: list[dict] = []   # tracks per-chapter length enforcement results
+    total_words_generated: int = 0            # running total for global length tracking
 
     # Format voice seed for prompt injection
     from novelforge.voice import format_voice_prompt
@@ -423,6 +425,7 @@ def _set_step(step_label: str) -> None:
                 degraded_passes=degraded_passes,
                 chapter_rhythm_shape=chapter_rhythm_shape,
                 chapter_rhythm_reason=chapter_rhythm_reason,
+                target_words=target_per_chapter,
             )
             summaries.append(summary)
 
@@ -494,6 +497,22 @@ def _set_step(step_label: str) -> None:
             chapter_usage = get_llm_usage()
             chapter_word_count = len(text.split())
 
+            # Track length enforcement results
+            from novelforge.agents.chapter._helpers import check_chapter_length
+            min_pct = config.CHAPTER_MIN_LENGTH_PCT
+            _, min_threshold, meets_min = check_chapter_length(
+                text, target_per_chapter, min_pct,
+            )
+            total_words_generated += chapter_word_count
+            length_enforcement_log.append({
+                "chapter_num": chapter_num,
+                "target": target_per_chapter,
+                "min_threshold": min_threshold,
+                "actual": chapter_word_count,
+                "meets_min_threshold": meets_min,
+                "total_words_so_far": total_words_generated,
+            })
+
             chapters_done.append({
                 "number": chapter_num,
                 "title": chapter_title,
@@ -516,6 +535,7 @@ def _set_step(step_label: str) -> None:
                 "chapters_done": list(chapters_done),
                 "character_state_log": list(character_state_log),
                 "degraded_passes": list(degraded_passes),
+                "length_enforcement": list(length_enforcement_log),
             })
             _persist_progress(force=True)  # always persist on chapter completion