world_writer.py

# world_writer --title "Tracking the Dead Wax" --pov "third person perspective" --outline_file outline.txt --characters_file characters.txt --detailed
# python -B world_writer.py --outline_file outline.txt --detailed
# pip install anthropic
# tested with: anthropic 0.49.0 circa March 2025
import anthropic
import os
import argparse
import re
import sys
import time
from datetime import datetime

parser = argparse.ArgumentParser(description='Extract and develop characters and world elements from a novel outline.')
parser.add_argument('--request_timeout', type=int, default=600, help='Maximum timeout for output (default: 600 seconds)')
"""
request_timeout is for each chunk and not for the entire request time:
script                      API
  |                          |
  |------ Request ---------->|
  |                          |
  |<----- first chunk -------|  ← timeout applies here
  |                          |
  |<----- next chunk --------|  ← and here (timeout is reset for each chunk)
  |                          |
  |        ...etc...         |
"""
parser.add_argument('--thinking_budget',    type=int, default=32000, help='Maximum tokens for AI thinking (default: 32000)')
parser.add_argument('--max_tokens',         type=int, default=20000, help='Maximum tokens for output (default: 20000)')
parser.add_argument('--context_window',     type=int, default=204648, help='Context window for Claude 3.7 Sonnet (default: 204648)')

parser.add_argument('--lang',               type=str, default="English", help='Language for writing (default: English)')
parser.add_argument('--title',              type=str, required=True, default="", help="Title of story")
parser.add_argument('--pov',                type=str, required=True, default="third person perspective", help="Point of view")
parser.add_argument('--characters_file',    type=str, required=True, default="characters.txt", help="Characters")
parser.add_argument('--outline_file',       type=str, required=True, help="Path to the outline file generated by outline_writer.py")
parser.add_argument('--detailed',           action='store_true', help='Generate more detailed character and world profiles')
parser.add_argument('--save_dir',           type=str, default=".")
args = parser.parse_args()

def count_words(text):
    return len(re.sub(r'(\r\n|\r|\n)', ' ', text).split())

# load the characters file
characters_content = ""
try:
    with open(args.characters_file, 'r', encoding='utf-8') as file:
        characters_content = file.read()
    print(f"Loaded characters from: {args.characters_file}")
except FileNotFoundError:
    print(f"Error: Characters file not found: {args.characters_file}")
    sys.exit(1)
except Exception as e:
    print(f"Error: Could not read Characters file: {e}")
    sys.exit(1)

# load the outline file
outline_content = ""
try:
    with open(args.outline_file, 'r', encoding='utf-8') as file:
        outline_content = file.read()
    print(f"Loaded outline from: {args.outline_file}")
except FileNotFoundError:
    print(f"Error: Outline file not found: {args.outline_file}")
    sys.exit(1)
except Exception as e:
    print(f"Error: Could not read outline file: {e}")
    sys.exit(1)

# create integrated world prompt with character section
world_prompt = f"""You are a skilled novelist, worldbuilder, and character developer helping to create a comprehensive world document in fluent, authentic {args.lang}.
This document will include both the world elements and detailed character profiles for a novel based on the outline below.

=== OUTLINE ===
{outline_content}
=== END OUTLINE ===

Create a detailed world document with the following sections:

----------------------------------------------
WORLD: {args.title}
----------------------------------------------

1. SETTING OVERVIEW:
   - Time period and era
   - General geography and environment
   - Notable locations mentioned in the outline

2. SOCIAL STRUCTURE:
   - Government or ruling systems
   - Social classes or hierarchies
   - Cultural norms and values

3. HISTORY:
   - Major historical events that impact the story
   - Historical figures relevant to the plot
   - Timeline of important developments

4. TECHNOLOGY AND MAGIC:
   - Level of technological development
   - Technological systems or devices crucial to the plot
   - If applicable: magic systems, supernatural elements, or fantastic creatures

5. ECONOMY:
   - Economic systems
   - Resources and trade
   - Economic conflicts relevant to the story

6. THEMES AND SYMBOLS:
   - Recurring motifs and symbols
   - Philosophical or moral questions explored
   - Cultural or religious symbolism

7. RULES OF THE WORLD:
   - Laws (both legal and natural/supernatural)
   - Limitations and constraints
   - Unique aspects of how this world functions

8. CHARACTER PROFILES:
"""

