Update selfapi_writer.py

selfapi_writer.py

@@ -1,21 +1,20 @@
-# File: selfapi_writer.py
-
 from anthropic import Anthropic
 import streamlit as st
 import json
 import os
-from typing import Dict, Any, Optional
+import tiktoken
+from typing import Dict, Any, Optional, List, Tuple
 
 class SelfApiWriter:
     def __init__(self):
         """Initialize the Self.api writer"""
         # Try to get API key from environment variables first, then from secrets
-        
         ANTHROPIC_API_KEY = os.getenv('api_key')
                     
-        """Initialize Claude service with API key from HuggingFace secrets"""
+        """Initialize Claude service with API key"""
         if not ANTHROPIC_API_KEY:
-            raise ValueError("Anthropic API key not found in HuggingFace secrets. Please ensure ANTHROPIC_API_KEY is set in your space's secrets.")
+            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 = {}
@@ -23,9 +22,73 @@ class SelfApiWriter:
         self.writing_guidelines = None
         self.initialized = False
 
+        # Configuration for generation
+        self.pages_per_chapter = 70  # Approximately 35,000 words
+        self.words_per_page = 500
+        self.max_iterations = 20
+        self.max_tokens = 40000
+
+        # Token encoding
+        self.tokenizer = tiktoken.encoding_for_model("gpt-4")
+
+    def _truncate_blueprint(self, blueprint: str, max_tokens: int = 15000) -> Tuple[str, str]:
+        """
+        Intelligently truncate the blueprint to fit within token limits
+        
+        :param blueprint: Full blueprint text
+        :param max_tokens: Maximum tokens to keep
+        :return: Tuple of (truncated_blueprint, overview_summary)
+        """
+        # Tokenize the blueprint
+        tokens = self.tokenizer.encode(blueprint)
+        
+        # If within token limit, return full blueprint
+        if len(tokens) <= max_tokens:
+            return blueprint, ""
+        
+        # Try to preserve key sections intelligently
+        truncated_tokens = tokens[:max_tokens]
+        truncated_blueprint = self.tokenizer.decode(truncated_tokens)
+        
+        # Generate a summary of the truncated sections
+        try:
+            overview_response = self.client.messages.create(
+                model=self.model,
+                max_tokens=1000,
+                messages=[
+                    {
+                        "role": "system",
+                        "content": "You are an expert at creating concise summaries of book blueprints."
+                    },
+                    {
+                        "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 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)
+
             system_prompt = """You are an expert book planner analyzing a blueprint.
             Extract ALL relevant information and return it in a structured format.
             Include:
@@ -60,17 +123,19 @@ class SelfApiWriter:
                     "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:
 
-            {blueprint}
+            {truncated_blueprint}
+
+            {overview_summary}
 
             Return only the JSON structure without any additional text."""
 
             response = self.client.messages.create(
-                #model="claude-3-sonnet-20240229",
                 model=self.model,
                 max_tokens=4000,
                 temperature=0,
@@ -80,6 +145,9 @@ class SelfApiWriter:
             
             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"]
@@ -93,42 +161,63 @@ class SelfApiWriter:
             return None
 
     def write_introduction(self) -> str:
-        """Generate the book's introduction based on extracted guidelines"""
+        """Generate the book's introduction"""
         if not self.initialized:
             raise ValueError("Writer not initialized. Process blueprint first.")
             
         try:
-            system_prompt = f"""You are writing the introduction for '{self.book_info["title"]}'
-            Core Vision: {self.book_info["vision"]}
-            Target Audience: {self.book_info["target_audience"]}
-            
-            Writing Style: {self.writing_guidelines["style"]}
-            Tone: {self.writing_guidelines["tone"]}
-            
-            Content Requirements:
-            {', '.join(self.writing_guidelines["content_requirements"])}"""
+            # Retrieve full original blueprint
+            full_blueprint = self.context.get('full_original_blueprint', '')
 
-            intro_prompt = f"""Write the introduction: "{self.book_structure['introduction']}"
-            
-            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."""
-
-            response = self.client.messages.create(
-                model=self.model,
-                max_tokens=20000,
-                temperature=0.7,
-                system=system_prompt,
-                messages=[{"role": "user", "content": intro_prompt}]
-            )
+            # Iterative generation for introduction
+            full_intro_content = ""
+            for iteration in range(1, self.max_iterations + 1):
+                response = self.client.messages.create(
+                    model=self.model,
+                    max_tokens=2000,
+                    temperature=0.7,
+                    messages=[
+                        {
+                            "role": "system",
+                            "content": 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')}
+                            """
+                        },
+                        {
+                            "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
             
-            intro_content = response.content[0].text
-            self.context['introduction'] = intro_content
-            return intro_content
+            # Store and return
+            self.context['introduction'] = full_intro_content
+            return full_intro_content
             
         except Exception as e:
             st.error(f"Error generating introduction: {str(e)}")
@@ -140,49 +229,81 @@ class SelfApiWriter:
             raise ValueError("Writer not initialized. Process blueprint first.")
             
         try:
+            # 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"]
-            previous_chapter = self.context.get(f'part_{part_idx}_chapter_{chapter_idx-1}', '')
-
-            system_prompt = f"""You are writing '{self.book_info["title"]}'
-            Target Audience: {self.book_info["target_audience"]}
-            Writing Style: {self.writing_guidelines["style"]}
-            Tone: {self.writing_guidelines["tone"]}
             
-            Chapter Structure Requirements:
-            {', '.join(self.writing_guidelines["chapter_structure"])}
+            # Iterative generation for chapter
+            full_chapter_content = ""
+            for iteration in range(1, self.max_iterations + 1):
+                response = self.client.messages.create(
+                    model=self.model,
+                    max_tokens=2000,
+                    temperature=0.7,
+                    messages=[
+                        {
+                            "role": "system",
+                            "content": 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', []))}"""
+                        },
+                        {
+                            "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
             
-            Content Requirements:
-            {', '.join(self.writing_guidelines["content_requirements"])}
+            # Store in context
+            if 'parts' not in self.context:
+                self.context['parts'] = []
             
-            Practical Elements to Include:
-            {', '.join(self.writing_guidelines["practical_elements"])}"""
-
-            chapter_prompt = f"""
-            Write Chapter: "{chapter_title}" in Part {part_idx + 1}: "{part_title}"
+            while len(self.context['parts']) <= part_idx:
+                self.context['parts'].append({'chapters': []})
             
-            Previous Chapter Context: {previous_chapter[:1000] if previous_chapter else 'Starting new part'}
-
-            Follow ALL provided guidelines for:
-            1. Structure and organization
-            2. Style and tone
-            3. Practical elements and exercises
-            4. Content depth and requirements
+            while len(self.context['parts'][part_idx]['chapters']) <= chapter_idx:
+                self.context['parts'][part_idx]['chapters'].append({})
             
-            Begin writing the complete chapter now."""
-
-            response = self.client.messages.create(
-                model=self.model,
-                max_tokens=50000,
-                temperature=0.7,
-                system=system_prompt,
-                messages=[{"role": "user", "content": chapter_prompt}]
-            )
+            self.context['parts'][part_idx]['chapters'][chapter_idx] = {
+                'title': chapter_title,
+                'content': full_chapter_content
+            }
             
-            chapter_content = response.content[0].text
-            self.context[f'part_{part_idx}_chapter_{chapter_idx}'] = chapter_content
-            return chapter_content
+            return full_chapter_content
             
         except Exception as e:
             st.error(f"Error generating chapter: {str(e)}")