Create app.py

app.py

@@ -0,0 +1,290 @@
+import gradio as gr
+import random
+import openai
+from openai import APIError, APIConnectionError, RateLimitError
+import os
+from PIL import Image  # This is the corrected import
+import io
+import base64
+import asyncio
+from queue import Queue
+from threading import Thread
+import time
+
+# Get the current script's directory
+current_dir = os.path.dirname(os.path.abspath(__file__))
+avatars_dir = os.path.join(current_dir, "avatars")
+
+# Dictionary mapping characters to their avatar image filenames
+character_avatars = {
+    "Harry Potter": "harry.png",
+    "Hermione Granger": "hermione.png",
+    "poor Ph.D. student": "phd.png",
+    "Donald Trump": "trump.png",
+    "a super cute red panda": "red_panda.png"
+}
+
+
+predefined_characters = ["Harry Potter", "Hermione Granger",  "poor Ph.D. student", "Donald Trump", "a super cute red panda"]
+
+def get_character(dropdown_value, custom_value):
+    return custom_value if dropdown_value == "Custom" else dropdown_value
+
+def resize_image(image_path, size=(100, 100)):
+    if not os.path.exists(image_path):
+        return None
+    with Image.open(image_path) as img:
+        img.thumbnail(size)
+        buffered = io.BytesIO()
+        img.save(buffered, format="PNG")
+        return base64.b64encode(buffered.getvalue()).decode()
+
+resized_avatars = {}
+for character, filename in character_avatars.items():
+    full_path = os.path.join(avatars_dir, filename)
+    if os.path.exists(full_path):
+        resized_avatars[character] = resize_image(full_path)
+    else:
+        pass
+
+async def generate_response_stream(messages, api_key):
+    client = openai.AsyncOpenAI(
+        api_key=api_key,
+        base_url="https://api.sambanova.ai/v1",
+    )
+    try:
+        if len(messages) >= 10:
+            # avoid hitting rate limit
+            time.sleep(0.5)
+        response = await client.chat.completions.create(
+            model='Meta-Llama-3.1-405B-Instruct',
+            messages=messages,
+            temperature=0.7,
+            top_p=0.9,
+            stream=True
+        )
+        full_response = ""
+        async for chunk in response:
+            if chunk.choices[0].delta.content is not None:
+                full_response += chunk.choices[0].delta.content
+                yield full_response
+    except Exception as e:
+        yield f"Error: {str(e)}"
+
+async def simulate_conversation_stream(character1, character2, initial_message, num_turns, api_key):
+    messages_character_1 = [{"role": "system", "content": f"Avoid overly verbose answer in your response. Act as {character1}."},
+                            {"role": "assistant", "content": initial_message}]
+    messages_character_2 = [{"role": "system", "content": f"Avoid overly verbose answer in your response. Act as {character2}."},
+                            {"role": "user", "content": initial_message}]
+    
+    conversation = [
+        {"character": character1, "content": initial_message},
+        {"character": character2, "content": ""}  # Initialize with an empty response for character2
+    ]
+    yield format_conversation_as_html(conversation)
+    num_turns *= 2
+    for turn_num in range(num_turns - 1):
+        current_character = character2 if turn_num % 2 == 0 else character1
+        messages = messages_character_2 if turn_num % 2 == 0 else messages_character_1
+
+        full_response = ""
+        async for response in generate_response_stream(messages, api_key):
+            full_response = response
+            conversation[-1]["content"] = full_response
+            yield format_conversation_as_html(conversation)
+
+        if turn_num % 2 == 0:
+            messages_character_1.append({"role": "user", "content": full_response})
+            messages_character_2.append({"role": "assistant", "content": full_response})
+        else:
+            messages_character_2.append({"role": "user", "content": full_response})
+            messages_character_1.append({"role": "assistant", "content": full_response})
+        
+        # Add a new empty message for the next turn, if it's not the last turn
+        if turn_num < num_turns - 2:
+            next_character = character1 if turn_num % 2 == 0 else character2
+            conversation.append({"character": next_character, "content": ""})
+
+def stream_conversation(character1, character2, initial_message, num_turns, api_key, queue):
+    async def run_simulation():
+        async for html in simulate_conversation_stream(character1, character2, initial_message, num_turns, api_key):
+            queue.put(html)
+        queue.put(None)  # Signal that the conversation is complete
+
+    asyncio.run(run_simulation())
+
+def chat_interface(character1_dropdown, character1_custom, character2_dropdown, character2_custom, 
+    initial_message, num_turns, api_key):
+
+    character1 = get_character(character1_dropdown, character1_custom)
+    character2 = get_character(character2_dropdown, character2_custom)
+
+    queue = Queue()
+    thread = Thread(target=stream_conversation, args=(character1, character2, initial_message, num_turns, api_key, queue))
+    thread.start()
+
+    while True:
+        result = queue.get()
+        if result is None:
+            break
+        yield result
+    
+    thread.join()
+
+def format_conversation_as_html(conversation):
+    html_output = """
+    <style>
+    .chat-container {
+        display: flex;
+        flex-direction: column;
+        gap: 10px;
+        font-family: Arial, sans-serif;
+    }
+    .message {
+        display: flex;
+        padding: 10px;
+        border-radius: 10px;
+        max-width: 80%;
+        align-items: flex-start;
+    }
+    .left {
+        align-self: flex-start;
+        background-color: #1565C0;
+        color: #FFFFFF;
+    }
+    .right {
+        align-self: flex-end;
+        background-color: #2E7D32;
+        color: #FFFFFF;
+        flex-direction: row-reverse;
+    }
+    .avatar-container {
+        flex-shrink: 0;
+        width: 40px;
+        height: 40px;
+        margin: 0 10px;
+    }
+    .avatar {
+        width: 100%;
+        height: 100%;
+        border-radius: 50%;
+        object-fit: cover;
+    }
+    .message-content {
+        display: flex;
+        flex-direction: column;
+        min-width: 150px;
+        flex-grow: 1;
+    }
+    .character-name {
+        font-weight: bold;
+        margin-bottom: 5px;
+    }
+    .message-text {
+        word-wrap: break-word;
+        overflow-wrap: break-word;
+    }
+    </style>
+    <div class="chat-container">
+    """
+    
+    for i, message in enumerate(conversation):
+        align = "left" if i % 2 == 0 else "right"
+        avatar_data = resized_avatars.get(message["character"])
+        
+        html_output += f'<div class="message {align}">'
+        
+        if avatar_data:
+            html_output += f'''
+            <div class="avatar-container">
+                <img src="data:image/png;base64,{avatar_data}" class="avatar" alt="{message["character"]} avatar">
+            </div>
+            '''
+        
+        html_output += f'''
+        <div class="message-content">
+            <div class="character-name">{message["character"]}</div>
+            <div class="message-text">{message["content"]}</div>
+        </div>
+        </div>
+        '''
+    
+    html_output += "</div>"
+    return html_output
+
+
+def format_chat_for_download(html_chat):
+    # Extract text content from HTML
+    import re
+    chat_text = re.findall(r'<div class="character-name">(.*?)</div>.*?<div class="message-text">(.*?)</div>', html_chat, re.DOTALL)
+    return "\n".join([f"{speaker.strip()}: {message.strip()}" for speaker, message in chat_text])
+
+def save_chat_to_file(chat_content):
+    # Create a downloads directory if it doesn't exist
+    downloads_dir = os.path.join(os.getcwd(), "downloads")
+    os.makedirs(downloads_dir, exist_ok=True)
+    
+    # Generate a unique filename
+    import datetime
+    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+    filename = f"chat_{timestamp}.txt"
+    file_path = os.path.join(downloads_dir, filename)
+    
+    # Save the chat content to the file
+    with open(file_path, "w", encoding="utf-8") as f:
+        f.write(chat_content)
+    
+    return file_path
+    
+
+with gr.Blocks() as app:
+    gr.Markdown("# Character Chat Generator")
+    
+    gr.Markdown("Powerd by [LLama3.1-405B-Instruct](https://huggingface.co/meta-llama/Meta-Llama-3.1-405B-Instruct) on [SambaNova Cloud](https://cloud.sambanova.ai/apis)")
+    api_key = gr.Textbox(label="Enter your Sambanova Cloud API Key\n(To get one, go to https://cloud.sambanova.ai/apis)", type="password")
+    
+    with gr.Column():
+            character1_dropdown = gr.Dropdown(choices=predefined_characters + ["Custom"], label="Select Character 1")
+            character1_custom = gr.Textbox(label="Custom Character 1 (if selected above)", visible=False)
+    with gr.Column():
+        character2_dropdown = gr.Dropdown(choices=predefined_characters + ["Custom"], label="Select Character 2")
+        character2_custom = gr.Textbox(label="Custom Character 2 (if selected above)", visible=False)
+    
+    initial_message = gr.Textbox(label="Initial message (for Character 1)")
+    num_turns = gr.Slider(minimum=1, maximum=10, step=1, value=5, label="Number of conversation turns")
+    
+    generate_btn = gr.Button("Generate Conversation")
+    output = gr.HTML(label="Generated Conversation")
+    
+
+    def show_custom_input(choice):
+        return gr.update(visible=choice == "Custom")
+
+    character1_dropdown.change(show_custom_input, inputs=character1_dropdown, outputs=character1_custom)
+    character2_dropdown.change(show_custom_input, inputs=character2_dropdown, outputs=character2_custom)
+    
+    generate_btn.click(
+        chat_interface,
+        inputs=[character1_dropdown, character1_custom, character2_dropdown, 
+        character2_custom, initial_message, num_turns, api_key],
+        outputs=output,
+    )
+    
+    gr.Markdown("## Download Chat History")
+    
+    download_btn = gr.Button("Download Conversation")
+    download_output = gr.File(label="Download")
+
+    def download_conversation(html_chat):
+        chat_content = format_chat_for_download(html_chat)
+        file_path = save_chat_to_file(chat_content)
+        return file_path
+
+    download_btn.click(
+        download_conversation,
+        inputs=output,
+        outputs=download_output
+    )
+
+if __name__ == "__main__":
+    app.launch()