Update bot.py to use asynchronous calls for get_history and get_user_model functions

Add slash command /user_stat to display user statistics

.github/workflows/main.yml

@@ -6,9 +6,11 @@ on:
       - main
 
 jobs:
-  # Run unit tests for the project
-  CI:
+  tests:
     runs-on: ubuntu-latest
+    env:
+      MONGODB_URI: ${{ secrets.MONGODB_URI }}
+    environment: Private Server Deploy 
     steps:
       - name: Checkout repository
         uses: actions/checkout@v4
@@ -17,7 +19,15 @@ jobs:
         uses: actions/setup-python@v5.3.0
         with:
           python-version: '3.12.3'
-        
+
+      - name: Cache Python dependencies
+        uses: actions/cache@v3
+        with:
+          path: ~/.cache/pip
+          key: ${{ runner.os }}-pip-${{ hashFiles('requirements.txt') }}
+          restore-keys: |
+            ${{ runner.os }}-pip-
+
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
@@ -27,23 +37,23 @@ jobs:
       - name: Run unit tests
         run: |
           python -m pytest tests/
-      
-      # Run security check
+
       - name: pyupio/safety-action
         uses: pyupio/safety-action@v1.0.1
         with:
           api-key: ${{ secrets.SAFETY_API_KEY }}
-     
-  
-  # Build and push package to GitHub Container Registry (GHCR)
-  build-and-push-to-ghcr:
-    runs-on: ubuntu-latest 
-    environment: Private Server Deploy
-    needs: CI  # This job depends on the CI job
+
+  build-and-push:
+    runs-on: ubuntu-latest
+    environment: Private Server Deploy 
+    needs: tests
     steps:
-      - name: Check out the repository
+      - name: Checkout repository
         uses: actions/checkout@v4
 
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v3
+
       - name: Log in to GitHub Container Registry
         uses: docker/login-action@v3
         with:
@@ -51,31 +61,22 @@ jobs:
           username: ${{ github.actor }}
           password: ${{ secrets.GITHUB_TOKEN }}
 
-      - name: Build the Docker image
-        run: |
-          IMAGE_NAME=ghcr.io/coder-vippro/chatgpt-discord-bot
-          IMAGE_TAG=latest
-          docker build -t $IMAGE_NAME:$IMAGE_TAG .
-
-      - name: Push the Docker image
-        run: |
-          IMAGE_NAME=ghcr.io/coder-vippro/chatgpt-discord-bot
-          IMAGE_TAG=latest
-          docker push $IMAGE_NAME:$IMAGE_TAG
+      - name: Build and push Docker image with cache
+        uses: docker/build-push-action@v6
+        with:
+          context: .
+          push: true
+          tags: "ghcr.io/coder-vippro/chatgpt-discord-bot:latest"
+          cache-from: type=gha
+          cache-to: type=gha,mode=max
+          github-token: ${{ secrets.GITHUB_TOKEN }}
 
   # Deploy from GHCR to the main server
-  deploy-to-main-server:
+  deploy:
     runs-on: self-hosted
     environment: Private Server Deploy  # Specify the deployment environment
-    needs: build-and-push-to-ghcr  # This job depends on the GHCR push job
+    needs: build-and-push  # This job depends on the GHCR push job
     steps:
-      # Step 1: Log in to GitHub Container Registry
-      - name: Log in to GitHub Container Registry
-        uses: docker/login-action@v3
-        with:
-          registry: ghcr.io
-          username: ${{ github.actor }}
-          password: ${{ secrets.GITHUB_TOKEN }}
 
       # Step 2: Stop and remove the previous running container
       - name: Remove old running container
@@ -109,4 +110,4 @@ jobs:
             -e GOOGLE_CX="${{ secrets.GOOGLE_CX }}" \
             -e OPENAI_BASE_URL="${{ secrets.OPENAI_BASE_URL }}" \
             -e MONGODB_URI="${{ secrets.MONGODB_URI }}" \
-            $IMAGE_NAME:$IMAGE_TAG
+            $IMAGE_NAME:$IMAGE_TAG

.github/workflows/pull.yml

@@ -9,6 +9,9 @@ jobs:
   # Run unit tests for the project
   test:
     runs-on: ubuntu-latest
+    env:
+      MONGODB_URI: ${{ secrets.MONGODB_URI }}
+    environment: Private Server Deploy 
     steps:
       - name: Checkout repository
         uses: actions/checkout@v4