# add character profile instructions
world_prompt += f"""
   For each of the following characters, create a detailed profile but do NOT change the character names:

    === CHARACTERS ===
    {characters_content}
    === END CHARACTERS ===

   Include for each character:

   a) CHARACTER NAME: [Full name]
   b) ROLE: [Protagonist, Antagonist, Supporting Character, etc.]
   c) AGE: [Age or age range]
   d) PHYSICAL DESCRIPTION: [Detailed physical appearance]
   e) BACKGROUND: [Personal history relevant to the story]
   f) PERSONALITY: [Core personality traits, strengths, and flaws]
   g) MOTIVATIONS: [What drives this character? What do they want?]
   h) CONFLICTS: [Internal struggles and external conflicts]
   i) RELATIONSHIPS: [Important relationships with other characters]
   j) ARC: [How this character changes throughout the story]
   k) NOTABLE QUOTES: [3-5 examples of how this character might speak]"""

if args.detailed:
    world_prompt += """
   l) SKILLS & ABILITIES: [Special skills, knowledge, or supernatural abilities]
   m) HABITS & QUIRKS: [Distinctive behaviors and mannerisms]
   n) SECRETS: [What this character is hiding]
   o) FEARS & WEAKNESSES: [What makes this character vulnerable]
   p) SYMBOLIC ELEMENTS: [Any symbolic elements associated with this character]
   q) NARRATIVE FUNCTION: [How this character serves the themes and plot]

"""

# add formatting instructions
world_prompt += """
IMPORTANT FORMATTING INSTRUCTIONS:
- Write in {pov}
- Make the existing character profiles deep and psychologically nuanced
- Ensure the existing character motivations are complex and realistic
- Ensure the existing characters have traits and backgrounds that naturally arise from the world of the story
- Ensure the existing characters will create interesting dynamics and conflicts with each other
- Keep all details consistent with the outline and list of characters
- Focus on elements that directly impact the characters and plot
- Provide enough detail to give the world depth while leaving room for creative development
- Ensure the world elements support and enhance the narrative
- Separate each major section with a line of dashes (------)
- Separate each character profile with a line of dashes (------)
- Be consistent in formatting throughout the document
- Use plain text formatting with NO markdown in your outputs
- Do NOT change nor add to character names
"""

