homeros.py

from adaptors.db import get_story, save_story
from models.story import Story
from loguru import logger
import json
from pprint import pprint
from helpers import gen_unique_id
import prompts
from adaptors.llm import answer

MAX_STORY_LEN = 2 #after how many chunks we force the story to end

'''
initiates a new story and saves in DB
'''
def init_story(story_data):
    story_data["uuid"] = gen_unique_id()
    pprint(story_data)
    story = Story(
        uuid=story_data['uuid'],
        status=story_data['status']
    )
    save_story(story)
    return story.to_dict()


'''
defines a field of metadata (world, style, plot etc)
'''
def define_metadata(user_input, field, story_data):
    story = get_story(story_data["uuid"])
    setattr(story, field, user_input)
    save_story(story)
    return story.to_dict()


'''
creates the first chunk of the story
'''
def start_story(story_data):
    return continue_story("Please begin", story_data)

'''
main function that manages adding the next chunk to a story (first text, then audio)
'''
def continue_story(user_input, story_data):
    story = get_story(story_data["uuid"])
    chunks = json.loads(story.chunks)
    messages = json.loads(story.messages)

    messages.append({
        "role":"user",
        "content": user_input
    })
    next_chunk_text = create_next_chunk_text(user_input, story)
    next_chunk_audio = create_next_chunk_audio(next_chunk_text)
    messages.append({
        "role":"assistant",
        "content": next_chunk_text
    })
    chunks.append({
        "text" : next_chunk_text,
        "audio_url" : next_chunk_audio
    })
    story.chunks = json.dumps(chunks)
    story.messages = json.dumps(messages)
    story.status = "ongoing"
    save_story(story)
    return story.to_dict()

'''
generates the last part of the story and changes status to "finished"
'''
def finish_story(user_input, story_data):
    story = get_story(story_data["uuid"])
    chunks = json.loads(story.chunks)
    messages = json.loads(story.messages)

    user_input = user_input + "\n\nPlease finish the story now."

    messages.append({
        "role":"user",
        "content": user_input
    })
    next_chunk_text = create_next_chunk_text(user_input, story)
    next_chunk_audio = create_next_chunk_audio(next_chunk_text)
    chunks.append({
        "text" : next_chunk_text,
        "audio_url" : next_chunk_audio
    })
    story.chunks = json.dumps(chunks)
    story.status = "finished"
    save_story(story)
    return story.to_dict()


'''
generates the next chunk of the story
'''
def create_next_chunk_text(user_input, story):

    next_chunk = ""

    # get system message from prompts
    system_message = prompts.get(
        prompt_name = "storyteller_general",
        substitutions = {
            "WORLD" : story.world,
            "HERO" : story.hero,
            "PLOT" : story.plot,
            "ENDING" : story.ending,
            "STYLE" : story.style,
        }
    )

    # get history of messages
    messages = json.loads(story.messages)

    # add user message to history of messages
    messages.append({
        "role" : "user",
        "content": user_input
    })

    # get llm to answer
    next_chunk = answer(system_message, messages)

    # return the answer
    return next_chunk


'''
evaluates the story up until now and returns a dict with the result of the evaluation
'''
def evaluate_story(story):
    evaluation = {}
    story_len = len(story["chunks"])
    logger.debug(story_len)
    evaluation["is_time_to_end"] = story_len > MAX_STORY_LEN

    return evaluation


'''
turns next story chunk into audio and returns a URL
'''
def create_next_chunk_audio(text):
    return "url.com"