Skip SSL certificate verification in web scraping and improve text sorting logic

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,5 @@ 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
+            -e ADMIN_ID="${{ secrets.ADMIN_ID }}" \
+            $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,33 @@
 # Use an official Python runtime as a parent image
 FROM python:3.11.10-slim
 
-# Install curl and other dependencies
-RUN apt-get update && apt-get install -y curl && apt-get clean
+# Set environment variables to reduce Python buffer and logs
+ENV PYTHONDONTWRITEBYTECODE=1 \
+    PYTHONUNBUFFERED=1
+
+# Install curl, g++ compiler, and other dependencies
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    curl \
+    g++ \
+    build-essential \
+    make \
+  && 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 +35,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)
@@ -40,8 +39,9 @@ To get started, ensure you have:
      RUNWARE_API_KEY=your_runware_api_key
      GOOGLE_API_KEY=your_google_api_key
      GOOGLE_CX=your_google_cx
-     OPENAI_BASE_URL=https://models.inference.ai.azure.com or https://api.openai.com/v1/models or any api else you want
+     OPENAI_BASE_URL=https://models.inference.ai.azure.com or https://api.openai.com/v1/models or any openai compatible api else you want
      MONGODB_URI=mongodb://localhost:27017/
+     ADMIN_ID=your_discord_user_id
      ```
    - Use the following `docker-compose.yml`:
      ```yaml
@@ -72,8 +72,9 @@ To get started, ensure you have:
      RUNWARE_API_KEY=your_runware_api_key
      GOOGLE_API_KEY=your_google_api_key
      GOOGLE_CX=your_google_cx
-     OPENAI_BASE_URL=https://models.inference.ai.azure.com or https://api.openai.com/v1/models or any api else you want
+     OPENAI_BASE_URL=https://models.inference.ai.azure.com or https://api.openai.com/v1/models or any openai compatible api else you want
      MONGODB_URI=mongodb://localhost:27017/
+     ADMIN_ID=your_discord_user_id
      ```
    - Install the dependencies:
      ```bash
@@ -119,6 +120,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
 

requirements.txt

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

tests/test_bot.py

@@ -1,63 +1,368 @@
+import sys
+import os
 import unittest
-from unittest.mock import AsyncMock, MagicMock
-from bot import bot, search, generate_image, web
+import asyncio
+import pytest
+from unittest.mock import patch, MagicMock, AsyncMock
 
