162 lines
5.5 KiB
Python
162 lines
5.5 KiB
Python
#
|
|
# Module: step_2
|
|
#
|
|
# This module implements the step_2 function for creating a new season for an agent.
|
|
#
|
|
# Title: Step 2
|
|
# Summary: Step 2 implementation.
|
|
# Authors:
|
|
# - @TheBlockRhino
|
|
# Created: 2025-01-02
|
|
# Last edited by: @TheBlockRhino
|
|
# Last edited date: 2025-01-09
|
|
# URLs:
|
|
# - https://arai-ai.io
|
|
# - https://github.com/ARAI-DevHub/arai-ai-agents
|
|
# - https://x.com/TheBlockRhino
|
|
|
|
# standard imports
|
|
import os
|
|
import json
|
|
import json
|
|
import sys
|
|
|
|
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
# custom ARAI imports
|
|
from utils.content_generator import ContentGenerator
|
|
from utils.template_types import TemplateType
|
|
import prompt_chaining.step_3_create_posts as next_step
|
|
|
|
def create_seasons_and_episodes(ai_model, master_file_path, number_of_episodes):
|
|
'''
|
|
Description:
|
|
Create a new season for the agent
|
|
|
|
Args:
|
|
ai_model: The AI model to use for generating responses
|
|
master_file_path: The path to the agent master json file
|
|
|
|
Returns:
|
|
season_file_path: The path to the season json file
|
|
|
|
Raises:
|
|
Exception: If there's an error creating the season
|
|
|
|
Example:
|
|
>>> ai_model = GeminiModel()
|
|
>>> master_file_path = "agent_template.json"
|
|
>>> step_2(ai_model, master_file_path)
|
|
'''
|
|
print("Step 2: Create a new season")
|
|
|
|
# step 2.1: load the season template json file
|
|
manager = ContentGenerator()
|
|
#season_template = manager.create_new_template_yaml(TemplateType.SEASON)
|
|
print(f"Creating a new season template")
|
|
season_template = manager.create_new_template_json(TemplateType.SEASON)
|
|
|
|
# step 2.2: load the agent json file
|
|
print(f"Loading agent master json file")
|
|
agent_master_json = None
|
|
with open(master_file_path, 'r', encoding='utf-8') as file:
|
|
agent_master_json = json.load(file)
|
|
|
|
# extract agent from master json
|
|
print(f"Extracting agent details from master json")
|
|
agent_details = agent_master_json['agent']['agent_details']
|
|
|
|
|
|
# if the seasons array is empty or None, we need to initialize it
|
|
if not agent_master_json['agent']['seasons'] or agent_master_json['agent']['seasons'] == []:
|
|
agent_master_json['agent']['seasons'] = season_template['seasons']
|
|
|
|
# step 2.3: find the previous season
|
|
# get last season file
|
|
previous_season = None
|
|
# Get the last season from the seasons array if it exists
|
|
seasons = agent_master_json['agent']['seasons']
|
|
previous_season = seasons[-1] if seasons else None
|
|
|
|
if previous_season:
|
|
print(f"Found previous season: {previous_season['season_number']}")
|
|
else:
|
|
print("No previous seasons found")
|
|
|
|
# step 2.3: Generate a new season name, topic, and communication style with the prompt_2 template
|
|
# prompt 2 Season Creation:
|
|
# note that emojis will be output as unicode characters due to the json dump
|
|
print("Crafting prompt for AI to create a new season")
|
|
prompt_2_vars = {
|
|
"agent_name": agent_details["name"],
|
|
"agent_json": json.dumps(agent_details),
|
|
"season_json": json.dumps(season_template),
|
|
"previous_season": json.dumps(previous_season),
|
|
"number_of_episodes": number_of_episodes
|
|
}
|
|
|
|
# Constants for retry configuration
|
|
max_retries = 3
|
|
delay = 2 # seconds between retries
|
|
|
|
# step 3.11: Run the prompt with retry logic for LLM failures
|
|
success = False
|
|
for attempt in range(max_retries):
|
|
try:
|
|
# step 2.4: Run the prompt
|
|
print("Sending prompt to AI to create a new season")
|
|
season_data = manager.run_prompt(
|
|
# prompt_key="prompt_1 (Character Creation)",
|
|
prompt_key="prompt_2 (Season Creation)",
|
|
template_vars=prompt_2_vars,
|
|
ai_model=ai_model,
|
|
)
|
|
|
|
# Validate that posts_data is valid JSON and has expected structure
|
|
if isinstance(season_data, dict) and 'seasons' in season_data:
|
|
success = True
|
|
break
|
|
else:
|
|
print(f"Invalid response format. Attempt {attempt + 1}/{max_retries}")
|
|
|
|
except (json.JSONDecodeError, Exception) as e:
|
|
print(f"Error on attempt {attempt + 1}/{max_retries}: {str(e)}")
|
|
if attempt < max_retries - 1:
|
|
print(f"Retrying in {delay} seconds...")
|
|
time.sleep(delay)
|
|
continue
|
|
|
|
if not success:
|
|
print("Max retries reached. Failed to get valid response from LLM.")
|
|
return None
|
|
|
|
#print(f"season data is: {season_data}")
|
|
|
|
# step 2.8: append the season data to the master data
|
|
if len(seasons) <= 1 and seasons[0]["season_number"] == 0:
|
|
print("First season, so we need to initialize the seasons array")
|
|
agent_master_json = manager.initialize_seasons(
|
|
master_data=agent_master_json,
|
|
seasons_data=season_data
|
|
)
|
|
else:
|
|
print("Appending the season data to the master data")
|
|
agent_master_json = manager.append_seasons(
|
|
master_data=agent_master_json,
|
|
seasons_data=season_data,
|
|
)
|
|
|
|
# step 2.9: save the master data to a file
|
|
print("Saving the master data to a file")
|
|
manager.save_json_file(
|
|
save_path=master_file_path,
|
|
json_data=agent_master_json
|
|
)
|
|
|
|
print("Step 2 complete")
|
|
#return season_file_path
|
|
|
|
import models.gemini_model as gemini_model
|
|
if __name__ == "__main__":
|
|
ai_model = gemini_model.GeminiModel()
|
|
create_seasons_and_episodes(ai_model, "configs/CipherCat/CipherCat_master.json", 3) |