Update selfapi_writer.py

selfapi_writer.py

@@ -3,7 +3,7 @@ import streamlit as st
 import json
 import os
 import tiktoken
-from typing import Dict, Any, Optional, List, Tuple
+from typing import Dict, Any, Optional, List, Tuple, Generator
 
 class SelfApiWriter:
     def __init__(self):
@@ -31,6 +31,55 @@ class SelfApiWriter:
         # Token encoding
         self.tokenizer = tiktoken.encoding_for_model("gpt-4")
 
+    def _generate_with_progress(self, 
+                                generate_func: callable, 
+                                title: str, 
+                                total_steps: int = 20) -> str:
+        """
+        Generate content with progress tracking
+        
+        :param generate_func: Function to generate content
+        :param title: Title for progress display
+        :param total_steps: Total number of generation steps
+        :return: Generated content
+        """
+        # Create Streamlit progress bar
+        progress_bar = st.progress(0, text=f"Generating {title}...")
+        
+        # Tracking variables
+        full_content = ""
+        
+        try:
+            for iteration in range(1, total_steps + 1):
+                # Update progress
+                progress = iteration / total_steps
+                progress_bar.progress(
+                    min(int(progress * 100), 100), 
+                    text=f"Generating {title}... (Iteration {iteration}/{total_steps})"
+                )
+                
+                # Generate content
+                new_content = generate_func(iteration)
+                full_content += new_content
+                
+                # Stopping criteria
+                if (len(full_content.split()) > self.pages_per_chapter * self.words_per_page 
+                    or iteration == total_steps):
+                    break
+            
+            # Complete progress
+            progress_bar.progress(100, text=f"Finished generating {title}")
+            
+            return full_content
+        
+        except Exception as e:
+            st.error(f"Error generating {title}: {e}")
+            progress_bar.empty()
+            return f"Error generating {title}: {e}"
+        finally:
+            # Ensure progress bar is cleared
+            progress_bar.empty()
+
     def _truncate_blueprint(self, blueprint: str, max_tokens: int = 15000) -> Tuple[str, str]:
         """
         Intelligently truncate the blueprint to fit within token limits
@@ -83,79 +132,81 @@ class SelfApiWriter:
     def process_blueprint(self, blueprint: str) -> Dict[str, Any]:
         """Process blueprint to extract complete writing guidelines and structure"""
         try:
-            # Truncate blueprint if too long
-            truncated_blueprint, overview_summary = self._truncate_blueprint(blueprint)
+            # Add a spinner during blueprint processing
+            with st.spinner("Processing blueprint..."):
+                # Truncate blueprint if too long
+                truncated_blueprint, overview_summary = self._truncate_blueprint(blueprint)
 
-            # Prepare system prompt for blueprint processing
-            system_prompt = """You are an expert book planner analyzing a blueprint.
-            Extract ALL relevant information and return it in a structured format.
-            Include:
-            1. Book title and high-level information
-            2. Complete structure (introduction, parts, chapters)
-            3. All writing style guidelines
-            4. Content requirements and constraints
-            5. Target audience details
-            6. Chapter structure requirements
-            7. Tone and voice requirements
-            8. Any other relevant guidelines or requirements
-            
-            Return a JSON structure with the following format:
-            {
-                "book_info": {
-                    "title": "Book title",
-                    "vision": "Core vision/purpose",
-                    "target_audience": "Detailed audience description"
-                },
-                "structure": {
-                    "introduction": "Introduction title",
-                    "parts": [
-                        {
-                            "title": "Part title",
-                            "chapters": ["Chapter 1 title", "Chapter 2 title", ...]
-                        }
-                    ]
-                },
-                "guidelines": {
-                    "style": "Writing style description",
-                    "tone": "Tone requirements",
-                    "chapter_structure": ["Required chapter components"],
-                    "content_requirements": ["Specific content requirements"],
-                    "practical_elements": ["Required practical elements"]
-                },
-                "overview_summary": "Summary of truncated sections"
-            }"""
+                # Prepare system prompt for blueprint processing
+                system_prompt = """You are an expert book planner analyzing a blueprint.
+                Extract ALL relevant information and return it in a structured format.
+                Include:
+                1. Book title and high-level information
+                2. Complete structure (introduction, parts, chapters)
+                3. All writing style guidelines
+                4. Content requirements and constraints
+                5. Target audience details
+                6. Chapter structure requirements
+                7. Tone and voice requirements
+                8. Any other relevant guidelines or requirements
+                
+                Return a JSON structure with the following format:
+                {
+                    "book_info": {
+                        "title": "Book title",
+                        "vision": "Core vision/purpose",
+                        "target_audience": "Detailed audience description"
+                    },
+                    "structure": {
+                        "introduction": "Introduction title",
+                        "parts": [
+                            {
+                                "title": "Part title",
+                                "chapters": ["Chapter 1 title", "Chapter 2 title", ...]
+                            }
+                        ]
+                    },
+                    "guidelines": {
+                        "style": "Writing style description",
+                        "tone": "Tone requirements",
+                        "chapter_structure": ["Required chapter components"],
+                        "content_requirements": ["Specific content requirements"],
+                        "practical_elements": ["Required practical elements"]
+                    },
+                    "overview_summary": "Summary of truncated sections"
+                }"""
 
-            prompt = f"""Analyze this book blueprint and extract ALL information:
+                prompt = f"""Analyze this book blueprint and extract ALL information:
 
-            {truncated_blueprint}
+                {truncated_blueprint}
 
-            {overview_summary}
+                {overview_summary}
 
-            Return only the JSON structure without any additional text."""
+                Return only the JSON structure without any additional text."""
 
-            response = self.client.messages.create(
-                model=self.model,
-                max_tokens=4000,
-                temperature=0,
-                system=system_prompt,
-                messages=[{"role": "user", "content": prompt}]
-            )
-            
-            extracted_info = json.loads(response.content[0].text)
-            
-            # Store full original blueprint for reference
-            extracted_info['full_original_blueprint'] = blueprint
-            
-            # Store extracted information
-            self.book_info = extracted_info["book_info"]
-            self.book_structure = extracted_info["structure"]
-            self.writing_guidelines = extracted_info["guidelines"]
-            self.initialized = True
-            
-            # Store full blueprint in context
-            self.context['full_original_blueprint'] = blueprint
-            
-            return extracted_info
+                response = self.client.messages.create(
+                    model=self.model,
+                    max_tokens=4000,
+                    temperature=0,
+                    system=system_prompt,
+                    messages=[{"role": "user", "content": prompt}]
+                )
+                
+                extracted_info = json.loads(response.content[0].text)
+                
+                # Store full original blueprint for reference
+                extracted_info['full_original_blueprint'] = blueprint
+                
+                # Store extracted information
+                self.book_info = extracted_info["book_info"]
+                self.book_structure = extracted_info["structure"]
+                self.writing_guidelines = extracted_info["guidelines"]
+                self.initialized = True
+                
+                # Store full blueprint in context
+                self.context['full_original_blueprint'] = blueprint
+                
+                return extracted_info
             
         except Exception as e:
             st.error(f"Error processing blueprint: {str(e)}")
@@ -166,149 +217,137 @@ class SelfApiWriter:
         if not self.initialized:
             raise ValueError("Writer not initialized. Process blueprint first.")
             
-        try:
+        def generate_intro_iteration(iteration: int) -> str:
+            """Generate a single iteration of the introduction"""
             # Retrieve full original blueprint
             full_blueprint = self.context.get('full_original_blueprint', '')
 
-            # Iterative generation for introduction
-            full_intro_content = ""
-            for iteration in range(1, self.max_iterations + 1):
-                # Prepare system prompt with full context
-                system_prompt = f"""You are writing the introduction for '{self.book_info.get('title', 'Untitled Book')}'
-                Full Blueprint Context:
-                {full_blueprint}
-                
-                Core Vision: {self.book_info.get('vision', '')}
-                Target Audience: {self.book_info.get('target_audience', '')}
-                
-                Writing Style: {self.writing_guidelines.get('style', 'Academic and clear')}
-                Tone: {self.writing_guidelines.get('tone', 'Professional')}
-                """
-
-                response = self.client.messages.create(
-                    model=self.model,
-                    max_tokens=2000,
-                    temperature=0.7,
-                    system=system_prompt,
-                    messages=[
-                        {
-                            "role": "user",
-                            "content": f"""Write the introduction: "{self.book_structure.get('introduction', 'Book Introduction')}"
-                            
-                            Iteration {iteration} of {self.max_iterations}
-                            
-                            Create an engaging opening that:
-                            1. Introduces the book's core concept
-                            2. Speaks directly to the target audience
-                            3. Outlines the book's approach and structure
-                            4. Sets the tone for the entire book
-                            
-                            Follow ALL provided guidelines for style, tone, and content."""
-                        }
-                    ]
-                )
-                
-                new_content = response.content[0].text
-                full_intro_content += new_content
-                
-                # Stopping criteria
-                if (len(full_intro_content.split()) > self.pages_per_chapter * self.words_per_page 
-                    or iteration == self.max_iterations):
-                    break
+            # Prepare system prompt with full context
+            system_prompt = f"""You are writing the introduction for '{self.book_info.get('title', 'Untitled Book')}'
+            Full Blueprint Context:
+            {full_blueprint}
             
