├── .gitignore ├── README.md ├── agents.py ├── book_generator.py ├── config.py ├── main.py ├── outline_generator.py └── requirements.txt /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | share/python-wheels/ 24 | *.egg-info/ 25 | .installed.cfg 26 | *.egg 27 | MANIFEST 28 | 29 | # PyInstaller 30 | # Usually these files are written by a python script from a template 31 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 32 | *.manifest 33 | *.spec 34 | 35 | # Installer logs 36 | pip-log.txt 37 | pip-delete-this-directory.txt 38 | 39 | # Unit test / coverage reports 40 | htmlcov/ 41 | .tox/ 42 | .nox/ 43 | .coverage 44 | .coverage.* 45 | .cache 46 | nosetests.xml 47 | coverage.xml 48 | *.cover 49 | *.py,cover 50 | .hypothesis/ 51 | .pytest_cache/ 52 | cover/ 53 | 54 | # Translations 55 | *.mo 56 | *.pot 57 | 58 | # Django stuff: 59 | *.log 60 | local_settings.py 61 | db.sqlite3 62 | db.sqlite3-journal 63 | 64 | # Flask stuff: 65 | instance/ 66 | .webassets-cache 67 | 68 | # Scrapy stuff: 69 | .scrapy 70 | 71 | # Sphinx documentation 72 | docs/_build/ 73 | 74 | # PyBuilder 75 | .pybuilder/ 76 | target/ 77 | 78 | # Jupyter Notebook 79 | .ipynb_checkpoints 80 | 81 | # IPython 82 | profile_default/ 83 | ipython_config.py 84 | 85 | # pyenv 86 | # For a library or package, you might want to ignore these files since the code is 87 | # intended to run in multiple environments; otherwise, check them in: 88 | # .python-version 89 | 90 | # pipenv 91 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 92 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 93 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 94 | # install all needed dependencies. 95 | #Pipfile.lock 96 | 97 | # UV 98 | # Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control. 99 | # This is especially recommended for binary packages to ensure reproducibility, and is more 100 | # commonly ignored for libraries. 101 | #uv.lock 102 | 103 | # poetry 104 | # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. 105 | # This is especially recommended for binary packages to ensure reproducibility, and is more 106 | # commonly ignored for libraries. 107 | # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control 108 | #poetry.lock 109 | 110 | # pdm 111 | # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. 112 | #pdm.lock 113 | # pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it 114 | # in version control. 115 | # https://pdm.fming.dev/latest/usage/project/#working-with-version-control 116 | .pdm.toml 117 | .pdm-python 118 | .pdm-build/ 119 | 120 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm 121 | __pypackages__/ 122 | 123 | # Celery stuff 124 | celerybeat-schedule 125 | celerybeat.pid 126 | 127 | # SageMath parsed files 128 | *.sage.py 129 | 130 | # Environments 131 | .env 132 | .venv 133 | env/ 134 | venv/ 135 | ENV/ 136 | env.bak/ 137 | venv.bak/ 138 | 139 | # Spyder project settings 140 | .spyderproject 141 | .spyproject 142 | 143 | # Rope project settings 144 | .ropeproject 145 | 146 | # mkdocs documentation 147 | /site 148 | 149 | # mypy 150 | .mypy_cache/ 151 | .dmypy.json 152 | dmypy.json 153 | 154 | # Pyre type checker 155 | .pyre/ 156 | 157 | # pytype static type analyzer 158 | .pytype/ 159 | 160 | # Cython debug symbols 161 | cython_debug/ 162 | 163 | # PyCharm 164 | # JetBrains specific template is maintained in a separate JetBrains.gitignore that can 165 | # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore 166 | # and can be added to the global gitignore or merged into this file. For a more nuclear 167 | # option (not recommended) you can uncomment the following to ignore the entire idea folder. 168 | #.idea/ 169 | 170 | # PyPI configuration file 171 | .pypirc 172 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # AutoGen Book Generator 2 | 3 | A Python-based system that uses AutoGen to generate complete books through collaborative AI agents. The system employs multiple specialized agents working together to create coherent, structured narratives from initial prompts. 4 | 5 | ## Features 6 | 7 | - Multi-agent collaborative writing system 8 | - Structured chapter generation with consistent formatting 9 | - Maintains story continuity and character development 10 | - Automated world-building and setting management 11 | - Support for complex, multi-chapter narratives 12 | - Built-in validation and error handling 13 | 14 | ## Architecture 15 | 16 | The system uses several specialized agents: 17 | 18 | - **Story Planner**: Creates high-level story arcs and plot points 19 | - **World Builder**: Establishes and maintains consistent settings 20 | - **Memory Keeper**: Tracks continuity and context 21 | - **Writer**: Generates the actual prose 22 | - **Editor**: Reviews and improves content 23 | - **Outline Creator**: Creates detailed chapter outlines 24 | 25 | ## Installation 26 | 27 | 1. Clone the repository: 28 | ```bash 29 | git clone https://github.com/yourusername/autogen-book-generator.git 30 | cd autogen-book-generator 31 | ``` 32 | 33 | 2. Create a virtual environment: 34 | ```bash 35 | python -m venv venv 36 | source venv/bin/activate # On Windows: venv\Scripts\activate 37 | ``` 38 | 39 | 3. Install dependencies: 40 | ```bash 41 | pip install -r requirements.txt 42 | ``` 43 | 44 | ## Usage 45 | 46 | 1. Basic usage: 47 | ```python 48 | from main import main 49 | 50 | if __name__ == "__main__": 51 | main() 52 | ``` 53 | 54 | 2. Custom initial prompt: 55 | ```python 56 | from config import get_config 57 | from agents import BookAgents 58 | from book_generator import BookGenerator 59 | from outline_generator import OutlineGenerator 60 | 61 | # Get configuration 62 | agent_config = get_config() 63 | 64 | # Create agents 65 | outline_agents = BookAgents(agent_config) 66 | agents = outline_agents.create_agents() 67 | 68 | # Generate outline 69 | outline_gen = OutlineGenerator(agents, agent_config) 70 | outline = outline_gen.generate_outline(your_prompt, num_chapters=25) 71 | 72 | # Initialize book generator 73 | book_agents = BookAgents(agent_config, outline) 74 | agents_with_context = book_agents.create_agents() 75 | book_gen = BookGenerator(agents_with_context, agent_config, outline) 76 | 77 | # Generate book 78 | book_gen.generate_book(outline) 79 | ``` 80 | 81 | ## Configuration 82 | 83 | The system can be configured through `config.py`. Key configurations include: 84 | 85 | - LLM endpoint URL 86 | - Number of chapters 87 | - Agent parameters 88 | - Output directory settings 89 | 90 | ## Output Structure 91 | 92 | Generated content is saved in the `book_output` directory: 93 | ``` 94 | book_output/ 95 | ├── outline.txt 96 | ├── chapter_01.txt 97 | ├── chapter_02.txt 98 | └── ... 99 | ``` 100 | 101 | ## Requirements 102 | 103 | - Python 3.8+ 104 | - AutoGen 0.2.0+ 105 | - Other dependencies listed in requirements.txt 106 | 107 | ## Development 108 | 109 | To contribute to the project: 110 | 111 | 1. Fork the repository 112 | 2. Create a new branch for your feature 113 | 3. Install development dependencies: 114 | ```bash 115 | pip install -r requirements.txt 116 | ``` 117 | 4. Make your changes 118 | 5. Run tests: 119 | ```bash 120 | pytest 121 | ``` 122 | 6. Submit a pull request 123 | 124 | ## Error Handling 125 | 126 | The system includes robust error handling: 127 | - Validates chapter completeness 128 | - Ensures proper formatting 129 | - Maintains backup copies of generated content 130 | - Implements retry logic for failed generations 131 | 132 | ## Limitations 133 | 134 | - Requires significant computational resources 135 | - Generation time increases with chapter count 136 | - Quality depends on the underlying LLM model 137 | - May require manual review for final polish 138 | 139 | ## Contributing 140 | 141 | Contributions are welcome! Please feel free to submit a Pull Request. 142 | 143 | ## License 144 | 145 | This project is licensed under the MIT License - see the LICENSE file for details. 146 | 147 | ## Acknowledgments 148 | 149 | - Built using the [AutoGen](https://github.com/microsoft/autogen) framework 150 | - Inspired by collaborative writing systems -------------------------------------------------------------------------------- /agents.py: -------------------------------------------------------------------------------- 1 | """Define the agents used in the book generation system with improved context management""" 2 | import autogen 3 | from typing import Dict, List, Optional 4 | 5 | class BookAgents: 6 | def __init__(self, agent_config: Dict, outline: Optional[List[Dict]] = None): 7 | """Initialize agents with book outline context""" 8 | self.agent_config = agent_config 9 | self.outline = outline 10 | self.world_elements = {} # Track described locations/elements 11 | self.character_developments = {} # Track character arcs 12 | 13 | def _format_outline_context(self) -> str: 14 | """Format the book outline into a readable context""" 15 | if not self.outline: 16 | return "" 17 | 18 | context_parts = ["Complete Book Outline:"] 19 | for chapter in self.outline: 20 | context_parts.extend([ 21 | f"\nChapter {chapter['chapter_number']}: {chapter['title']}", 22 | chapter['prompt'] 23 | ]) 24 | return "\n".join(context_parts) 25 | 26 | def create_agents(self, initial_prompt, num_chapters) -> Dict: 27 | """Create and return all agents needed for book generation""" 28 | outline_context = self._format_outline_context() 29 | 30 | # Memory Keeper: Maintains story continuity and context 31 | memory_keeper = autogen.AssistantAgent( 32 | name="memory_keeper", 33 | system_message=f"""You are the keeper of the story's continuity and context. 34 | Your responsibilities: 35 | 1. Track and summarize each chapter's key events 36 | 2. Monitor character development and relationships 37 | 3. Maintain world-building consistency 38 | 4. Flag any continuity issues 39 | 40 | Book Overview: 41 | {outline_context} 42 | 43 | Format your responses as follows: 44 | - Start updates with 'MEMORY UPDATE:' 45 | - List key events with 'EVENT:' 46 | - List character developments with 'CHARACTER:' 47 | - List world details with 'WORLD:' 48 | - Flag issues with 'CONTINUITY ALERT:'""", 49 | llm_config=self.agent_config, 50 | ) 51 | 52 | # Story Planner - Focuses on high-level story structure 53 | story_planner = autogen.AssistantAgent( 54 | name="story_planner", 55 | system_message=f"""You are an expert story arc planner focused on overall narrative structure. 56 | 57 | Your sole responsibility is creating the high-level story arc. 58 | When given an initial story premise: 59 | 1. Identify major plot points and story beats 60 | 2. Map character arcs and development 61 | 3. Note major story transitions 62 | 4. Plan narrative pacing 63 | 64 | Format your output EXACTLY as: 65 | STORY_ARC: 66 | - Major Plot Points: 67 | [List each major event that drives the story] 68 | 69 | - Character Arcs: 70 | [For each main character, describe their development path] 71 | 72 | - Story Beats: 73 | [List key emotional and narrative moments in sequence] 74 | 75 | - Key Transitions: 76 | [Describe major shifts in story direction or tone] 77 | 78 | Always provide specific, detailed content - never use placeholders.""", 79 | llm_config=self.agent_config, 80 | ) 81 | 82 | # Outline Creator - Creates detailed chapter outlines 83 | outline_creator = autogen.AssistantAgent( 84 | name="outline_creator", 85 | system_message=f"""Generate a detailed {num_chapters}-chapter outline. 86 | 87 | YOU MUST USE EXACTLY THIS FORMAT FOR EACH CHAPTER - NO DEVIATIONS: 88 | 89 | Chapter 1: [Title] 90 | Chapter Title: [Same title as above] 91 | Key Events: 92 | - [Event 1] 93 | - [Event 2] 94 | - [Event 3] 95 | Character Developments: [Specific character moments and changes] 96 | Setting: [Specific location and atmosphere] 97 | Tone: [Specific emotional and narrative tone] 98 | 99 | [REPEAT THIS EXACT FORMAT FOR ALL {num_chapters} CHAPTERS] 100 | 101 | Requirements: 102 | 1. EVERY field must be present for EVERY chapter 103 | 2. EVERY chapter must have AT LEAST 3 specific Key Events 104 | 3. ALL chapters must be detailed - no placeholders 105 | 4. Format must match EXACTLY - including all headings and bullet points 106 | 107 | Initial Premise: 108 | {initial_prompt} 109 | 110 | START WITH 'OUTLINE:' AND END WITH 'END OF OUTLINE' 111 | """, 112 | llm_config=self.agent_config, 113 | ) 114 | 115 | # World Builder: Creates and maintains the story setting 116 | world_builder = autogen.AssistantAgent( 117 | name="world_builder", 118 | system_message=f"""You are an expert in world-building who creates rich, consistent settings. 119 | 120 | Your role is to establish ALL settings and locations needed for the entire story based on a provided story arc. 121 | 122 | Book Overview: 123 | {outline_context} 124 | 125 | Your responsibilities: 126 | 1. Review the story arc to identify every location and setting needed 127 | 2. Create detailed descriptions for each setting, including: 128 | - Physical layout and appearance 129 | - Atmosphere and environmental details 130 | - Important objects or features 131 | - Sensory details (sights, sounds, smells) 132 | 3. Identify recurring locations that appear multiple times 133 | 4. Note how settings might change over time 134 | 5. Create a cohesive world that supports the story's themes 135 | 136 | Format your response as: 137 | WORLD_ELEMENTS: 138 | 139 | [LOCATION NAME]: 140 | - Physical Description: [detailed description] 141 | - Atmosphere: [mood, time of day, lighting, etc.] 142 | - Key Features: [important objects, layout elements] 143 | - Sensory Details: [what characters would experience] 144 | 145 | [RECURRING ELEMENTS]: 146 | - List any settings that appear multiple times 147 | - Note any changes to settings over time 148 | 149 | [TRANSITIONS]: 150 | - How settings connect to each other 151 | - How characters move between locations""", 152 | llm_config=self.agent_config, 153 | ) 154 | 155 | # Writer: Generates the actual prose 156 | writer = autogen.AssistantAgent( 157 | name="writer", 158 | system_message=f"""You are an expert creative writer who brings scenes to life. 159 | 160 | Book Context: 161 | {outline_context} 162 | 163 | Your focus: 164 | 1. Write according to the outlined plot points 165 | 2. Maintain consistent character voices 166 | 3. Incorporate world-building details 167 | 4. Create engaging prose 168 | 5. Please make sure that you write the complete scene, do not leave it incomplete 169 | 6. Each chapter MUST be at least 5000 words (approximately 30,000 characters). Consider this a hard requirement. If your output is shorter, continue writing until you reach this minimum length 170 | 7. Ensure transitions are smooth and logical 171 | 8. Do not cut off the scene, make sure it has a proper ending 172 | 9. Add a lot of details, and describe the environment and characters where it makes sense 173 | 174 | Always reference the outline and previous content. 175 | Mark drafts with 'SCENE:' and final versions with 'SCENE FINAL:'""", 176 | llm_config=self.agent_config, 177 | ) 178 | 179 | # Editor: Reviews and improves content 180 | editor = autogen.AssistantAgent( 181 | name="editor", 182 | system_message=f"""You are an expert editor ensuring quality and consistency. 183 | 184 | Book Overview: 185 | {outline_context} 186 | 187 | Your focus: 188 | 1. Check alignment with outline 189 | 2. Verify character consistency 190 | 3. Maintain world-building rules 191 | 4. Improve prose quality 192 | 5. Return complete edited chapter 193 | 6. Never ask to start the next chapter, as the next step is finalizing this chapter 194 | 7. Each chapter MUST be at least 5000 words. If the content is shorter, return it to the writer for expansion. This is a hard requirement - do not approve chapters shorter than 5000 words 195 | 196 | Format your responses: 197 | 1. Start critiques with 'FEEDBACK:' 198 | 2. Provide suggestions with 'SUGGEST:' 199 | 3. Return full edited chapter with 'EDITED_SCENE:' 200 | 201 | Reference specific outline elements in your feedback.""", 202 | llm_config=self.agent_config, 203 | ) 204 | 205 | # User Proxy: Manages the interaction 206 | user_proxy = autogen.UserProxyAgent( 207 | name="user_proxy", 208 | human_input_mode="TERMINATE", 209 | code_execution_config={ 210 | "work_dir": "book_output", 211 | "use_docker": False 212 | } 213 | ) 214 | 215 | return { 216 | "story_planner": story_planner, 217 | "world_builder": world_builder, 218 | "memory_keeper": memory_keeper, 219 | "writer": writer, 220 | "editor": editor, 221 | "user_proxy": user_proxy, 222 | "outline_creator": outline_creator 223 | } 224 | 225 | def update_world_element(self, element_name: str, description: str) -> None: 226 | """Track a new or updated world element""" 227 | self.world_elements[element_name] = description 228 | 229 | def update_character_development(self, character_name: str, development: str) -> None: 230 | """Track character development""" 231 | if character_name not in self.character_developments: 232 | self.character_developments[character_name] = [] 233 | self.character_developments[character_name].append(development) 234 | 235 | def get_world_context(self) -> str: 236 | """Get formatted world-building context""" 237 | if not self.world_elements: 238 | return "No established world elements yet." 239 | 240 | return "\n".join([ 241 | "Established World Elements:", 242 | *[f"- {name}: {desc}" for name, desc in self.world_elements.items()] 243 | ]) 244 | 245 | def get_character_context(self) -> str: 246 | """Get formatted character development context""" 247 | if not self.character_developments: 248 | return "No character developments tracked yet." 249 | 250 | return "\n".join([ 251 | "Character Development History:", 252 | *[f"- {name}:\n " + "\n ".join(devs) 253 | for name, devs in self.character_developments.items()] 254 | ]) -------------------------------------------------------------------------------- /book_generator.py: -------------------------------------------------------------------------------- 1 | """Main class for generating books using AutoGen with improved iteration control""" 2 | import autogen 3 | from typing import Dict, List, Optional 4 | import os 5 | import time 6 | import re 7 | 8 | class BookGenerator: 9 | def __init__(self, agents: Dict[str, autogen.ConversableAgent], agent_config: Dict, outline: List[Dict]): 10 | """Initialize with outline to maintain chapter count context""" 11 | self.agents = agents 12 | self.agent_config = agent_config 13 | self.output_dir = "book_output" 14 | self.chapters_memory = [] # Store chapter summaries 15 | self.max_iterations = 3 # Limit editor-writer iterations 16 | self.outline = outline # Store the outline 17 | os.makedirs(self.output_dir, exist_ok=True) 18 | 19 | def _clean_chapter_content(self, content: str) -> str: 20 | """Clean up chapter content by removing artifacts and chapter numbers""" 21 | # Remove chapter number references 22 | content = re.sub(r'\*?\s*\(Chapter \d+.*?\)', '', content) 23 | content = re.sub(r'\*?\s*Chapter \d+.*?\n', '', content, count=1) 24 | 25 | # Clean up any remaining markdown artifacts 26 | content = content.replace('*', '') 27 | content = content.strip() 28 | 29 | return content 30 | 31 | 32 | def initiate_group_chat(self) -> autogen.GroupChat: 33 | """Create a new group chat for the agents with improved speaking order""" 34 | outline_context = "\n".join([ 35 | f"\nChapter {ch['chapter_number']}: {ch['title']}\n{ch['prompt']}" 36 | for ch in sorted(self.outline, key=lambda x: x['chapter_number']) 37 | ]) 38 | 39 | messages = [{ 40 | "role": "system", 41 | "content": f"Complete Book Outline:\n{outline_context}" 42 | }] 43 | 44 | writer_final = autogen.AssistantAgent( 45 | name="writer_final", 46 | system_message=self.agents["writer"].system_message, 47 | llm_config=self.agent_config 48 | ) 49 | 50 | return autogen.GroupChat( 51 | agents=[ 52 | self.agents["user_proxy"], 53 | self.agents["memory_keeper"], 54 | self.agents["writer"], 55 | self.agents["editor"], 56 | writer_final 57 | ], 58 | messages=messages, 59 | max_round=5, 60 | speaker_selection_method="round_robin" 61 | ) 62 | 63 | def _get_sender(self, msg: Dict) -> str: 64 | """Helper to get sender from message regardless of format""" 65 | return msg.get("sender") or msg.get("name", "") 66 | 67 | def _verify_chapter_complete(self, messages: List[Dict]) -> bool: 68 | """Verify chapter completion by analyzing entire conversation context""" 69 | print("******************** VERIFYING CHAPTER COMPLETION ****************") 70 | current_chapter = None 71 | chapter_content = None 72 | sequence_complete = { 73 | 'memory_update': False, 74 | 'plan': False, 75 | 'setting': False, 76 | 'scene': False, 77 | 'feedback': False, 78 | 'scene_final': False, 79 | 'confirmation': False 80 | } 81 | 82 | # Analyze full conversation 83 | for msg in messages: 84 | content = msg.get("content", "") 85 | 86 | # Track chapter number 87 | if not current_chapter: 88 | num_match = re.search(r"Chapter (\d+):", content) 89 | if num_match: 90 | current_chapter = int(num_match.group(1)) 91 | 92 | # Track completion sequence 93 | if "MEMORY UPDATE:" in content: sequence_complete['memory_update'] = True 94 | if "PLAN:" in content: sequence_complete['plan'] = True 95 | if "SETTING:" in content: sequence_complete['setting'] = True 96 | if "SCENE:" in content: sequence_complete['scene'] = True 97 | if "FEEDBACK:" in content: sequence_complete['feedback'] = True 98 | if "SCENE FINAL:" in content: 99 | sequence_complete['scene_final'] = True 100 | chapter_content = content.split("SCENE FINAL:")[1].strip() 101 | if "**Confirmation:**" in content and "successfully" in content: 102 | sequence_complete['confirmation'] = True 103 | 104 | #print all sequence_complete flags 105 | print("******************** SEQUENCE COMPLETE **************", sequence_complete) 106 | print("******************** CURRENT_CHAPTER ****************", current_chapter) 107 | print("******************** CHAPTER_CONTENT ****************", chapter_content) 108 | 109 | # Verify all steps completed and content exists 110 | if all(sequence_complete.values()) and current_chapter and chapter_content: 111 | self._save_chapter(current_chapter, chapter_content) 112 | return True 113 | 114 | return False 115 | 116 | def _prepare_chapter_context(self, chapter_number: int, prompt: str) -> str: 117 | """Prepare context for chapter generation""" 118 | if chapter_number == 1: 119 | return f"Initial Chapter\nRequirements:\n{prompt}" 120 | 121 | context_parts = [ 122 | "Previous Chapter Summaries:", 123 | *[f"Chapter {i+1}: {summary}" for i, summary in enumerate(self.chapters_memory)], 124 | "\nCurrent Chapter Requirements:", 125 | prompt 126 | ] 127 | return "\n".join(context_parts) 128 | 129 | def generate_chapter(self, chapter_number: int, prompt: str) -> None: 130 | """Generate a single chapter with completion verification""" 131 | print(f"\nGenerating Chapter {chapter_number}...") 132 | 133 | try: 134 | # Create group chat with reduced rounds 135 | groupchat = self.initiate_group_chat() 136 | manager = autogen.GroupChatManager( 137 | groupchat=groupchat, 138 | llm_config=self.agent_config 139 | ) 140 | 141 | # Prepare context 142 | context = self._prepare_chapter_context(chapter_number, prompt) 143 | chapter_prompt = f""" 144 | IMPORTANT: Wait for confirmation before proceeding. 145 | IMPORTANT: This is Chapter {chapter_number}. Do not proceed to next chapter until explicitly instructed. 146 | DO NOT END THE STORY HERE unless this is actually the final chapter ({self.outline[-1]['chapter_number']}). 147 | 148 | Current Task: Generate Chapter {chapter_number} content only. 149 | 150 | Chapter Outline: 151 | Title: {self.outline[chapter_number - 1]['title']} 152 | 153 | Chapter Requirements: 154 | {prompt} 155 | 156 | Previous Context for Reference: 157 | {context} 158 | 159 | Follow this exact sequence for Chapter {chapter_number} only: 160 | 161 | 1. Memory Keeper: Context (MEMORY UPDATE) 162 | 2. Writer: Draft (CHAPTER) 163 | 3. Editor: Review (FEEDBACK) 164 | 4. Writer Final: Revision (CHAPTER FINAL) 165 | 166 | Wait for each step to complete before proceeding.""" 167 | 168 | # Start generation 169 | self.agents["user_proxy"].initiate_chat( 170 | manager, 171 | message=chapter_prompt 172 | ) 173 | 174 | if not self._verify_chapter_complete(groupchat.messages): 175 | raise ValueError(f"Chapter {chapter_number} generation incomplete") 176 | 177 | self._process_chapter_results(chapter_number, groupchat.messages) 178 | chapter_file = os.path.join(self.output_dir, f"chapter_{chapter_number:02d}.txt") 179 | if not os.path.exists(chapter_file): 180 | raise FileNotFoundError(f"Chapter {chapter_number} file not created") 181 | 182 | completion_msg = f"Chapter {chapter_number} is complete. Proceed with next chapter." 183 | self.agents["user_proxy"].send(completion_msg, manager) 184 | 185 | except Exception as e: 186 | print(f"Error in chapter {chapter_number}: {str(e)}") 187 | self._handle_chapter_generation_failure(chapter_number, prompt) 188 | 189 | def _extract_final_scene(self, messages: List[Dict]) -> Optional[str]: 190 | """Extract chapter content with improved content detection""" 191 | for msg in reversed(messages): 192 | content = msg.get("content", "") 193 | sender = self._get_sender(msg) 194 | 195 | if sender in ["writer", "writer_final"]: 196 | # Handle complete scene content 197 | if "SCENE FINAL:" in content: 198 | scene_text = content.split("SCENE FINAL:")[1].strip() 199 | if scene_text: 200 | return scene_text 201 | 202 | # Fallback to scene content 203 | if "SCENE:" in content: 204 | scene_text = content.split("SCENE:")[1].strip() 205 | if scene_text: 206 | return scene_text 207 | 208 | # Handle raw content 209 | if len(content.strip()) > 100: # Minimum content threshold 210 | return content.strip() 211 | 212 | return None 213 | 214 | def _handle_chapter_generation_failure(self, chapter_number: int, prompt: str) -> None: 215 | """Handle failed chapter generation with simplified retry""" 216 | print(f"Attempting simplified retry for Chapter {chapter_number}...") 217 | 218 | try: 219 | # Create a new group chat with just essential agents 220 | retry_groupchat = autogen.GroupChat( 221 | agents=[ 222 | self.agents["user_proxy"], 223 | self.agents["story_planner"], 224 | self.agents["writer"] 225 | ], 226 | messages=[], 227 | max_round=3 228 | ) 229 | 230 | manager = autogen.GroupChatManager( 231 | groupchat=retry_groupchat, 232 | llm_config=self.agent_config 233 | ) 234 | 235 | retry_prompt = f"""Emergency chapter generation for Chapter {chapter_number}. 236 | 237 | {prompt} 238 | 239 | Please generate this chapter in two steps: 240 | 1. Story Planner: Create a basic outline (tag: PLAN) 241 | 2. Writer: Write the complete chapter (tag: SCENE FINAL) 242 | 243 | Keep it simple and direct.""" 244 | 245 | self.agents["user_proxy"].initiate_chat( 246 | manager, 247 | message=retry_prompt 248 | ) 249 | 250 | # Save the retry results 251 | self._process_chapter_results(chapter_number, retry_groupchat.messages) 252 | 253 | except Exception as e: 254 | print(f"Error in retry attempt for Chapter {chapter_number}: {str(e)}") 255 | print("Unable to generate chapter content after retry") 256 | 257 | def _process_chapter_results(self, chapter_number: int, messages: List[Dict]) -> None: 258 | """Process and save chapter results, updating memory""" 259 | try: 260 | # Extract the Memory Keeper's final summary 261 | memory_updates = [] 262 | for msg in reversed(messages): 263 | sender = self._get_sender(msg) 264 | content = msg.get("content", "") 265 | 266 | if sender == "memory_keeper" and "MEMORY UPDATE:" in content: 267 | update_start = content.find("MEMORY UPDATE:") + 14 268 | memory_updates.append(content[update_start:].strip()) 269 | break 270 | 271 | # Add to memory even if no explicit update (use basic content summary) 272 | if memory_updates: 273 | self.chapters_memory.append(memory_updates[0]) 274 | else: 275 | # Create basic memory from chapter content 276 | chapter_content = self._extract_final_scene(messages) 277 | if chapter_content: 278 | basic_summary = f"Chapter {chapter_number} Summary: {chapter_content[:200]}..." 279 | self.chapters_memory.append(basic_summary) 280 | 281 | # Extract and save the chapter content 282 | self._save_chapter(chapter_number, messages) 283 | 284 | except Exception as e: 285 | print(f"Error processing chapter results: {str(e)}") 286 | raise 287 | 288 | def _save_chapter(self, chapter_number: int, messages: List[Dict]) -> None: 289 | print(f"\nSaving Chapter {chapter_number}") 290 | try: 291 | chapter_content = self._extract_final_scene(messages) 292 | if not chapter_content: 293 | raise ValueError(f"No content found for Chapter {chapter_number}") 294 | 295 | chapter_content = self._clean_chapter_content(chapter_content) 296 | 297 | filename = os.path.join(self.output_dir, f"chapter_{chapter_number:02d}.txt") 298 | 299 | # Create backup if file exists 300 | if os.path.exists(filename): 301 | backup_filename = f"{filename}.backup" 302 | import shutil 303 | shutil.copy2(filename, backup_filename) 304 | 305 | with open(filename, "w", encoding='utf-8') as f: 306 | f.write(f"Chapter {chapter_number}\n\n{chapter_content}") 307 | 308 | # Verify file 309 | with open(filename, "r", encoding='utf-8') as f: 310 | saved_content = f.read() 311 | if len(saved_content.strip()) == 0: 312 | raise IOError(f"File {filename} is empty") 313 | 314 | print(f"✓ Saved to: {filename}") 315 | 316 | except Exception as e: 317 | print(f"Error saving chapter: {str(e)}") 318 | raise 319 | 320 | def generate_book(self, outline: List[Dict]) -> None: 321 | """Generate the book with strict chapter sequencing""" 322 | print("\nStarting Book Generation...") 323 | print(f"Total chapters: {len(outline)}") 324 | 325 | # Sort outline by chapter number 326 | sorted_outline = sorted(outline, key=lambda x: x["chapter_number"]) 327 | 328 | for chapter in sorted_outline: 329 | chapter_number = chapter["chapter_number"] 330 | 331 | # Verify previous chapter exists and is valid 332 | if chapter_number > 1: 333 | prev_file = os.path.join(self.output_dir, f"chapter_{chapter_number-1:02d}.txt") 334 | if not os.path.exists(prev_file): 335 | print(f"Previous chapter {chapter_number-1} not found. Stopping.") 336 | break 337 | 338 | # Verify previous chapter content 339 | with open(prev_file, 'r', encoding='utf-8') as f: 340 | content = f.read() 341 | if not self._verify_chapter_content(content, chapter_number-1): 342 | print(f"Previous chapter {chapter_number-1} content invalid. Stopping.") 343 | break 344 | 345 | # Generate current chapter 346 | print(f"\n{'='*20} Chapter {chapter_number} {'='*20}") 347 | self.generate_chapter(chapter_number, chapter["prompt"]) 348 | 349 | # Verify current chapter 350 | chapter_file = os.path.join(self.output_dir, f"chapter_{chapter_number:02d}.txt") 351 | if not os.path.exists(chapter_file): 352 | print(f"Failed to generate chapter {chapter_number}") 353 | break 354 | 355 | with open(chapter_file, 'r', encoding='utf-8') as f: 356 | content = f.read() 357 | if not self._verify_chapter_content(content, chapter_number): 358 | print(f"Chapter {chapter_number} content invalid") 359 | break 360 | 361 | print(f"✓ Chapter {chapter_number} complete") 362 | time.sleep(5) 363 | 364 | def _verify_chapter_content(self, content: str, chapter_number: int) -> bool: 365 | """Verify chapter content is valid""" 366 | if not content: 367 | return False 368 | 369 | # Check for chapter header 370 | if f"Chapter {chapter_number}" not in content: 371 | return False 372 | 373 | # Ensure content isn't just metadata 374 | lines = content.split('\n') 375 | content_lines = [line for line in lines if line.strip() and 'MEMORY UPDATE:' not in line] 376 | 377 | return len(content_lines) >= 3 # At least chapter header + 2 content lines -------------------------------------------------------------------------------- /config.py: -------------------------------------------------------------------------------- 1 | """Configuration for the book generation system""" 2 | import os 3 | from typing import Dict, List 4 | 5 | def get_config(local_url: str = "http://localhost:1234/v1") -> Dict: 6 | """Get the configuration for the agents""" 7 | 8 | # Basic config for local LLM 9 | config_list = [{ 10 | 'model': 'Mistral-Nemo-Instruct-2407', 11 | 'base_url': local_url, 12 | 'api_key': "not-needed" 13 | }] 14 | 15 | # Common configuration for all agents 16 | agent_config = { 17 | "seed": 42, 18 | "temperature": 0.7, 19 | "config_list": config_list, 20 | "timeout": 600, 21 | "cache_seed": None 22 | } 23 | 24 | return agent_config -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | """Main script for running the book generation system""" 2 | from config import get_config 3 | from agents import BookAgents 4 | from book_generator import BookGenerator 5 | from outline_generator import OutlineGenerator 6 | 7 | def main(): 8 | # Get configuration 9 | agent_config = get_config() 10 | 11 | 12 | # Initial prompt for the book 13 | initial_prompt = """ 14 | Create a story in my established writing style with these key elements: 15 | Its important that it has several key storylines that intersect and influence each other. The story should be set in a modern corporate environment, with a focus on technology and finance. The protagonist is a software engineer named Dane who has just completed a groundbreaking stock prediction algorithm. The algorithm predicts a catastrophic market crash, but Dane oversleeps and must rush to an important presentation to share his findings with executives. The tension arises from the questioning of whether his "error" might actually be correct. 16 | 17 | The piece is written in third-person limited perspective, following Dane's thoughts and experiences. The prose is direct and technical when describing the protagonist's work, but becomes more introspective during personal moments. The author employs a mix of dialogue and internal monologue, with particular attention to time progression and technical details around the algorithm and stock predictions. 18 | Story Arch: 19 | 20 | Setup: Dane completes a groundbreaking stock prediction algorithm late at night 21 | Initial Conflict: The algorithm predicts a catastrophic market crash 22 | Rising Action: Dane oversleeps and must rush to an important presentation 23 | Climax: The presentation to executives where he must explain his findings 24 | Tension Point: The questioning of whether his "error" might actually be correct 25 | 26 | Characters: 27 | 28 | Dane: The protagonist; a dedicated software engineer who prioritizes work over personal life. Wears grey polo shirts on Thursdays, tends to get lost in his work, and struggles with work-life balance. More comfortable with code than public speaking. 29 | Gary: Dane's nervous boss who seems caught between supporting Dane and managing upper management's expectations 30 | Jonathan Morego: Senior VP of Investor Relations who raises pointed questions about the validity of Dane's predictions 31 | Silence: Brief mention as an Uber driver 32 | C-Level Executives: Present as an audience during the presentation 33 | 34 | World Description: 35 | The story takes place in a contemporary corporate setting, likely a financial technology company. The world appears to be our modern one, with familiar elements like: 36 | 37 | Major tech companies (Tesla, Google, Apple, Microsoft) 38 | Stock market and financial systems 39 | Modern technology (neural networks, predictive analytics) 40 | Urban environment with rideshare services like Uber 41 | Corporate hierarchy and office culture 42 | 43 | The story creates tension between the familiar corporate world and the potential for an unprecedented financial catastrophe, blending elements of technical thriller with workplace drama. The setting feels grounded in reality but hints at potentially apocalyptic economic consequences. 44 | """ 45 | 46 | num_chapters = 25 47 | # Create agents 48 | outline_agents = BookAgents(agent_config) 49 | agents = outline_agents.create_agents(initial_prompt, num_chapters) 50 | 51 | # Generate the outline 52 | outline_gen = OutlineGenerator(agents, agent_config) 53 | print("Generating book outline...") 54 | outline = outline_gen.generate_outline(initial_prompt, num_chapters) 55 | 56 | # Create new agents with outline context 57 | book_agents = BookAgents(agent_config, outline) 58 | agents_with_context = book_agents.create_agents(initial_prompt, num_chapters) 59 | 60 | # Initialize book generator with contextual agents 61 | book_gen = BookGenerator(agents_with_context, agent_config, outline) 62 | 63 | # Print the generated outline 64 | print("\nGenerated Outline:") 65 | for chapter in outline: 66 | print(f"\nChapter {chapter['chapter_number']}: {chapter['title']}") 67 | print("-" * 50) 68 | print(chapter['prompt']) 69 | 70 | # Save the outline for reference 71 | print("\nSaving outline to file...") 72 | with open("book_output/outline.txt", "w") as f: 73 | for chapter in outline: 74 | f.write(f"\nChapter {chapter['chapter_number']}: {chapter['title']}\n") 75 | f.write("-" * 50 + "\n") 76 | f.write(chapter['prompt'] + "\n") 77 | 78 | # Generate the book using the outline 79 | print("\nGenerating book chapters...") 80 | if outline: 81 | book_gen.generate_book(outline) 82 | else: 83 | print("Error: No outline was generated.") 84 | 85 | if __name__ == "__main__": 86 | main() -------------------------------------------------------------------------------- /outline_generator.py: -------------------------------------------------------------------------------- 1 | """Generate book outlines using AutoGen agents with improved error handling""" 2 | import autogen 3 | from typing import Dict, List 4 | import re 5 | 6 | class OutlineGenerator: 7 | def __init__(self, agents: Dict[str, autogen.ConversableAgent], agent_config: Dict): 8 | self.agents = agents 9 | self.agent_config = agent_config 10 | 11 | def generate_outline(self, initial_prompt: str, num_chapters: int = 25) -> List[Dict]: 12 | """Generate a book outline based on initial prompt""" 13 | print("\nGenerating outline...") 14 | 15 | 16 | groupchat = autogen.GroupChat( 17 | agents=[ 18 | self.agents["user_proxy"], 19 | self.agents["story_planner"], 20 | self.agents["world_builder"], 21 | self.agents["outline_creator"] 22 | ], 23 | messages=[], 24 | max_round=4, 25 | speaker_selection_method="round_robin" 26 | ) 27 | 28 | manager = autogen.GroupChatManager(groupchat=groupchat, llm_config=self.agent_config) 29 | 30 | outline_prompt = f"""Let's create a {num_chapters}-chapter outline for a book with the following premise: 31 | 32 | {initial_prompt} 33 | 34 | Process: 35 | 1. Story Planner: Create a high-level story arc and major plot points 36 | 2. World Builder: Suggest key settings and world elements needed 37 | 3. Outline Creator: Generate a detailed outline with chapter titles and prompts 38 | 39 | Start with Chapter 1 and number chapters sequentially. 40 | 41 | Make sure there are at least 3 scenes in each chapter. 42 | 43 | [Continue with remaining chapters] 44 | 45 | Please output all chapters, do not leave out any chapters. Think through every chapter carefully, none should be to be determined later 46 | It is of utmost importance that you detail out every chapter, do not combine chapters, or leave any out 47 | There should be clear content for each chapter. There should be a total of {num_chapters} chapters. 48 | 49 | End the outline with 'END OF OUTLINE'""" 50 | 51 | try: 52 | # Initiate the chat 53 | self.agents["user_proxy"].initiate_chat( 54 | manager, 55 | message=outline_prompt 56 | ) 57 | 58 | # Extract the outline from the chat messages 59 | return self._process_outline_results(groupchat.messages, num_chapters) 60 | 61 | except Exception as e: 62 | print(f"Error generating outline: {str(e)}") 63 | # Try to salvage any outline content we can find 64 | return self._emergency_outline_processing(groupchat.messages, num_chapters) 65 | 66 | def _get_sender(self, msg: Dict) -> str: 67 | """Helper to get sender from message regardless of format""" 68 | return msg.get("sender") or msg.get("name", "") 69 | 70 | def _extract_outline_content(self, messages: List[Dict]) -> str: 71 | """Extract outline content from messages with better error handling""" 72 | print("Searching for outline content in messages...") 73 | 74 | # Look for content between "OUTLINE:" and "END OF OUTLINE" 75 | for msg in reversed(messages): 76 | content = msg.get("content", "") 77 | if "OUTLINE:" in content: 78 | # Extract content between OUTLINE: and END OF OUTLINE 79 | start_idx = content.find("OUTLINE:") 80 | end_idx = content.find("END OF OUTLINE") 81 | 82 | if start_idx != -1: 83 | if end_idx != -1: 84 | return content[start_idx:end_idx].strip() 85 | else: 86 | # If no END OF OUTLINE marker, take everything after OUTLINE: 87 | return content[start_idx:].strip() 88 | 89 | # Fallback: look for content with chapter markers 90 | for msg in reversed(messages): 91 | content = msg.get("content", "") 92 | if "Chapter 1:" in content or "**Chapter 1:**" in content: 93 | return content 94 | 95 | return "" 96 | 97 | def _process_outline_results(self, messages: List[Dict], num_chapters: int) -> List[Dict]: 98 | """Extract and process the outline with strict format requirements""" 99 | outline_content = self._extract_outline_content(messages) 100 | 101 | if not outline_content: 102 | print("No structured outline found, attempting emergency processing...") 103 | return self._emergency_outline_processing(messages, num_chapters) 104 | 105 | chapters = [] 106 | chapter_sections = re.split(r'Chapter \d+:', outline_content) 107 | 108 | for i, section in enumerate(chapter_sections[1:], 1): # Skip first empty section 109 | try: 110 | # Extract required components 111 | title_match = re.search(r'\*?\*?Title:\*?\*?\s*(.+?)(?=\n|$)', section, re.IGNORECASE) 112 | events_match = re.search(r'\*?\*?Key Events:\*?\*?\s*(.*?)(?=\*?\*?Character Developments:|$)', section, re.DOTALL | re.IGNORECASE) 113 | character_match = re.search(r'\*?\*?Character Developments:\*?\*?\s*(.*?)(?=\*?\*?Setting:|$)', section, re.DOTALL | re.IGNORECASE) 114 | setting_match = re.search(r'\*?\*?Setting:\*?\*?\s*(.*?)(?=\*?\*?Tone:|$)', section, re.DOTALL | re.IGNORECASE) 115 | tone_match = re.search(r'\*?\*?Tone:\*?\*?\s*(.*?)(?=\*?\*?Chapter \d+:|$)', section, re.DOTALL | re.IGNORECASE) 116 | 117 | # If no explicit title match, try to get it from the chapter header 118 | if not title_match: 119 | title_match = re.search(r'\*?\*?Chapter \d+:\s*(.+?)(?=\n|$)', section) 120 | 121 | # Verify all components exist 122 | if not all([title_match, events_match, character_match, setting_match, tone_match]): 123 | print(f"Missing required components in Chapter {i}") 124 | missing = [] 125 | if not title_match: missing.append("Title") 126 | if not events_match: missing.append("Key Events") 127 | if not character_match: missing.append("Character Developments") 128 | if not setting_match: missing.append("Setting") 129 | if not tone_match: missing.append("Tone") 130 | print(f" Missing: {', '.join(missing)}") 131 | continue 132 | 133 | # Format chapter content 134 | chapter_info = { 135 | "chapter_number": i, 136 | "title": title_match.group(1).strip(), 137 | "prompt": "\n".join([ 138 | f"- Key Events: {events_match.group(1).strip()}", 139 | f"- Character Developments: {character_match.group(1).strip()}", 140 | f"- Setting: {setting_match.group(1).strip()}", 141 | f"- Tone: {tone_match.group(1).strip()}" 142 | ]) 143 | } 144 | 145 | # Verify events (at least 3) 146 | events = re.findall(r'-\s*(.+?)(?=\n|$)', events_match.group(1)) 147 | if len(events) < 3: 148 | print(f"Chapter {i} has fewer than 3 events") 149 | continue 150 | 151 | chapters.append(chapter_info) 152 | 153 | except Exception as e: 154 | print(f"Error processing Chapter {i}: {str(e)}") 155 | continue 156 | 157 | # If we don't have enough valid chapters, raise error to trigger retry 158 | if len(chapters) < num_chapters: 159 | raise ValueError(f"Only processed {len(chapters)} valid chapters out of {num_chapters} required") 160 | 161 | return chapters 162 | 163 | def _verify_chapter_sequence(self, chapters: List[Dict], num_chapters: int) -> List[Dict]: 164 | """Verify and fix chapter numbering""" 165 | # Sort chapters by their current number 166 | chapters.sort(key=lambda x: x['chapter_number']) 167 | 168 | # Renumber chapters sequentially starting from 1 169 | for i, chapter in enumerate(chapters, 1): 170 | chapter['chapter_number'] = i 171 | 172 | # Add placeholder chapters if needed 173 | while len(chapters) < num_chapters: 174 | next_num = len(chapters) + 1 175 | chapters.append({ 176 | 'chapter_number': next_num, 177 | 'title': f'Chapter {next_num}', 178 | 'prompt': '- Key events: [To be determined]\n- Character developments: [To be determined]\n- Setting: [To be determined]\n- Tone: [To be determined]' 179 | }) 180 | 181 | # Trim excess chapters if needed 182 | chapters = chapters[:num_chapters] 183 | 184 | return chapters 185 | 186 | def _emergency_outline_processing(self, messages: List[Dict], num_chapters: int) -> List[Dict]: 187 | """Emergency processing when normal outline extraction fails""" 188 | print("Attempting emergency outline processing...") 189 | 190 | chapters = [] 191 | current_chapter = None 192 | 193 | # Look through all messages for any chapter content 194 | for msg in messages: 195 | content = msg.get("content", "") 196 | lines = content.split('\n') 197 | 198 | for line in lines: 199 | # Look for chapter markers 200 | chapter_match = re.search(r'Chapter (\d+)', line) 201 | if chapter_match and "Key events:" in content: 202 | if current_chapter: 203 | chapters.append(current_chapter) 204 | 205 | current_chapter = { 206 | 'chapter_number': int(chapter_match.group(1)), 207 | 'title': line.split(':')[-1].strip() if ':' in line else f"Chapter {chapter_match.group(1)}", 208 | 'prompt': [] 209 | } 210 | 211 | # Collect bullet points 212 | if current_chapter and line.strip().startswith('-'): 213 | current_chapter['prompt'].append(line.strip()) 214 | 215 | # Add the last chapter if it exists 216 | if current_chapter and current_chapter.get('prompt'): 217 | current_chapter['prompt'] = '\n'.join(current_chapter['prompt']) 218 | chapters.append(current_chapter) 219 | current_chapter = None 220 | 221 | if not chapters: 222 | print("Emergency processing failed to find any chapters") 223 | # Create a basic outline structure 224 | chapters = [ 225 | { 226 | 'chapter_number': i, 227 | 'title': f'Chapter {i}', 228 | 'prompt': '- Key events: [To be determined]\n- Character developments: [To be determined]\n- Setting: [To be determined]\n- Tone: [To be determined]' 229 | } 230 | for i in range(1, num_chapters + 1) 231 | ] 232 | 233 | # Ensure proper sequence and number of chapters 234 | return self._verify_chapter_sequence(chapters, num_chapters) -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | # Core dependencies 2 | autogen>=0.2.0 3 | typing>=3.7.4 4 | 5 | # Development dependencies 6 | pytest>=7.0.0 # For testing 7 | black>=22.0.0 # For code formatting 8 | flake8>=4.0.0 # For linting 9 | mypy>=0.910 # For type checking 10 | 11 | # Optional dependencies 12 | tqdm>=4.65.0 # For progress bars 13 | python-dotenv>=0.19.0 # For environment variable management --------------------------------------------------------------------------------