2wawdw / creat_stor.py
Nouman112233's picture
Upload creat_stor.py
da09fa8 verified
from concurrent.futures import ThreadPoolExecutor, as_completed
import os
from openai import AzureOpenAI, OpenAI
import json
from story_schema import Story, StoryPage
client = AzureOpenAI(
azure_endpoint = 'https://storymii-openai.openai.azure.com/',
api_key="8u2PbSBg9U29jGIUJVYy67pWaKU2UYe9l1TnAZvtXKbTg62xr5KlJQQJ99ALACYeBjFXJ3w3AAABACOG2gJ2",
api_version="2024-08-01-preview",
)
image_client = OpenAI(api_key = "sk-storymii-prod-images-R53ukBLUScAtCUoX2TAkT3BlbkFJO9JzB9NTYBUddahQj4Iw")
def generate_story(ageRange, character_count, style_story, story_theme, language, plot, main_character_name, main_character_attr, other_details, tone, diverse_learner) -> Story:
try:
if diverse_learner == "Dyslexic":
system_prompt = system_dislecxic(ageRange, character_count)
user_prompt = user_dislecxic(ageRange, character_count, style_story, story_theme, language, plot, main_character_name, main_character_attr, other_details, tone)
elif diverse_learner == "ADHD":
system_prompt = system_adhdprompt(ageRange, character_count)
user_prompt = user_adhdprompt(ageRange, character_count, style_story, story_theme, language, plot, main_character_name, main_character_attr, other_details, tone)
else:
system_prompt = other_system_prompt(ageRange, character_count, diverse_learner)
user_prompt = other_user_prompt(ageRange, character_count, style_story, story_theme, language, plot, main_character_name, main_character_attr, other_details, tone)
response = client.chat.completions.create(
model="gpt-4",
messages=[
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_prompt}
]
)
story_content = response.choices[0].message.content
story_data = json.loads(story_content)
# Validate and structure the story data according to the schema
story: Story = {
'title': story_data['title'],
'cover_image_prompt': story_data['cover_image_prompt'],
'pages': [
{
'content': page['content'],
'image_prompt': page['image_prompt']
} for page in story_data['pages']
]
}
# Ensure exactly 4 pages
if len(story['pages']) != 4:
raise ValueError("Story must contain exactly 4 pages")
return story
except json.JSONDecodeError:
raise ValueError("Invalid JSON response from story generation")
except KeyError as e:
raise ValueError(f"Missing required field in story response: {e}")
except Exception as e:
raise Exception(f"Error generating story: {str(e)}")
def system_dislecxic(ageRange, character_count):
return f"""
[Your existing system_dislecxic function content]
"""
def user_dislecxic(ageRange, character_count, style_story, story_theme, language, plot, main_character_name, main_character_attr, other_details, tone):
return f"""
[Your existing user_dislecxic function content]
"""
def system_adhdprompt(ageRange, character_count):
return f"""
[Your existing system_adhdprompt function content]
"""
def user_adhdprompt(ageRange, character_count, style_story, story_theme, language, plot, main_character_name, main_character_attr, other_details, tone):
return f"""
[Your existing user_adhdprompt function content]
"""
def other_system_prompt(ageRange, character_count, diverse_learner):
return f"""
[Your existing other_system_prompt function content]
"""
def other_user_prompt(ageRange, character_count, style_story, story_theme, language, plot, main_character_name, main_character_attr, other_details, tone):
return f"""
[Your existing other_user_prompt function content]
"""