Code Extender | JavaScript

AI Book Generator Enhancements

This proposal details the extension of a Python-based Streamlit application that generates books using AI. It introduces a new feature allowing users to set a word limit for each section, enhancing content control and customization.

Prompt

# 1: Import libraries
import streamlit as st
from groq import Groq
import json

from infinite_bookshelf.agents import (
    generate_section,
    generate_book_structure,
    generate_book_title,
)
from infinite_bookshelf.inference import GenerationStatistics
from infinite_bookshelf.tools import create_markdown_file, create_pdf_file
from infinite_bookshelf.ui.components import (
    render_groq_form,
    render_advanced_groq_form,
    display_statistics,
    render_download_buttons,
)
from infinite_bookshelf.ui import Book, load_return_env, ensure_states


# 2: Initialize env variables and session states
GROQ_API_KEY = load_return_env(["GROQ_API_KEY"])["GROQ_API_KEY"]

states = {
    "api_key": GROQ_API_KEY,
    "button_disabled": False,
    "button_text": "Generate",
    "statistics_text": "",
    "book_title": "",
}

if GROQ_API_KEY:
    states["groq"] = (
        Groq()
    )  # Define Groq provider if API key provided. Otherwise defined later after API key is provided.

ensure_states(states)


# 3: Define Streamlit page structure and functionality
st.write(
    """
# Infinite Bookshelf: Write full books using llama3 (8b and 70b) on Groq
"""
)


def disable():
    st.session_state.button_disabled = True


def enable():
    st.session_state.button_disabled = False


def empty_st():
    st.empty()


try:
    if st.button("End Generation and Download Book"):
        if "book" in st.session_state:
            render_download_buttons(st.session_state.get("book"))

    (
        submitted,
        groq_input_key,
        topic_text,
        additional_instructions,
        writing_style,
        complexity_level,
        seed_content,
        uploaded_file,
        title_agent_model,
        structure_agent_model,
        section_agent_model,
    ) = render_advanced_groq_form(
        on_submit=disable,
        button_disabled=st.session_state.button_disabled,
        button_text=st.session_state.button_text,
    )

    # New content for advanced mode
    additional_section_writer_prompt = "The book chapters should be comprehensive. The writing should be: \nEngaging and tailored to the specified writing style, tone, and complexity level. \nWell-structured with clear subheadings, paragraphs, and transitions. \nRich in relevant examples, analogies, and explanations. \nConsistent with provided seed content and additional instructions. \nFocused on delivering value through insightful analysis and information. \nFactually accurate based on the latest available information. \nCreative, offering unique perspectives or thought-provoking ideas. \nEnsure each section flows logically, maintaining coherence throughout the chapter."
    advanced_settings_prompt = f"Use the following parameters:\nWriting Style: {writing_style}\nComplexity Level: {complexity_level}"
    total_seed_content = ""

    # Fill total_seed_content
    if seed_content:
        total_seed_content += seed_content
    if uploaded_file:
        total_seed_content += uploaded_file.read().decode("utf-8")
    if total_seed_content != "":
        total_seed_content = f"The user has provided seed content for context. Develop the structure and content around the provided seed: {total_seed_content}"

    if submitted:
        if len(topic_text) < 10:
            raise ValueError("Book topic must be at least 10 characters long")

        st.session_state.button_disabled = True
        st.session_state.statistics_text = (
            "Generating book title and structure in background...."
        )

        placeholder = st.empty()
        display_statistics(
            placeholder=placeholder, statistics_text=st.session_state.statistics_text
        )

        if not GROQ_API_KEY:
            st.session_state.groq = Groq(api_key=groq_input_key)

        # Step 1: Generate book structure using structure_writer agent
        additional_instructions_prompt = (
            additional_instructions + advanced_settings_prompt
        )
        if total_seed_content != "":
            additional_instructions_prompt += "\n" + total_seed_content

        large_model_generation_statistics, book_structure = generate_book_structure(
            prompt=topic_text,
            additional_instructions=additional_instructions_prompt,
            model=structure_agent_model,
            groq_provider=st.session_state.groq,
            long=True # Use longer version in advanced
        )

        # Step 2: Generate book title using title_writer agent
        st.session_state.book_title = generate_book_title(
            prompt=topic_text,
            model=title_agent_model,
            groq_provider=st.session_state.groq,
        )

        st.write(f"## {st.session_state.book_title}")

        total_generation_statistics = GenerationStatistics(
            model_name=section_agent_model
        )

        # Step 3: Generate book section content using section_writer agent
        try:
            book_structure_json = json.loads(book_structure)
            book = Book(st.session_state.book_title, book_structure_json)

            if "book" not in st.session_state:
                st.session_state.book = book

            # Print the book structure to the terminal to show structure
            print(json.dumps(book_structure_json, indent=2))

            st.session_state.book.display_structure()

            def stream_section_content(sections):
                for title, content in sections.items():
                    if isinstance(content, str):
                        additional_instructions_prompt = f"{additional_section_writer_prompt}\n{additional_instructions}\n{advanced_settings_prompt}"
                        if total_seed_content != "":
                            additional_instructions_prompt += "\n" + total_seed_content

                        content_stream = generate_section(
                            prompt=(title + ": " + content),
                            additional_instructions=additional_instructions_prompt,
                            model=section_agent_model,
                            groq_provider=st.session_state.groq,
                        )
                        for chunk in content_stream:
                            # Check if GenerationStatistics data is returned instead of str tokens
                            chunk_data = chunk
                            if type(chunk_data) == GenerationStatistics:
                                total_generation_statistics.add(chunk_data)

                                st.session_state.statistics_text = str(
                                    total_generation_statistics
                                )
                                display_statistics(
                                    placeholder=placeholder,
                                    statistics_text=st.session_state.statistics_text,
                                )

                            elif chunk != None:
                                st.session_state.book.update_content(title, chunk)
                    elif isinstance(content, dict):
                        stream_section_content(content)

            stream_section_content(book_structure_json)

        except json.JSONDecodeError:
            st.error("Failed to decode the book structure. Please try again.")