@@ -32,4 +35,4 @@ jobs:
       - name: pyupio/safety-action
         uses: pyupio/safety-action@v1.0.1
         with:
-          api-key: ${{ secrets.SAFETY_API_KEY }}
+          api-key: ${{ secrets.SAFETY_API_KEY }}

.gitignore

@@ -2,3 +2,8 @@ test.py
 .env
 chat_history.db
 bot_copy.py
+__pycache__/bot.cpython-312.pyc
+tests/__pycache__/test_bot.cpython-312.pyc
+.vscode/settings.json
+chatgpt.zip
+response.txt

Dockerfile

@@ -1,22 +1,30 @@
 # Use an official Python runtime as a parent image
 FROM python:3.11.10-slim
 
+# Set environment variables to reduce Python buffer and logs
+ENV PYTHONDONTWRITEBYTECODE=1 \
+    PYTHONUNBUFFERED=1
+
 # Install curl and other dependencies
-RUN apt-get update && apt-get install -y curl && apt-get clean
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    curl \
+  && apt-get clean \
+  && rm -rf /var/lib/apt/lists/*
 
 # Set the working directory in the container
 WORKDIR /usr/src/discordbot
 
 # Copy the requirements file first to leverage Docker cache
-COPY requirements.txt ./
+COPY requirements.txt .
 
-# Install any needed packages specified in requirements.txt
-RUN pip install --no-cache-dir -r requirements.txt
+# Install Python dependencies
+RUN --mount=type=cache,target=/root/.cache/pip \
+    pip install --no-cache-dir -r requirements.txt
 
-# Expose port for health check
+# Expose port (optional, only if needed for a web server)
 EXPOSE 5000
 
-# Add health check
+# Add health check (update endpoint if needed)
 HEALTHCHECK --interval=30s --timeout=10s --start-period=10s --retries=3 \
   CMD curl --fail http://localhost:5000/health || exit 1
 
@@ -24,4 +32,4 @@ HEALTHCHECK --interval=30s --timeout=10s --start-period=10s --retries=3 \
 COPY . .
 
 # Command to run the application
-CMD ["python3", "bot.py"]
+CMD ["python3", "bot.py"]

README.md

@@ -1,4 +1,3 @@
-
 # ChatGPT Discord Bot
 
 ![Build and Push](https://github.com/coder-vippro/ChatGPT-Discord-Bot/actions/workflows/main.yml/badge.svg)
@@ -119,6 +118,7 @@ Once the bot is running, it connects to Discord using credentials from `.env`. C
 - **Scrape Web Content**: `/web url: "https://example.com"`
 - **Search Google**: `/search prompt: "latest news in Vietnam"`
 - **Normal chat**: `Ping the bot with a question or send a dms to the bot to start`
+- **User Statistics**: `/user_stat` - Get your current input token, output token, and model.
 
 ## CI/CD
 

bot.py

@@ -1,24 +1,24 @@
 import os
 import discord
 import io
-import pymongo
-from discord.ext import commands, tasks
-from discord import app_commands
+import threading
+import tiktoken
+import asyncio
 import requests
-from bs4 import BeautifulSoup
 import logging
 import sys
-from openai import OpenAI, RateLimitError
 import aiohttp
+from discord.ext import commands, tasks
+from discord import app_commands
+from motor.motor_asyncio import AsyncIOMotorClient
+from bs4 import BeautifulSoup
+from openai import OpenAI, RateLimitError
 from runware import Runware, IImageInference
 from collections import defaultdict
-import asyncio
-from PIL import Image
-from io import BytesIO
 from dotenv import load_dotenv
-from pymongo import MongoClient
 from flask import Flask, jsonify
-import threading
+
+# Load environment variables
 load_dotenv()
 
 # Flask app for health-check
@@ -34,6 +34,8 @@ def health():
         return jsonify(status="unhealthy", error="Bot is disconnected"), 500
     elif not bot.is_ready():  # Bot is not ready yet
         return jsonify(status="unhealthy", error="Bot is not ready"), 500
+    elif bot.latency > 151:  # Bot heartbeat is blocked for more than 151 seconds
+        return jsonify(status="unhealthy", error=f"Heartbeat to websocket blocked for {bot.latency:.2f} seconds"), 500
     else:
         return jsonify(status="healthy"), 200
 
@@ -110,13 +112,15 @@ MODEL_OPTIONS = [
     "gpt-4o",
     "gpt-4o-mini",
     "o1-preview",
-    "o1-mini"
+    "o1-mini",
+    "o1",
+    "o3-mini"
 ]
 
 # Prompt for different plugins
 WEB_SCRAPING_PROMPT = "You are using the Web Scraping Plugin, gathering information from given url. Respond accurately and combine data to provide a clear, insightful summary. "
 NORMAL_CHAT_PROMPT = "You're ChatGPT for Discord! You can chat, generate images, and perform searches. Craft responses that are easy to copy directly into Discord chats, without using markdown, code blocks, or extra formatting. When you solving any problems you must remember that: Let's solve this step-by-step. What information do we need to find? What operation might help us solve this? Explain your reasoning and provide the answer."
-SEARCH_PROMPT = "You are using the Google Search Plugin, accessing information from the top 3 Google results. Summarize these findings clearly, adding relevant insights to answer the users question."
+SEARCH_PROMPT = "You are using the Google Search Plugin, accessing information from the top 3 Google results link which is the scraped content from these 3 website. Summarize these findings clearly, adding relevant insights to answer the users question."
 
 # Google API details
 GOOGLE_API_KEY = str(os.getenv("GOOGLE_API_KEY"))  # Google API Key
@@ -132,7 +136,7 @@ MONGODB_URI = str(os.getenv("MONGODB_URI"))
 runware = Runware(api_key=RUNWARE_API_KEY)
 
 # MongoDB client initialization
-mongo_client = MongoClient(MONGODB_URI)
+mongo_client = AsyncIOMotorClient(MONGODB_URI)
 db = mongo_client['chatgpt_discord_bot']  # Database name
 
 # Dictionary to keep track of user requests and their cooldowns
@@ -146,42 +150,36 @@ TOKEN = str(os.getenv("DISCORD_TOKEN"))
 
 # --- Database functions ---
 
-def get_history(user_id):
-    user_data = db.user_histories.find_one({'user_id': user_id})
-    if user_data and 'history' in user_data:
-        return user_data['history']
-    else:
-        return [{"role": "system", "content": NORMAL_CHAT_PROMPT}]
-    
-def save_history(user_id, history):
-    db.user_histories.update_one(
+async def get_history(user_id):
+    user_data = await db.user_histories.find_one({'user_id': user_id})
+    return user_data['history'] if user_data and 'history' in user_data else [{"role": "system", "content": NORMAL_CHAT_PROMPT}]
+
+async def save_history(user_id, history):
+    await db.user_histories.update_one(
         {'user_id': user_id},
         {'$set': {'history': history}},
         upsert=True
     )
 
-# New function to get the user's model preference
-def get_user_model(user_id):
-    user_pref = db.user_preferences.find_one({'user_id': user_id})
-    if user_pref and 'model' in user_pref:
-        return user_pref['model']
-    else:
-        return "gpt-4o"  # Default to "gpt-4o" if no preference
+async def get_user_model(user_id):
+    user_pref = await db.user_preferences.find_one({'user_id': user_id})
+    return user_pref['model'] if user_pref and 'model' in user_pref else "gpt-4o"
 
-def save_user_model(user_id, model):
-    db.user_preferences.update_one(
+async def save_user_model(user_id, model):
+    await db.user_preferences.update_one(
         {'user_id': user_id},
         {'$set': {'model': model}},
         upsert=True
     )
 
+# --- End of Database functions ---
 
 # Intents and bot initialization
 intents = discord.Intents.default()
 intents.message_content = True
 
 # Bot initialization
-bot = commands.Bot(command_prefix="!", intents=intents, heartbeat_timeout=120)
+bot = commands.Bot(command_prefix="//quocanhvu", intents=intents, heartbeat_timeout=120)
 tree = bot.tree  # For slash commands
 
 # Function to perform a Google search and return results
@@ -194,7 +192,7 @@ def google_custom_search(query: str, num_results: int = 3) -> list:
         "num": num_results
     }
     try:
-        response = requests.get(search_url, params=params, timeout=15)  # Add timeout
+        response = requests.get(search_url, params=params, timeout=30)  # Add timeout
         response.raise_for_status()  # Check for any errors in the response
         data = response.json()
 
@@ -284,7 +282,7 @@ async def choose_model(interaction: discord.Interaction):
         user_id = interaction.user.id
         
         # Save the model selection to the database
-        save_user_model(user_id, selected_model)
+        await save_user_model(user_id, selected_model)
         await interaction.response.send_message(
             f"Model set to `{selected_model}` for your responses.", ephemeral=True
         )
@@ -301,48 +299,44 @@ async def search(interaction: discord.Interaction, query: str):
     """Searches Google and sends results to the AI model."""
     await interaction.response.defer(thinking=True)
     user_id = interaction.user.id
-    history = get_history(user_id)
+    history = await get_history(user_id)
 
     history.append({"role": "user", "content": query})
 
     try:
         # Perform Google search
-        search_results = google_custom_search(query, num_results=5)
+        search_results = google_custom_search(query, num_results=2)
         if not search_results:
             await interaction.followup.send("No search results found.")
             return
 
-        # Prepare the search results for the AI model
-        combined_input = f"{SEARCH_PROMPT}\nUser query: {query}\nGoogle search results:\n"
-        
-        # Extract URLs and prepare the message
-        links = []
+        # Scrape content from the first 5 links
+        scraped_contents = []
         for result in search_results:
-            url = result.split('\n')[1].split('Link: ')[1]  # Extract URL from the result string
-            links.append(url)
-            combined_input += f"{result}\n"
+            url = result.split('\n')[1].split('Link: ')[1]
+            content = scrape_web_content(url)
+            scraped_contents.append(content)
 
-        # Add links at the end of the combined input
-        combined_input += "\nLinks:\n" + "\n".join(links)
+        # Prepare the combined input for the AI model
+        combined_input = f"{SEARCH_PROMPT}\nUser query: {query}\nScraped Contents:\n" + "\n".join(scraped_contents)
 
         history.append({"role": "system", "content": combined_input})
 
         # Send the history to the AI model
         response = client.chat.completions.create(
-            model="gpt-4o-mini",
+            model="gpt-4o",
             messages=history,
-            temperature=0.3,
+            temperature=0.4,
             max_tokens=4096,
-            top_p=0.7
+            top_p=1
         )
 
         reply = response.choices[0].message.content
         history.append({"role": "assistant", "content": reply})
         save_history(user_id, history)
 
-        # Prepare the final response including the links
-        links_message = "\nLinks:\n" + "\n".join(links)
-        await interaction.followup.send(reply + links_message)
+        # Send the final response to the user
+        await interaction.followup.send(reply)
 
     except Exception as e:
         await interaction.followup.send(f"Error: {str(e)}", ephemeral=True)
@@ -354,7 +348,7 @@ async def web(interaction: discord.Interaction, url: str):
     """Scrapes a webpage and sends data to the AI model."""
     await interaction.response.defer(thinking=True)
     user_id = interaction.user.id
-    history = get_history(user_id)
+    history = await get_history(user_id)
 
     try:
         content = scrape_web_content(url)
@@ -390,25 +384,77 @@ async def reset(interaction: discord.Interaction):
     db.user_histories.delete_one({'user_id': user_id})
     await interaction.response.send_message("Your data has been cleared and reset!", ephemeral=True)
 
+# Slash command for user statistics (/user_stat)
+@tree.command(name="user_stat", description="Get your current input token, output token, and model.")
+async def user_stat(interaction: discord.Interaction):
+    """Fetches and displays the current input token, output token, and model for the user."""
+    user_id = interaction.user.id
+    history = await get_history(user_id)
+    model = await get_user_model(user_id)
+
+    # Handle cases where user model is not found
+    if not model:
+        model = "gpt-4o"  # Default model
+
+    # Adjust model for encoding purposes
+    if model in ["gpt-4o", "o1", "o1-preview", "o1-mini", "o3-mini"]:
+        encoding_model = "gpt-4o"
+    else:
+        encoding_model = model
+
+    # Retrieve the appropriate encoding for the selected model
+    encoding = tiktoken.encoding_for_model(encoding_model)
+
+    # Initialize token counts
+    input_tokens = 0
+    output_tokens = 0
+
+    # Calculate input and output tokens
+    if history:
+        for item in history:
+            content = item.get('content')  # Safely access 'content'
+
+            # Handle case where content is a list or other type
+            if isinstance(content, list):
+                # Convert list of objects to a single string (e.g., join texts with a space)
+                content = " ".join(
+                    sub_item.get('text', '') for sub_item in content if isinstance(sub_item, dict)
+                )
+
+            # Ensure content is a string before processing
+            if isinstance(content, str):
+                token_count = len(encoding.encode(content))
+                if item['role'] == 'user':
+                    input_tokens += token_count
+                elif item['role'] in ['assistant', 'developer']:
+                    # Treat 'developer' as 'assistant' for token counting
+                    output_tokens += token_count
+
+    # Create the statistics message
+    stat_message = (
+        f"**User Statistics:**\n"
+        f"Model: `{model}`\n"
+        f"Input Tokens: `{input_tokens}`\n"
+        f"Output Tokens: `{output_tokens}`\n"
+    )
+
+    # Send the response
+    await interaction.response.send_message(stat_message, ephemeral=True)
+
 
 # Slash command for help (/help)
 @tree.command(name="help", description="Display a list of available commands.")
 async def help_command(interaction: discord.Interaction):
     """Sends a list of available commands to the user."""
     help_message = (
-        "**Available Commands:**\n"
-        "/choose_model - Select the AI model to use for responses (gpt-4o, gpt-4o-mini, o1-preview, o1-mini).\n"
-        "/search `<query>` - Search on Google and send results to AI model.\n"
-        "/web `<url>` - Scrape a webpage and send data to AI model.\n"
-        "/generate `<prompt>` - Generate an image from a text prompt.\n"
-        "/reset - Reset your conversation history.\n"
-        "/help - Display this help message.\n"
         "**Các lệnh có sẵn:**\n"
         "/choose_model - Chọn mô hình AI để sử dụng cho phản hồi (gpt-4o, gpt-4o-mini, o1-preview, o1-mini).\n"
         "/search `<truy vấn>` - Tìm kiếm trên Google và gửi kết quả đến mô hình AI.\n"
         "/web `<url>` - Thu thập dữ liệu từ trang web và gửi đến mô hình AI.\n"
         "/generate `<gợi ý>` - Tạo hình ảnh từ gợi ý văn bản.\n"
         "/reset - Đặt lại lịch sử trò chuyện của bạn.\n"
+        "/remaining_turns - Kiểm tra số lượt trò chuyện còn lại cho mỗi mô hình.\n"
+        "/user_stat - Nhận thông tin về token đầu vào, token đầu ra và mô hình hiện tại của bạn.\n"
         "/help - Hiển thị tin nhắn trợ giúp này.\n"
     )
     await interaction.response.send_message(help_message, ephemeral=True)
@@ -434,6 +480,7 @@ async def send_response(interaction: discord.Interaction, reply: str):
     else:
         await interaction.followup.send(reply)
 
+
 # Event to handle incoming messages
 @bot.event
 async def on_message(message: discord.Message):
@@ -446,10 +493,16 @@ async def on_message(message: discord.Message):
     else:
         await bot.process_commands(message)
 
+# Function to handle user messages
 async def handle_user_message(message: discord.Message):
+    # Offload processing to a non-blocking task
+    asyncio.create_task(process_user_message(message)) 
+
+# Function to process user messages
+async def process_user_message(message: discord.Message):
     user_id = message.author.id
-    history = get_history(user_id)
-    model = get_user_model(user_id)
+    history = await get_history(user_id)
+    model = await get_user_model(user_id)
 
     # Initialize content list for the current message
     content = []
@@ -507,7 +560,17 @@ async def handle_user_message(message: discord.Message):
     # Prepare messages to send to API
     messages_to_send = history.copy()
 
-    if model in ["gpt-4o", "gpt-4o-mini"]:
+    if model in ["gpt-4o", "gpt-4o-mini", "o1", "o3-mini"]:
+        # If the model is "o1", rename "system" role to "developer"
+        if model == "o1" or model == "o3-mini":
+            for msg in messages_to_send:
+                if msg["role"] == "system":
+                    msg["role"] = "developer"
+        elif model != "o1":
+            for msg in messages_to_send:
+                if msg["role"] == "developer":
+                    msg["role"] = "system"
+
         # Include up to 10 previous images
         def get_last_n_images(history, n=10):
             images = []
@@ -515,9 +578,10 @@ async def handle_user_message(message: discord.Message):
                 if msg["role"] == "user" and isinstance(msg["content"], list):
                     for part in reversed(msg["content"]):
                         if part["type"] == "image_url":
+                            part["details"] = "high"
                             images.append(part)
                             if len(images) == n:
-                                return images[::-1]  # Reverse to maintain order
+                                return images[::-1]
             return images[::-1]
 
         # Get the last 10 images
@@ -532,20 +596,18 @@ async def handle_user_message(message: discord.Message):
                 ]
                 last_message["content"].extend(latest_images)
             else:
-                # Ensure content is a list
                 last_message["content"] = [{"type": "text", "text": last_message["content"]}]
                 last_message["content"].extend(latest_images)
             messages_to_send[-1] = last_message
 
         # Fix the 431 error by limiting the number of images
-        max_images = 10  # Adjust the limit as needed
+        max_images = 10
         total_images = 0
         for msg in messages_to_send:
             if msg["role"] == "user" and isinstance(msg["content"], list):
                 image_parts = [part for part in msg["content"] if part.get("type") == "image_url"]
                 total_images += len(image_parts)
         if total_images > max_images:
-            # Remove older images to keep total_images <= max_images
             images_removed = 0
             for msg in messages_to_send:
                 if msg["role"] == "user" and isinstance(msg["content"], list):
@@ -553,19 +615,17 @@ async def handle_user_message(message: discord.Message):
                     for part in msg["content"]:
                         if part.get("type") == "image_url" and images_removed < (total_images - max_images):
                             images_removed += 1
-                            continue  # Skip this image
+                            continue
                         new_content.append(part)
                     msg["content"] = new_content
 
     else:
-        # Exclude image URLs and system prompts for 'o1' model family
-        # Remove 'image_url' content from messages
+        # Exclude image URLs and system prompts for other models
         for msg in messages_to_send:
             if msg["role"] == "user" and isinstance(msg["content"], list):
                 msg["content"] = [
-                    part for part in msg["content"] if part.get("type") != "image_url"
+                    part for part in msg["content"] if part["type"] != "image_url"
                 ]
-        # Remove system prompts from messages
         messages_to_send = [
             msg for msg in messages_to_send if msg.get("role") != "system"
         ]
@@ -578,45 +638,46 @@ async def handle_user_message(message: discord.Message):
         }
 
         if model in ["gpt-4o", "gpt-4o-mini"]:
-            # Include parameters for 'gpt-4o' models
             api_params.update({
                 "temperature": 0.3,
-                "max_tokens": 4096,
+                "max_tokens": 8096,
                 "top_p": 0.7,
             })
 
-        # Send messages to the API
-        response = client.chat.completions.create(**api_params)
-
+        # The non-blocking call, done in a background thread
+        response = await asyncio.to_thread(client.chat.completions.create, **api_params)
         reply = response.choices[0].message.content
         history.append({"role": "assistant", "content": reply})
-        save_history(user_id, history)
+        await save_history(user_id, history)
 
         await send_response(message.channel, reply)
 
+    # Handle rate limit errors
     except RateLimitError:
         error_message = (
-            "Error: Rate limit exceeded for your model"
+            "Error: Rate limit exceeded for your model. "
             "Please try again later or use /choose_model to change to any models else."
         )
         logging.error(f"Rate limit error: {error_message}")
         await message.channel.send(error_message)
 
+    # Handle other exceptions
     except Exception as e:
         error_message = f"Error: {str(e)}"
         logging.error(f"Error handling user message: {error_message}")
         await message.channel.send(error_message)
+        db.user_histories.delete_one({'user_id': user_id})
 
-# Function to trim the history to avoid exceeding token limits
+# Function to get the remaining turns for each model
 def trim_history(history):
-    """Trims the history to avoid exceeding token limits."""
+    """Trims the history to avoid exceeding token limits by removing older messages first."""
     tokens_used = sum(len(str(item['content'])) for item in history)
     max_tokens_allowed = 9000
-    while tokens_used > max_tokens_allowed:
-        removed_item = history.pop(1)
+    while tokens_used > max_tokens_allowed and len(history) > 1:
+        removed_item = history.pop(0)
         tokens_used -= len(str(removed_item['content']))
 
-# Function to send a response to the channel
+# Function to send response to the discord channel
 async def send_response(channel: discord.TextChannel, reply: str):
     """Sends the reply to the channel, handling long responses."""
     if len(reply) > 2000:
@@ -681,6 +742,7 @@ async def change_status():
             await bot.change_presence(activity=discord.Game(name=status))
             await asyncio.sleep(300)  # Change every 60 seconds
 
+# Event to run when the bot is ready
 @bot.event
 async def on_ready():
     """Bot startup event to sync slash commands and start status loop."""
@@ -688,6 +750,7 @@ async def on_ready():
     print(f"Logged in as {bot.user}")
     change_status.start() # Start the status changing loop   
 
+
 # Start Flask in a separate thread
 flask_thread = threading.Thread(target=run_flask)
 flask_thread.daemon = True  # Ensure it closes when the main program exits
@@ -696,4 +759,4 @@ flask_thread.start()
 # Main bot startup
 if __name__ == "__main__":
     logging.basicConfig(level=logging.INFO, stream=sys.stdout)
-    bot.run(TOKEN)
+    bot.run(TOKEN)

requirements.txt

@@ -7,4 +7,6 @@ runware
 Pillow              
 discord.py
 pymongo
-flask
+flask
+tiktoken
+motor

tests/test_bot.py

@@ -1,63 +1,164 @@
 import unittest
-from unittest.mock import AsyncMock, MagicMock
-from bot import bot, search, generate_image, web
+import unittest.mock as mock
+from unittest.mock import MagicMock, patch, AsyncMock
+import asyncio
+import requests
+from flask import Flask
+from bot import (
+    app,
+    run_flask,
+    client,
+    statuses,
+    MODEL_OPTIONS,
+    WEB_SCRAPING_PROMPT,
+    NORMAL_CHAT_PROMPT,
+    SEARCH_PROMPT,
+    google_custom_search,
+    scrape_web_content,
+    get_history,
+    save_history,
+    get_user_model,
+    save_user_model,
+    bot,
+    process_request,
+    process_queue,
+    choose_model,
+    search,
+    web,
+    reset,
+    help_command,
+    should_respond_to_message,
+    handle_user_message,
+    trim_history,
+    generate_image,
+    _generate_image_command,
+    change_status,
+    on_ready
+)
+
+class TestFullBot(unittest.TestCase):
 
-class TestDiscordBotCommands(unittest.TestCase):
     def setUp(self):
-        self.bot = bot
-        self.interaction = AsyncMock()
-        self.interaction.user.id = 123456789  # Mock user ID
+        # Sử dụng app của Flask để test endpoint
+        self.app = app.test_client()
 
-    async def test_search_command(self):
-        # Set up mocks for interaction methods
-        self.interaction.response.defer = AsyncMock()
-        self.interaction.followup.send = AsyncMock()
+    def test_flask_health_endpoint(self):
+        with patch("bot.bot.is_closed", return_value=False), \
+             patch("bot.bot.is_ready", return_value=True):
+            response = self.app.get('/health')
+            self.assertEqual(response.status_code, 200, "Health endpoint should return 200 if bot is ready.")
 
-        # Call the search command with a sample query
-        await search(self.interaction, query="Python")
+    def test_run_flask(self):
+        # Kiểm tra run_flask khởi động mà không báo lỗi
+        with patch.object(Flask, 'run') as mock_run:
+            run_flask()
+            mock_run.assert_called_once()
 
-        # Check if followup.send was called
-        self.interaction.followup.send.assert_called()
-        self.interaction.response.defer.assert_called_with(thinking=True)
+    @patch("requests.get")
+    def test_google_custom_search(self, mock_get):
+        mock_resp = MagicMock()
+        mock_resp.json.return_value = {"items": [{"title": "Result 1"}]}
+        mock_get.return_value = mock_resp
+        results = google_custom_search("test")
+        self.assertEqual(len(results), 1, "Should return 1 search result when JSON has items.")
 
-    async def test_generate_image_command(self):
-        # Mock the deferred response
-        self.interaction.response.defer = AsyncMock()
-        self.interaction.followup.send = AsyncMock()
+    @patch("requests.get")
+    def test_scrape_web_content(self, mock_get):
+        mock_response = MagicMock()
+        mock_response.status_code = 200
+        mock_response.content = b"<p>Some scraped text</p>"
+        mock_get.return_value = mock_response
+        result = scrape_web_content("https://example.com")
+        self.assertIn("Some scraped text", result, "Scraped content should include known text.")
 
-        # Patch Runware API to return a mock image URL
-        with unittest.mock.patch('bot.runware.imageInference', return_value=[MagicMock(imageURL="http://example.com/image.png")]):
-            await generate_image(self.interaction, prompt="Sunset over mountains")
+    async def async_test_get_history(self):
+        with patch("bot.db.user_histories.find_one", new_callable=AsyncMock) as mock_find_one:
+            mock_find_one.return_value = {"history": [{"role": "system", "content": NORMAL_CHAT_PROMPT}]}
+            history = await get_history(1234)
+            self.assertIsInstance(history, list, "History should be a list.")
 
-            # Check if defer and followup were called
-            self.interaction.response.defer.assert_called_with(thinking=True)
-            self.interaction.followup.send.assert_called()
+    def test_get_history(self):
+        asyncio.run(self.async_test_get_history())
 
-    async def test_web_scraping_command(self):
-        # Mock the interaction methods
-        self.interaction.response.defer = AsyncMock()
-        self.interaction.followup.send = AsyncMock()
+    async def async_test_get_user_model_default(self):
+        with patch("bot.db.user_preferences.find_one", new_callable=AsyncMock) as mock_find_one:
+            mock_find_one.return_value = None
+            model = await get_user_model(1234)
+            self.assertEqual(model, "gpt-4o")
 
-        # Call the web command with a mock URL
-        await web(self.interaction, url="https://vnexpress.net/nguon-con-khien-arm-huy-giay-phep-chip-voi-qualcomm-4807985.html")
 
-        # Ensure a followup message was sent
-        self.interaction.followup.send.assert_called()
-        self.interaction.response.defer.assert_called_with(thinking=True)
+    def test_trim_history_with_large_content(self):
+        sample_history = [
+            {"role": "user", "content": "x" * 5000},
+            {"role": "user", "content": "y" * 5000}
+        ]
+        trim_history(sample_history)
+        # Giả sử hàm trim_history không xóa hết nội dung mà vẫn giữ lại tối thiểu 1 message
+        self.assertGreaterEqual(len(sample_history), 1, "History should not be completely removed.")
 
-    async def test_message_processing(self):
-        # Mock a direct message
-        message = MagicMock()
-        message.author.id = 987654321
-        message.content = "Hello, bot!"
-        message.guild = None  # Simulate a DM
+    def test_trim_history_with_empty_history(self):
+        history = []
+        trim_history(history)
+        self.assertEqual(len(history), 0, "Empty history should remain empty.")
 
-        # Mock channel.send to test if the bot sends a message
-        message.channel.send = AsyncMock()
+    def test_trim_history_with_single_message(self):
+        history = [{"role": "user", "content": "test"}]
+        trim_history(history)
+        self.assertEqual(len(history), 1, "Single message history should remain unchanged.")
 
-        # Test the bot's response
-        await bot.on_message(message)
-        message.channel.send.assert_called()  # Check if the bot replied
+    async def async_test_process_message_with_attachment(self):
+        message = AsyncMock()
+        message.author.id = 1234
+        message.content = "Check this file"
+        message.attachments = [
+            AsyncMock(filename="test.txt", read=AsyncMock(return_value=b"File content"))
+        ]
+        # Patch bot.user để tránh lỗi khi gọi mentioned_in (bot.user có thể là None trong môi trường test)
+        with patch.object(bot.user, 'mentioned_in', return_value=False):
+            await bot.on_message(message)
 
-if __name__ == '__main__':
-    unittest.main()
+
+    async def async_test_search(self):
+        interaction = AsyncMock()
+        interaction.user.id = 1234
+        # Patch get_history để trả về list thay vì coroutine, tránh lỗi khi gọi append
+        with patch("bot.get_history", new=AsyncMock(return_value=[])):
+            await search.callback(interaction, query="Python")
+            interaction.response.defer.assert_called()
+            interaction.followup.send.assert_called()
+
+    def test_search_command(self):
+        asyncio.run(self.async_test_search())
+
+    async def async_test_web(self):
+        interaction = AsyncMock()
+        interaction.user.id = 1234
+        with patch("bot.get_history", new=AsyncMock(return_value=[])):
+            await web.callback(interaction, url="https://test.com")
+            interaction.response.defer.assert_called()
+            interaction.followup.send.assert_called()
+
+    def test_web_command(self):
+        asyncio.run(self.async_test_web())
+
+    async def async_test_reset(self):
+        interaction = AsyncMock()
+        interaction.user.id = 1234
+        await reset.callback(interaction)
+        interaction.response.send_message.assert_called()
+
+    async def test_reset_command(self):
+        await self.async_test_reset()
+
+    async def async_test_help_command(self):
+        interaction = AsyncMock()
+        # Nếu help_command được đăng ký dưới dạng command (Command object),
+        # bạn cần gọi .callback thay vì gọi trực tiếp đối tượng.
+        await help_command.callback(interaction)
+        interaction.response.send_message.assert_called()
+
+    def test_help_command(self):
+        asyncio.run(self.async_test_help_command())
+
+if __name__ == "__main__":
+    unittest.main()