app.py

# File: app.py
# Location: /app.py (root directory for Hugging Face Spaces)
# Description: Main Streamlit application for AI Book Writing Assistant

import streamlit as st
import os
import json
import uuid

from multi_agent_book_workflow import BookWritingOrchestrator
from persistence_manager import ProjectPersistenceManager

class BookWritingApp:
    def __init__(self):
        """
        Initialize the Streamlit application for book writing
        
        Manages session state and persistence
        """
        # Set page configuration
        st.set_page_config(
            page_title="AI Book Writing Assistant",
            page_icon="📖",
            layout="wide"
        )
        
        # Initialize persistence manager
        self.persistence_manager = ProjectPersistenceManager()
        
        # Initialize session state
        self._initialize_session_state()
        
        # Set up project orchestrator 
        self._setup_project_orchestrator()
    
    def _setup_project_orchestrator(self):
        """
        Initialize project orchestrator with proper error handling
        """
        try:
            if 'project_orchestrator' not in st.session_state or st.session_state.project_orchestrator is None:
                st.session_state.project_orchestrator = BookWritingOrchestrator()
        except ValueError as ve:
            st.error(f"API Key Error: {str(ve)}")
            st.info("Please ensure your API keys are properly set in the environment variables.")
            st.session_state.project_orchestrator = None
        except Exception as e:
            st.error(f"Error setting up project orchestrator: {e}")
            st.session_state.project_orchestrator = None
    
    def _initialize_session_state(self):
        """
        Initialize and manage Streamlit session state
        """
        # Project tracking
        if 'active_project_id' not in st.session_state:
            st.session_state.active_project_id = None
        
        # Restore last active project if exists
        if st.session_state.active_project_id:
            self._load_project(st.session_state.active_project_id)
        
        # Ensure core session state variables exist
        session_defaults = {
            'project_orchestrator': None,
            'book_concept': None,
            'generated_chapters': {},
            'total_chapters': 10,  # Default value
            'project_metadata': {
                'title': 'Untitled Project',
                'genre': 'Unspecified',
                'total_chapters': 10
            }
        }
        
        for key, default in session_defaults.items():
            if key not in st.session_state:
                st.session_state[key] = default
    
    def _load_project(self, project_id):
        """
        Load a project from persistent storage
        """
        try:
            loaded_project = self.persistence_manager.load_project(project_id)
            
            if loaded_project:
                # Update session state with loaded project
                st.session_state.active_project_id = project_id
                st.session_state.book_concept = loaded_project.get('book_concept')
                st.session_state.generated_chapters = loaded_project.get('chapters', {})
                st.session_state.total_chapters = loaded_project.get('total_chapters', 10)
                st.session_state.project_metadata = {
                    'title': loaded_project.get('title', 'Untitled Project'),
                    'genre': loaded_project.get('genre', 'Unspecified'),
                    'total_chapters': loaded_project.get('total_chapters', 10)
                }
                
                st.success(f"Loaded project: {loaded_project.get('title', 'Untitled')}")
            else:
                st.warning("Could not load project.")
        
        except Exception as e:
            st.error(f"Error loading project: {e}")
    
    def _save_current_project(self):
        """
        Save the current project to persistent storage
        """
        try:
            # Prepare project data for saving
            project_data = {
                'project_id': st.session_state.active_project_id or str(uuid.uuid4()),
                'title': st.session_state.project_metadata.get('title', 'Untitled Project'),
                'genre': st.session_state.project_metadata.get('genre', 'Unspecified'),
                'book_concept': st.session_state.book_concept,
                'chapters': st.session_state.generated_chapters,
                'total_chapters': st.session_state.total_chapters
            }
            
            # Save project
            project_id = self.persistence_manager.save_project(project_data)
            
            if project_id:
                # Update active project ID
                st.session_state.active_project_id = project_id
                st.success(f"Project saved successfully!")
            else:
                st.error("Failed to save project.")
        
        except Exception as e:
            st.error(f"Error saving project: {e}")
    
    def render_project_management(self):
        """
        Render project management interface
        """
        st.sidebar.header("Project Management")
        
        # New Project Button
        if st.sidebar.button("Start New Project"):
            # Reset session state
            st.session_state.active_project_id = None
            st.session_state.book_concept = None
            st.session_state.generated_chapters = {}
            st.session_state.total_chapters = 10
            st.session_state.project_metadata = {
                'title': 'Untitled Project',
                'genre': 'Unspecified',
                'total_chapters': 10
            }
            st.experimental_rerun()
        
        # Load Project
        project_list = self.persistence_manager.list_projects()
        if project_list:
            selected_project = st.sidebar.selectbox(
                "Load Existing Project",
                options=[p['title'] for p in project_list],
                index=None,
                placeholder="Select a project..."
            )
            
            if selected_project:
                # Find project ID
                project_id = next(
                    (p['project_id'] for p in project_list if p['title'] == selected_project),
                    None
                )
                
                if project_id and st.sidebar.button("Load Selected Project"):
                    self._load_project(project_id)
        
        # Save Current Project
        if st.sidebar.button("Save Current Project"):
            self._save_current_project()
    
    def render_book_generation_interface(self):
        """
        Render the main book generation interface
        """
        # Tabs for different stages of book writing
        tab1, tab2, tab3 = st.tabs([
            "Book Concept", 
            "Chapter Generation", 
            "Project Progress"
        ])
        
        with tab1:
            self._render_concept_development()
        
        with tab2:
            self._render_chapter_generation()
        
        with tab3:
            self._render_project_progress()
    
    def _render_concept_development(self):
        """
        Render the book concept development interface
        """
        st.header("📘 Book Concept Development")
        
        # Initial concept input
        initial_concept = st.text_area(
            "Describe Your Book Idea", 
            height=200, 
            placeholder="Enter a comprehensive description of your book concept..."
        )
        
        if st.button("Generate Book Concept"):
            if st.session_state.project_orchestrator is None:
                st.error("Project orchestrator is not properly initialized. Please check your API keys.")
                return
                
            with st.spinner("Developing Book Concept..."):
                try:
                    # Generate book concept using multi-agent approach
                    book_concept = st.session_state.project_orchestrator.generate_book_concept(
                        initial_concept
                    )
                    
                    if book_concept:
                        # Update session state
                        st.session_state.book_concept = book_concept
                        st.session_state.project_metadata.update({
                            'title': book_concept.get('title', 'Untitled Project'),
                            'genre': book_concept.get('genre', 'Unspecified')
                        })
                        
                        # Display generated concept
                        st.subheader("Generated Book Concept")
                        st.json(book_concept)
                    else:
                        st.error("Failed to generate book concept.")
                
                except Exception as e:
                    st.error(f"Error generating book concept: {e}")
    
    def _render_chapter_generation(self):
        """
        Render the chapter generation interface
        """
        st.header("✍️ Chapter Generation")
        
        # Check if book concept exists
        if not st.session_state.book_concept:
            st.warning("Please generate a book concept first.")
            return
        
        # Chapter generation controls
        col1, col2 = st.columns(2)
        
        with col1:
            # Chapter number selection
            max_chapters = 20  # Can be adjusted
            current_chapter = len(st.session_state.generated_chapters) + 1
            chapter_number = st.number_input(
                "Select Chapter to Generate", 
                min_value=1, 
                max_value=max_chapters, 
                value=min(current_chapter, max_chapters)
            )
        
        with col2:
            # Total chapters planning
            total_chapters = st.number_input(
                "Total Planned Chapters", 
                min_value=1, 
                max_value=max_chapters, 
                value=st.session_state.total_chapters
            )
            
            # Update session state
            if total_chapters != st.session_state.total_chapters:
                st.session_state.total_chapters = total_chapters
                st.session_state.project_metadata['total_chapters'] = total_chapters
        
        # Generate Chapter Button
        if st.button("Generate Chapter"):
            if st.session_state.project_orchestrator is None:
                st.error("Project orchestrator is not properly initialized. Please check your API keys.")
                return
                
            with st.spinner(f"Generating Chapter {chapter_number}..."):
                try:
                    # Generate chapter content
                    chapter_content = st.session_state.project_orchestrator.generate_chapter_content(
                        st.session_state.book_concept, 
                        chapter_number
                    )
                    
                    if chapter_content:
                        # Store generated chapter
                        st.session_state.generated_chapters[chapter_number] = {
                            'content': chapter_content,
                            'status': 'Generated'
                        }
                        
                        # Display Chapter Content
                        st.subheader(f"Chapter {chapter_number}")
                        st.write(chapter_content)
                        
                        # Create a unique key for the text area
                        edit_key = f"edit_chapter_{chapter_number}"
                        
                        # Optional: Edit Chapter
                        edited_content = st.text_area(
                            "Edit Chapter Content", 
                            value=chapter_content, 
                            height=400,
                            key=edit_key
                        )
                        
                        # Save button with unique key
                        save_key = f"save_chapter_{chapter_number}"
                        if st.button(f"Save Chapter {chapter_number}", key=save_key):
                            st.session_state.generated_chapters[chapter_number]['content'] = edited_content
                            st.success(f"Chapter {chapter_number} saved!")
                    else:
                        st.error("Failed to generate chapter content.")
                
                except Exception as e:
                    st.error(f"Error generating chapter: {e}")
    
    def _render_project_progress(self):
        """
        Render project progress and tracking interface
        """
        st.header("📊 Project Progress")
        
        # Project Metadata Overview
        st.subheader("Project Overview")
        col1, col2 = st.columns(2)
        
        with col1:
            st.metric("Book Title", st.session_state.project_metadata.get('title', 'Untitled'))
            st.metric("Genre", st.session_state.project_metadata.get('genre', 'Unspecified'))
        
        with col2:
            total_chapters = st.session_state.total_chapters
            generated_chapters = len(st.session_state.generated_chapters)
            
            st.metric("Total Planned Chapters", total_chapters)
            st.metric("Chapters Generated", generated_chapters)
            
            # Progress Bar
            progress = generated_chapters / total_chapters if total_chapters > 0 else 0
            st.progress(progress)
        
        # Chapter Navigation and Details
        st.subheader("Generated Chapters")
        
        # Create tabs for each generated chapter
        if st.session_state.generated_chapters:
            chapter_tabs = st.tabs([
                f"Chapter {ch_num}" for ch_num in sorted(st.session_state.generated_chapters.keys())
            ])
            
            for i, ch_num in enumerate(sorted(st.session_state.generated_chapters.keys())):
                with chapter_tabs[i]:
                    chapter_data = st.session_state.generated_chapters[ch_num]
                    st.write(chapter_data['content'])
                    
                    # Chapter status and actions
                    col1, col2 = st.columns(2)
                    with col1:
                        status = st.selectbox(
                            "Chapter Status", 
                            ["Generated", "In Review", "Completed"],
                            key=f"status_{ch_num}"
                        )
                        if status != chapter_data['status']:
                            chapter_data['status'] = status
                    
                    with col2:
                        if st.button(f"Export Chapter {ch_num}", key=f"export_{ch_num}"):
                            try:
                                # Export chapter functionality
                                export_path = os.path.join('data', f"chapter_{ch_num}.txt")
                                os.makedirs('data', exist_ok=True)
                                with open(export_path, "w") as f:
                                    f.write(chapter_data['content'])
                                st.success(f"Chapter {ch_num} exported to {export_path}")
                            except Exception as e:
                                st.error(f"Error exporting chapter: {e}")
        else:
            st.info("No chapters generated yet. Start writing your book!")
    
    def run(self):
        """
        Run the main Streamlit application
        """
        # Render project management sidebar
        self.render_project_management()
        
        # Main title
        st.title("🚀 AI-Powered Book Writing Assistant")
        
        # Check if project orchestrator is initialized
        if st.session_state.project_orchestrator is None:
            st.warning("Please ensure your API keys are properly set in the environment variables.")
        
        # Render the main interface
        self.render_book_generation_interface()

def main():
    """
    Initialize and run the application
    """
    app = BookWritingApp()
    app.run()

if __name__ == "__main__":
    main()