-class TestDiscordBotCommands(unittest.TestCase):
+# Add parent directory to path for imports
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
+import bot
+
+class TestBotUtils(unittest.TestCase):
+    
+    def test_count_tokens(self):
+        """Test token counting function"""
+        # Basic test
+        self.assertGreater(bot.count_tokens("Hello world"), 0)
+        # Empty string should return 0 or small value
+        self.assertLessEqual(bot.count_tokens(""), 3)
+        # Longer text should have more tokens
+        short_text = "Hello"
+        long_text = "Hello " * 100
+        self.assertLess(bot.count_tokens(short_text), bot.count_tokens(long_text))
+    
+
+    def test_prepare_messages_for_api(self):
+        """Test message preparation for API"""
+        messages = [
+            {"role": "system", "content": "You are a helpful assistant"},
+            {"role": "user", "content": "Hello"},
+            {"role": "assistant", "content": "Hi there"},
+            {"role": "user", "content": "Help me with Python"}
+        ]
+        prepared = bot.prepare_messages_for_api(messages, max_tokens=1000)
+        self.assertIsInstance(prepared, list)
+        # Should have role and content
+        for msg in prepared:
+            self.assertIn("role", msg)
+            self.assertIn("content", msg)
+
+class TestCodeSanitization(unittest.TestCase):
+    
+    def test_python_safe_code(self):
+        """Test Python code sanitization with safe code"""
+        code = """
+def factorial(n):
+    if n == 0:
+        return 1
+    return n * factorial(n-1)
+        
+print(factorial(5))
+"""
+        is_safe, sanitized = bot.sanitize_code(code, "python")
+        self.assertTrue(is_safe)
+        self.assertIn("signal.alarm(10)", sanitized)
+        
+    def test_python_unsafe_code(self):
+        """Test Python code sanitization with unsafe imports"""
+        code = """
+import os
+print(os.system('ls'))
+"""
+        is_safe, message = bot.sanitize_code(code, "python")
+        self.assertFalse(is_safe)
+        self.assertIn("Forbidden module import", message)
+        
+    def test_cpp_safe_code(self):
+        """Test C++ code sanitization with safe code"""
+        code = """
+#include <iostream>
+using namespace std;
+
+int main() {
+    cout << "Hello World" << endl;
+    return 0;
+}
+"""
+        is_safe, sanitized = bot.sanitize_code(code, "cpp")
+        self.assertTrue(is_safe)
+        self.assertIn("userMain(", sanitized)
+        
+    def test_cpp_unsafe_code(self):
+        """Test C++ code sanitization with unsafe includes"""
+        code = """
+#include <iostream>
+#include <fstream>
+int main() {
+    ofstream file("test.txt");
+    file << "Hello World";
+    return 0;
+}
+"""
+        is_safe, message = bot.sanitize_code(code, "cpp")
+        self.assertFalse(is_safe)
+        self.assertIn("Forbidden header", message)
+
+    def test_python_eval_exec_detection(self):
+        """Test detection of eval/exec in Python code"""
+        code = """
+print("Hello")
+eval("print('This is dangerous')")
+"""
+        is_safe, message = bot.sanitize_code(code, "python")
+        self.assertFalse(is_safe)
+        self.assertIn("Forbidden", message)
+
+    def test_python_add_missing_structure(self):
+        """Test Python code gets proper safety structure"""
+        code = "print('Hello world')"
+        is_safe, sanitized = bot.sanitize_code(code, "python")
+        self.assertTrue(is_safe)
+        self.assertIn("try:", sanitized)
+        self.assertIn("except Exception as e:", sanitized)
+        self.assertIn("finally:", sanitized)
+        self.assertIn("signal.alarm(0)", sanitized)
+
+    def test_cpp_add_missing_iostream(self):
+        """Test C++ code gets iostream added when using cout"""
+        code = """
+int main() {
+    cout << "Hello" << endl;
+    return 0;
+}
+"""
+        is_safe, sanitized = bot.sanitize_code(code, "cpp")
+        self.assertTrue(is_safe)
+        self.assertIn("#include <iostream>", sanitized)
+
+    def test_cpp_add_missing_namespace(self):
+        """Test C++ code gets namespace std added when needed"""
+        code = """
+#include <iostream>
+int main() {
+    cout << "Hello" << endl;
+    return 0;
+}
+"""
+        is_safe, sanitized = bot.sanitize_code(code, "cpp")
+        self.assertTrue(is_safe)
+        self.assertIn("using namespace std;", sanitized)
+
+    def test_edge_case_empty_code(self):
+        """Test sanitization with empty code"""
+        code = ""
+        is_safe, sanitized = bot.sanitize_code(code, "python")
+        self.assertTrue(is_safe)
+        self.assertNotEqual(sanitized, "")  # Should add safety structure
+
+    def test_multiple_forbidden_imports(self):
+        """Test code with multiple forbidden imports"""
+        code = """
+import os
+import sys
+import subprocess
+print("This is malicious")
+"""
+        is_safe, message = bot.sanitize_code(code, "python")
+        self.assertFalse(is_safe)
+        self.assertIn("Forbidden", message)
+
+
+@pytest.mark.asyncio
+class TestAsyncFunctions:
+    
+    @pytest.fixture
+    def mock_channel(self):
+        mock = AsyncMock()
+        mock.send = AsyncMock()
+        return mock
+        
+    @patch('bot.client.chat.completions.create')
+    async def test_process_batch(self, mock_create, mock_channel):
+        """Test batch processing"""
+        mock_response = MagicMock()
+        mock_response.choices = [MagicMock()]
+        mock_response.choices[0].message.content = "Processed content"
+        mock_create.return_value = mock_response
+        
+        result = await bot.process_batch(
+            model="gpt-4o-mini",
+            user_prompt="Analyze this",
+            batch_content="Test content",
+            current_batch=1,
+            total_batches=1,
+            channel=mock_channel()
+        )
+        
+        # Should call the API
+        mock_create.assert_called_once()
+        self.assertTrue(result)
+        
+    @patch('asyncio.create_subprocess_exec')
+    async def test_execute_code(self, mock_subprocess):
+        """Test code execution"""
+        # Configure the mock
+        mock_proc = AsyncMock()
+        mock_proc.communicate = AsyncMock(return_value=(b"Hello World", b""))
+        mock_subprocess.return_value = mock_proc
+        
+        result = await bot.execute_code("print('Hello World')", "python")
+        self.assertIn("Hello World", result)
+        
+
+
+class TestToolFunctions(unittest.TestCase):
+    
+    @patch('bot.google_custom_search')
+    def test_tool_function_google_search(self, mock_search):
+        """Test Google search tool function"""
+        mock_search.return_value = {"results": [{"title": "Test", "link": "https://example.com"}]}
+        
+        result = bot.tool_functions["google_search"]({"query": "test", "num_results": 1})
+        mock_search.assert_called_once_with("test", 1)
+        self.assertIsInstance(result, dict)
+    
+    @patch('bot.scrape_web_content')
+    def test_tool_function_scrape_webpage(self, mock_scrape):
+        """Test web scraping tool function"""
+        mock_scrape.return_value = "Scraped content"
+        
+        result = bot.tool_functions["scrape_webpage"]({"url": "https://example.com"})
+        mock_scrape.assert_called_once_with("https://example.com")
+        self.assertEqual(result, "Scraped content")
+
+    @patch('bot.requests.get')
+    def test_scrape_web_content_detailed(self, mock_requests_get):
+        """Test web scraping function with different scenarios"""
+        # Setup mock response for successful HTML page with article content
+        mock_response = MagicMock()
+        mock_response.status_code = 200
+        mock_response.headers = {'content-type': 'text/html'}
+        mock_response.content = """
+        <html><body>
+            <article>
+                <h1>Test Article</h1>
+                <p>This is a test paragraph with meaningful content.</p>
+            </article>
+        </body></html>
+        """
+        mock_requests_get.return_value = mock_response
+        
+        # Test successful scraping with article tag
+        content = bot.scrape_web_content("https://example.com")
+        self.assertIn("Test Article", content)
+        self.assertIn("test paragraph", content)
+        
+        # Test with HTTP error
+        mock_response.status_code = 404
+        mock_response.headers = {'content-type': 'text/html'}
+        content = bot.scrape_web_content("https://example.com/not-found")
+        self.assertIn("Error: Received status code 404", content)
+        
+        # Test with request exception
+        mock_requests_get.side_effect = Exception("Connection error")
+        content = bot.scrape_web_content("https://example.com")
+        self.assertIn("An error occurred", content)
+
+
+@pytest.mark.asyncio
+class TestErrorHandling:
+    
+    @pytest.fixture
+    def mock_channel(self):
+        mock = AsyncMock()
+        mock.send = AsyncMock()
+        return mock
+        
+    @patch('bot.client.chat.completions.create')
+    async def test_handle_api_error(self, mock_create, mock_channel):
+        """Test handling of API errors"""
+        mock_create.side_effect = Exception("API Error")
+        
+        channel = mock_channel()
+        # Call function that would use the API
+        await bot.send_chatgpt_response(
+            prompt="Test prompt",
+            channel=channel,
+            conversation_history=[],
+            user_id="12345"
+        )
+        
+        # Verify error was handled and communicated
+        channel.send.assert_called_once()
+        args, _ = channel.send.call_args
+        self.assertIn("error", args[0].lower())
+
+    @patch('bot.execute_code')
+    async def test_code_execution_timeout(self, mock_execute):
+        """Test handling of code execution timeout"""
+        mock_execute.side_effect = asyncio.TimeoutError()
+        
+        result = await bot.process_tool_calls(
+            {"tool_calls": [{"function": {"name": "code_interpreter", "arguments": '{"code": "print(\\"test\\")", "language": "python"}'}}]},
+            []
+        )
+        
+        self.assertIn("execution timed out", result[0]["content"].lower())
+
+class TestCommandHandling(unittest.TestCase):
+    
     def setUp(self):
