add poll answer support

4 files changed, 100 insertions, 4 deletions

diff --git a/helpers.py b/helpers.py
index e2450f0..d42b315 100644
--- a/helpers.py
+++ b/helpers.py
@@ -1,6 +1,6 @@
 import time
 
-from aiogram.methods import AnswerGuestQuery
+from aiogram.methods import AnswerGuestQuery, SendPoll
 from aiogram.types import InlineQueryResultArticle, InputTextMessageContent, Message
 from aiogram.enums import ChatType
 
@@ -88,3 +88,27 @@ async def g_answer(data, message, text, title="da", id="1", parse_mode="Markdown
             ),
         )
     )
+
+
+async def answer_poll(
+    data,
+    message,
+    question,
+    options,
+    is_anonymous=True,
+    poll_type="regular",
+    allows_multiple_answers=False,
+    correct_option_id=None,
+    user_id=None
+):
+    return await data["bot"](
+        SendPoll(
+            chat_id=user_id,
+            question=question,
+            options=options,
+            is_anonymous=is_anonymous,
+            type=poll_type,
+            allows_multiple_answers=allows_multiple_answers,
+            correct_option_id=correct_option_id,
+        )
+    )
diff --git a/main.py b/main.py
index c9d1cfa..14c1ba2 100644
--- a/main.py
+++ b/main.py
@@ -33,7 +33,7 @@ from colors import *
 from helpers import *
 from commands import *
 from dotenv import load_dotenv
-from random import uniform
+from random import uniform, shuffle
 
 try:
     from config import *
@@ -331,6 +331,68 @@ async def guest_middleware(handler, event, data):
 
     last_command_time[user_id][1] = False
     last_command_time[user_id][2] = []
+
+    if result.startswith("[") and ":" in result:
+        rl = result.split("\n")
+        poll_type = rl[0].split(":")[0].replace("[", "").replace("]", "")
+        poll_name = rl[0].split(":")[1]
+        if len(rl) > 1:
+            options = rl[1].split(" =| ")
+        else:
+            options = ["варианты не указаны"]
+        correct_option = options[0]
+        correct_option_id = None
+        is_anon = True
+        allow_multiple_answers = False
+        
+        if event.guest_message:
+            mid = user_id        
+        elif message.chat.type == "private":
+            mid = user_id
+        elif message.chat.type in ("group", "supergroup"):
+            mid = message.chat.id
+        else:
+            return
+
+        if poll_type == "quiz":
+            shuffle(options)
+            correct_option_id = options.index(correct_option)
+        elif poll_type == "poll":
+            poll_type = "regular"
+            allow_multiple_answers = False
+        elif poll_type == "multipoll":
+            poll_type = "regular"
+            allow_multiple_answers = True
+        else:
+            poll_type = "regular"
+
+        try:
+            result = await answer_poll(
+                data=data,
+                message=message,
+                question=poll_name,
+                options=options,
+                is_anonymous=is_anon,
+                poll_type=poll_type,
+                correct_option_id=correct_option_id,
+                user_id=mid
+            )
+        except Exception as e:
+            if event.guest_message:
+                await g_answer(data, message, "гостевой режим не работает с опросами. либо добавь меня " \
+                    "в группу, либо напиши мне в личку чтобы я выдал тебе опрос. ✅"
+                )
+            else:
+                print(f"{RED} -- err poll {e} {RESET}")
+            return
+
+        if result and event.guest_message:
+            await g_answer(data, message, "отправил опрос в личку, так как гостевой режим не работает с " \
+                "опросами✅✅"
+            )
+
+    
+        return
     
     try:
         if event.guest_message:
diff --git a/prompts.py b/prompts.py
index 27275ab..02b14eb 100644
--- a/prompts.py
+++ b/prompts.py
@@ -36,6 +36,7 @@ SYSTEM_PROMPTS = {
 Не добавляй других эмодзи кроме галочек.
 Пиши с сука огромным количеством сарказма. 
 Не ставь знаки препинания в ответе и отвечай с маленькой буквы.
+не будь настолько тупым и хотя бы пиши правильно структуру опросов
     """,
     
     "тру узбек (февраль 2026)": """
@@ -66,6 +67,15 @@ uzmoji (вставляй эмодзи всегда и в большом коли
     """
 }
 
+ADD_TO_PROMPT = """
+типы опросов (указывается в квадратных скобках): poll (обычный опрос), multipoll (опрос с множественным ответом), quiz (квиз с выбор правильного ответа)
+чтобы отправить опрос (если попросит пользователь) нужно ответить в следующем формате без лишнего текста:
+[тип опроса]:название опроса
+вариант 1 =| вариант 2 =| и так далее
+
+(p.s ДЛЯ quiz: 1-й вариант - верный)
+"""
+
 IMAGE_PROMPT = """
 опиши это фото максимально подробно на ~500 символов.
 перепиши весь текст на изображении.
diff --git a/supergenerator.py b/supergenerator.py
index 89b0902..17c4c03 100644
--- a/supergenerator.py
+++ b/supergenerator.py
@@ -204,7 +204,7 @@ async def generate(
         stream_cf = config["stream"]
         system_prompt_name = config["prompt"]
         if system_prompt_name in SYSTEM_PROMPTS:
-            system_prompt = SYSTEM_PROMPTS[system_prompt_name]
+            system_prompt = SYSTEM_PROMPTS[system_prompt_name] + ADD_TO_PROMPT
         else:
             text = "❌произошла ошибка, возможные проблемы:\n" \
                 f"1. у вас выбран удалённый сист промпт (`{system_prompt_name}`)\n" \
@@ -259,7 +259,7 @@ async def generate(
                         print(f"{RED} -- err photo - {e} : {model}{RESET}")
 
         content += [{"type": "text", "text": prompt}]
-        messages = [{"role": "system", "content": SYSTEM_PROMPTS[DEFAULT_PROMPT]}] + [
+        messages = [{"role": "system", "content": SYSTEM_PROMPTS[DEFAULT_PROMPT]+ADD_TO_PROMPT}] + [
             {"role": "user", "content": content}
         ]