-            # Store and return
-            self.context['introduction'] = full_intro_content
-            return full_intro_content
+            Core Vision: {self.book_info.get('vision', '')}
+            Target Audience: {self.book_info.get('target_audience', '')}
             
-        except Exception as e:
-            st.error(f"Error generating introduction: {str(e)}")
-            return f"Error generating introduction: {str(e)}"
+            Writing Style: {self.writing_guidelines.get('style', 'Academic and clear')}
+            Tone: {self.writing_guidelines.get('tone', 'Professional')}
+            """
+
+            response = self.client.messages.create(
+                model=self.model,
+                max_tokens=2000,
+                temperature=0.7,
+                system=system_prompt,
+                messages=[
+                    {
+                        "role": "user",
+                        "content": f"""Write the introduction: "{self.book_structure.get('introduction', 'Book Introduction')}"
+                        
+                        Iteration {iteration} of {self.max_iterations}
+                        
+                        Create an engaging opening that:
+                        1. Introduces the book's core concept
+                        2. Speaks directly to the target audience
+                        3. Outlines the book's approach and structure
+                        4. Sets the tone for the entire book
+                        
+                        Follow ALL provided guidelines for style, tone, and content."""
+                    }
+                ]
+            )
+            
+            return response.content[0].text
+
+        # Generate with progress tracking
+        full_intro_content = self._generate_with_progress(
+            generate_intro_iteration, 
+            "Introduction"
+        )
+        
+        # Store and return
+        self.context['introduction'] = full_intro_content
+        return full_intro_content
 
     def write_chapter(self, part_idx: int, chapter_idx: int) -> str:
         """Generate a chapter using extracted guidelines"""
         if not self.initialized:
             raise ValueError("Writer not initialized. Process blueprint first.")
             
-        try:
+        def generate_chapter_iteration(iteration: int) -> str:
+            """Generate a single iteration of the chapter"""
             # Retrieve full original blueprint
             full_blueprint = self.context.get('full_original_blueprint', '')
 
             part = self.book_structure["parts"][part_idx]
             chapter_title = part["chapters"][chapter_idx]
             part_title = part["title"]
-            
-            # Iterative generation for chapter
-            full_chapter_content = ""
-            for iteration in range(1, self.max_iterations + 1):
-                # Prepare system prompt with full context
-                system_prompt = f"""You are writing '{self.book_info.get('title', 'Untitled Book')}'
-                Full Blueprint Context:
-                {full_blueprint}
-                
-                Chapter: {chapter_title}
-                Part: {part_title}
-                
-                Target Audience: {self.book_info.get('target_audience', '')}
-                Writing Style: {self.writing_guidelines.get('style', 'Academic and clear')}
-                Tone: {self.writing_guidelines.get('tone', 'Professional')}
-                
-                Chapter Structure Requirements:
-                {', '.join(self.writing_guidelines.get('chapter_structure', []))}
-                
-                Content Requirements:
-                {', '.join(self.writing_guidelines.get('content_requirements', []))}"""
 
-                response = self.client.messages.create(
-                    model=self.model,
-                    max_tokens=2000,
-                    temperature=0.7,
-                    system=system_prompt,
-                    messages=[
-                        {
-                            "role": "user",
-                            "content": f"""Write Chapter: "{chapter_title}" in Part {part_idx + 1}: "{part_title}"
-                            
-                            Iteration {iteration} of {self.max_iterations}
-                            
-                            Follow ALL provided guidelines for:
-                            1. Structure and organization
-                            2. Style and tone
-                            3. Practical elements and exercises
-                            4. Content depth and requirements
-                            
-                            Begin writing the complete chapter now."""
-                        }
-                    ]
-                )
-                
-                new_content = response.content[0].text
-                full_chapter_content += new_content
-                
-                # Stopping criteria
-                if (len(full_chapter_content.split()) > self.pages_per_chapter * self.words_per_page 
-                    or iteration == self.max_iterations):
-                    break
-            
-            # Store in context
-            if 'parts' not in self.context:
-                self.context['parts'] = []
+            # Prepare system prompt with full context
+            system_prompt = f"""You are writing '{self.book_info.get('title', 'Untitled Book')}'
+            Full Blueprint Context:
+            {full_blueprint}
             
