diff --git a/README.md b/README.md
index f430c61..0017cbd 100644
--- a/README.md
+++ b/README.md
@@ -1,4 +1,4 @@
-# Video Downloader Bot Setup Guide
+# Video Downloader Bot
[](LICENSE)
@@ -7,219 +7,437 @@
[](https://github.com/ovchynnikov/load-bot-linux/actions/workflows/github-actions-push-image.yml)
[](https://github.com/ovchynnikov/load-bot-linux/actions/workflows/github-action-push-to-remote.yml)
-This guide provides step-by-step instructions on installation and running the Video Downloader bot on a Linux system.
-- Backend code uses [yt-dlp](https://github.com/yt-dlp/yt-dlp) which is released under The [Unlicense](https://unlicense.org/). All rights for yt-dlp belong to their respective authors.
----
+A Telegram bot that downloads videos from 1000+ platforms (YouTube, Instagram, TikTok, Reddit, X, Facebook, etc.) with automatic compression and optional AI chat capabilities.
-## Deploy with Docker
+## Contents
+- [Quick Start](#quick-start)
+- [Features](#features)
+- [Setup](#setup)
+- [Usage](#usage)
+- [Troubleshooting](#troubleshooting)
+- [Contributing](#contributing)
-Prerequisites:
- 1. Create `.env` file in the project root folder with your token (mandatory) and access configuration (optional). Use `.env.example` as a reference.
- 2. Clone the repo
- ```sh
- git clone https://github.com/ovchynnikov/load-bot-linux.git
- ```
-Build and run the container
-```
-docker build . -t downloader-bot:latest
-```
-```
-docker run -d --name downloader-bot --restart always --env-file .env downloader-bot:latest
-```
-To persist user data (conversation history, rate limits) between restarts, add a volume:
-```
-docker run -d --name downloader-bot --restart always --env-file .env -v bot-data:/bot/data downloader-bot:latest
-```
-or use a built image from **Docker hub**
-```
-docker run -d --name downloader-bot --restart always --env-file .env ovchynnikov/load-bot-linux:latest
-```
-With persistent data:
-```
-docker run -d --name downloader-bot --restart always --env-file .env -v bot-data:/bot/data ovchynnikov/load-bot-linux:latest
-```
-or if you use instagram cookies
+## Quick Start
+
+### Docker (recommended)
+
+```bash
+# Create .env with your bot token
+echo "BOT_TOKEN=your_token_here" > .env
+
+# Run
+docker run -d --name downloader-bot --restart always --env-file .env \
+ -v bot-data:/bot/data \
+ ovchynnikov/load-bot-linux:latest
```
-docker run -d --name downloader-bot --restart always --env-file .env -v bot-data:/bot/data -v /absolute/path/to/instagram_cookies.txt:/bot/instagram_cookies.txt ovchynnikov/load-bot-linux:latest
+
+### Systemd
+
+```bash
+git clone https://github.com/ovchynnikov/load-bot-linux.git
+cd load-bot-linux
+pip install -r src/requirements.txt
+sudo apt install ffmpeg
+
+# Create service (see Deploy with Linux Service section below)
```
-or if you want use GPU power of intel chip and set USE_GPU_COMPRESSING=True variable
+
+### Docker Compose
+
+```bash
+git clone https://github.com/ovchynnikov/load-bot-linux.git
+cd load-bot-linux
+docker-compose up -d
```
-docker run --rm --device /dev/dri:/dev/dri --group-add video downloader-bot .....
+
+Get a Telegram bot token from [@BotFather](https://t.me/botfather), then send `bot_health` to test.
+
+## Features
+
+- Downloads from 1000+ video platforms
+- Automatic compression to fit Telegram's 50 MB limit
+- GPU acceleration (Intel VAAPI)
+- Instagram Stories/Carousels with automatic fallback
+- Access control via allowlist (by username or chat ID)
+- Error reporting to admin chats
+- Multi-language support (Ukrainian, English)
+- Optional AI chat (Grok or Google Gemini)
+- Conversation context prompt history per user for AI
+
+## Setup
+
+### Prerequisites
+
+- Python 3.9+
+- FFmpeg
+- Linux OS
+
+### Get Bot Token
+
+1. Chat with [@BotFather](https://t.me/botfather) on Telegram
+2. Create a bot and copy the token
+3. Add to `.env` file
+
+### Environment Variables
+
+**Required:**
+- `BOT_TOKEN` - Your Telegram bot token
+
+
+ Click to expand
+
+**Optional - Basic:**
+- `LANGUAGE` - `en` or `uk` (default: uk)
+- `LOG_LEVEL` - DEBUG, INFO, WARNING, ERROR (default: INFO)
+
+**Optional - Video Processing:**
+- `H_CODEC` - `libx265` (smaller) or `libx264` (default: libx265)
+- `USE_GPU_COMPRESSING` - Enable Intel VAAPI (default: false)
+- `INSTACOOKIES` - Use Instagram cookies file (default: false)
+
+**Optional - Access Control:**
+- `LIMIT_BOT_ACCESS` - Restrict to allowlist (default: false)
+- `ALLOWED_USERNAMES` - Comma-separated usernames
+- `ALLOWED_CHAT_IDS` - Comma-separated chat IDs
+
+**Optional - Error Reporting:**
+- `SEND_ERROR_TO_ADMIN` - Forward errors to admin (default: false)
+- `ADMINS_CHAT_IDS` - Comma-separated admin chat IDs
+- `SEND_USER_INFO_WITH_HEALTHCHECK` - Send user info when bot_health command is triggered
+
+**Optional - AI/LLM (Grok or Gemini):**
+- `USE_LLM` - Enable AI chat (default: false)
+- `LLM_PROVIDER` - `grok` or `gemini` (default: grok)
+- `GROK_API_KEY` - xAI API key (get from https://console.grok.ai)
+- `GEMINI_API_KEY` - Google API key (get from https://aistudio.google.com)
+- `LLM_RPM_LIMIT` - LLM Requests (prompts) per minute (default: 50)
+- `LLM_RPD_LIMIT` - LLM Requests per day (default: 500)
+- `GROK_IMG_MODEL` - default: grok-imagine-image
+- `IMG_GEN_RPM_LIMIT` - Image generations per minute (default: 1)
+- `IMG_GEN_RPD_LIMIT` - Image generations per day (default: 25)
+- `MAX_CONTEXT_MESSAGES` - LLM Messages (prompts) to remember per user (default: 3)
+- `MAX_CONTEXT_CHARS` - Max chars per message (default: 500)
+
+**Optional - Cleanup:**
+- `USER_CLEANUP_TTL_DAYS` - Remove LLM context messages (prompts) for inactive users after N days (default: 3)
+- `USER_CLEANUP_INTERVAL_HOURS` - Cleanup interval (default: 24)
+
+
+
+### Example .env
+
+```ini
+BOT_TOKEN=123456789:ABCDEFghijklmnopqrstuvwxyz
+LANGUAGE=en
+LIMIT_BOT_ACCESS=false
+ALLOWED_USERNAMES=
+ALLOWED_CHAT_IDS=
+H_CODEC=libx265
+USE_GPU_COMPRESSING=false
+INSTACOOKIES=false
+SEND_ERROR_TO_ADMIN=false
+ADMINS_CHAT_IDS=
+USE_LLM=false
+LLM_PROVIDER=grok
```
-Alternatively, you can use **docker-compose**
+
+## Deploy with Docker
+
+### Basic
+
+```bash
+docker run -d --name downloader-bot --restart always --env-file .env \
+ -v bot-data:/bot/data \
+ ovchynnikov/load-bot-linux:latest
```
-docker-compose build
+
+### With Instagram Cookies
+
+```bash
+docker run -d --name downloader-bot --restart always --env-file .env \
+ -v bot-data:/bot/data \
+ -v /path/to/instagram_cookies.txt:/bot/instagram_cookies.txt \
+ ovchynnikov/load-bot-linux:latest
```
+
+Enable `INSTACOOKIES=true` in `.env`.
+
+### With GPU (Intel)
+
+```bash
+docker run -d --name downloader-bot --restart always --env-file .env \
+ -v bot-data:/bot/data \
+ --device /dev/dri:/dev/dri \
+ --group-add video \
+ ovchynnikov/load-bot-linux:latest
```
-docker-compose up
+
+Set `USE_GPU_COMPRESSING=true` in `.env`.
+
+### Build Custom Image
+
+```bash
+git clone https://github.com/ovchynnikov/load-bot-linux.git
+cd load-bot-linux
+docker build . -t downloader-bot:latest
+docker run -d --name downloader-bot --restart always --env-file .env \
+ -v bot-data:/bot/data \
+ downloader-bot:latest
```
----
-## Deploy with Linux Service (daemon)
+## Deploy with Linux Service (Systemd)
+
Click to expand
- 1. Clone and Install
-Clone the repo
-```sh
-git clone https://github.com/ovchynnikov/load-bot-linux.git
-```
- 2. Install dependencies
+### Install
+
```bash
-pip install -r scr/requirements.txt
-```
-```sh
-sudo apt update && sudo apt install ffmpeg -y
-```
- 3. Change permissions for the yt-dlp
-```
+git clone https://github.com/ovchynnikov/load-bot-linux.git
+cd load-bot-linux
+pip install -r src/requirements.txt
+sudo apt install ffmpeg
sudo chmod a+rx $(which yt-dlp)
```
- 4. Create and configure Linux service
-```sh
+### Create Service
+
+```bash
sudo nano /etc/systemd/system/downloader-bot.service
```
-Add the following configuration to the file:
```ini
[Unit]
Description=Video Downloader Bot Service
After=network.target
[Service]
-User=your_linux_user # <====== REPLACE `your_linux_user` with the username that will run the bot.
-WorkingDirectory=/path/to/your/bot # <====== REPLACE THIS with the absolute path to your bot's folder.
-ExecStart=/usr/bin/python3 /path/to/your/bot/main.py # <====== REPLACE THIS with the command to start your bot. Adjust if you're using a virtual environment.
-Restart=always # Ensures the bot restarts automatically if it crashes.
+User=your_user
+WorkingDirectory=/path/to/bot
+ExecStart=/usr/bin/python3 /path/to/bot/main.py
+Restart=always
RestartSec=5
-Environment="BOT_TOKEN=your_bot_token" # <====== REPLACE THIS with your bot token.
+Environment="BOT_TOKEN=your_token_here"
+Environment="LANGUAGE=en"
Environment="LOG_LEVEL=INFO"
-Environment="LIMIT_BOT_ACCESS=False" # <====== REPLACE THIS (value is optional. False by default) Type: Boolean
-Environment="ALLOWED_USERNAMES=" # <====== REPLACE THIS (value is optional) Type: string separated by commas. Example: ALLOWED_USERNAMES=username1,username2,username3
-Environment="ALLOWED_CHAT_IDS=" # <====== REPLACE THIS (value is optional) Type: string separated by commas. Example: ALLOWED_CHAT_IDS=12349,12345,123456
-Environment="INSTACOOKIES=False" # <====== REPLACE THIS (value is optional) Type: Boolean. False by default.
-Environment="ADMINS_CHAT_IDS=" # <====== REPLACE THIS (value is optional) Type: string separated by commas. IDS to send Exceptions errors to private messages. Get this from bot health check
-Environment="SEND_ERROR_TO_ADMIN=True" # <====== REPLACE THIS (value is optional) Type: Boolean. Send errors to admins in private messages
-Environment="H_CODEC=libx265" # <====== REPLACE THIS (value is optional) Type: String. libx265 or libx264
-Environment="USE_GPU_COMPRESSING=False" # <====== Enable to use GPU for video compression using Intel chip and VAAPI. False by default
[Install]
WantedBy=multi-user.target
```
- 5. Start the Bot Service
-
-Reload the systemd daemon and start the bot service:
+### Start
```bash
sudo systemctl daemon-reload
-```
-```bash
sudo systemctl enable downloader-bot.service
-```
-```bash
sudo systemctl start downloader-bot.service
-```
-```bash
sudo systemctl status downloader-bot.service
```
- 6. Troubleshooting
+### View Logs
+
+```bash
+journalctl -u downloader-bot.service -f
+```
-- Check the status of the service:
- ```sh
- sudo systemctl status downloader-bot.service
- ```
-- View logs for more details:
- ```sh
- journalctl -u downloader-bot.service
- ```
-## How to use the bot
-
- 1. Create Your Token for the Telegram Bot
-- Follow this guide to create your Telegram bot and obtain the bot token:
- [How to Get Your Bot Token](https://www.freecodecamp.org/news/how-to-create-a-telegram-bot-using-python/).
- Make sure you put token in `.env` file
-
- 2. Health Check
-- Verify the bot is running by sending a message with the trigger word:
- ```sh
- bot_health
- ```
- or
- ```sh
- ботяра
- ```
-
- If the bot is active, it will respond accordingly.
-
- 3. Once the bot is created and the Linux service or Docker image is running:
- Send a URL from **YouTube Shorts**, **Instagram Reels**, or similar platforms to the bot.
- Example:
- ```
- https://youtube.com/shorts/kaTxVLGd6IE?si=YaUM3gYjr1kcXqTm
- ```
- Wait for the bot to process the URL and respond.
-
-## Supported platforms by default
-```
-instagram reels
-facebook reels
-tiktok
-reddit
-x.com
-youtube shorts
-```
-
-### Download videos from other sources.
-Videos shorter than 10 minutes usually work fine. The Telegram limitation for a video is 50 MB.
-- To download the full video from YouTube add two asterisks before the url address.
-Example:
-```
- **https://www.youtube.com/watch?v=rxdu3whDVSM or with a space ** https://www.youtube.com/watch?v=rxdu3whDVSM
-```
-The expected waiting time for videos up to 10 minutes is 3-10 minutes depending on the internet speed.
-- Full list of supported sites here: [yt-dlp Supported Sites](https://github.com/yt-dlp/yt-dlp/blob/master/supportedsites.md)
-
-### Instagram Stories and Reels credentials
-- To download Instagram stories and reels you need to create a cookies file `instagram_cookies.txt` in the `bot` folder and set env var `INSTACOOKIES` to `True`.
-- You can use the `instagram_cookies_example.txt` file as a reference from the `src` folder of the repo.
-- Suggestion on how to get the file: easy export with [chrome extension](https://chromewebstore.google.com/detail/get-cookiestxt-locally/cclelndahbckbenkjhflpdbgdldlbecc)
-- When you run the bot with Docker, place `instagram_cookies.txt` to the folder with your `.env` file and add `-v instagram_cookies.txt:/bot/instagram_cookies.txt` to the start command
-- The bot supports downloading Instagram stories and carousels with pictures using `gallery-dl` when `yt-dlp` fails.
-- The bot will automatically fall back to `gallery-dl` for Instagram URLs without `reels` when `yt-dlp` fails, multiple media files is supported.
-- The same cookies file used for `yt-dlp` can be used for `gallery-dl`
-
-## Access Control with Safe List
-The bot can use 'Safelist' to restrict access for users or groups.
-Ensure these variables are set in your `.env` file, without them or with the chat ID and username.
-You can get your `chat_id` and `username` by setting `LIMIT_BOT_ACCESS=True` first. Then, send the word `bot_health` or `ботяра`, and the bot will answer you with the chat ID and username.
-The priority for allowed Group Chat is highest. All users in the Group Chat can use the bot even if they do not have access to it in private chat.
-- When `LIMIT_BOT_ACCESS=True` to use the bot in private messages add the username to the `ALLOWED_USERNAMES` variable or chat ID to `ALLOWED_CHAT_IDS`.
-- If you want a bot in your Group Chat with restrictions, leave `ALLOWED_CHAT_IDS` empty and define the `ALLOWED_USERNAMES` variable list.
-```ini
-LIMIT_BOT_ACCESS=False # If True, the bot will only work for users in ALLOWED_USERNAMES or ALLOWED_CHAT_IDS
-ALLOWED_USERNAMES= # a list of allowed usernames as strings separated by commas. Example: ALLOWED_USERNAMES=username1,username2,username3
-ALLOWED_CHAT_IDS= # a list of allowed chat IDs as strings separated by commas. Example: ALLOWED_CHAT_IDS=-412349,12345,123456
+## Usage
+
+### Send a Video URL
+
+Simply send any supported platform URL:
+
+```
+https://youtube.com/shorts/video_id
+https://www.instagram.com/reel/ABC123/
+https://www.tiktok.com/@user/video/123456
```
-## Troubleshooting
-If you sent a link to the bot and got no response, send the word `bot_health` or `ботяра` to the bot to check if it's working.
+### Download Full YouTube Video
+
+Prefix with `**`:
+
+```
+**https://www.youtube.com/watch?v=video_id
+```
+
+### Check Bot Status
+
+Send `bot_health` to the bot. It will respond with status.
+
+### AI Chat (if enabled)
+
+Send `ботяра, ` any message and the bot will respond using Grok or Gemini.
+
+### Generate Image (Grok only)
+
+```
+ботяра, image: a sunset over mountains
+```
+
+## Supported Platforms
+
+- Instagram (Reels, Stories, Carousels)
+- Facebook Reels
+- TikTok
+- YouTube (Shorts and full videos)
+- Reddit
+- X.com (Twitter)
+- 1000+ others via yt-dlp
+
+See the [full list](https://github.com/yt-dlp/yt-dlp/blob/master/supportedsites.md).
+
+## Instagram Stories and Carousels
-- in .env file set IDS to send Exceptions errors to in private messages to Admins. Get these ids from bot healthcheck.
-Works only for Exceptions errors that are not handled by the bot code.
+To download private/age-restricted content:
+
+1. Export cookies using a browser extension
+2. Save as `instagram_cookies.txt`
+3. Mount in Docker or place in working directory
+4. Set `INSTACOOKIES=true`
+
+The bot automatically falls back to gallery-dl if yt-dlp fails.
+
+## Access Control
+
+Restrict bot access to specific users or groups:
```ini
-ADMINS_CHAT_IDS="your_admins_chat_id_here" # ADMINS_CHAT_IDS=chatid_1,chatid_2,chatid_3
+LIMIT_BOT_ACCESS=true
+ALLOWED_USERNAMES=username1,username2
+ALLOWED_CHAT_IDS=12345,67890
```
-- in .env file set SEND_ERROR_TO_ADMIN=True to send errors to admins in private messages
+To get your IDs, send `bot_health` to the bot.
+
+## Error Reporting
+
+Forward errors to admin chats:
```ini
-SEND_ERROR_TO_ADMIN=True
+SEND_ERROR_TO_ADMIN=true
+ADMINS_CHAT_IDS=12345,67890
+```
+
+## AI/LLM Chat
+
+Optional integration with language models.
+
+### Setup Grok (xAI)
+
+1. Sign up at https://console.grok.ai
+2. Get API key
+3. Set in `.env`:
+ ```ini
+ USE_LLM=true
+ LLM_PROVIDER=grok
+ GROK_API_KEY=xai-your-key
+ ```
+
+### Setup Gemini (Google)
+
+1. Get API key at https://aistudio.google.com
+2. Set in `.env`:
+ ```ini
+ USE_LLM=true
+ LLM_PROVIDER=gemini
+ GEMINI_API_KEY=your-key
+ ```
+
+### Usage
+
+- Send a message: bot responds with AI
+- Image generation (Grok): `ботяра, image: prompt`
+- Bot remembers conversation history (configurable)
+
+## Troubleshooting
+
+### Bot not responding
+
+```bash
+# Check if running
+docker ps | grep downloader-bot
+
+# View logs
+docker logs downloader-bot
+
+# Systemd logs
+journalctl -u downloader-bot.service -n 50
+```
+
+Send `bot_health` to test.
+
+### Video download fails
+
+- Check if platform is supported
+- For YouTube, use `**` prefix for full videos
+- Check available disk space
+- Enable debug logging: `LOG_LEVEL=DEBUG`
+
+### Instagram downloads don't work
+
+- Set up cookies file (see Instagram section)
+- Enable `INSTACOOKIES=true`
+- Ensure cookies are valid
+
+### GPU not working
+
+Check if Intel GPU is present:
+
+```bash
+vainfo
+```
+
+If not found, install drivers:
+
+```bash
+sudo apt install intel-media-va-driver-non-free
```
+
+### Database locked
+
+```bash
+docker restart downloader-bot
+```
+
+Or clear database (WARNING: loses user data):
+
+```bash
+docker exec downloader-bot rm /bot/data/bot.db
+docker restart downloader-bot
+```
+
+## Contributing
+
+Contributions welcome. Please:
+
+1. Check existing issues
+2. Open an issue or fork and submit a PR
+3. Follow code style (black, type hints)
+
+To set up development:
+
+```bash
+git clone https://github.com/yourusername/load-bot-linux.git
+cd load-bot-linux
+python3 -m venv venv
+source venv/bin/activate
+pip install -r src/requirements.txt
+```
+
+## License
+
+MIT License - see [LICENSE](LICENSE) file.
+
+## Credits
+
+- [yt-dlp](https://github.com/yt-dlp/yt-dlp) - Video downloader
+- [python-telegram-bot](https://github.com/python-telegram-bot/python-telegram-bot) - Telegram API
+- [gallery-dl](https://github.com/mikf/gallery-dl) - Media gallery downloader
+- [FFmpeg](https://ffmpeg.org) - Video processing
+
---
+
+Backend code uses [yt-dlp](https://github.com/yt-dlp/yt-dlp) which is released under The [Unlicense](https://unlicense.org/). All rights for yt-dlp belong to their respective authors.
diff --git a/src/db_storage.py b/src/db_storage.py
index 2ecfdb4..651407c 100644
--- a/src/db_storage.py
+++ b/src/db_storage.py
@@ -28,16 +28,38 @@ def _create_tables(self):
rate_limit_timestamps TEXT,
daily_count INTEGER DEFAULT 0,
daily_date TEXT,
- last_seen REAL
+ last_seen REAL,
+ img_gen_rate_limit_timestamps TEXT,
+ img_gen_daily_count INTEGER DEFAULT 0,
+ img_gen_daily_date TEXT
)
""")
cursor.execute("CREATE INDEX IF NOT EXISTS idx_user_data_last_seen ON user_data(last_seen)")
+ # Migrate existing tables that don't have image gen columns
+ for col, definition in [
+ ("img_gen_rate_limit_timestamps", "TEXT"),
+ ("img_gen_daily_count", "INTEGER DEFAULT 0"),
+ ("img_gen_daily_date", "TEXT"),
+ ]:
+ try:
+ cursor.execute(f"ALTER TABLE user_data ADD COLUMN {col} {definition}")
+ except sqlite3.OperationalError as e:
+ # Only suppress the error if the column already exists
+ error_msg = str(e).lower()
+ if "duplicate" not in error_msg and "already exists" not in error_msg:
+ raise
+ # Column already exists, continue
self.conn.commit()
def load_user_data(self, user_id):
"""Load user data from database."""
cursor = self.conn.cursor()
- cursor.execute("SELECT * FROM user_data WHERE user_id = ?", (user_id,))
+ cursor.execute(
+ "SELECT user_id, conversation_context, rate_limit_timestamps, daily_count, daily_date, last_seen,"
+ " img_gen_rate_limit_timestamps, img_gen_daily_count, img_gen_daily_date"
+ " FROM user_data WHERE user_id = ?",
+ (user_id,),
+ )
row = cursor.fetchone()
if row:
return {
@@ -46,17 +68,32 @@ def load_user_data(self, user_id):
"daily_count": row[3],
"daily_date": row[4],
"last_seen": row[5],
+ "img_gen_rate_limit_timestamps": json.loads(row[6]) if row[6] else [],
+ "img_gen_daily_count": row[7] or 0,
+ "img_gen_daily_date": row[8] or "",
}
return None
- def save_user_data(self, user_id, conversation_context, rate_limit_timestamps, daily_count, daily_date, last_seen):
+ def save_user_data(
+ self,
+ user_id,
+ conversation_context,
+ rate_limit_timestamps,
+ daily_count,
+ daily_date,
+ last_seen,
+ img_gen_rate_limit_timestamps=None,
+ img_gen_daily_count=0,
+ img_gen_daily_date="",
+ ):
"""Save user data to database."""
cursor = self.conn.cursor()
cursor.execute(
"""
- INSERT OR REPLACE INTO user_data
- (user_id, conversation_context, rate_limit_timestamps, daily_count, daily_date, last_seen)
- VALUES (?, ?, ?, ?, ?, ?)
+ INSERT OR REPLACE INTO user_data
+ (user_id, conversation_context, rate_limit_timestamps, daily_count, daily_date, last_seen,
+ img_gen_rate_limit_timestamps, img_gen_daily_count, img_gen_daily_date)
+ VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
""",
(
user_id,
@@ -65,6 +102,9 @@ def save_user_data(self, user_id, conversation_context, rate_limit_timestamps, d
daily_count,
daily_date,
last_seen,
+ json.dumps(img_gen_rate_limit_timestamps or []),
+ img_gen_daily_count,
+ img_gen_daily_date,
),
)
self.conn.commit()
@@ -75,6 +115,31 @@ def delete_user_data(self, user_id):
cursor.execute("DELETE FROM user_data WHERE user_id = ?", (user_id,))
self.conn.commit()
+ def update_user_image_limits(self, user_id, img_gen_rate_limit_timestamps, img_gen_daily_count, img_gen_daily_date):
+ """Update or insert image generation limit fields.
+
+ Ensures the user row exists with image limit fields set, either by updating
+ an existing row or creating a new one with defaults for other fields.
+ """
+ cursor = self.conn.cursor()
+ cursor.execute(
+ """
+ INSERT INTO user_data (user_id, img_gen_rate_limit_timestamps, img_gen_daily_count, img_gen_daily_date)
+ VALUES (?, ?, ?, ?)
+ ON CONFLICT(user_id) DO UPDATE SET
+ img_gen_rate_limit_timestamps = excluded.img_gen_rate_limit_timestamps,
+ img_gen_daily_count = excluded.img_gen_daily_count,
+ img_gen_daily_date = excluded.img_gen_daily_date
+ """,
+ (
+ user_id,
+ json.dumps(img_gen_rate_limit_timestamps or []),
+ img_gen_daily_count,
+ img_gen_daily_date,
+ ),
+ )
+ self.conn.commit()
+
def get_stale_users(self, ttl_seconds):
"""Get list of user IDs that haven't been seen within TTL."""
current_time = time.time()
diff --git a/src/main.py b/src/main.py
index ab72811..53ba02b 100644
--- a/src/main.py
+++ b/src/main.py
@@ -1,5 +1,6 @@
"""Download videos from tiktok, x(twitter), reddit, youtube shorts, instagram reels and many more"""
+import base64
import os
import random
import json
@@ -8,8 +9,14 @@
import time
import traceback
from datetime import datetime
+from typing import Optional
import google.generativeai as genai
from openai import AsyncOpenAI
+
+try:
+ import xai_sdk
+except ImportError:
+ xai_sdk = None
from functools import lru_cache
from collections import defaultdict
from dotenv import load_dotenv
@@ -48,10 +55,24 @@
GEMINI_MODEL = os.getenv("GEMINI_MODEL", "gemini-flash-latest")
GROK_API_KEY = os.getenv("GROK_API_KEY")
GROK_MODEL = os.getenv("GROK_MODEL", "grok-4-latest")
+GROK_IMG_MODEL = os.getenv("GROK_IMG_MODEL", "grok-imagine-image")
TELEGRAM_CONNECT_TIMEOUT = 60
TELEGRAM_POOL_TIMEOUT = 30
TELEGRAM_READ_TIMEOUT = 120
TELEGRAM_WRITE_TIMEOUT = 120
+MAX_PROMPT_LEN = 1000
+
+
+def get_image_caption():
+ """Get localized image caption."""
+ if language == "uk":
+ return "Ось ваше зображення 🖼️"
+ else:
+ return "Here's your image 🖼️"
+
+
+IMAGE_CAPTION_STUB = get_image_caption() # Legacy reference for compatibility
+IMAGE_TIMEOUT_SEC = 30.0
# Configure Gemini API
if GEMINI_API_KEY:
@@ -62,11 +83,26 @@
if GROK_API_KEY:
grok_client = AsyncOpenAI(api_key=GROK_API_KEY, base_url="https://api.x.ai/v1")
+# Configure xAI image API client (grok-imagine-image)
+xai_client = None
+if GROK_API_KEY and xai_sdk is not None:
+ try:
+ xai_client = xai_sdk.Client(api_key=GROK_API_KEY)
+ except Exception as e: # pylint: disable=broad-except
+ error("Failed to initialize xai_sdk.Client: %s", e)
+ xai_client = None
+
# Rate limiting for LLM APIs
llm_rate_limit = defaultdict(list) # {user_id: [timestamp1, timestamp2, ...]}
llm_daily_limit = defaultdict(lambda: {"count": 0, "date": ""}) # {user_id: {count, date}}
-LLM_RPM_LIMIT = int(os.getenv("LLM_RPM_LIMIT", "50")) # Requests per minute per user
-LLM_RPD_LIMIT = int(os.getenv("LLM_RPD_LIMIT", "500")) # Requests per day per user
+
+# Rate limiting for Image Generation
+img_gen_rate_limit = defaultdict(list) # {user_id: [timestamp1, timestamp2, ...]}
+img_gen_daily_limit = defaultdict(lambda: {"count": 0, "date": ""}) # {user_id: {count, date}}
+LLM_RPM_LIMIT = int(os.getenv("LLM_RPM_LIMIT", "50")) # LLM Requests per minute per user
+LLM_RPD_LIMIT = int(os.getenv("LLM_RPD_LIMIT", "500")) # LLM Requests per day per user
+IMG_GEN_RPM_LIMIT = int(os.getenv("IMG_GEN_RPM_LIMIT", "1")) # Image Generation Requests per minute per user
+IMG_GEN_RPD_LIMIT = int(os.getenv("IMG_GEN_RPD_LIMIT", "25")) # Image Generation Requests per day per user
# Conversation context storage: {user_id: [(user_msg, bot_response), ...]}
conversation_context = defaultdict(list)
@@ -160,6 +196,181 @@ def is_bot_mentioned(message_text: str) -> bool:
return False
+def extract_image_prompt(message_text: str) -> Optional[str]:
+ """Extract image generation prompt for commands like 'ботяра, image: ...'."""
+ if not message_text:
+ return None
+
+ lower = message_text.lower()
+ # Match bot command for image generation: ботяра, image: prompt
+ match = re.search(r"ботяра[^\w\d]*image\s*:\s*(.+)", lower)
+ if match:
+ prompt = match.group(1).strip()
+ return prompt or None
+
+ # Fallback for english trigger
+ match = re.search(r"bot\s*:\s*image\s*:\s*(.+)", lower)
+ if match:
+ prompt = match.group(1).strip()
+ return prompt or None
+
+ return None
+
+
+async def generate_image_and_send(update: Update, prompt: str) -> None:
+ """Generate image through Grok image API and send to Telegram."""
+ if not prompt:
+ await update.message.reply_text(
+ (
+ "Вкажіть, що саме потрібно згенерувати після 'botyara, image:'"
+ if language == "uk"
+ else "Please specify what you want to generate after 'bot, image:'"
+ ),
+ reply_to_message_id=update.message.message_id,
+ )
+ return
+
+ if not GROK_API_KEY:
+ await update.message.reply_text(
+ (
+ "Grok API key не налаштовано. Будь ласка, встановіть GROK_API_KEY."
+ if language == "uk"
+ else "Grok API key is not configured. Please set GROK_API_KEY."
+ ),
+ reply_to_message_id=update.message.message_id,
+ )
+ return
+
+ if not xai_sdk or not xai_client:
+ await update.message.reply_text(
+ (
+ "xAI клієнт недоступний. Перевірте встановлення xai-sdk та GROK_API_KEY."
+ if language == "uk"
+ else "xAI client is unavailable. Please check xai-sdk installation and GROK_API_KEY."
+ ),
+ reply_to_message_id=update.message.message_id,
+ )
+ return
+
+ # Rate limiting for image generation
+ user_id = update.effective_user.id
+ current_time = time.time()
+
+ # Load img_gen data from DB on first access
+ if user_id not in img_gen_daily_limit:
+ user_data = await asyncio.to_thread(db_storage.load_user_data, user_id)
+ if user_data:
+ img_gen_rate_limit[user_id] = user_data["img_gen_rate_limit_timestamps"]
+ img_gen_daily_limit[user_id] = {
+ "count": user_data["img_gen_daily_count"],
+ "date": user_data["img_gen_daily_date"],
+ }
+
+ # Clean old timestamps (older than 60 seconds)
+ img_gen_rate_limit[user_id] = [t for t in img_gen_rate_limit[user_id] if current_time - t < 60]
+
+ if len(img_gen_rate_limit[user_id]) >= IMG_GEN_RPM_LIMIT:
+ debug("Image gen RPM limit hit for user %s", user_id)
+ await update.message.reply_text(
+ (
+ "Вибачте, забагато запитів на генерацію зображень. Почекайте хвилину."
+ if language == "uk"
+ else "Sorry, too many image generation requests. Please wait a minute."
+ ),
+ reply_to_message_id=update.message.message_id,
+ )
+ return
+
+ # Check daily image gen limit
+ today = datetime.now().strftime("%Y-%m-%d")
+ if img_gen_daily_limit[user_id]["date"] != today:
+ img_gen_daily_limit[user_id] = {"count": 0, "date": today}
+
+ if img_gen_daily_limit[user_id]["count"] >= IMG_GEN_RPD_LIMIT:
+ debug("Image gen RPD limit hit for user %s", user_id)
+ await update.message.reply_text(
+ (
+ "Вибачте, денний ліміт генерації зображень вичерпано. Спробуйте завтра."
+ if language == "uk"
+ else "Sorry, daily image generation limit reached. Try again tomorrow."
+ ),
+ reply_to_message_id=update.message.message_id,
+ )
+ return
+
+ # Tentatively add current request timestamp (will be removed on failure)
+ img_gen_rate_limit[user_id].append(current_time)
+
+ prompt = prompt[:MAX_PROMPT_LEN].strip()
+
+ try:
+ image_response = await asyncio.wait_for(
+ asyncio.to_thread(
+ xai_client.image.sample,
+ prompt=prompt,
+ model=GROK_IMG_MODEL,
+ ),
+ timeout=IMAGE_TIMEOUT_SEC,
+ )
+
+ image_url = getattr(image_response, "url", None)
+ image_b64 = getattr(image_response, "image", None) if not image_url else None
+
+ if not image_url and not image_b64:
+ raise ValueError("Не вдалося отримати результат з xAI API")
+
+ if image_url:
+ await update.message.reply_photo(photo=image_url, caption=get_image_caption())
+ else:
+ file_bytes = base64.b64decode(image_b64)
+ await update.message.reply_photo(photo=file_bytes, caption=get_image_caption())
+
+ # Increment daily limit only after successful generation
+ img_gen_daily_limit[user_id]["count"] += 1
+
+ # Save img_gen rate limit data to DB (best-effort, targeted update only)
+ async def save_img_gen_to_db():
+ try:
+ await asyncio.to_thread(
+ db_storage.update_user_image_limits,
+ user_id,
+ img_gen_rate_limit[user_id],
+ img_gen_daily_limit[user_id]["count"],
+ img_gen_daily_limit[user_id]["date"],
+ )
+ except Exception as db_error: # pylint: disable=broad-except
+ error("Failed to save img_gen data to database: %s", db_error)
+
+ asyncio.create_task(save_img_gen_to_db())
+
+ except asyncio.TimeoutError:
+ error("Image generation timed out for prompt: %.100s", prompt)
+ # Remove tentative timestamp on failure
+ if img_gen_rate_limit[user_id] and img_gen_rate_limit[user_id][-1] == current_time:
+ img_gen_rate_limit[user_id].pop()
+ await update.message.reply_text(
+ (
+ "Генерація зайняла надто багато часу. Спробуйте пізніше."
+ if language == "uk"
+ else "Image generation took too long. Please try again later."
+ ),
+ reply_to_message_id=update.message.message_id,
+ )
+ except Exception as e: # pylint: disable=broad-except
+ error("Image generation failed: %s", e)
+ # Remove tentative timestamp on failure
+ if img_gen_rate_limit[user_id] and img_gen_rate_limit[user_id][-1] == current_time:
+ img_gen_rate_limit[user_id].pop()
+ await update.message.reply_text(
+ (
+ "Вибачте, не вдалося згенерувати зображення. Спробуйте пізніше."
+ if language == "uk"
+ else "Sorry, I couldn't generate the image. Please try again later."
+ ),
+ reply_to_message_id=update.message.message_id,
+ )
+
+
def clean_url(message_text: str) -> str:
"""
Cleans the URL from the message text by removing unwanted characters and usernames.
@@ -246,6 +457,22 @@ async def handle_message(update: Update, context: ContextTypes.DEFAULT_TYPE): #
debug("LLM_PROVIDER: %s", LLM_PROVIDER)
if bot_mentioned:
+ cleaned_text = message_text.strip().lower()
+
+ # Health check always takes priority, even with LLM enabled
+ if cleaned_text.startswith("bot_health"):
+ # Check if it's a pure health check command (no additional parameters like 'image:')
+ if "image:" not in cleaned_text:
+ debug("Health check command detected")
+ await respond_with_bot_message(update)
+ return
+
+ image_prompt = extract_image_prompt(message_text)
+ if image_prompt:
+ debug("Bot image command detected with prompt: %s", image_prompt)
+ await generate_image_and_send(update, image_prompt)
+ return
+
if USE_LLM:
debug("Calling LLM response function")
await respond_with_llm_message(update)
@@ -615,19 +842,19 @@ async def respond_with_llm_message(update):
try:
# Check if user is asking for image generation and modify prompt
image_keywords = [
- 'картинку',
- 'картинка',
- 'зображення',
- 'image',
- 'фото',
- 'picture',
- 'згенеруй',
- 'generate',
- 'створи',
- 'create',
- 'покажи',
- 'покажи мне',
- 'покажи мені',
+ # 'картинку',
+ # 'картинка',
+ # 'зображення',
+ # 'image',
+ # 'фото',
+ # 'picture',
+ # 'згенеруй',
+ # 'generate',
+ # 'створи',
+ # 'create',
+ # 'покажи',
+ # 'покажи мне',
+ # 'покажи мені',
]
# Check both original message and processed prompt
original_text = message_text.lower()
@@ -720,15 +947,25 @@ async def save_to_db():
llm_daily_limit[user_id]["count"],
llm_daily_limit[user_id]["date"],
)
- await asyncio.to_thread(
- db_storage.save_user_data,
- user_id,
- conversation_context[user_id],
- llm_rate_limit[user_id],
- llm_daily_limit[user_id]["count"],
- llm_daily_limit[user_id]["date"],
- user_last_seen[user_id],
- )
+
+ # Build save arguments, only including image gen data if explicitly set for this user
+ save_kwargs = {
+ "user_id": user_id,
+ "conversation_context": conversation_context[user_id],
+ "rate_limit_timestamps": llm_rate_limit[user_id],
+ "daily_count": llm_daily_limit[user_id]["count"],
+ "daily_date": llm_daily_limit[user_id]["date"],
+ "last_seen": user_last_seen[user_id],
+ }
+
+ # Only include image gen data if user has actually interacted with image generation
+ if user_id in img_gen_rate_limit:
+ save_kwargs["img_gen_rate_limit_timestamps"] = img_gen_rate_limit[user_id]
+ if user_id in img_gen_daily_limit:
+ save_kwargs["img_gen_daily_count"] = img_gen_daily_limit[user_id]["count"]
+ save_kwargs["img_gen_daily_date"] = img_gen_daily_limit[user_id]["date"]
+
+ await asyncio.to_thread(db_storage.save_user_data, **save_kwargs)
except Exception as db_error: # pylint: disable=broad-except
error("Failed to save user data to database: %s", db_error)
@@ -914,6 +1151,10 @@ async def cleanup_stale_users():
del llm_rate_limit[user_id]
if user_id in llm_daily_limit:
del llm_daily_limit[user_id]
+ if user_id in img_gen_rate_limit:
+ del img_gen_rate_limit[user_id]
+ if user_id in img_gen_daily_limit:
+ del img_gen_daily_limit[user_id]
if user_id in user_last_seen:
del user_last_seen[user_id]
# Remove from database
diff --git a/src/requirements.txt b/src/requirements.txt
index 244b45a..d615d20 100644
--- a/src/requirements.txt
+++ b/src/requirements.txt
@@ -5,3 +5,4 @@ gallery-dl>=1.31.7
aiohttp>=3.13.3
google-generativeai>=0.8.6
openai>=2.24.0
+xai-sdk>=1.8.1