151 lines
5.2 KiB
Python
151 lines
5.2 KiB
Python
![]() |
#
|
||
|
# Module: step_1
|
||
|
#
|
||
|
# This module implements the step_1 function for creating a new agent.
|
||
|
#
|
||
|
# Title: Step 1
|
||
|
# Summary: Step 1 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 yaml
|
||
|
import json
|
||
|
import sys
|
||
|
import os
|
||
|
import time
|
||
|
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||
|
|
||
|
# custom ARAI imports
|
||
|
from utils.content_generator import ContentGenerator
|
||
|
import prompt_chaining.step_2_create_content as next_step
|
||
|
from utils.template_types import TemplateType
|
||
|
|
||
|
# -------------------------------------------------------------------
|
||
|
# Step 1: Create a new agent
|
||
|
# -------------------------------------------------------------------
|
||
|
def create_agent(ai_model, concept: str):
|
||
|
'''
|
||
|
Description:
|
||
|
Create a new agent
|
||
|
|
||
|
Args:
|
||
|
ai_model: The AI model to use for generating responses
|
||
|
debug (bool, optional): whether to print debug information. Defaults to False.
|
||
|
|
||
|
Returns:
|
||
|
agent_file_path: The path to the agent yaml file
|
||
|
|
||
|
Raises:
|
||
|
Exception: If there's an error creating the agent
|
||
|
|
||
|
Example:
|
||
|
>>> ai_model = GeminiModel()
|
||
|
>>> step_1(ai_model, "alien drone pilot who is a sarcastic asshole visiting earth to report back his findings to his home planet")
|
||
|
'''
|
||
|
|
||
|
# Step 1.1: Create a new agent
|
||
|
manager = ContentGenerator()
|
||
|
# agent_template = manager.create_new_template_yaml(TemplateType.AGENT)
|
||
|
agent_template = manager.create_new_template_json(TemplateType.AGENT)
|
||
|
agent_master_template = manager.create_new_template_json(TemplateType.MASTER)
|
||
|
|
||
|
# Gather a list of agent names based on folder names in config folder
|
||
|
agent_names_blacklist = manager.get_agent_names_blacklist()
|
||
|
# remove _ from names
|
||
|
# print(f"agent_names_blacklist: {agent_names_blacklist}")
|
||
|
# agent_names_blacklist = [name.replace("_", " ") for name in agent_names_blacklist]
|
||
|
print(f"agent_names_blacklist: {agent_names_blacklist}")
|
||
|
|
||
|
# step 1.2: Generate a new agent name, topic, personality, and communication style with the prompt_1 template
|
||
|
# prompt 1 Character Creation:
|
||
|
print("Crafting prompt for AI to create a new agent")
|
||
|
prompt_1_vars = {
|
||
|
# "agent_name": "",
|
||
|
# "personality": "",
|
||
|
# "communication_style": "",
|
||
|
# "topic": "",
|
||
|
"concept": concept,
|
||
|
"agent_names_blacklist": agent_names_blacklist,
|
||
|
"agent_json": json.dumps(agent_template)
|
||
|
}
|
||
|
|
||
|
# 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 1.3: Run the prompt
|
||
|
print("Sending prompt to AI to create a new agent")
|
||
|
agent_data = manager.run_prompt(
|
||
|
prompt_key="prompt_1 (Character Sheet Creation)",
|
||
|
template_vars=prompt_1_vars,
|
||
|
ai_model=ai_model,
|
||
|
debug=False
|
||
|
)
|
||
|
|
||
|
# Validate that posts_data is valid JSON and has expected structure
|
||
|
if isinstance(agent_data, dict) and 'agent' in agent_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
|
||
|
|
||
|
# step 1.4: Merge agent details into the master template
|
||
|
print("Merging agent details into the master template")
|
||
|
agent_master_template = manager.merge_agent_details(
|
||
|
master_data=agent_master_template,
|
||
|
agent_data=agent_data
|
||
|
)
|
||
|
|
||
|
# step 1.5: store the concept in the agent template
|
||
|
print("Storing the concept in the agent template")
|
||
|
agent_master_template["agent"]["concept"] = prompt_1_vars["concept"]
|
||
|
|
||
|
# step 1.6: create the file path for master file
|
||
|
print("Creating the file path for the master file")
|
||
|
agent_master_file_path = manager.create_filepath(
|
||
|
agent_name=agent_master_template["agent"]["agent_details"]["name"],
|
||
|
season_number=0,
|
||
|
episode_number=0,
|
||
|
template_type=TemplateType.MASTER
|
||
|
)
|
||
|
|
||
|
# step 1.7: Save the agent data to a file
|
||
|
print("Saving the agent data to a file")
|
||
|
manager.save_json_file(
|
||
|
save_path=agent_master_file_path,
|
||
|
json_data=agent_master_template
|
||
|
)
|
||
|
|
||
|
print("Step 1 complete")
|
||
|
return agent_master_file_path
|
||
|
|
||
|
|
||
|
import models.gemini_model as gemini_model
|
||
|
if __name__ == "__main__":
|
||
|
ai_model = gemini_model.GeminiModel()
|
||
|
create_agent(ai_model, "time-traveling historian documenting the evolution of human technology through the ages")
|
||
|
|