Generating and Managing Short Stories

Introduction and Context Setting

Welcome to the third unit of our course on building a short story generation service. In this lesson, we will focus on generating short stories using AI. We will explore the StoryGeneratorService class, which is a key component of our application. This class will help us take user input, interact with the model, and manage the generated content effectively. By the end of this lesson, you will understand how to integrate these components to create a functional story generation service.

Before we dive into the new material, let's briefly revisit some key concepts from our previous lessons. We have already covered the PromptManager class, which is responsible for formatting user input into a structured prompt. Additionally, we have discussed the StoryManager class, which helps us manage and store generated stories. These components are essential for the functionality of our story generation service, and we will build upon them in this lesson.

Exploring the StoryGeneratorService Class

Let's start by understanding the structure and purpose of the StoryGeneratorService class. This class is responsible for generating stories based on user input and managing the interaction with the model.

First, we initialize the StoryGeneratorService class:

import os
from anthropic import Anthropic
from .models.story_manager import StoryManager
from .models.prompt_manager import PromptManager

class StoryGeneratorService:
    def __init__(self):
        self.story_manager = StoryManager()
        self.claude_client = Anthropic(api_key=os.environ.get("ANTHROPIC_API_KEY"))
  • Here, we import the necessary modules and classes. The Anthropic class is used to interact with the model.
  • The StoryGeneratorService class initializes with two main components: StoryManager for managing stories and Anthropic for interacting with the AI model.
  • The api_key is retrieved from environment variables, ensuring secure access to the model. In the CodeSignal environment, the ANTHROPIC_API_KEY environment variable is already set up for you. This means you do not need to manually configure the API key.

Generating a Story with User Input

Now, let's walk through the process of generating a story using the generate_story method. This method takes user input, formats it, and sends it to the model.

def generate_story(self, user_input: str, tone: str):
    prompt = PromptManager.format_prompt(user_input, tone)
    try:
        response = self.claude_client.messages.create(
            model='claude-sonnet-4-5-20250929',
            max_tokens=1500,
            messages=[
                {
                    "role": "user",
                    "content": prompt,
                }
            ]
        )
        # Note: The structure of the response may vary depending on the model and API version.
        # For example, the generated story might be in response.content[0].text or response.content.
        # Always check the actual response structure and update this line if necessary.
        story = response.content[0].text
        return self.story_manager.add_story(prompt, story)
    except Exception as e:
        print(str(e))
        raise RuntimeError(f"Error generating story: {str(e)}") from e
  • The generate_story method starts by formatting the user input using PromptManager.format_prompt(user_input, tone).
  • It then sends a request to the model using self.claude_client.messages.create(), specifying the model and the formatted prompt.
  • The response from the model contains the generated story. In this example, we access it with response.content[0].text, but you should always check the actual response structure and update this line if necessary.
  • The generated story is then stored using self.story_manager.add_story(prompt, story).
  • If an error occurs during this process, it is caught, and a RuntimeError is raised with a descriptive message.
Sign up

Join the 1M+ learners on CodeSignal

Be a part of our community of 1M+ users who develop and demonstrate their skills on CodeSignal