recognition + graphs + schemes
This commit is contained in:
parent
f7cf147e86
commit
b6274f7106
6 changed files with 1680 additions and 2 deletions
419
src/bot.py
Normal file
419
src/bot.py
Normal file
|
|
@ -0,0 +1,419 @@
|
|||
"""
|
||||
Telegram-бот с четырьмя функциями:
|
||||
1) получает изображение → возвращает его описание (vision-модель Qwen);
|
||||
2) /plot <описание> → matplotlib-график;
|
||||
3) /arch <описание системы> → архитектурная схема через Graphviz;
|
||||
4) /repro (в подписи или ответе к фото) → воспроизводит график/схему с картинки.
|
||||
"""
|
||||
|
||||
import os
|
||||
import io
|
||||
import re
|
||||
import json
|
||||
import base64
|
||||
import logging
|
||||
import asyncio
|
||||
|
||||
import matplotlib
|
||||
matplotlib.use("Agg") # без GUI-бэкенда, мы только сохраняем в буфер
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
import graphviz
|
||||
|
||||
from dotenv import load_dotenv
|
||||
from openai import AsyncOpenAI
|
||||
from telegram import Update
|
||||
from telegram.ext import (
|
||||
Application,
|
||||
CommandHandler,
|
||||
MessageHandler,
|
||||
ContextTypes,
|
||||
filters,
|
||||
)
|
||||
|
||||
# ---------- конфиг ----------
|
||||
|
||||
load_dotenv()
|
||||
|
||||
TELEGRAM_TOKEN = os.getenv("TELEGRAM_BOT_TOKEN")
|
||||
LLM_API_KEY = os.getenv("LLM_API_KEY")
|
||||
LLM_BASE_URL = os.getenv("LLM_BASE_URL", "https://llm.lambda.coredump.ru/v1")
|
||||
LLM_MODEL = os.getenv("LLM_MODEL", "qwen3.5-122b")
|
||||
# Отключить thinking у reasoning-моделей (актуально для Qwen3/DeepSeek-R1 и т.п.),
|
||||
# чтобы все токены шли в финальный ответ, а не в <think>-блок.
|
||||
DISABLE_THINKING = os.getenv("LLM_DISABLE_THINKING", "false").lower() in ("1", "true", "yes")
|
||||
|
||||
if not TELEGRAM_TOKEN or not LLM_API_KEY:
|
||||
raise RuntimeError("В .env должны быть заданы TELEGRAM_BOT_TOKEN и LLM_API_KEY")
|
||||
|
||||
logging.basicConfig(
|
||||
format="%(asctime)s | %(name)s | %(levelname)s | %(message)s",
|
||||
level=logging.INFO,
|
||||
)
|
||||
logger = logging.getLogger("qwen-bot")
|
||||
|
||||
client = AsyncOpenAI(api_key=LLM_API_KEY, base_url=LLM_BASE_URL)
|
||||
|
||||
|
||||
def _chat_kwargs(max_tokens: int) -> dict:
|
||||
"""Собирает общие kwargs для chat.completions.create, включая отключение
|
||||
reasoning если попросили. Параметры передаются через extra_body и игнорируются
|
||||
провайдерами, которые их не поддерживают."""
|
||||
kwargs: dict = {"max_tokens": max_tokens}
|
||||
if DISABLE_THINKING:
|
||||
# OpenRouter: reasoning.exclude. DashScope/vLLM с Qwen3: enable_thinking=False
|
||||
kwargs["extra_body"] = {
|
||||
"reasoning": {"exclude": True},
|
||||
"enable_thinking": False,
|
||||
"chat_template_kwargs": {"enable_thinking": False},
|
||||
}
|
||||
return kwargs
|
||||
|
||||
|
||||
# ---------- 1. Описание изображения ----------
|
||||
|
||||
async def describe_image(image_bytes: bytes) -> str:
|
||||
"""Отправляет картинку в vision-модель и возвращает текстовое описание."""
|
||||
b64 = base64.b64encode(image_bytes).decode("utf-8")
|
||||
response = await client.chat.completions.create(
|
||||
model=LLM_MODEL,
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": (
|
||||
"Опиши подробно, что изображено на картинке: "
|
||||
"объекты, их расположение, цвета, контекст и общее настроение."
|
||||
),
|
||||
},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": f"data:image/jpeg;base64,{b64}"},
|
||||
},
|
||||
],
|
||||
}
|
||||
],
|
||||
**_chat_kwargs(max_tokens=1024),
|
||||
)
|
||||
return _extract_content(response).strip()
|
||||
|
||||
|
||||
# ---------- 2. Генерация графиков ----------
|
||||
|
||||
PLOT_SYSTEM_PROMPT = """Ты генерируешь Python-код для построения графиков matplotlib.
|
||||
|
||||
Строгие правила:
|
||||
- в окружении доступны только: matplotlib.pyplot как `plt` и numpy как `np`;
|
||||
- фигура уже создана заранее, просто рисуй через plt;
|
||||
- НЕ вызывай plt.show() и ничего не сохраняй в файлы;
|
||||
- НЕ используй import — модули уже импортированы;
|
||||
- возвращай ТОЛЬКО исполняемый python-код, без markdown-обрамления, без комментариев и пояснений."""
|
||||
|
||||
|
||||
def _strip_code_fences(text: str) -> str:
|
||||
"""Убирает ```python ... ``` если модель всё же обернула код."""
|
||||
match = re.search(r"```(?:[a-zA-Z0-9_+-]*)?\s*(.*?)```", text, re.DOTALL)
|
||||
return (match.group(1) if match else text).strip()
|
||||
|
||||
|
||||
def _extract_content(response) -> str:
|
||||
"""Достаёт текстовый ответ из choices[0], кидая понятную ошибку на пустой ответ."""
|
||||
choice = response.choices[0]
|
||||
content = choice.message.content
|
||||
if not content:
|
||||
finish = getattr(choice, "finish_reason", "unknown")
|
||||
raise RuntimeError(
|
||||
f"Модель вернула пустой ответ (finish_reason={finish}). "
|
||||
"Возможные причины: thinking-модель упёрлась в max_tokens и не дошла до финального ответа; "
|
||||
"провайдер отфильтровал запрос; кончилась квота."
|
||||
)
|
||||
return content
|
||||
|
||||
|
||||
async def generate_plot_code(request: str) -> str:
|
||||
response = await client.chat.completions.create(
|
||||
model=LLM_MODEL,
|
||||
messages=[
|
||||
{"role": "system", "content": PLOT_SYSTEM_PROMPT},
|
||||
{"role": "user", "content": request},
|
||||
],
|
||||
**_chat_kwargs(max_tokens=8192),
|
||||
)
|
||||
return _strip_code_fences(_extract_content(response))
|
||||
|
||||
|
||||
def render_plot(code: str) -> bytes:
|
||||
"""
|
||||
Выполняет сгенерированный код и возвращает PNG-байты.
|
||||
|
||||
ВНИМАНИЕ: exec используется в ограниченном неймспейсе, но это не настоящая
|
||||
песочница. Для публичного бота лучше изолировать исполнение (отдельный
|
||||
процесс с resource-лимитами, Docker, nsjail и т. п.).
|
||||
"""
|
||||
plt.close("all")
|
||||
plt.figure(figsize=(8, 5))
|
||||
|
||||
safe_builtins = {
|
||||
"range": range, "len": len, "abs": abs, "min": min, "max": max,
|
||||
"sum": sum, "enumerate": enumerate, "zip": zip, "map": map, "filter": filter,
|
||||
"sorted": sorted, "reversed": reversed, "any": any, "all": all,
|
||||
"list": list, "tuple": tuple, "dict": dict, "set": set,
|
||||
"float": float, "int": int, "str": str, "bool": bool, "round": round,
|
||||
"isinstance": isinstance, "type": type, "print": print,
|
||||
}
|
||||
safe_globals = {"plt": plt, "np": np, "__builtins__": safe_builtins}
|
||||
|
||||
exec(code, safe_globals)
|
||||
|
||||
buf = io.BytesIO()
|
||||
plt.savefig(buf, format="png", dpi=120, bbox_inches="tight")
|
||||
plt.close("all")
|
||||
buf.seek(0)
|
||||
return buf.read()
|
||||
|
||||
|
||||
# ---------- 3. Генерация архитектурных схем ----------
|
||||
|
||||
ARCH_SYSTEM_PROMPT = """Ты генерируешь архитектурные схемы на языке Graphviz DOT.
|
||||
|
||||
Правила:
|
||||
- пиши корректный DOT-код для ориентированного графа (digraph);
|
||||
- давай узлам осмысленные метки; подписывай стрелки (протокол, формат данных, действие) где уместно;
|
||||
- группируй связанные компоненты в subgraph cluster_* с подписями;
|
||||
- используй разные shape и цвета заливки для разных типов компонентов:
|
||||
* клиент/пользователь — shape=actor или oval;
|
||||
* сервисы/бэкенды — shape=box, style=filled;
|
||||
* базы данных — shape=cylinder;
|
||||
* очереди/брокеры — shape=box3d;
|
||||
* внешние API и облачные сервисы — shape=cloud;
|
||||
* хранилища файлов — shape=folder;
|
||||
- задавай rankdir=LR для горизонтальной компоновки, если компонентов больше 4;
|
||||
- верни ТОЛЬКО валидный DOT-код без markdown-обрамления, без комментариев и пояснений."""
|
||||
|
||||
|
||||
async def generate_arch_dot(description: str) -> str:
|
||||
response = await client.chat.completions.create(
|
||||
model=LLM_MODEL,
|
||||
messages=[
|
||||
{"role": "system", "content": ARCH_SYSTEM_PROMPT},
|
||||
{"role": "user", "content": description},
|
||||
],
|
||||
**_chat_kwargs(max_tokens=8192), # DOT + возможный reasoning могут съесть много
|
||||
)
|
||||
return _strip_code_fences(_extract_content(response))
|
||||
|
||||
|
||||
def render_architecture(dot_source: str) -> bytes:
|
||||
"""Рендерит DOT-описание в PNG через локальный бинарь graphviz."""
|
||||
src = graphviz.Source(dot_source, engine="dot", format="png")
|
||||
return src.pipe()
|
||||
|
||||
|
||||
# ---------- 4. Воспроизведение графика/схемы с фото ----------
|
||||
|
||||
REPRO_SYSTEM_PROMPT = """Ты анализируешь изображение и воспроизводишь его кодом.
|
||||
|
||||
Определи тип:
|
||||
- "plot" — график любого вида (line, bar, scatter, histogram, pie, heatmap и т.п.);
|
||||
- "architecture" — архитектурная схема, блок-схема, граф компонентов, диаграмма потоков;
|
||||
- "other" — всё остальное.
|
||||
|
||||
Если "plot" — сгенерируй Python-код на matplotlib, максимально близко воспроизводящий оригинал:
|
||||
- тот же тип графика, подписи осей, заголовок, легенда (если есть), примерная форма данных;
|
||||
- доступны только `plt` и `np`; БЕЗ import, БЕЗ plt.show() и savefig — фигура создана заранее.
|
||||
|
||||
Если "architecture" — сгенерируй валидный DOT-код для Graphviz digraph:
|
||||
- те же узлы с теми же метками, те же связи и направления, та же группировка в subgraph cluster_*;
|
||||
- подходящие shape (cylinder для БД, cloud для внешних API, box3d для очередей и т.п.).
|
||||
|
||||
Верни СТРОГО JSON-объект без markdown-обрамления, без пояснений до или после:
|
||||
{"type": "plot|architecture|other", "notes": "<краткое описание>", "code": "<исполняемый код>"}
|
||||
|
||||
Поле `code` должно быть одной строкой JSON с экранированными переводами строк (\\n)."""
|
||||
|
||||
|
||||
def _extract_json(text: str) -> dict:
|
||||
"""Парсит JSON из ответа модели, терпимо снимая markdown-обёртку и мусор по краям."""
|
||||
cleaned = _strip_code_fences(text)
|
||||
try:
|
||||
return json.loads(cleaned)
|
||||
except json.JSONDecodeError:
|
||||
# резервный план: вытащить первый JSON-объект из текста
|
||||
match = re.search(r"\{.*\}", cleaned, re.DOTALL)
|
||||
if not match:
|
||||
raise RuntimeError(f"Модель вернула не-JSON:\n{text[:500]}")
|
||||
return json.loads(match.group(0))
|
||||
|
||||
|
||||
async def analyze_and_reproduce(image_bytes: bytes) -> tuple[str, str, str]:
|
||||
"""Возвращает (type, code, notes)."""
|
||||
b64 = base64.b64encode(image_bytes).decode("utf-8")
|
||||
response = await client.chat.completions.create(
|
||||
model=LLM_MODEL,
|
||||
messages=[
|
||||
{"role": "system", "content": REPRO_SYSTEM_PROMPT},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "Проанализируй изображение и воспроизведи его кодом."},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": f"data:image/jpeg;base64,{b64}"},
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
**_chat_kwargs(max_tokens=8192),
|
||||
)
|
||||
data = _extract_json(_extract_content(response))
|
||||
return (
|
||||
data.get("type", "other"),
|
||||
data.get("code", ""),
|
||||
data.get("notes", ""),
|
||||
)
|
||||
|
||||
|
||||
# ---------- Telegram-хэндлеры ----------
|
||||
|
||||
async def cmd_start(update: Update, _: ContextTypes.DEFAULT_TYPE) -> None:
|
||||
await update.message.reply_text(
|
||||
"Привет! Я умею:\n"
|
||||
"📷 описывать фото — просто пришли картинку\n"
|
||||
"📊 строить графики — команда /plot <описание>\n"
|
||||
"🏛 рисовать архитектурные схемы — команда /arch <описание системы>\n"
|
||||
"♻️ воспроизводить график/схему с фото — пришли фото с подписью /repro "
|
||||
"или ответь /repro на сообщение с фото\n\n"
|
||||
"Примеры:\n"
|
||||
" /plot sin(x) и cos(x) на отрезке [0, 2π]\n"
|
||||
" /arch React-фронт, FastAPI-бэкенд, Postgres и Redis за nginx"
|
||||
)
|
||||
|
||||
|
||||
async def on_photo(update: Update, _: ContextTypes.DEFAULT_TYPE) -> None:
|
||||
await update.message.chat.send_action("typing")
|
||||
photo = update.message.photo[-1] # берём самый большой размер
|
||||
tg_file = await photo.get_file()
|
||||
image_bytes = bytes(await tg_file.download_as_bytearray())
|
||||
try:
|
||||
text = await describe_image(image_bytes)
|
||||
await update.message.reply_text(text)
|
||||
except Exception as e:
|
||||
logger.exception("Ошибка при описании изображения")
|
||||
await update.message.reply_text(f"Не смог разобрать изображение: {e}")
|
||||
|
||||
|
||||
async def cmd_plot(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
|
||||
request = " ".join(context.args).strip()
|
||||
if not request:
|
||||
await update.message.reply_text(
|
||||
"Опиши, какой график построить. Пример:\n/plot затухающая синусоида на [0, 10]"
|
||||
)
|
||||
return
|
||||
|
||||
await update.message.chat.send_action("upload_photo")
|
||||
try:
|
||||
code = await generate_plot_code(request)
|
||||
logger.info("Сгенерирован код:\n%s", code)
|
||||
png = await asyncio.to_thread(render_plot, code)
|
||||
await update.message.reply_photo(photo=png, caption=request[:1024])
|
||||
except Exception as e:
|
||||
logger.exception("Ошибка при построении графика")
|
||||
await update.message.reply_text(f"Не получилось построить график: {e}")
|
||||
|
||||
|
||||
async def cmd_arch(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
|
||||
description = " ".join(context.args).strip()
|
||||
if not description:
|
||||
await update.message.reply_text(
|
||||
"Опиши архитектуру системы. Пример:\n"
|
||||
"/arch React-фронт, FastAPI-бэкенд с воркерами Celery, Postgres, Redis и S3"
|
||||
)
|
||||
return
|
||||
|
||||
await update.message.chat.send_action("upload_photo")
|
||||
try:
|
||||
dot = await generate_arch_dot(description)
|
||||
logger.info("Сгенерирован DOT:\n%s", dot)
|
||||
png = await asyncio.to_thread(render_architecture, dot)
|
||||
await update.message.reply_photo(photo=png, caption=description[:1024])
|
||||
except graphviz.ExecutableNotFound:
|
||||
logger.exception("Нет бинаря graphviz")
|
||||
await update.message.reply_text(
|
||||
"На сервере не установлен Graphviz. Поставь его: `apt install graphviz` или `brew install graphviz`."
|
||||
)
|
||||
except Exception as e:
|
||||
logger.exception("Ошибка при построении архитектуры")
|
||||
await update.message.reply_text(f"Не получилось построить схему: {e}")
|
||||
|
||||
|
||||
async def cmd_repro(update: Update, _: ContextTypes.DEFAULT_TYPE) -> None:
|
||||
"""
|
||||
Воспроизводит график/схему с приложенного фото.
|
||||
Работает двумя способами:
|
||||
1) фото с подписью "/repro"
|
||||
2) команда "/repro" в ответ (reply) на сообщение с фото
|
||||
"""
|
||||
msg = update.message
|
||||
if msg.photo:
|
||||
source_msg = msg
|
||||
elif msg.reply_to_message and msg.reply_to_message.photo:
|
||||
source_msg = msg.reply_to_message
|
||||
else:
|
||||
await msg.reply_text(
|
||||
"Пришли фото с подписью /repro или ответь /repro на сообщение с фото."
|
||||
)
|
||||
return
|
||||
|
||||
await msg.chat.send_action("upload_photo")
|
||||
tg_file = await source_msg.photo[-1].get_file()
|
||||
image_bytes = bytes(await tg_file.download_as_bytearray())
|
||||
|
||||
try:
|
||||
kind, code, notes = await analyze_and_reproduce(image_bytes)
|
||||
logger.info("repro type=%s notes=%s\ncode:\n%s", kind, notes, code)
|
||||
|
||||
if kind == "plot":
|
||||
png = await asyncio.to_thread(render_plot, code)
|
||||
caption = f"📊 {notes}" if notes else "📊 Воспроизведённый график"
|
||||
elif kind == "architecture":
|
||||
png = await asyncio.to_thread(render_architecture, code)
|
||||
caption = f"🏛 {notes}" if notes else "🏛 Воспроизведённая схема"
|
||||
else:
|
||||
await msg.reply_text(
|
||||
"Это не похоже ни на график, ни на архитектурную схему."
|
||||
+ (f"\n\n{notes}" if notes else "")
|
||||
)
|
||||
return
|
||||
|
||||
await msg.reply_photo(photo=png, caption=caption[:1024])
|
||||
except graphviz.ExecutableNotFound:
|
||||
logger.exception("Нет бинаря graphviz")
|
||||
await msg.reply_text(
|
||||
"На сервере не установлен Graphviz. Поставь его: `apt install graphviz` / "
|
||||
"`brew install graphviz` / `winget install graphviz`."
|
||||
)
|
||||
except Exception as e:
|
||||
logger.exception("Ошибка при воспроизведении")
|
||||
await msg.reply_text(f"Не получилось воспроизвести: {e}")
|
||||
|
||||
|
||||
# ---------- запуск ----------
|
||||
|
||||
def main() -> None:
|
||||
app = Application.builder().token(TELEGRAM_TOKEN).build()
|
||||
app.add_handler(CommandHandler(["start", "help"], cmd_start))
|
||||
app.add_handler(CommandHandler("plot", cmd_plot))
|
||||
app.add_handler(CommandHandler("arch", cmd_arch))
|
||||
app.add_handler(CommandHandler("repro", cmd_repro))
|
||||
# фото без подписи-команды → описание (иначе команда уже обработана выше)
|
||||
app.add_handler(
|
||||
MessageHandler(filters.PHOTO & ~filters.CaptionRegex(r"^/"), on_photo)
|
||||
)
|
||||
logger.info("Бот запущен")
|
||||
app.run_polling(allowed_updates=Update.ALL_TYPES)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Loading…
Add table
Add a link
Reference in a new issue