feat: Abstract more functionality from main to library

.gitignore

@@ -1,4 +1,5 @@
 .env
 venv
 .venv
-*.pyc
+*.pyc
+__pycache__

infinite_bookshelf/tools/pdf.py

@@ -4,6 +4,7 @@
 
 from io import BytesIO
 from markdown import markdown
+from weasyprint import HTML, CSS
 
 def create_pdf_file(content: str) -> BytesIO:
     """

infinite_bookshelf/ui/__init__.py

@@ -0,0 +1,2 @@
+from .book import Book
+from .initialization import load_return_env, ensure_states

infinite_bookshelf/ui/book.py

@@ -0,0 +1,76 @@
+"""
+Class to store and display structured book
+"""
+
+import streamlit as st
+
+class Book:
+    def __init__(self, book_title, structure):
+        self.book_title = book_title
+        self.structure = structure
+        self.contents = {title: "" for title in self.flatten_structure(structure)}
+        self.placeholders = {title: st.empty() for title in self.flatten_structure(structure)}
+        st.markdown(f"# {self.book_title}")
+        st.markdown("## Generating the following:")
+        toc_columns = st.columns(4)
+        self.display_toc(self.structure, toc_columns)
+        st.markdown("---")
+
+    def flatten_structure(self, structure):
+        sections = []
+        for title, content in structure.items():
+            sections.append(title)
+            if isinstance(content, dict):
+                sections.extend(self.flatten_structure(content))
+        return sections
+
+    def update_content(self, title, new_content):
+        try:
+            self.contents[title] += new_content
+            self.display_content(title)
+        except TypeError as e:
+            pass
+
+    def display_content(self, title):
+        if self.contents[title].strip():
+            self.placeholders[title].markdown(f"## {title}\n{self.contents[title]}")
+
+    def display_structure(self, structure=None, level=1):
+        if structure is None:
+            structure = self.structure
+
+        for title, content in structure.items():
+            if self.contents[title].strip():  # Only display title if there is content
+                st.markdown(f"{'#' * level} {title}")
+                self.placeholders[title].markdown(self.contents[title])
+            if isinstance(content, dict):
+                self.display_structure(content, level + 1)
+
+    def display_toc(self, structure, columns, level=1, col_index=0):
+        for title, content in structure.items():
+            with columns[col_index % len(columns)]:
+                st.markdown(f"{' ' * (level-1) * 2}- {title}")
+            col_index += 1
+            if isinstance(content, dict):
+                col_index = self.display_toc(content, columns, level + 1, col_index)
+        return col_index
+
+    def get_markdown_content(self, structure=None, level=1):
+        """
+        Returns the markdown styled pure string with the contents.
+        """
+        if structure is None:
+            structure = self.structure
+
+        if level==1:
+            markdown_content = f"# {self.book_title}\n\n"
+
+        else:
+            markdown_content = ""
+
+        for title, content in structure.items():
+            if self.contents[title].strip():  # Only include title if there is content
+                markdown_content += f"{'#' * level} {title}\n{self.contents[title]}\n\n"
+            if isinstance(content, dict):
+                markdown_content += self.get_markdown_content(content, level + 1)
+        return markdown_content

infinite_bookshelf/ui/components/__init__.py

@@ -0,0 +1,3 @@
+from .basic_form import render_groq_form
+from .statistics import display_statistics
+from .download import render_download_buttons

infinite_bookshelf/ui/components/basic_form.py

@@ -0,0 +1,32 @@
+"""
+Component function to render basic prompt input form
+"""
+
+import streamlit as st
+
+def render_groq_form(on_submit, button_disabled=False, button_text="Generate"):
+    with st.form("groqform"):
+        groq_input_key = st.text_input(
+            "Enter your Groq API Key (gsk_yA...):", "", type="password"
+        ) if not st.session_state.get('api_key') else None
+
+        topic_text = st.text_input(
+            "What do you want the book to be about?",
+            value="",
+            help="Enter the main topic or title of your book",
+        )
+
+        additional_instructions = st.text_area(
+            "Additional Instructions (optional)",
+            help="Provide any specific guidelines or preferences for the book's content",
+            placeholder="E.g., 'Focus on beginner-friendly content', 'Include case studies', etc.",
+            value="",
+        )
+
+        submitted = st.form_submit_button(
+            button_text,
+            on_click=on_submit,
+            disabled=button_disabled,
+        )
+
+        return submitted, groq_input_key, topic_text, additional_instructions

infinite_bookshelf/ui/components/download.py

@@ -0,0 +1,28 @@
+"""
+Component function to download button
+"""
+
+import streamlit as st
+from ...tools import create_markdown_file, create_pdf_file
+
+def render_download_buttons(book):
+    if book:
+        # Create markdown file
+        markdown_file = create_markdown_file(book.get_markdown_content())
+        st.download_button(
+            label="Download Text",
+            data=markdown_file,
+            file_name=f'{book.book_title}.txt',
+            mime='text/plain'
+        )
+
+        # Create pdf file (styled)
+        pdf_file = create_pdf_file(book.get_markdown_content())
+        st.download_button(
+            label="Download PDF",
+            data=pdf_file,
+            file_name=f'{book.book_title}.pdf',
+            mime='application/pdf'
+        )
+    else:
+        st.error("Please generate content first before downloading the book.")

infinite_bookshelf/ui/components/statistics.py

@@ -0,0 +1,15 @@
+"""
+Component function to render inference statistics
+"""
+
+import streamlit as st
+
+def display_statistics(placeholder, statistics_text):
+    with placeholder.container():
+        if statistics_text:
+            if "Generating structure in background" not in statistics_text:
+                st.markdown(statistics_text + "\n\n---\n")
+            else:
+                st.markdown(statistics_text)
+        else:
+            placeholder.empty()

infinite_bookshelf/ui/initialization.py

@@ -0,0 +1,23 @@
+"""
+Function to initialize streamlit session states and environment variables
+"""
+
+import streamlit as st
+from dotenv import load_dotenv
+from typing import List, Dict, Any
+import os
+
+# load .env file to environment
+load_dotenv()
+
+def load_return_env(variables: List[str]) -> Dict[str, str]:
+    return {var: os.getenv(var, None) for var in variables}
+
+def ensure_states(state_dict: Dict[str, Any]) -> None:
+    """
+    Define key values in session state 
+    if key not already defined
+    """
+    for key, default_value in state_dict.items():
+        if key not in st.session_state:
+            st.session_state[key] = default_value