Create data_persistence.py

data_persistence.py

@@ -0,0 +1,261 @@
+# File: data_persistence.py
+# Location: /data_persistence.py
+# Description: Robust data persistence and serialization for book writing project
+
+import os
+import json
+import sqlite3
+from typing import Dict, Any
+import streamlit as st
+
+class ProjectPersistenceManager:
+    def __init__(self, project_id: str = None):
+        """
+        Initialize persistent storage for book writing projects
+        
+        Args:
+            project_id (str, optional): Unique identifier for the project
+        """
+        # Ensure data directory exists
+        os.makedirs('project_data', exist_ok=True)
+        
+        # SQLite Database Setup
+        self.db_path = 'project_data/book_projects.db'
+        self.conn = sqlite3.connect(self.db_path)
+        self._create_tables()
+    
+    def _create_tables(self):
+        """
+        Create necessary tables for project persistence
+        """
+        cursor = self.conn.cursor()
+        
+        # Project Metadata Table
+        cursor.execute('''
+            CREATE TABLE IF NOT EXISTS projects (
+                project_id TEXT PRIMARY KEY,
+                title TEXT,
+                genre TEXT,
+                total_chapters INTEGER,
+                created_at DATETIME DEFAULT CURRENT_TIMESTAMP
+            )
+        ''')
+        
+        # Chapters Table
+        cursor.execute('''
+            CREATE TABLE IF NOT EXISTS chapters (
+                project_id TEXT,
+                chapter_number INTEGER,
+                content TEXT,
+                status TEXT,
+                generated_at DATETIME DEFAULT CURRENT_TIMESTAMP,
+                PRIMARY KEY (project_id, chapter_number),
+                FOREIGN KEY (project_id) REFERENCES projects (project_id)
+            )
+        ''')
+        
+        # Book Concept Table
+        cursor.execute('''
+            CREATE TABLE IF NOT EXISTS book_concepts (
+                project_id TEXT PRIMARY KEY,
+                concept_data TEXT,
+                FOREIGN KEY (project_id) REFERENCES projects (project_id)
+            )
+        ''')
+        
+        self.conn.commit()
+    
+    def save_project_metadata(self, project_data: Dict[str, Any]):
+        """
+        Save project metadata to persistent storage
+        
+        Args:
+            project_data (Dict): Project metadata dictionary
+        """
+        cursor = self.conn.cursor()
+        
+        # Upsert project metadata
+        cursor.execute('''
+            INSERT OR REPLACE INTO projects 
+            (project_id, title, genre, total_chapters) 
+            VALUES (?, ?, ?, ?)
+        ''', (
+            project_data.get('project_id', 'unknown'),
+            project_data.get('title', 'Untitled'),
+            project_data.get('genre', 'Unspecified'),
+            project_data.get('total_chapters', 0)
+        ))
+        
+        self.conn.commit()
+    
+    def save_book_concept(self, project_id: str, concept_data: Dict[str, Any]):
+        """
+        Save book concept to persistent storage
+        
+        Args:
+            project_id (str): Unique project identifier
+            concept_data (Dict): Book concept details
+        """
+        cursor = self.conn.cursor()
+        
+        # Upsert book concept
+        cursor.execute('''
+            INSERT OR REPLACE INTO book_concepts 
+            (project_id, concept_data) 
+            VALUES (?, ?)
+        ''', (
+            project_id, 
+            json.dumps(concept_data)
+        ))
+        
+        self.conn.commit()
+    
+    def save_chapter(
+        self, 
+        project_id: str, 
+        chapter_number: int, 
+        content: str, 
+        status: str = 'Generated'
+    ):
+        """
+        Save chapter content to persistent storage
+        
+        Args:
+            project_id (str): Unique project identifier
+            chapter_number (int): Chapter number
+            content (str): Chapter content
+            status (str): Chapter status
+        """
+        cursor = self.conn.cursor()
+        
+        # Upsert chapter
+        cursor.execute('''
+            INSERT OR REPLACE INTO chapters 
+            (project_id, chapter_number, content, status) 
+            VALUES (?, ?, ?, ?)
+        ''', (
+            project_id, 
+            chapter_number, 
+            content, 
+            status
+        ))
+        
+        self.conn.commit()
+    
+    def load_project_metadata(self, project_id: str) -> Dict[str, Any]:
+        """
+        Load project metadata from persistent storage
+        
+        Args:
+            project_id (str): Unique project identifier
+        
+        Returns:
+            Dict containing project metadata
+        """
+        cursor = self.conn.cursor()
+        cursor.execute('SELECT * FROM projects WHERE project_id = ?', (project_id,))
+        
+        project = cursor.fetchone()
+        if project:
+            return {
+                'project_id': project[0],
+                'title': project[1],
+                'genre': project[2],
+                'total_chapters': project[3]
+            }
+        
+        return {}
+    
+    def load_book_concept(self, project_id: str) -> Dict[str, Any]:
+        """
+        Load book concept from persistent storage
+        
+        Args:
+            project_id (str): Unique project identifier
+        
+        Returns:
+            Dict containing book concept
+        """
+        cursor = self.conn.cursor()
+        cursor.execute('SELECT concept_data FROM book_concepts WHERE project_id = ?', (project_id,))
+        
+        concept = cursor.fetchone()
+        return json.loads(concept[0]) if concept else {}
+    
+    def load_chapters(self, project_id: str) -> Dict[int, Dict]:
+        """
+        Load all chapters for a project
+        
+        Args:
+            project_id (str): Unique project identifier
+        
+        Returns:
+            Dict of chapters with chapter number as key
+        """
+        cursor = self.conn.cursor()
+        cursor.execute('SELECT chapter_number, content, status FROM chapters WHERE project_id = ?', (project_id,))
+        
+        chapters = {}
+        for chapter_num, content, status in cursor.fetchall():
+            chapters[chapter_num] = {
+                'content': content,
+                'status': status
+            }
+        
+        return chapters
+    
+    def close(self):
+        """
+        Close database connection
+        """
+        self.conn.close()
+
+# Test functionality
+def test_persistence():
+    """
+    Run comprehensive tests for data persistence
+    """
+    # Initialize persistence manager
+    pm = ProjectPersistenceManager()
+    
+    # Test project metadata
+    test_project = {
+        'project_id': 'test_project_1',
+        'title': 'Test Book',
+        'genre': 'Science Fiction',
+        'total_chapters': 10
+    }
+    
+    # Save project metadata
+    pm.save_project_metadata(test_project)
+    
+    # Load and verify project metadata
+    loaded_project = pm.load_project_metadata('test_project_1')
+    assert loaded_project == test_project, "Project metadata not saved/loaded correctly"
+    
+    # Test book concept
+    test_concept = {
+        'narrative_premise': 'A story about AI and humanity',
+        'core_themes': ['Technology', 'Ethics']
+    }
+    
+    pm.save_book_concept('test_project_1', test_concept)
+    
+    # Load and verify book concept
+    loaded_concept = pm.load_book_concept('test_project_1')
+    assert loaded_concept == test_concept, "Book concept not saved/loaded correctly"
+    
+    # Test chapter saving and loading
+    pm.save_chapter('test_project_1', 1, "Chapter 1 content", "Generated")
+    pm.save_chapter('test_project_1', 2, "Chapter 2 content", "In Review")
+    
+    # Load chapters
+    loaded_chapters = pm.load_chapters('test_project_1')
+    assert len(loaded_chapters) == 2, "Chapters not saved/loaded correctly"
+    assert loaded_chapters[1]['content'] == "Chapter 1 content", "Chapter 1 content mismatch"
+    
+    print("All persistence tests passed successfully!")
+
+# Run tests if script is executed directly
+if __name__ == "__main__":
+    test_persistence()