From f87c5f05444b53f2fee033b99a400888538941a4 Mon Sep 17 00:00:00 2001 From: Ben Sima Date: Sun, 1 Dec 2024 20:33:09 -0500 Subject: Implement storybook prototype This paritally used gptme to create a storybook generator. The problem I ran into is that gptme doesn't do any architecting or considerations for maintainable code, or even readable code, so it just wrote a long script. I couldn't test it. Also, it didn't actually generate a 10-page story, it generated 10 separate stories. So, I ended up writing it myself and using gptme to fixup TODOs that I wrote along the way. --- Biz/Storybook.py | 291 +++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 291 insertions(+) create mode 100644 Biz/Storybook.py (limited to 'Biz/Storybook.py') diff --git a/Biz/Storybook.py b/Biz/Storybook.py new file mode 100644 index 0000000..2f362c4 --- /dev/null +++ b/Biz/Storybook.py @@ -0,0 +1,291 @@ +"""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"

""" + 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() -- cgit v1.2.3