From 2263fabfb9be152cee64331206f8c26fbd54d1a5 Mon Sep 17 00:00:00 2001 From: NikitolProject Date: Sun, 27 Apr 2025 06:32:34 +0300 Subject: [PATCH] Init commit --- .gitignore | 174 ++++++++++++++++++++ Dockerfile | 19 +++ README.md | 69 ++++++++ bot.py | 342 +++++++++++++++++++++++++++++++++++++++ docker-compose.yaml | 54 +++++++ docker/postgres/init.sql | 16 ++ requirements.txt | 4 + test-prompts.txt | 1 + wait-for-postgres.sh | 16 ++ 9 files changed, 695 insertions(+) create mode 100644 .gitignore create mode 100644 Dockerfile create mode 100644 README.md create mode 100644 bot.py create mode 100644 docker-compose.yaml create mode 100644 docker/postgres/init.sql create mode 100644 requirements.txt create mode 100644 test-prompts.txt create mode 100755 wait-for-postgres.sh diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..1800114 --- /dev/null +++ b/.gitignore @@ -0,0 +1,174 @@ +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ +cover/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +.pybuilder/ +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +# For a library or package, you might want to ignore these files since the code is +# intended to run in multiple environments; otherwise, check them in: +# .python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# UV +# Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +#uv.lock + +# poetry +# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control +#poetry.lock + +# pdm +# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. +#pdm.lock +# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it +# in version control. +# https://pdm.fming.dev/latest/usage/project/#working-with-version-control +.pdm.toml +.pdm-python +.pdm-build/ + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static type analyzer +.pytype/ + +# Cython debug symbols +cython_debug/ + +# PyCharm +# JetBrains specific template is maintained in a separate JetBrains.gitignore that can +# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore +# and can be added to the global gitignore or merged into this file. For a more nuclear +# option (not recommended) you can uncomment the following to ignore the entire idea folder. +#.idea/ + +# Ruff stuff: +.ruff_cache/ + +# PyPI configuration file +.pypirc \ No newline at end of file diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..10d0e4b --- /dev/null +++ b/Dockerfile @@ -0,0 +1,19 @@ +FROM python:3.12-slim + +WORKDIR /app + +# Установка необходимых пакетов +RUN apt-get update && apt-get install -y postgresql-client && rm -rf /var/lib/apt/lists/* + +# Установка зависимостей +COPY requirements.txt . +RUN pip install --no-cache-dir -r requirements.txt + +# Копирование файлов проекта +COPY . . + +# Делаем скрипт исполняемым +RUN chmod +x wait-for-postgres.sh + +# Запуск бота с ожиданием PostgreSQL +CMD ["./wait-for-postgres.sh", "postgres", "python", "bot.py"] \ No newline at end of file diff --git a/README.md b/README.md new file mode 100644 index 0000000..07aa649 --- /dev/null +++ b/README.md @@ -0,0 +1,69 @@ +# Telegram бот с Gemma + +Телеграм бот, использующий модель Gemma через Ollama для обработки сообщений и изображений. Поддерживает сохранение контекста в PostgreSQL. + +## Функциональность + +- Обработка текстовых сообщений +- Обработка фотографий +- Поддержка команды `/clear` для очистки контекста диалога +- Поддержка команды `/start` для начала работы с ботом +- Игнорирование стикеров, GIF, видео и документов +- Сохранение контекста диалогов в PostgreSQL + +## Запуск с помощью Docker + +1. Убедитесь, что Docker и Docker Compose установлены +2. Настройте файл `.env`: + ``` + BOT_TOKEN=your_telegram_bot_token_here + OLLAMA_HOST=http://192.168.50.168:11434 + OLLAMA_MODEL=gemma3:12b-it-qat + POSTGRES_DB=gemma_bot + POSTGRES_USER=postgres + POSTGRES_PASSWORD=postgres + ``` +3. Запустите контейнеры: + ``` + docker compose up -d + ``` + +## Установка и запуск без Docker + +1. Установите зависимости: + ``` + pip install aiogram ollama python-dotenv asyncpg + ``` + +2. Настройте файл `.env`: + ``` + BOT_TOKEN=your_telegram_bot_token_here + OLLAMA_HOST=http://192.168.50.168:11434 + OLLAMA_MODEL=gemma3:12b-it-qat + POSTGRES_HOST=localhost + POSTGRES_PORT=5432 + POSTGRES_DB=gemma_bot + POSTGRES_USER=postgres + POSTGRES_PASSWORD=postgres + ``` + +3. Убедитесь, что Ollama запущена и модель Gemma доступна. +4. Убедитесь, что PostgreSQL запущена и создана база данных. + +5. Запустите бота: + ``` + python bot.py + ``` + +## Как использовать + +1. Отправьте боту текстовое сообщение, и он ответит на него, используя модель Gemma. +2. Отправьте боту фотографию (с подписью или без), и он проанализирует ее. +3. Используйте команду `/clear` для очистки истории диалога. +4. Используйте команду `/start` для начала нового диалога. + +## Примечания + +- Бот поддерживает только одну фотографию за раз. +- Стикеры, GIF, видео и документы игнорируются. +- Контекст диалогов сохраняется в PostgreSQL, поэтому при перезапуске бота история сообщений сохраняется. \ No newline at end of file diff --git a/bot.py b/bot.py new file mode 100644 index 0000000..c4a3848 --- /dev/null +++ b/bot.py @@ -0,0 +1,342 @@ +import os +import asyncio +import logging +import json +from io import BytesIO +from typing import Dict, List, Any, Optional + +from aiogram import Bot, Dispatcher, types, F +from aiogram.filters import Command +from aiogram.types import Message +from dotenv import load_dotenv +from ollama import AsyncClient +import asyncpg + +# Настройка логирования +logging.basicConfig(level=logging.INFO) + +# Загрузка переменных окружения +load_dotenv() + +# Инициализация бота и диспетчера +bot = Bot(token=os.getenv("BOT_TOKEN")) +dp = Dispatcher() + +# Инициализация клиента Ollama +ollama_client = AsyncClient(host=os.getenv("OLLAMA_HOST")) +model_name = os.getenv("OLLAMA_MODEL") + +# Соединение с PostgreSQL +postgres_pool: Optional[asyncpg.Pool] = None + +# Словарь для хранения контекста диалога для каждого пользователя в памяти (кэш) +user_messages: Dict[int, List[Dict[str, Any]]] = {} +# Словарь для хранения задач обновления статуса печати +typing_tasks: Dict[int, asyncio.Task] = {} + + +async def init_db(): + """Инициализация базы данных""" + global postgres_pool + + # Создаем пул соединений с базой данных + postgres_pool = await asyncpg.create_pool( + host=os.getenv("POSTGRES_HOST"), + port=os.getenv("POSTGRES_PORT"), + database=os.getenv("POSTGRES_DB"), + user=os.getenv("POSTGRES_USER"), + password=os.getenv("POSTGRES_PASSWORD") + ) + + # Создаем таблицу, если она не существует + async with postgres_pool.acquire() as conn: + await conn.execute(''' + CREATE TABLE IF NOT EXISTS messages ( + id SERIAL PRIMARY KEY, + user_id BIGINT NOT NULL, + message_data JSONB NOT NULL, + created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW() + ) + ''') + + # Создаем индекс для быстрого поиска по user_id + await conn.execute(''' + CREATE INDEX IF NOT EXISTS idx_messages_user_id ON messages (user_id) + ''') + + +async def save_message_to_db(user_id: int, message_data: Dict[str, Any]): + """Сохранение сообщения в базу данных""" + if not postgres_pool: + logging.error("Пул соединений PostgreSQL не инициализирован") + return + + # Преобразуем изображения в строку для хранения + if "images" in message_data: + # Обходим проблему хранения бинарных данных - просто удаляем изображения + # Для реального приложения можно реализовать хранение изображений в отдельной таблице + message_data_copy = message_data.copy() + message_data_copy["has_image"] = True + message_data_copy.pop("images", None) + else: + message_data_copy = message_data + + try: + async with postgres_pool.acquire() as conn: + await conn.execute( + ''' + INSERT INTO messages (user_id, message_data) + VALUES ($1, $2) + ''', + user_id, json.dumps(message_data_copy) + ) + except Exception as e: + logging.error(f"Ошибка сохранения сообщения в БД: {e}") + + +async def load_messages_from_db(user_id: int) -> List[Dict[str, Any]]: + """Загрузка сообщений из базы данных""" + if not postgres_pool: + logging.error("Пул соединений PostgreSQL не инициализирован") + return [] + + try: + async with postgres_pool.acquire() as conn: + rows = await conn.fetch( + ''' + SELECT message_data + FROM messages + WHERE user_id = $1 + ORDER BY created_at ASC + ''', + user_id + ) + + # Преобразуем строки JSON обратно в словари + return [json.loads(row['message_data']) for row in rows] + except Exception as e: + logging.error(f"Ошибка загрузки сообщений из БД: {e}") + return [] + + +async def clear_messages_from_db(user_id: int): + """Очистка сообщений пользователя из базы данных""" + if not postgres_pool: + logging.error("Пул соединений PostgreSQL не инициализирован") + return + + try: + async with postgres_pool.acquire() as conn: + await conn.execute( + ''' + DELETE FROM messages + WHERE user_id = $1 + ''', + user_id + ) + except Exception as e: + logging.error(f"Ошибка очистки сообщений из БД: {e}") + + +async def keep_typing(chat_id: int): + """Функция для поддержания статуса 'печатает...' до окончания генерации ответа""" + try: + while True: + await bot.send_chat_action(chat_id=chat_id, action="typing") + await asyncio.sleep(4) # Обновляем каждые 4 секунды (статус активен 5 секунд) + except asyncio.CancelledError: + # Задача отменена, завершаем работу + pass + except Exception as e: + logging.error(f"Ошибка в keep_typing: {e}") + + +@dp.message(Command("start")) +async def cmd_start(message: Message): + """Обработчик команды /start""" + user_id = message.from_user.id + + # Загружаем контекст из БД при первом запуске + messages = await load_messages_from_db(user_id) + if not messages: + # Если сообщений нет, создаем новый список + user_messages[user_id] = [] + else: + # Если есть, сохраняем их в кэше + user_messages[user_id] = messages + + await message.answer("Привет! Я бот на базе Gemma. Отправьте мне сообщение или фото, и я отвечу вам.") + + +@dp.message(Command("clear")) +async def cmd_clear(message: Message): + """Очистка контекста диалога""" + user_id = message.from_user.id + + # Очищаем контекст в памяти + user_messages[user_id] = [] + + # Очищаем контекст в БД + await clear_messages_from_db(user_id) + + await message.answer("Контекст диалога очищен.") + + +@dp.message(F.photo) +async def handle_photo(message: Message): + """Обработчик фотографий""" + user_id = message.from_user.id + + # Инициализация контекста пользователя, если он еще не существует + if user_id not in user_messages: + # Пытаемся загрузить контекст из БД + messages = await load_messages_from_db(user_id) + user_messages[user_id] = messages or [] + + # Получаем фото наилучшего качества + photo = message.photo[-1] + file_info = await bot.get_file(photo.file_id) + file_content = await bot.download_file(file_info.file_path) + + # Конвертируем фото в Base64 + image_data = file_content.read() + image_io = BytesIO(image_data) + + # Получаем подпись к фото (если есть) + caption = message.caption or "Проанализируй эту фотографию" + + # Создаем сообщение с изображением + user_message = { + "role": "user", + "content": caption, + "images": [image_io.getvalue()] + } + + # Добавляем сообщение в контекст + user_messages[user_id].append(user_message) + + # Сохраняем сообщение в БД + await save_message_to_db(user_id, user_message) + + # Запускаем задачу обновления статуса печати + typing_tasks[user_id] = asyncio.create_task(keep_typing(user_id)) + + try: + # Получаем ответ от Gemma + answer = "" + async for part in await ollama_client.chat( + model=model_name, + messages=user_messages[user_id], + stream=True + ): + answer += part['message']['content'] + + # Останавливаем задачу обновления статуса печати + if user_id in typing_tasks: + typing_tasks[user_id].cancel() + del typing_tasks[user_id] + + # Создаем ответ ассистента + assistant_message = {"role": "assistant", "content": answer} + + # Добавляем ответ в историю + user_messages[user_id].append(assistant_message) + + # Сохраняем ответ в БД + await save_message_to_db(user_id, assistant_message) + + # Отправляем ответ пользователю + await message.answer(answer) + + except Exception as e: + # Останавливаем задачу обновления статуса печати в случае ошибки + if user_id in typing_tasks: + typing_tasks[user_id].cancel() + del typing_tasks[user_id] + + logging.error(f"Ошибка при обработке фото: {e}") + await message.answer(f"Произошла ошибка при обработке фото: {e}") + + +@dp.message(F.text) +async def handle_text(message: Message): + """Обработчик текстовых сообщений""" + user_id = message.from_user.id + + # Инициализация контекста пользователя, если он еще не существует + if user_id not in user_messages: + # Пытаемся загрузить контекст из БД + messages = await load_messages_from_db(user_id) + user_messages[user_id] = messages or [] + + # Создаем сообщение пользователя + user_message = {"role": "user", "content": message.text} + + # Добавляем сообщение пользователя в контекст + user_messages[user_id].append(user_message) + + # Сохраняем сообщение в БД + await save_message_to_db(user_id, user_message) + + # Запускаем задачу обновления статуса печати + typing_tasks[user_id] = asyncio.create_task(keep_typing(user_id)) + + try: + # Получаем ответ от Gemma + answer = "" + async for part in await ollama_client.chat( + model=model_name, + messages=user_messages[user_id], + stream=True + ): + answer += part['message']['content'] + + # Останавливаем задачу обновления статуса печати + if user_id in typing_tasks: + typing_tasks[user_id].cancel() + del typing_tasks[user_id] + + # Создаем ответ ассистента + assistant_message = {"role": "assistant", "content": answer} + + # Добавляем ответ в историю + user_messages[user_id].append(assistant_message) + + # Сохраняем ответ в БД + await save_message_to_db(user_id, assistant_message) + + # Отправляем ответ пользователю + await message.answer(answer) + + except Exception as e: + # Останавливаем задачу обновления статуса печати в случае ошибки + if user_id in typing_tasks: + typing_tasks[user_id].cancel() + del typing_tasks[user_id] + + logging.error(f"Ошибка при обработке сообщения: {e}") + await message.answer(f"Произошла ошибка при обработке сообщения: {e}") + + +@dp.message() +async def handle_other(message: Message): + """Обработчик всех остальных типов сообщений""" + await message.answer("Я могу обрабатывать только текст и фотографии. Пожалуйста, отправьте текст или одну фотографию.") + + +async def main(): + """Запуск бота""" + # Инициализируем подключение к базе данных + await init_db() + + try: + # Запускаем бота + await dp.start_polling(bot) + finally: + # Закрываем пул соединений при завершении + if postgres_pool: + await postgres_pool.close() + + +if __name__ == "__main__": + asyncio.run(main()) \ No newline at end of file diff --git a/docker-compose.yaml b/docker-compose.yaml new file mode 100644 index 0000000..001a6c6 --- /dev/null +++ b/docker-compose.yaml @@ -0,0 +1,54 @@ +version: '3.8' + +services: + # Сервис бота + bot: + build: + context: . + dockerfile: Dockerfile + restart: always + depends_on: + postgres: + condition: service_healthy + environment: + - BOT_TOKEN=${BOT_TOKEN} + - OLLAMA_HOST=${OLLAMA_HOST} + - OLLAMA_MODEL=${OLLAMA_MODEL} + - POSTGRES_HOST=postgres + - POSTGRES_PORT=5432 + - POSTGRES_DB=${POSTGRES_DB} + - POSTGRES_USER=${POSTGRES_USER} + - POSTGRES_PASSWORD=${POSTGRES_PASSWORD} + volumes: + - ./:/app + networks: + - bot-network + + # Сервис базы данных + postgres: + image: postgres:15-alpine + restart: always + environment: + - POSTGRES_DB=${POSTGRES_DB} + - POSTGRES_USER=${POSTGRES_USER} + - POSTGRES_PASSWORD=${POSTGRES_PASSWORD} + volumes: + - postgres_data:/var/lib/postgresql/data + - ./docker/postgres/init.sql:/docker-entrypoint-initdb.d/init.sql + ports: + - "5432:5432" + healthcheck: + test: ["CMD-SHELL", "pg_isready -U ${POSTGRES_USER} -d ${POSTGRES_DB}"] + interval: 10s + timeout: 5s + retries: 5 + start_period: 10s + networks: + - bot-network + +networks: + bot-network: + driver: bridge + +volumes: + postgres_data: \ No newline at end of file diff --git a/docker/postgres/init.sql b/docker/postgres/init.sql new file mode 100644 index 0000000..9e1fdb3 --- /dev/null +++ b/docker/postgres/init.sql @@ -0,0 +1,16 @@ +-- Создаем базу данных для бота, если она еще не существует +CREATE DATABASE gemma_bot WITH OWNER postgres ENCODING 'UTF8' LC_COLLATE = 'en_US.utf8' LC_CTYPE = 'en_US.utf8'; + +-- Подключаемся к базе данных +\c gemma_bot; + +-- Создаем таблицу для хранения сообщений +CREATE TABLE IF NOT EXISTS messages ( + id SERIAL PRIMARY KEY, + user_id BIGINT NOT NULL, + message_data JSONB NOT NULL, + created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW() +); + +-- Создаем индекс для быстрого поиска по user_id +CREATE INDEX IF NOT EXISTS idx_messages_user_id ON messages (user_id); \ No newline at end of file diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..9934afd --- /dev/null +++ b/requirements.txt @@ -0,0 +1,4 @@ +aiogram==3.20.0 +python-dotenv==1.1.0 +ollama==0.4.8 +asyncpg==0.30.0 \ No newline at end of file diff --git a/test-prompts.txt b/test-prompts.txt new file mode 100644 index 0000000..ca25c85 --- /dev/null +++ b/test-prompts.txt @@ -0,0 +1 @@ +Ты - JSON-парсер. Твоя задача - принимать сообщение от пользователя и отвечать {\"root\": \"<корень слова\"} \ No newline at end of file diff --git a/wait-for-postgres.sh b/wait-for-postgres.sh new file mode 100755 index 0000000..ed33058 --- /dev/null +++ b/wait-for-postgres.sh @@ -0,0 +1,16 @@ +#!/bin/sh +# wait-for-postgres.sh + +set -e + +host="$1" +shift +cmd="$@" + +until PGPASSWORD=$POSTGRES_PASSWORD psql -h "$host" -U "$POSTGRES_USER" -d "$POSTGRES_DB" -c '\q'; do + >&2 echo "Postgres is unavailable - sleeping" + sleep 1 +done + +>&2 echo "Postgres is up - executing command" +exec $cmd \ No newline at end of file