selfapi_writer.py

from anthropic import Anthropic
import streamlit as st
import json
import os
import tiktoken
from typing import Dict, Any, Optional, List, Tuple, Generator
from content_state import ContentState

class SelfApiWriter:
    def __init__(self):
        """Initialize the Self.api writer with enhanced content tracking"""
        ANTHROPIC_API_KEY = os.getenv('api_key')
            
        if not ANTHROPIC_API_KEY:
            raise ValueError("Anthropic API key not found. Please ensure ANTHROPIC_API_KEY is set.")
        
        self.client = Anthropic(api_key=ANTHROPIC_API_KEY)
        self.model = "claude-3-opus-20240229"
        self.context = {}
        self.book_structure = None
        self.writing_guidelines = None
        self.initialized = False

        # Configuration for generation
        self.pages_per_chapter = 25
        self.words_per_page = 250
        self.max_iterations = 10
        self.max_tokens = 15000

        # Token encoding
        self.tokenizer = tiktoken.encoding_for_model("gpt-4")
        
        # Add content state tracking
        self.content_states = {}

    def _initialize_content_state(self, content_id: str) -> None:
        """Initialize a new content state tracker"""
        if content_id not in self.content_states:
            self.content_states[content_id] = ContentState()
    
    def set_manual_content(self, content_id: str, content: str) -> None:
        """Set manual content for a specific section"""
        if content_id not in self.content_states:
            self._initialize_content_state(content_id)
        self.content_states[content_id].set_manual_content(content)
        
        # If the content is for introduction, also store in context
        if content_id == 'introduction':
            if 'manual_content' not in self.context:
                self.context['manual_content'] = {}
            self.context['manual_content']['introduction'] = content

    def _truncate_blueprint(self, blueprint: str, max_tokens: int = 15000) -> Tuple[str, str]:
        """Intelligently truncate the blueprint to fit within token limits"""
        tokens = self.tokenizer.encode(blueprint)
        
        if len(tokens) <= max_tokens:
            return blueprint, ""
        
        truncated_tokens = tokens[:max_tokens]
        truncated_blueprint = self.tokenizer.decode(truncated_tokens)
        
        try:
            overview_response = self.client.messages.create(
                model=self.model,
                max_tokens=1000,
                system="You are an expert at creating concise summaries of book blueprints.",
                messages=[{
                    "role": "user",
                    "content": f"""The following blueprint was truncated due to length constraints. 
                    Please create a comprehensive overview that captures the essence of the 
                    truncated sections:

                    Truncated Blueprint Ending:
                    {blueprint[len(truncated_blueprint):]}

                    Provide a summary that:
                    1. Captures key themes and intentions
                    2. Highlights main sections that were cut off
                    3. Ensures no critical information is lost
                    4. Is concise but comprehensive"""
                }]
            )
            overview_summary = overview_response.content[0].text
        except Exception as e:
            overview_summary = f"Note: Some blueprint content was truncated. Original blueprint exceeded {max_tokens} tokens."
        
        return truncated_blueprint, overview_summary

    def _generate_section_outline(self, content_id: str, section_type: str, title: str) -> List[str]:
        """Generate detailed outline for a section before writing"""
        state = self.content_states[content_id]
        
        outline_prompt = f"""Based on the current progress:
        Previous Summary: {state.current_summary}
        Key Points Covered: {', '.join(state.key_points_covered)}
        
        Create a detailed outline for {section_type}: "{title}" that:
        1. Builds on previously covered material
        2. Introduces new concepts progressively
        3. Maintains narrative continuity
        4. Plans clear transitions between subsections
        
        Return the outline as a list of specific points to cover."""

        response = self.client.messages.create(
            model=self.model,
            max_tokens=1000,
            temperature=0.5,
            messages=[{"role": "user", "content": outline_prompt}]
        )
        
        outline = [point.strip() for point in response.content[0].text.split('\n') if point.strip()]
        state.section_outlines = outline
        return outline

    def process_blueprint(self, blueprint: str) -> Dict[str, Any]:
        """Process blueprint to extract structure and guidelines"""
        try:
            with st.spinner("Processing blueprint..."):
                truncated_blueprint, overview_summary = self._truncate_blueprint(blueprint)

                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"]
                    }
                }"""

                response = self.client.messages.create(
                    model=self.model,
                    max_tokens=4000,
                    temperature=0,
                    system=system_prompt,
                    messages=[{
                        "role": "user", 
                        "content": f"""Analyze this book blueprint and extract ALL information:

                        {truncated_blueprint}

                        {overview_summary}

                        Return only the JSON structure without any additional text."""
                    }]
                )
                
                extracted_info = json.loads(response.content[0].text)
                extracted_info['full_original_blueprint'] = blueprint
                
                self.book_info = extracted_info["book_info"]
                self.book_structure = extracted_info["structure"]
                self.writing_guidelines = extracted_info["guidelines"]
                self.initialized = True
                self.context['full_original_blueprint'] = blueprint
                
                return extracted_info
            
        except Exception as e:
            st.error(f"Error processing blueprint: {str(e)}")
            return None


    def write_introduction(self, additional_prompt: str = "") -> str:
        """Generate the book's introduction with enhanced continuity"""
        if not self.initialized:
            raise ValueError("Writer not initialized. Process blueprint first.")
            
        content_id = "introduction"
        self._initialize_content_state(content_id)
            
        def generate_intro_iteration(iteration: int,
                                        previous_summary: str,
                                        points_to_cover: List[str],
                                        narrative_threads: List[str]) -> str:
            """Generate a single iteration of the introduction"""
            full_blueprint = self.context.get('full_original_blueprint', '')

            system_prompt = f"""You are writing the introduction for '{self.book_info.get('title', 'Untitled Book')}'
            
            Previous Content Summary: {previous_summary}
            Points to Cover in This Section: {', '.join(points_to_cover)}
            Active Narrative Threads: {', '.join(narrative_threads)}
            
            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')}
            
            Additional Instructions: {additional_prompt}
            
            Write the introduction following these guidelines."""

            response = self.client.messages.create(
                model=self.model,
                max_tokens=2000,
                temperature=0.7,
                system=system_prompt,
                messages=[{
                    "role": "user",
                    "content": f"""Write the next section of the introduction, building on:
                    Previous Summary: {previous_summary}
                    Points to Cover: {', '.join(points_to_cover)}"""
                }]
            )
            
            return response.content[0].text

        full_intro_content = self._generate_with_continuity(
            generate_intro_iteration,
            content_id,
            "Introduction"
        )
        
        self.context['introduction'] = full_intro_content
        return full_intro_content

    def write_chapter(self, part_idx: int, chapter_idx: int, additional_prompt: str = "") -> str:
        """Generate a chapter using enhanced content continuity and additional prompts"""
        if not self.initialized:
            raise ValueError("Writer not initialized. Process blueprint first.")
            
        content_id = f"part_{part_idx}_chapter_{chapter_idx}"
        self._initialize_content_state(content_id)
        
        # Add any additional prompts to the content state
        if additional_prompt:
            self.content_states[content_id].add_custom_prompt(content_id, additional_prompt)
        
        def generate_chapter_iteration(iteration: int,
                                        previous_summary: str,
                                        points_to_cover: List[str],
                                        narrative_threads: List[str]) -> str:
            """Generate a single chapter iteration with enhanced context"""
            part = self.book_structure["parts"][part_idx]
            chapter_title = part["chapters"][chapter_idx]
            part_title = part["title"]
            
            # Get complete context including custom prompts
            section_context = self.content_states[content_id].get_section_context(content_id)
            
            # Enhanced system prompt with additional context
            system_prompt = f"""You are writing '{self.book_info.get('title', 'Untitled Book')}'
            Chapter: {chapter_title}
            Part: {part_title}
            
            Blueprint Context: {self.context.get('full_original_blueprint', '')}
            
            Additional Instructions: {additional_prompt}
            Custom Guidelines: {section_context.get('custom_instructions', '')}
            
            Previous Content Summary: {previous_summary}
            Points to Cover in This Section: {', '.join(points_to_cover)}
            Active Narrative Threads: {', '.join(narrative_threads)}
            
            Writing Guidelines: {json.dumps(self.writing_guidelines, indent=2)}
            
            Create content that:
            1. Builds naturally on previous sections
            2. Incorporates the additional instructions and custom guidelines
            3. Maintains consistent narrative threads
            4. Creates smooth transitions
            5. Follows all style and structure guidelines
            
            If additional instructions are provided, ensure they are seamlessly integrated
            into the content while maintaining the overall style and structure."""

            response = self.client.messages.create(
                model=self.model,
                max_tokens=2000,
                temperature=0.7,
                system=system_prompt,
                messages=[{
                    "role": "user", 
                    "content": f"Write the next section of Chapter: {chapter_title}, incorporating any additional instructions provided."
                }]
            )
            
            return response.content[0].text

        full_chapter_content = self._generate_with_continuity(
            generate_chapter_iteration,
            content_id,
            f"Chapter: {self.book_structure['parts'][part_idx]['chapters'][chapter_idx]}"
        )
        
        # Store context history
        self.content_states[content_id].update_context_history(
            content_id,
            self.content_states[content_id].get_section_context(content_id)
        )
        
        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 _generate_with_continuity(self, 
                                    generate_func: callable,
                                    content_id: str,
                                    title: str,
                                    total_steps: int = 10) -> str:
        """Enhanced generation with content continuity tracking"""
        progress_bar = st.progress(0, text=f"Generating {title}...")
        full_content = ""
        state = self.content_states[content_id]
        
        try:
            # If manual content exists, use it as a starting point
            if state.manual_content:
                full_content = state.manual_content + "\n\n"
                # Generate initial summary from manual content
                state.current_summary = self._generate_progressive_summary(
                    content_id,
                    full_content
                )
            
            # Generate initial outline
            outline = self._generate_section_outline(content_id, "section", title)
            points_per_iteration = max(1, len(outline) // total_steps)
            
            for iteration in range(1, total_steps + 1):
                progress = iteration / total_steps
                progress_bar.progress(
                    min(int(progress * 100), 100),
                    text=f"Generating {title}... (Iteration {iteration}/{total_steps})"
                )
                
                start_idx = (iteration - 1) * points_per_iteration
                end_idx = min(start_idx + points_per_iteration, len(outline))
                current_points = outline[start_idx:end_idx]
                
                new_content = generate_func(
                    iteration=iteration,
                    previous_summary=state.current_summary,
                    points_to_cover=current_points,
                    narrative_threads=state.narrative_threads
                )
                
                state.generated_sections.append(new_content)
                
                if iteration > 1:
                    transition = self._generate_transition(
                        content_id,
                        state.generated_sections[-2],
                        new_content
                    )
                    full_content += transition
                
                full_content += new_content
                
                state.current_summary = self._generate_progressive_summary(
                    content_id,
                    full_content
                )
                state.key_points_covered.update(current_points)
                
                if len(full_content.split()) > self.pages_per_chapter * self.words_per_page:
                    break
            
            conclusion = self._generate_conclusion(content_id, full_content)
            full_content += conclusion
            
            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:
            progress_bar.empty()

    def _generate_transition(self, content_id: str, prev_content: str, next_content: str) -> str:
        """Generate smooth transition between sections"""
        state = self.content_states[content_id]
        
        transition_prompt = f"""Create a smooth transition between these sections:
        
        Previous Section Summary: {self._summarize_text(prev_content)}
        Next Section Key Points: {self._summarize_text(next_content)}
        
        Create a natural bridge that:
        1. References relevant previous points
        2. Introduces upcoming concepts
        3. Maintains narrative flow
        4. Feels organic and not forced"""

        response = self.client.messages.create(
            model=self.model,
            max_tokens=300,
        temperature=0.7,
            messages=[{"role": "user", "content": transition_prompt}]
        )
        
        transition = response.content[0].text
        state.transition_points.append(transition)
        return transition

    def _generate_progressive_summary(self, content_id: str, content: str) -> str:
        """Generate a running summary of content progress"""
        summary_prompt = f"""Summarize the key points and narrative progression of:

        {content}

        Focus on:
        1. Main concepts introduced
        2. Key arguments developed
        3. Narrative threads established
        4. Important conclusions reached
        
        Keep the summary concise but comprehensive."""

        response = self.client.messages.create(
            model=self.model,
            max_tokens=500,
            temperature=0.3,
            messages=[{"role": "user", "content": summary_prompt}]
        )
        
        return response.content[0].text

    def _summarize_text(self, text: str) -> str:
        """Generate a concise summary of text"""
        response = self.client.messages.create(
            model=self.model,
            max_tokens=300,
            temperature=0.3,
            messages=[{
                "role": "user",
                "content": f"Summarize the key points from this text:\n\n{text}"
            }]
        )
        return response.content[0].text

    def _generate_conclusion(self, content_id: str, full_content: str) -> str:
        """Generate a conclusion that ties everything together"""
        state = self.content_states[content_id]
        
        conclusion_prompt = f"""Create a conclusion that ties together:

        Content Summary: {state.current_summary}
        Key Points Covered: {', '.join(state.key_points_covered)}
        Narrative Threads: {', '.join(state.narrative_threads)}
        
        The conclusion should:
        1. Summarize main arguments
        2. Connect key themes
        3. Reinforce core messages
        4. Provide closure while maintaining interest"""

        response = self.client.messages.create(
            model=self.model,
            max_tokens=500,
            temperature=0.7,
            messages=[{"role": "user", "content": conclusion_prompt}]
        )
        
        return response.content[0].text

    def get_current_structure(self) -> Optional[Dict[str, Any]]:
        """Get current book structure and guidelines"""
        if not self.initialized:
            return None
            
        return {
            "book_info": self.book_info,
            "structure": self.book_structure,
            "guidelines": self.writing_guidelines
        }