except Exception as e:
    st.session_state.button_disabled = False
    st.error(e)

    if st.button("Clear"):
        st.rerun()

Answer

Code Extension Proposal

Existing Code Analysis

The provided code snippet is written in Python and is part of a Streamlit application designed for writing and generating books using AI models from Groq. The main functionality includes user input handling, book structure generation, and content creation based on various parameters defined by the user.

Key Features of the Existing Code:

Imports: It imports necessary libraries, including Streamlit, Groq, and custom modules for generating book content and handling user interactions.
Environment Variables: It loads the Groq API key and initializes session states to manage UI states and book data.
User Interaction: It provides a form where users can specify topics, writing styles, and other parameters for book generation.
Error Handling: It includes basic error handling for cases like invalid topic length and JSON decoding issues.

Objective: Extend Functionality

Instruction

The objective is to add a feature that allows users to specify a word limit for each section generated in the book. This word limit should be an input parameter in the render_advanced_groq_form and should be enforced when generating each section in the stream_section_content function.

Extended Code Snippet

The following modifications were made to extend the existing functionality:

New Input Parameter: Added word_limit as a parameter in the render_advanced_groq_form.
Word Limit Enforcement: Updated the stream_section_content function to respect the specified word limit during section generation.

# 1: Import libraries remains unchanged

# 2: Initialize env variables and session states remains unchanged

# 3: Define Streamlit page structure and functionality remains unchanged

try:
    if st.button("End Generation and Download Book"):
        if "book" in st.session_state:
            render_download_buttons(st.session_state.get("book"))

    (
        submitted,
        groq_input_key,
        topic_text,
        additional_instructions,
        writing_style,
        complexity_level,
        seed_content,
        uploaded_file,
        title_agent_model,
        structure_agent_model,
        section_agent_model,
        word_limit  # New word limit parameter
    ) = render_advanced_groq_form(
        on_submit=disable,
        button_disabled=st.session_state.button_disabled,
        button_text=st.session_state.button_text,
    )

    # Remaining code before section generation remains unchanged

    def stream_section_content(sections):
        for title, content in sections.items():
            if isinstance(content, str):
                # Enforce word limit on generated content
                additional_instructions_prompt = f"{additional_section_writer_prompt}\n{additional_instructions}\n{advanced_settings_prompt}\nWord Limit: {word_limit}"

                # Check if word_limit is specified and valid
                if word_limit is not None and word_limit > 0:
                    content_stream = generate_section(
                        prompt=(title + ": " + content),
                        additional_instructions=additional_instructions_prompt,
                        model=section_agent_model,
                        groq_provider=st.session_state.groq,
                        word_limit=word_limit  # Include word limit in the request
                    )
                else:
                    # Generate without limit if invalid word limit
                    content_stream = generate_section(
                        prompt=(title + ": " + content),
                        additional_instructions=additional_instructions_prompt,
                        model=section_agent_model,
                        groq_provider=st.session_state.groq,
                    )

                for chunk in content_stream:
                    # Check GenerationStatistics and update UI as before
                    chunk_data = chunk
                    if isinstance(chunk_data, GenerationStatistics):
                        total_generation_statistics.add(chunk_data)
                        st.session_state.statistics_text = str(total_generation_statistics)
                        display_statistics(placeholder=placeholder, statistics_text=st.session_state.statistics_text)
                    elif chunk is not None:
                        st.session_state.book.update_content(title, chunk)
            elif isinstance(content, dict):
                stream_section_content(content)
            
    stream_section_content(book_structure_json)

except Exception as e:
    st.session_state.button_disabled = False
    st.error(e)

    if st.button("Clear"):
        st.rerun()

Summary of Changes

User Input: Added word_limit to the form inputs allowing users to specify maximum word count for each section.
Content Generation Logic: Integrated word_limit within the stream_section_content method, ensuring that it is applied during content generation.

This modification adheres to best practices, ensuring clarity and maintainability, and does not disrupt the existing functionality of the application.