# calculate a safe max_tokens value
estimated_input_tokens = int(len(world_prompt) // 5.5)
max_safe_tokens = max(5000, args.context_window - estimated_input_tokens - 1000)  # 1000 token buffer
max_tokens = int(min(args.max_tokens, max_safe_tokens))

# ensure max_tokens is always greater than thinking budget
if max_tokens <= args.thinking_budget:
    max_tokens = args.thinking_budget + args.max_tokens
    print(f"Adjusted max_tokens to {max_tokens} to exceed thinking budget of {args.thinking_budget} (room for thinking/writing)")

absolute_path = os.path.abspath(args.save_dir)

print(f"Max request timeout: {args.request_timeout} seconds")
print(f"Max retries: 0 (anthropic's default was 2)")
print(f"Max AI model context window: {args.context_window} tokens")
print(f"AI model thinking budget: {args.thinking_budget} tokens")
print(f"Generating world document (including characters) for novel: {args.title}")
print(f"Setting max_tokens to: {max_tokens}")

client = anthropic.Anthropic(
    timeout=args.request_timeout,
    max_retries=0  # default is 2
)

# count tokens for the prompt
prompt_token_count = 0
try:
    response = client.beta.messages.count_tokens(
        model="claude-3-7-sonnet-20250219",
        messages=[{"role": "user", "content": world_prompt}],
        thinking={
            "type": "enabled",
            "budget_tokens": args.thinking_budget
        },
        betas=["output-128k-2025-02-19"]
    )
    prompt_token_count = response.input_tokens
    print(f"Prompt tokens: {prompt_token_count}")
except Exception as e:
    print(f"Error counting tokens:\n{e}\n")

world_response = ""
world_thinking = ""

start_time = time.time()

dt = datetime.fromtimestamp(start_time)
formatted_time = dt.strftime("%A %B %d, %Y %I:%M:%S %p").replace(" 0", " ").lower()
print(f"****************************************************************************")
print(f"*  Generating world document with character profiles...")
print(f"*  sending to API at: {formatted_time}")
print(f"*  This process typically takes several minutes.")
print(f"*  ")
print(f"*  It's recommended to keep the Terminal or command line the sole 'focus'")
print(f"*  and to avoid browsing online or running other apps, as these API")
print(f"*  network connections are often flakey, like delicate echoes of whispers.")
print(f"*  ")
print(f"*  So breathe, remove eye glasses, stretch, relax, and be like water 🥋 🧘🏽‍♀️")
print(f"****************************************************************************")

# print(world_prompt)
# sys.exit(1)

try:
    with client.beta.messages.stream(
        model="claude-3-7-sonnet-20250219",
        max_tokens=max_tokens,
        messages=[{"role": "user", "content": world_prompt}],
        thinking={
            "type": "enabled",
            "budget_tokens": args.thinking_budget
        },
        betas=["output-128k-2025-02-19"]
    ) as stream:
        for event in stream:
            if event.type == "content_block_delta":
                if event.delta.type == "thinking_delta":
                    world_thinking += event.delta.thinking
                elif event.delta.type == "text_delta":
                    world_response += event.delta.text
except Exception as e:
    print(f"Error generating world document:\n{e}\n")

elapsed = time.time() - start_time
minutes = int(elapsed // 60)
seconds = elapsed % 60

timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
world_filename = f"{args.save_dir}/world_{timestamp}.txt"
with open(world_filename, 'w', encoding='utf-8') as file:
    file.write(world_response)

word_count = count_words(world_response)

print(f"\nWorld document completed in: {minutes} minutes, {seconds:.2f} seconds.")
print(f"World document has {word_count} words.")
print(f"World document saved to: {world_filename}")

# calculate token count of output
output_token_count = 0
try:
    response = client.beta.messages.count_tokens(
        model="claude-3-7-sonnet-20250219",
        messages=[{"role": "user", "content": world_response}],
        thinking={
            "type": "enabled",
            "budget_tokens": args.thinking_budget
        },
        betas=["output-128k-2025-02-19"]
    )
    output_token_count = response.input_tokens
    print(f"World document is {output_token_count} tokens")
except Exception as e:
    print(f"Error counting output tokens:\n{e}\n")

# save thinking content if available
if world_thinking:
    thinking_filename = f"{args.save_dir}/world_thinking_{timestamp}.txt"
    with open(thinking_filename, 'w', encoding='utf-8') as file:
        file.write("=== PROMPT USED (EXCLUDING REFERENCE CONTENT) ===\n")
        file.write(f"Generating world document (including characters) for novel: {args.title}\n")
        file.write("\n\n=== AI'S THINKING PROCESS ===\n\n")
        file.write(world_thinking)
        file.write("\n=== END AI'S THINKING PROCESS ===\n")
        file.write(f"\nStats:\n")
        file.write(f"Prompt tokens: {prompt_token_count}\n")
        file.write(f"Elapsed time: {minutes} minutes, {seconds:.2f} seconds\n")
        file.write(f"Word count: {word_count}\n")
        file.write(f"\nFiles saved to: {absolute_path}")
        file.write(f"###\n")
    print(f"AI thinking process saved to: {thinking_filename}")

print(f"\nFiles saved to: {absolute_path}")
print(f"###\n")

# clean up variables
world_prompt = None
world_response = None
world_thinking = None
client = None