"""Storybook Generator Application. This application generates a children's storybook using the OpenAI API. The user can select a theme, specify the main character's name, and choose a setting. The app generates a 10-page storybook with images. The tech stack is: Python, Flask, HTMX, and bootstrap. All of the code is in this single file. """ # : out storybook # : dep flask # : dep openai import flask import json import openai import os import pydantic import sys import unittest app = flask.Flask(__name__) app.secret_key = os.urandom(24) def main() -> None: """Run the Flask application.""" if sys.argv[1] == "test": test() else: move() def move() -> None: """Run the application.""" app.run() def test() -> None: """Run the unittest suite manually.""" suite = unittest.TestSuite() tests = [StorybookTestCase, StoryTestCase] suite.addTests([ unittest.defaultTestLoader.loadTestsFromTestCase(tc) for tc in tests ]) unittest.TextTestRunner(verbosity=2).run(suite) @app.route("/") def index() -> str: """Render the main page.""" return flask.render_template_string(f""" Storybook Generator

Storybook Generator

""") class Page(pydantic.BaseModel): """Represents a single page in the storybook.""" model_config = pydantic.ConfigDict(arbitrary_types_allowed=True) text: str image_prompt: str image_url: str | None def load_image(page: Page) -> Page: """Load an image for a given page using the OpenAI API.""" if page.image_url is not None: return page prompt = page.image_prompt client = openai.OpenAI(api_key=os.getenv("OPENAI_API_KEY")) image_response = client.images.generate( prompt=prompt, n=1, size="256x256", ) page.image_url = image_response.data[0].url # Handle if image is None return page class Story(pydantic.BaseModel): """Represents a story with multiple pages.""" model_config = pydantic.ConfigDict(arbitrary_types_allowed=True) theme: str character: str setting: str moral: str pages: list[Page] | None = None def system_prompt() -> str: """Generate the system prompt for the OpenAI API.""" return ( "You are an author and illustrator of childrens books. " "Each book is 10 pages long. " "All output must be in valid JSON. " "Don't add explanation or markdown formatting beyond the JSON. " "In your output, include the text on the page and a description of the " "image to be generated with an AI image generator." ) def user_prompt(story: Story) -> str: """Generate the user prompt based on the story details.""" return " ".join([ "Write a story with the following details.", "Output must be in valid JSON where each page is an array of text and" "image like the following example:", """[{"text": "",""", """"image": ""}...],""", f"Character: {story.character}\n", f"Setting: {story.setting}\n", f"Theme: {story.theme}\n", f"Moral: {story.moral}\n", ]) def _openai_generate_text(story: Story) -> openai.types.chat.ChatCompletion: """Generate story text using the OpenAI API.""" messages: list[ openai.types.chat.ChatCompletionUserMessageParam | openai.types.chat.ChatCompletionSystemMessageParam ] = [ {"role": "system", "content": system_prompt()}, {"role": "user", "content": user_prompt(story)}, ] client = openai.OpenAI(api_key=os.getenv("OPENAI_API_KEY")) return client.chat.completions.create( model="gpt-4o-mini", messages=messages, max_tokens=1500, ) def generate_text(story: Story) -> Story: """Generate the text for a story and update its pages. Raises: ValueError: If the content is None or JSON parsing fails. """ response = _openai_generate_text(story) content = response.choices[0].message.content if content is None: error_message = "No content in response" raise ValueError(error_message) try: response_messages = json.loads(content) except (json.JSONDecodeError, ValueError) as e: error_message = f"Failed to parse story JSON: {e}" raise ValueError(error_message) from e story.pages = [ Page(text=msg["text"], image_prompt=msg["image"], image_url=None) for msg in response_messages ] return story @app.route("/generate/story", methods=["POST"]) def generate_story() -> str: """Generate a story based on user input.""" story = Story( theme=flask.request.form["theme"], character=flask.request.form["character"], setting=flask.request.form["setting"], moral="Honor thy mother and father", # request.form["moral"], ) story = generate_text(story) if story.pages is None: return "

error: no story pages found

" flask.session["story"] = story.model_dump_json() return "".join( f"
" f"Loading...""" f"
" f"

{page.text}

" for i, page in enumerate(story.pages) ) @app.route("/generate/image/", methods=["GET"]) def generate_image(n: int) -> str: """Generate an image for a specific page in the story.""" story_data = flask.session.get("story") if story_data is None: return "

error: no story data found

" try: story = Story.model_validate_json(story_data) except pydantic.ValidationError as e: return f"

error: story validation failed: {e}

" if story.pages is not None and 0 <= n < len(story.pages): page = load_image(story.pages[n]) return f"" return "

Image not found

" class StoryTestCase(unittest.TestCase): """Unit test case for the Story class and related functions.""" def test_story_creation(self) -> None: """Test the creation of a story and its text generation.""" s = Story( theme="Christian", character="Lia and her pet bunny", setting="A suburban park", moral="Honor thy mother and father", ) s = generate_text(s) self.assertIsNotNone(s.pages) self.assertEqual(len(s.pages), 10) # type: ignore[arg-type] class StorybookTestCase(unittest.TestCase): """Unit test case for the Storybook application.""" def setUp(self) -> None: """Set up the test client.""" self.app = app.test_client() def test_index_page(self) -> None: """Test that the index page loads successfully.""" response = self.app.get("/") self.assertEqual(response.status_code, 200) self.assertIn(b"Storybook Generator", response.data) def test_generate_story(self) -> None: """Test the story generation endpoint.""" response = self.app.post( "/generate/story", data={ "theme": "Christian", "character": "Alice", "setting": "forest", }, ) self.assertEqual(response.status_code, 200) self.assertIn(b"
", response.data) if __name__ == "__main__": main()