-        self.bot = bot
-        self.interaction = AsyncMock()
-        self.interaction.user.id = 123456789  # Mock user ID
+        self.ctx = MagicMock()
+        self.ctx.author = MagicMock()
+        self.ctx.author.id = "12345"
+        self.ctx.send = AsyncMock()
+    
+    @patch('bot.prepare_messages_for_api')
+    @patch('bot.client.chat.completions.create')
+    async def test_chat_command(self, mock_create, mock_prepare):
+        """Test chat command"""
+        mock_response = MagicMock()
+        mock_response.choices = [MagicMock()]
+        mock_response.choices[0].message.content = "Hello, I'm ChatGPT"
+        mock_create.return_value = mock_response
+        mock_prepare.return_value = []
+        
+        # Call chat command
+        await bot.chat(self.ctx, "Tell me about Python")
+        
+        # Verify
+        self.ctx.send.assert_called()
+        mock_create.assert_called_once()
 
-    async def test_search_command(self):
-        # Set up mocks for interaction methods
-        self.interaction.response.defer = AsyncMock()
-        self.interaction.followup.send = AsyncMock()
-
-        # Call the search command with a sample query
-        await search(self.interaction, query="Python")
-
-        # Check if followup.send was called
-        self.interaction.followup.send.assert_called()
-        self.interaction.response.defer.assert_called_with(thinking=True)
-
-    async def test_generate_image_command(self):
-        # Mock the deferred response
-        self.interaction.response.defer = AsyncMock()
-        self.interaction.followup.send = AsyncMock()
-
-        # 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")
-
-            # Check if defer and followup were called
-            self.interaction.response.defer.assert_called_with(thinking=True)
-            self.interaction.followup.send.assert_called()
-
-    async def test_web_scraping_command(self):
-        # Mock the interaction methods
-        self.interaction.response.defer = AsyncMock()
-        self.interaction.followup.send = AsyncMock()
-
-        # 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)
-
-    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
-
-        # Mock channel.send to test if the bot sends a message
-        message.channel.send = AsyncMock()
-
-        # Test the bot's response
-        await bot.on_message(message)
-        message.channel.send.assert_called()  # Check if the bot replied
+@patch('bot.HISTORY_DB', {})  # Mocking the database dictionary
+class TestDatabaseFunctions(unittest.TestCase):
+    
+    async def test_history_functions(self):
+        """Test user history saving and retrieval"""
+        user_id = "12345"
+        test_history = [
+            {"role": "user", "content": "Hello"},
+            {"role": "assistant", "content": "Hi there!"}
+        ]
+        
+        # Test saving history
+        await bot.save_history(user_id, test_history)
+        
+        # Test retrieving history
+        retrieved = await bot.get_history(user_id)
+        self.assertEqual(retrieved, test_history)
+        
+        # Test history for new user
+        new_user = "67890"
+        new_history = await bot.get_history(new_user)
+        self.assertEqual(len(new_history), 1)  # Should have system message
+        self.assertEqual(new_history[0]["role"], "system")
+    
+    @patch('bot.DEFAULT_MODEL', 'gpt-4o-mini')
+    async def test_user_settings(self):
+        """Test user settings management"""
+        user_id = "12345"
+        
+        # Test default settings
+        model = await bot.get_user_model(user_id)
+        self.assertEqual(model, "gpt-4o-mini")
+        
+        # Test updating settings
+        await bot.set_user_model(user_id, "gpt-4o")
+        updated_model = await bot.get_user_model(user_id)
+        self.assertEqual(updated_model, "gpt-4o")
+        
+        # Test invalid model handling
+        with self.assertRaises(ValueError):
+            await bot.set_user_model(user_id, "invalid-model")
 
 if __name__ == '__main__':
     unittest.main()