Upgrading Your AI Friend: Building a Gradio GUI for AutoGen and Mem0 Chatbots

Lennex Zinyando

10 Sep 2024 • 6 min read

In a previous post, I walked you through creating an AI-powered friend chatbot built with AutoGen and Mem0. You can read the blog post here and view the code on GitHub. In this post, we are going to evolve the AI friend chatbot to have a GUI powered by Gradio. This will make it a lot more user-friendly and easier to get started with. Whether you're a beginner looking to dip your toes into AI development or an experienced programmer seeking to expand your skills, this tutorial has something for everyone.

If you are in a rush and just want to see the code head over to GitHub.

Core Components

Our AI friend chatbot relies on several key components:

AutoGen: A framework for building conversational AI agents.
Gradio: A library for creating web-based interfaces for machine learning models.
Mem0: A memory management system for AI applications.
Chroma: A vector database for efficient storage and retrieval of conversational context.

Setting Up the Environment

Follow the setup instructions from the previous post with one addition: we are going to be adding Gradio as well.

pip install python-dotenv pyautogen mem0ai gradio chromadb groq

Getting the `OPENAI_API_KEY`

For embeddings, Mem0 uses text-embedding-3-small from OpenAI by default. To use it, we need to obtain an OPENAI_API_KEY. This will be used to create the memory embeddings that will be stored in the chromadb vector store. Head over to https://platform.openai.com/api-keys and grab an API key. If you don't have an account, you might need to register first. You might need to buy credits worth $5 to get started, and if your account is new, you might qualify for some free credits to get started.

Getting the `GROQ_API_KEY`

Groq will provide our language model. Follow these steps:

Go to https://console.groq.com/ and create an account
Navigate to the API keys section
Create a new API key for the project

With these steps completed, your environment should be ready for building our AI friend.

Building your AI friend

Create a folder for the application; you can name it gradio-ai-friend. Inside the folder, create a .env file where you will put your secret keys:

GROQ_API_KEY=gsk_secret_key_xxxxxxxxxxxxxxxxxxx
OPENAI_API_KEY=sk-proj-secret-key-xxxxxxxxxxxxx

Import Required Libraries

Create a file named ai-friend.py in your project directory. We'll build our AI friend in this file.

First, let's import the necessary libraries:

import os
from dotenv import load_dotenv
from autogen import ConversableAgent
from mem0 import Memory
import gradio as gr
import logging

load_dotenv()

These imports set us up with AutoGen for creating our chatbot, Mem0 for managing the agent's memory, Gradio for the user interface, and ChromaDB for our vector store, as well as loading environment variables using the dotenv library.

Configuring the Language Model

We're using the Groq API with the llama-3.1-8b-instant model. Here's how we set it up:

llm_config = {
    "config_list": [
        {
            "model": "llama-3.1-8b-instant",
            "api_key": os.getenv("GROQ_API_KEY"),
            "api_type": "groq",
        }
    ]
}

This configuration allows us to leverage a powerful language model for generating human-like responses.

Creating the AI Agent

We create our AI friend agent using AutoGen's ConversableAgent:

ai_friend_agent = ConversableAgent(
    name="Hazel",
    system_message="You are an AI friend chatbot",
    llm_config=llm_config,
    code_execution_config=False,
    human_input_mode="NEVER",
)

This agent, named Hazel, serves as the core of our chatbot's personality and conversational abilities.

Implementing Memory with Chroma

To give our chatbot a sense of context and history, we use Chroma as a vector store:

config = {
    "vector_store": {
        "provider": "chroma",
        "config": {
            "collection_name": "ai_friend_chatbot_memory",
            "path": "./chroma_db",
        },
    },
}

memory = Memory.from_config(config)

This setup enables Hazel to remember past interactions and provide more contextually relevant responses.

Generating Dynamic Prompts

The create_prompt function is crucial for generating context-aware responses:

def create_prompt(user_input, user_id, chat_history):
    """Generate the prompt based on user input, conversation history, and memory context."""
    memories = memory.search(user_input, user_id=user_id)
    context = "\\n".join([m["memory"] for m in memories])

    latest_messages = chat_history[-50:]
    history = "\\n".join([f"{name}: {message}" for name, message in latest_messages])

    prompt = f"""
    You are a warm, empathetic listener with a knack for understanding and responding to your user's needs. 
    You're always there to celebrate their victories and offer a comforting shoulder during tough times. 
    With your vast knowledge and emotional intelligence, Hazel is your go-to companion for insightful advice, 
    collaborative problem-solving, and meaningful conversations.

    Your goal is to be a positive, reassuring force in the user's life – a trusted companion they can rely on. 
    By building a rewarding, authentic friendship with the user, Hazel strives to be a source of support, 
    encouragement, and meaningful connection.

    Through contextual awareness and personalized responses, you adapt your communication style to the user's unique personality 
    and preferences, creating a tailored, immersive experience.

    Remember:

    1. Always strive to be helpful, supportive, and understanding.
    2. Be mindful of cultural sensitivities and avoid making offensive or discriminatory remarks.
    3. Use a conversational tone that is natural and engaging.
    4. Try to maintain a sense of humor and positivity.
    5. Be open to learning new things and adapting to different situations.

    IMPORTANT:

    DO NOT GREET THE USER ON EVERY INTERACTION UNLESS IT'S BEEN A SIGNIFICANT 
    AMOUNT OF TIME SINCE THE LAST INTERACTION.

    DO NOT SAY YOUR NAME ON EVERY INTERACTION UNLESS IT'S IMPORTANT FOR THE CONTEXT.

    Memories:
    {context}

    Chat History:
    {history}

    User's name is {user_id}
    User's input: {user_input}
    """
    return prompt

