import gradio as gr import yaml import os import datetime from llm import AVAILABLE_LLMS, create_client, get_response_from_llm, get_batch_responses_from_llm from character import load_character_config, get_character_response, build_prompt, build_system_prompt # Directory containing character YAML files CHARACTER_DIR = "characters" # Directory to save conversations CONVERSATION_DIR = "conversations" def list_character_files(): """List all YAML files in the character directory.""" return [f for f in os.listdir(CHARACTER_DIR) if f.endswith('.yaml')] def load_all_characters(): """Load all character configurations from YAML files.""" characters = {} character_files = list_character_files() for filename in character_files: filepath = os.path.join(CHARACTER_DIR, filename) config = load_character_config(filepath) character_name = config['character']['name'] characters[character_name] = {'config': config, 'file': filepath} return characters # Load all characters all_characters = load_all_characters() class ChatApp: """A class to encapsulate the chat application logic.""" def __init__(self): self.available_llms = AVAILABLE_LLMS self.api_key_dict = {} # To store API keys for different LLMs self.all_characters = all_characters self.character_config = None self.load_default_character() # Ensure the conversation directory exists if not os.path.exists(CONVERSATION_DIR): os.makedirs(CONVERSATION_DIR) def load_default_character(self): """Load the default character configuration.""" if self.all_characters: default_character_name = list(self.all_characters.keys())[0] self.character_config = self.all_characters[default_character_name]['config'] else: self.character_config = None def set_api_key(self, api_key, selected_llm): """Set the API key based on the selected LLM.""" api_key = api_key.strip() self.api_key_dict[selected_llm] = api_key # Store the API key return gr.update(value='', placeholder='API Key Set!') def select_character(self, character_name): """Update the selected character.""" if character_name == "New Character": # Initialize a new character configuration with empty fields self.character_config = { 'character': { 'name': '', 'pronouns': '', 'alternate_names': [], 'age': '', 'core_description': '', 'motivations': [], 'flaws': [], 'dialogue_style': '', 'example_dialogue': [] }, 'personality': { 'traits': [], 'mood': {}, 'personality': {} }, 'knowledge_and_cognition': { 'details': [], 'dementia_facts': [], 'loneliness_info': [] } } # Clear the UI components return ('', '', '', '', '', '', '', '', '', '', '', '', '', '', '', f"Creating a new character.") elif character_name in self.all_characters: self.character_config = self.all_characters[character_name]['config'] # Update the UI components with the character's data return ( self.character_config['character']['name'], self.character_config['character']['pronouns'], ', '.join(self.character_config['character'].get('alternate_names', [])), self.character_config['character']['age'], self.character_config['character']['core_description'], '\n'.join(self.character_config['character']['motivations']), '\n'.join(self.character_config['character']['flaws']), self.character_config['character']['dialogue_style'], '\n'.join([f"{list(d.keys())[0]}: {list(d.values())[0]}" for d in self.character_config['character']['example_dialogue']]), ', '.join(self.character_config['personality']['traits']), ', '.join([f"{k}: {v}" for k, v in self.character_config['personality']['mood'].items()]), ', '.join([f"{k}: {v}" for k, v in self.character_config['personality']['personality'].items()]), '\n'.join(self.character_config['knowledge_and_cognition']['details']), '\n'.join(self.character_config['knowledge_and_cognition']['dementia_facts']), '\n'.join(self.character_config['knowledge_and_cognition']['loneliness_info']), f"Character {character_name} loaded successfully." ) else: return (gr.update(),) * 15 + (f"Character {character_name} not found.",) def update_and_save_character_config( self, name, pronouns, alternate_names, age, core_description, motivations, flaws, dialogue_style, example_dialogue, traits, mood, personality, details, dementia_facts, loneliness_info, selected_character_name ): """Update the character configuration based on user input and save it.""" if self.character_config is None: return "No character selected.", gr.update() # Update character details self.character_config['character']['name'] = name self.character_config['character']['pronouns'] = pronouns self.character_config['character']['alternate_names'] = [n.strip() for n in alternate_names.split(',') if n.strip()] self.character_config['character']['age'] = age self.character_config['character']['core_description'] = core_description self.character_config['character']['motivations'] = [m.strip() for m in motivations.split('\n') if m.strip()] self.character_config['character']['flaws'] = [f.strip() for f in flaws.split('\n') if f.strip()] self.character_config['character']['dialogue_style'] = dialogue_style # Process example dialogue example_dialogue_list = [] if example_dialogue.strip(): lines = example_dialogue.strip().split('\n') for line in lines: if ':' in line: speaker, text = line.split(':', 1) example_dialogue_list.append({speaker.strip(): text.strip()}) self.character_config['character']['example_dialogue'] = example_dialogue_list # Update personality traits self.character_config['personality']['traits'] = [trait.strip() for trait in traits.split(',') if trait.strip()] # Update mood mood_dict = {} for item in mood.split(','): if ':' in item: key, value = item.split(':', 1) mood_dict[key.strip()] = float(value.strip()) self.character_config['personality']['mood'] = mood_dict # Update personality metrics personality_dict = {} for item in personality.split(','): if ':' in item: key, value = item.split(':', 1) personality_dict[key.strip()] = float(value.strip()) self.character_config['personality']['personality'] = personality_dict # Update knowledge and cognition self.character_config['knowledge_and_cognition']['details'] = [d.strip() for d in details.split('\n') if d.strip()] self.character_config['knowledge_and_cognition']['dementia_facts'] = [d.strip() for d in dementia_facts.split('\n') if d.strip()] self.character_config['knowledge_and_cognition']['loneliness_info'] = [d.strip() for d in loneliness_info.split('\n') if d.strip()] if selected_character_name == "New Character": # Create a new character file new_character_name = name.strip() if not new_character_name: return "Please enter a name for the new character.", gr.update() # Generate a filename based on the character's name filename = new_character_name.lower().replace(' ', '_') + '.yaml' filepath = os.path.join(CHARACTER_DIR, filename) if os.path.exists(filepath): return f"A character with the name '{new_character_name}' already exists.", gr.update() else: try: with open(filepath, "w") as file: yaml.dump(self.character_config, file) # Update the all_characters dictionary and the dropdown options self.all_characters[new_character_name] = {'config': self.character_config, 'file': filepath} # Update the character dropdown options character_names = list(self.all_characters.keys()) return ( f"New character '{new_character_name}' created and saved successfully.", gr.update(choices=["New Character"] + character_names, value=new_character_name) ) except Exception as e: return f"Error saving new character: {e}", gr.update() else: # Save the updated character configuration try: filepath = self.all_characters[selected_character_name]['file'] with open(filepath, "w") as file: yaml.dump(self.character_config, file) return f"Character configuration for '{selected_character_name}' updated and saved successfully.", gr.update() except Exception as e: return f"Error saving character configuration: {e}", gr.update() def chat(self, user_input, chat_history, selected_llm): """Handle the chat interaction.""" if not user_input: return "", [] if self.character_config is None: return "", [{"role": "assistant", "content": "No character selected."}] # Set the API key based on the selected LLM api_key = self.api_key_dict.get(selected_llm, None) if api_key: self.set_environment_api_key(selected_llm, api_key) else: return "", [{"role": "assistant", "content": "Please set the API key for the selected LLM."}] # Get the character's response using the provided function try: response_content = get_character_response(user_input, self.character_config, llm_model=selected_llm) # Create new messages list with proper format messages = chat_history + [ {"role": "user", "content": user_input}, {"role": "assistant", "content": response_content} ] return "", messages except Exception as e: error_message = [{"role": "assistant", "content": f"Error during LLM processing: {e}"}] return "", error_message def set_environment_api_key(self, llm, api_key): """Set the environment variable for the API key based on the LLM.""" if llm.startswith('gpt') or llm in ["o1-preview-2024-09-12", "o1-mini-2024-09-12"]: os.environ["OPENAI_API_KEY"] = api_key elif llm.startswith('deepseek'): os.environ["DEEPSEEK_API_KEY"] = api_key elif llm in ["llama3.1-405b", "llama3.1-405b-instruct"]: os.environ["OPENROUTER_API_KEY"] = api_key def save_conversation(self, chat_history): """Save the conversation to a file in the /data/conversations directory.""" # Use the persistent storage path for Hugging Face Spaces folder = os.path.join("/data/conversations") if not os.path.exists(folder): os.makedirs(folder, exist_ok=True) timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S_%f") character_name = self.character_config['character']['name'] if self.character_config else "AI" safe_char_name = "".join(c for c in character_name if c.isalnum() or c in ('-', '_')).lower() filename = f"chat_{timestamp}_{safe_char_name}.txt" filepath = os.path.join(folder, filename) try: with open(filepath, 'w', encoding='utf-8') as f: f.write(f"Conversation with {character_name}\n") f.write(f"Timestamp: {timestamp}\n") f.write("-" * 50 + "\n\n") for message in chat_history: role = message["role"] content = message["content"] if role == "user": f.write(f"User: {content}\n") else: f.write(f"{character_name}: {content}\n") f.write("\n") return f"Conversation saved as {filename}" except Exception as e: return f"Error saving conversation: {e}" def new_conversation(self): """Clear the chat history.""" return [], [] def collect_data(self, user_question, k, selected_llm): """Have LLM respond K times to a user question for data collection.""" if not user_question: return "Please enter a question.", None if self.character_config is None: return "No character selected.", None api_key = self.api_key_dict.get(selected_llm, None) if api_key: self.set_environment_api_key(selected_llm, api_key) else: return "Please set the API key for the selected LLM.", None prompt = build_prompt(user_question, self.character_config) system_prompt = build_system_prompt(self.character_config['character']['name']) client, model = create_client(selected_llm) try: responses, _ = get_batch_responses_from_llm( msg=prompt, client=client, model=model, system_message=system_prompt, temperature=0.9, n_responses=int(k) ) # Format responses for the Dataframe with separate up/down columns responses_data = [[str(i+1), resp, "⬆️", "⬇️"] for i, resp in enumerate(responses)] return "", responses_data except Exception as e: return f"Error during LLM processing: {e}", None def save_data_collection(self, user_question, ranked_responses): """Save the question, responses, and their rankings in a format optimized for DPO training.""" # Use the persistent storage path for Hugging Face Spaces folder = os.path.join("/data/dpo_training") if not os.path.exists(folder): os.makedirs(folder, exist_ok=True) try: timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S_%f") character_config = self.character_config or {} character_name = character_config.get('character', {}).get('name', "Unknown") safe_char_name = "".join(c for c in character_name if c.isalnum() or c in ('-', '_')).lower() jsonl_filename = f"dpo_{timestamp}_{safe_char_name}.jsonl" jsonl_file = os.path.join(folder, jsonl_filename) system_prompt = build_system_prompt(character_name) if self.character_config else "" import json with open(jsonl_file, 'w', encoding='utf-8') as f: for i, (rank1, chosen_response, _, _) in enumerate(ranked_responses[:-1]): for j, (rank2, rejected_response, _, _) in enumerate(ranked_responses[i+1:], i+1): messages = [ { "role": "system", "content": system_prompt }, { "role": "user", "content": user_question } ] dpo_example = { "messages": messages, "chosen": chosen_response, "rejected": rejected_response, "prompt": user_question, "system_prompt": system_prompt, "character_name": character_name, "ranking_info": { "chosen_rank": int(rank1), "rejected_rank": int(rank2) }, "metadata": { "timestamp": timestamp, "session_id": timestamp } } f.write(json.dumps(dpo_example, ensure_ascii=False) + '\n') return f"Rankings saved to {jsonl_filename}" except Exception as e: return f"Error saving rankings: {e}" def move_row(self, data, evt: gr.SelectData): """Move a row up or down based on which column was clicked.""" if not data or not evt: return data row_idx = evt.index[0] col_idx = evt.index[1] # Get column index instead of name if col_idx == 2 and row_idx > 0: # Up column (index 2) # Swap with row above data[row_idx], data[row_idx-1] = data[row_idx-1], data[row_idx] elif col_idx == 3 and row_idx < len(data) - 1: # Down column (index 3) # Swap with row below data[row_idx], data[row_idx+1] = data[row_idx+1], data[row_idx] # Update ranks for i, row in enumerate(data): row[0] = str(i + 1) return data # Instantiate the ChatApp app = ChatApp() with gr.Blocks() as demo: with gr.Tab("Settings"): gr.Markdown("## LLM Selection & Configuration") llm_dropdown = gr.Dropdown( label="Select LLM", choices=app.available_llms ) api_key_input = gr.Textbox( label="Enter API Key (if required)", type="password", placeholder="Your API Key" ) set_api_button = gr.Button("Set API Key") set_api_button.click( app.set_api_key, inputs=[api_key_input, llm_dropdown], outputs=api_key_input ) gr.Markdown("## AI Character Configuration") # Character Selection Dropdown character_names = list(app.all_characters.keys()) character_dropdown = gr.Dropdown( label="Select Character", choices=["New Character"] + character_names, value=character_names[0] if character_names else "New Character" ) # Character Details with gr.Accordion("Character Details", open=True): name_input = gr.Textbox(label="Name") pronouns_input = gr.Textbox(label="Pronouns") alternate_names_input = gr.Textbox(label="Alternate Names (comma-separated)") age_input = gr.Textbox(label="Age") core_description_input = gr.Textbox(label="Core Description", lines=5) motivations_input = gr.Textbox(label="Motivations (one per line)", lines=3) flaws_input = gr.Textbox(label="Flaws (one per line)", lines=3) dialogue_style_input = gr.Textbox(label="Dialogue Style") example_dialogue_input = gr.Textbox(label="Example Dialogue (format: speaker: text)", lines=5) # Personality Traits with gr.Accordion("Personality", open=False): traits_input = gr.Textbox(label="Traits (comma-separated)") mood_input = gr.Textbox(label="Mood (format: emotion: value, comma-separated)") personality_input = gr.Textbox(label="Personality Metrics (format: trait: value, comma-separated)") # Knowledge and Cognition with gr.Accordion("Knowledge and Cognition", open=False): details_input = gr.Textbox(label="Details (one per line)", lines=5) dementia_facts_input = gr.Textbox(label="Dementia Facts (one per line)", lines=5) loneliness_info_input = gr.Textbox(label="Loneliness Info (one per line)", lines=5) # Update and Save Button update_and_save_button = gr.Button("Update and Save Character Configuration") update_status = gr.Textbox(label="Status", interactive=False) update_and_save_button.click( app.update_and_save_character_config, inputs=[ name_input, pronouns_input, alternate_names_input, age_input, core_description_input, motivations_input, flaws_input, dialogue_style_input, example_dialogue_input, traits_input, mood_input, personality_input, details_input, dementia_facts_input, loneliness_info_input, character_dropdown ], outputs=[update_status, character_dropdown], ) # Character Selection Handling character_dropdown.change( app.select_character, inputs=[character_dropdown], outputs=[ name_input, pronouns_input, alternate_names_input, age_input, core_description_input, motivations_input, flaws_input, dialogue_style_input, example_dialogue_input, traits_input, mood_input, personality_input, details_input, dementia_facts_input, loneliness_info_input, update_status ] ) with gr.Tab("Chat"): gr.Markdown("## Chat Interface") chatbot = gr.Chatbot(type="messages") user_input = gr.Textbox(label="Your Message:", placeholder="Type your message here...") send_button = gr.Button("Send") def submit_message(user_input, chat_history, selected_llm): new_user_input, chat_history = app.chat(user_input, chat_history, selected_llm) return "", chat_history send_button.click( submit_message, inputs=[user_input, chatbot, llm_dropdown], outputs=[user_input, chatbot] ) user_input.submit( submit_message, inputs=[user_input, chatbot, llm_dropdown], outputs=[user_input, chatbot] ) new_conversation_button = gr.Button("🆕 New Conversation") new_conversation_button.click(app.new_conversation, outputs=[chatbot, chatbot]) save_button = gr.Button("💾 Save Conversation") save_status = gr.Textbox(label="Save Status", interactive=False) save_button.click(app.save_conversation, inputs=[chatbot], outputs=save_status) with gr.Tab("Data Collection"): gr.Markdown(""" ## Data Collection Interface This interface helps collect multiple AI responses for the same question to evaluate response quality. ### How to use: 1. Enter your question 2. Choose how many responses you want 3. Generate responses 4. Use ⬆️ and ⬇️ buttons to reorder responses (top = best) 5. Save the rankings """) with gr.Row(): with gr.Column(scale=3): data_question_input = gr.Textbox( label="Question for the AI Character", placeholder="Type your question here...", lines=3 ) with gr.Column(scale=1): k_input = gr.Slider( minimum=2, maximum=10, value=5, step=1, label="Number of Responses to Generate" ) llm_dropdown_data = gr.Dropdown( label="Select Language Model", choices=app.available_llms, value=app.available_llms[0] if app.available_llms else None ) generate_button = gr.Button("🔄 Generate Responses", variant="primary") collection_status = gr.Textbox( label="Generation Status", interactive=False, visible=False ) # New interface for ranking responses responses_df = gr.Dataframe( headers=["Rank", "Response", "Up", "Down"], datatype=["str", "str", "str", "str"], col_count=(4, "fixed"), interactive=True, wrap=True, row_count=10, label="Click ⬆️ or ⬇️ to reorder responses (top = best)", type="array" ) def move_row(data, evt: gr.SelectData): """Move a row up or down based on which column was clicked.""" if not data or not evt: return data row_idx = evt.index[0] col_idx = evt.index[1] # Get column index instead of name if col_idx == 2 and row_idx > 0: # Up column (index 2) # Swap with row above data[row_idx], data[row_idx-1] = data[row_idx-1], data[row_idx] elif col_idx == 3 and row_idx < len(data) - 1: # Down column (index 3) # Swap with row below data[row_idx], data[row_idx+1] = data[row_idx+1], data[row_idx] # Update ranks for i, row in enumerate(data): row[0] = str(i + 1) return data # Add click handler for both Up and Down columns responses_df.select( move_row, inputs=[responses_df], outputs=[responses_df] ) submit_ranking_button = gr.Button("💾 Save Rankings", variant="secondary") data_save_status = gr.Textbox( label="Save Status", interactive=False, visible=False ) # Show status messages when they contain content collection_status.change( lambda x: gr.update(visible=bool(x.strip())), inputs=[collection_status], outputs=[collection_status] ) data_save_status.change( lambda x: gr.update(visible=bool(x.strip())), inputs=[data_save_status], outputs=[data_save_status] ) generate_button.click( app.collect_data, inputs=[data_question_input, k_input, llm_dropdown_data], outputs=[collection_status, responses_df] ) submit_ranking_button.click( app.save_data_collection, inputs=[data_question_input, responses_df], outputs=[data_save_status] ) # Initialize UI components with default character data if app.character_config: character_dropdown.value = app.character_config['character']['name'] name_input.value = app.character_config['character']['name'] pronouns_input.value = app.character_config['character']['pronouns'] alternate_names_input.value = ', '.join(app.character_config['character'].get('alternate_names', [])) age_input.value = app.character_config['character']['age'] core_description_input.value = app.character_config['character']['core_description'] motivations_input.value = '\n'.join(app.character_config['character']['motivations']) flaws_input.value = '\n'.join(app.character_config['character']['flaws']) dialogue_style_input.value = app.character_config['character']['dialogue_style'] example_dialogue_input.value = '\n'.join([f"{list(d.keys())[0]}: {list(d.values())[0]}" for d in app.character_config['character']['example_dialogue']]) traits_input.value = ', '.join(app.character_config['personality']['traits']) mood_input.value = ', '.join([f"{k}: {v}" for k, v in app.character_config['personality']['mood'].items()]) personality_input.value = ', '.join([f"{k}: {v}" for k, v in app.character_config['personality']['personality'].items()]) details_input.value = '\n'.join(app.character_config['knowledge_and_cognition']['details']) dementia_facts_input.value = '\n'.join(app.character_config['knowledge_and_cognition']['dementia_facts']) loneliness_info_input.value = '\n'.join(app.character_config['knowledge_and_cognition']['loneliness_info']) demo.launch()