-            while len(self.context['parts']) <= part_idx:
-                self.context['parts'].append({'chapters': []})
+            Chapter: {chapter_title}
+            Part: {part_title}
             
-            while len(self.context['parts'][part_idx]['chapters']) <= chapter_idx:
-                self.context['parts'][part_idx]['chapters'].append({})
+            Target Audience: {self.book_info.get('target_audience', '')}
+            Writing Style: {self.writing_guidelines.get('style', 'Academic and clear')}
+            Tone: {self.writing_guidelines.get('tone', 'Professional')}
             
-            self.context['parts'][part_idx]['chapters'][chapter_idx] = {
-                'title': chapter_title,
-                'content': full_chapter_content
-            }
+            Chapter Structure Requirements:
+            {', '.join(self.writing_guidelines.get('chapter_structure', []))}
             
-            return full_chapter_content
+            Content Requirements:
+            {', '.join(self.writing_guidelines.get('content_requirements', []))}"""
+
+            response = self.client.messages.create(
+                model=self.model,
+                max_tokens=2000,
+                temperature=0.7,
+                system=system_prompt,
+                messages=[
+                    {
+                        "role": "user",
+                        "content": f"""Write Chapter: "{chapter_title}" in Part {part_idx + 1}: "{part_title}"
+                        
+                        Iteration {iteration} of {self.max_iterations}
+                        
+                        Follow ALL provided guidelines for:
+                        1. Structure and organization
+                        2. Style and tone
+                        3. Practical elements and exercises
+                        4. Content depth and requirements
+                        
+                        Begin writing the complete chapter now."""
+                    }
+                ]
+            )
             
-        except Exception as e:
-            st.error(f"Error generating chapter: {str(e)}")
-            return f"Error generating chapter {chapter_title}: {str(e)}"
+            return response.content[0].text
+
+        # Generate with progress tracking
+        full_chapter_content = self._generate_with_progress(
+            generate_chapter_iteration, 
+            f"Chapter: {self.book_structure['parts'][part_idx]['chapters'][chapter_idx]}"
+        )
+        
+        # Store in context
+        if 'parts' not in self.context:
+            self.context['parts'] = []
+        
+        while len(self.context['parts']) <= part_idx:
+            self.context['parts'].append({'chapters': []})
+        
+        while len(self.context['parts'][part_idx]['chapters']) <= chapter_idx:
+            self.context['parts'][part_idx]['chapters'].append({})
+        
+        self.context['parts'][part_idx]['chapters'][chapter_idx] = {
+            'title': self.book_structure['parts'][part_idx]['chapters'][chapter_idx],
+            'content': full_chapter_content
+        }
+        
+        return full_chapter_content
 
     def get_current_structure(self) -> Optional[Dict[str, Any]]:
         """Get current book structure and guidelines"""