This function retrieves relevant memories and constructs a prompt that guides Hazel's responses, ensuring they are personalized and contextually appropriate. There are 3 important components to the prompt:

Personality: personality of the AI friend. This guides how responses will be formed and be consistent.
Memories: Memories from Mem0 will give the AI friend context of past conversations.
Chat history: Current conversation history; this gives the AI friend a record of the current conversation between the user and the AI friend. This guides the AI friend on how to respond to the next user input.

Handling Chat Interactions

The chatbot_response function manages the flow of conversation:

def chatbot_response(user_input, chat_history, user_id):
    """Handle chat interaction and update the UI."""
    if user_input:
        # Immediately display the user's message in the chat
        chat_history.append((f"{user_id}", user_input))

        # Generate a prompt and get a response from the chatbot
        prompt = create_prompt(user_input, user_id, chat_history)
        reply = ai_friend_agent.generate_reply(
            messages=[{"content": prompt, "role": "user"}]
        )

        # Store the conversation in Mem0
        memory.add(f"{user_id}: {user_input}", user_id=user_id)

        # Add the bot's reply to the chat history
        chat_history.append(("Hazel", reply["content"]))

    return chat_history, ""

This function processes user input, generates a response, and updates the chat history and memory.

Building the Gradio Interface

We use Gradio to create an intuitive web interface for our chatbot:

def start_chat(name):
    """Function to start the chat by getting the user's name."""
    if name:
        return (
            gr.update(visible=False),
            gr.update(visible=False),
            gr.update(visible=True),
            gr.update(visible=True),
            name,
        )

# Create Gradio interface
with gr.Blocks() as ai_friend:
    # Input section for the user to enter their name
    with gr.Row() as name_group:
        name_input = gr.Textbox(
            label="Enter your name to start the chat",
            placeholder="Your name here",
            interactive=True,
        )
        start_button = gr.Button("Start Chat")

    # Chat interface elements that will appear after name is entered
    with gr.Group(visible=False) as chat_group:
        chat_interface = gr.Chatbot(
            label="Chat with Hazel",
            height="80vh",
            bubble_full_width=False,
            show_label=False,
        )
        user_input = gr.Textbox(
            label="Your message",
            placeholder="Type your message here...",
            interactive=True,
        )
        send_button = gr.Button("Send")

    # State to store chat history and user name
    chat_history = gr.State([])
    user_id = gr.State("")

    # When "Start Chat" is pressed, hide the name input and show the chat interface
    start_button.click(
        start_chat,
        inputs=name_input,
        outputs=[name_group, start_button, chat_group, user_input, user_id],
    )

    # When "Send" is pressed, clear the input box and update chat history
    send_button.click(
        chatbot_response,
        inputs=[user_input, chat_history, user_id],
        outputs=[chat_interface, user_input],
    )

This code creates a simple chat interface using Gradio. It defines a chat application with two main sections: an initial name input screen and the main chat interface. It uses Gradio's Blocks to structure the UI components. The start_chat function handles the transition from the name input to the chat interface. The chat interface includes a chatbot display, a text input for user messages, and a send button. The application maintains the chat history and user ID as state variables. When the user sends a message, the chatbot_response function is called to process the input and update the chat. This setup allows for a smooth, interactive chatbot experience within a web browser, demonstrating how to evolve a console-based AI assistant into a graphical user interface.

Launching the Chatbot

To run the chatbot, we simply call:

if __name__ == "__main__":
    ai_friend.launch()

This launches a local web server where users can interact with Hazel.

Talking to your AI friend

To start the chatbot, navigate to your project directory in the terminal and run:

python ai-friend.py

You can view your application locally on http://127.0.0.1:7860

Conclusion

We've built an AI friend that has a functional GUI, which is a huge improvement from the previous console-based UX. Gradio is perfect for creating user interfaces quickly and prototyping. It's rich in functionality, and you need to take time to study all it has to offer. For the next steps, I suggest that you play around with your AI friend and try to make improvements. A few suggestions on where to start are below:

Personality Tuning: Modify the system message and prompt template to adjust Hazel's personality and conversation style.
Error Handling: Implement try-except blocks to gracefully handle potential API errors or unexpected inputs.
Multi-Modal Capabilities: Consider integrating image or audio processing to enhance the chatbot's abilities.

Remember, the code is available on GitHub. If you have any questions, you can reach out to me on X (formerly Twitter) or LinkedIn. Happy coding!

AI should drive results, not complexity. AgentemAI helps businesses build scalable, efficient, and secure AI solutions. See how we can help.