├── bot
├── handlers
│ ├── __init__.py
│ ├── start.py
│ └── chat.py
├── __init__.py
├── __main__.py
├── database.py
├── PaLM.py
└── helpers.py
├── .pylintrc
├── requirements.txt
├── config.py
├── .vscode
└── launch.json
├── README.md
├── transfer.py
└── .gitignore
/bot/handlers/__init__.py:
--------------------------------------------------------------------------------
1 |
--------------------------------------------------------------------------------
/.pylintrc:
--------------------------------------------------------------------------------
https://raw.githubusercontent.com/EverythingSuckz/PaLM-Bot/HEAD/.pylintrc
--------------------------------------------------------------------------------
/requirements.txt:
--------------------------------------------------------------------------------
1 | python-dotenv
2 | pyrogram
3 | tgcrypto
4 | google-generativeai
5 | sqlalchemy
6 | pony
7 | psycopg2-binary
--------------------------------------------------------------------------------
/config.py:
--------------------------------------------------------------------------------
1 | import os
2 | from dotenv import load_dotenv
3 |
4 | load_dotenv()
5 |
6 |
7 | class Config(object):
8 | BOT_TOKEN = os.getenv("BOT_TOKEN")
9 | API_ID = os.getenv("API_ID")
10 | API_HASH = os.getenv("API_HASH")
11 | PALM_API_KEY = os.getenv("PALM_API_KEY")
12 | DB_HOST = os.getenv("DB_HOST")
13 | DB_USER = os.getenv("DB_USER")
14 | DB_PASSWORD = os.getenv("DB_PASSWORD")
15 | DB_NAME = os.getenv("DB_NAME")
16 |
--------------------------------------------------------------------------------
/.vscode/launch.json:
--------------------------------------------------------------------------------
1 | {
2 | // Use IntelliSense to learn about possible attributes.
3 | // Hover to view descriptions of existing attributes.
4 | // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
5 | "version": "0.2.0",
6 | "configurations": [
7 | {
8 | "name": "Python: Module",
9 | "type": "python",
10 | "request": "launch",
11 | "module": "bot",
12 | "justMyCode": true
13 | }
14 | ]
15 | }
--------------------------------------------------------------------------------
/bot/__init__.py:
--------------------------------------------------------------------------------
1 | import logging
2 | import sys
3 |
4 | from pyrogram import Client
5 |
6 | from config import Config
7 | from bot.PaLM import PaLMChat
8 |
9 | logging.basicConfig(
10 | level=logging.INFO,
11 | datefmt="%Y/%m/%d %H:%M:%S",
12 | format="[%(asctime)s][%(name)s][%(levelname)s] ==> %(message)s",
13 | handlers=[
14 | logging.FileHandler("bot.log"),
15 | logging.StreamHandler(sys.stdout)
16 | ]
17 | )
18 |
19 | logging.getLogger("pyrogram").setLevel(logging.ERROR)
20 |
21 | Bot: Client = None
22 | palm = PaLMChat(Config.PALM_API_KEY)
23 |
--------------------------------------------------------------------------------
/bot/__main__.py:
--------------------------------------------------------------------------------
1 | import asyncio
2 | import logging
3 | from pyrogram import Client, idle, types
4 |
5 | from config import Config
6 |
7 | async def main():
8 | import bot
9 | bot.Bot = bot = Client(
10 | "PaLM-Bot",
11 | api_id=Config.API_ID,
12 | api_hash=Config.API_HASH,
13 | bot_token=Config.BOT_TOKEN,
14 | plugins=dict(root="bot/handlers"),
15 | )
16 | await bot.start()
17 | await bot.set_bot_commands(
18 | commands=[
19 | types.BotCommand("start", "Get the start message."),
20 | types.BotCommand("help", "Get this same message."),
21 | types.BotCommand("about", "Get more info about the bot."),
22 | types.BotCommand("clearhistory", "Remove all your message history."),
23 | ], scope=types.BotCommandScopeAllPrivateChats()
24 | )
25 | logging.info("Bot started as @%s.", bot.me.username)
26 | await idle()
27 | await bot.stop()
28 |
29 | if __name__ == "__main__":
30 | asyncio.run(main())
--------------------------------------------------------------------------------
/bot/handlers/start.py:
--------------------------------------------------------------------------------
1 | import asyncio
2 | from pyrogram import Client, filters, types, enums
3 |
4 | from bot.helpers import limiter
5 | from bot.database import add_user
6 |
7 | gfn = lambda x: x.first_name + (f" {x.last_name}" if x.last_name else "")
8 |
9 | @Client.on_message(filters.private, group=-1)
10 | async def log_users(_, message: types.Message):
11 | asyncio.create_task(add_user(message.from_user.id, gfn(message.from_user), message.from_user.username))
12 |
13 | @Client.on_message(filters.command("start") & filters.private)
14 | @limiter(5)
15 | async def start_cmd(_, message: types.Message):
16 | await message.reply(
17 | f"Hi {message.from_user.mention(style=enums.ParseMode.HTML)}. Send me something.",
18 | parse_mode=enums.ParseMode.HTML,
19 | )
20 |
21 | @Client.on_message(filters.command("help") & filters.private)
22 | @limiter(3)
23 | async def help_cmd(_, message: types.Message):
24 | await message.reply(
25 | f"You can start a conversation with me by texting me anything. If you are looking for chat commands, type\n\n- /start - Get the start message.\n- /help - Get this same message.\n- /about - Get more info about the bot.\n- /clearhistory - Remove all your message history.\n\nHave a great day!", parse_mode=enums.ParseMode.HTML
26 | )
27 |
28 |
29 | @Client.on_message(filters.command("about") & filters.private)
30 | @limiter(5)
31 | async def about_cmd(_, message: types.Message):
32 | markup = types.InlineKeyboardMarkup([[
33 | types.InlineKeyboardButton(text="What's PaLM?", url="https://t.me/wrenchies/226"),
34 | types.InlineKeyboardButton(text="Source Code", url="https://github.com/EverythingSuckz/PaLM-bot"),
35 | ]])
36 | await message.reply(
37 | "I'm a conversational bot powerd by Google's PaLM API.",
38 | parse_mode=enums.ParseMode.HTML,
39 | reply_markup=markup,
40 | )
41 |
--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
1 | # PaLM Telegram Bot
2 | A telegram bot that interacts with Google's PaLM Chat API and can be hosted on serverless functions.
3 |
4 | > Looking for serverless verson based on aiogram?
5 | > Check out [this](https://github.com/EverythingSuckz/PaLM-Bot/tree/aiogram) branch.
6 |
7 | # Demo
8 | A working demo can be found at [@NotAIChatBot](https://telegram.dog/NotAIChatBot).
9 |
10 | # Requirements
11 |
12 | #### Telegram API ID and Hash
13 | Get your API ID and Hash from [my.telegram.org](https://my.telegram.org).
14 | #### Telegram Bot API Token
15 | Create a bot and get the bot token from [@BotFather](https://telegram.dog/BotFather).
16 | #### Python 3.8+
17 | Install Python 3.8 or higher from [here](https://www.python.org/downloads/)
18 | #### PostgreSQL Database
19 | Install PostgreSQL from [here](https://www.postgresql.org/download/) or use a managed database service like [ElephantSQL](https://www.elephantsql.com/)
20 | #### PaLM API Key
21 | Get your PaLM API key from [here](https://makersuite.google.com/)
22 |
23 | # Environment Variables
24 |
25 | - `API_ID` - Your [API ID](#telegram-api-id-and-hash)
26 | - `API_HASH` - Your [API Hash](#telegram-api-id-and-hash)
27 | - `BOT_TOKEN` - Your [bot token](#telegram-bot-api-token)
28 | - `PALM_API_KEY` - Your [PaLM API key](#palm-api-key)
29 | - `DB_USER` - [Database](#postgresql-database) username
30 | - `DB_PASSWORD` - [Database](#postgresql-database) password
31 | - `DB_HOST` - [Database](#postgresql-database) host
32 | - `DB_NAME` - [Database](#postgresql-database) name
33 |
34 | # Hosting
35 | # Self Hosting
36 |
37 | ```bash
38 | git clone -b pyrogram https://github.com/EverythingSuckz/PaLM-Bot
39 | cd PaLM-Bot
40 | python3 -m venv venv
41 | source venv/bin/activate # Linux
42 | .\venv\Scripts\activate # Windows
43 | pip install -r requirements.txt
44 | python -m bot
45 | ```
46 |
47 | _Based on [pyrogram](https://github.com/pyrogram/pyrogram)._
48 |
49 | **Give a ⭐ if you like this project!**
--------------------------------------------------------------------------------
/bot/handlers/chat.py:
--------------------------------------------------------------------------------
1 | import re
2 | import logging
3 | from typing import List
4 |
5 | from pyrogram import enums, types, filters
6 | from google.generativeai.types.safety_types import ContentFilterDict
7 |
8 | from bot import Bot, palm
9 | from bot.database import clear_history
10 | from bot.helpers import limiter, mentioned
11 |
12 | logger = logging.getLogger(__name__)
13 |
14 | gfn = lambda x: f"{x.first_name} {x.last_name}" if x.last_name else ""
15 |
16 | @Bot.on_message(filters.command("clearhistory") & filters.private)
17 | @limiter(15)
18 | async def clearhistory_cmd(_, message: types.Message):
19 | status = await message.reply(
20 | "Please wait...", parse_mode=enums.ParseMode.HTML
21 | )
22 | await clear_history(message.from_user.id if message.from_user else message.sender_chat.id)
23 | await status.edit_text("Done.", parse_mode=enums.ParseMode.HTML)
24 |
25 |
26 |
27 | @Bot.on_message((filters.text & filters.private) | (filters.text & mentioned & filters.group) & filters.incoming, group=2)
28 | @limiter(3)
29 | async def send_handler(_, message: types.Message):
30 | text = re.sub(f"@{Bot.me.username}", "", message.text, flags=re.IGNORECASE).strip()
31 | text = message.text.strip()
32 | if text and text.startswith("/"):
33 | return
34 | if not text:
35 | return
36 | await message.reply_chat_action(enums.ChatAction.TYPING)
37 | user_id = message.from_user.id if message.from_user else message.sender_chat.id
38 | name = message.from_user.first_name if message.from_user else message.sender_chat.title
39 | resp = await palm.get_reponse(user_id=user_id, name=name, message=text)
40 | if not resp:
41 | await message.reply("*AI did not repond*")
42 | return logger.info("No reponse to %s's message", name)
43 | if not resp.last:
44 | filers: List[ContentFilterDict] = resp.filters
45 | await message.reply("*ignores you*")
46 | logger.info("No reponse to %s's message", name)
47 | if filers:
48 | logger.info("Due to the following filters:")
49 | for fil in filers:
50 | logger.info(f"\t[%s]: %s", fil.get("reason"), fil.get("message"))
51 | return
52 | try:
53 | await message.reply(resp.last, parse_mode=enums.ParseMode.MARKDOWN, disable_web_page_preview=True, quote=True)
54 | except Exception as e:
55 | logger.exception(e, stack_info=True)
56 | await message.reply(resp.last, disable_web_page_preview=True, quote=True)
57 |
--------------------------------------------------------------------------------
/transfer.py:
--------------------------------------------------------------------------------
1 | import logging
2 |
3 |
4 | from pony import orm
5 | from datetime import datetime
6 |
7 | from config import Config
8 |
9 | logger = logging.getLogger(__name__)
10 |
11 | remote_db = orm.Database()
12 |
13 |
14 | class RemoteUser(remote_db.Entity):
15 | __table__ = "users"
16 | id = orm.PrimaryKey(int, size=64)
17 | name = orm.Required(str)
18 | username = orm.Optional(str)
19 | started_at = orm.Required(datetime, default=datetime.utcnow)
20 |
21 |
22 | class RemoteHistory(remote_db.Entity):
23 | __table__ = "history"
24 | chat_id = orm.Required(int, size=64)
25 | author = orm.Required(str)
26 | name = orm.Required(str)
27 | message = orm.Required(str)
28 |
29 |
30 |
31 | remote_db.bind(
32 | provider="postgres",
33 | host=Config.DB_HOST,
34 | user=Config.DB_USER,
35 | password=Config.DB_PASSWORD,
36 | database=Config.DB_NAME,
37 | sslmode="require",
38 | )
39 | logger.info("Postgres Database configured.")
40 |
41 | remote_db.generate_mapping(create_tables=True)
42 |
43 | local_db = orm.Database()
44 |
45 | class LocalUser(local_db.Entity):
46 | __table__ = "users"
47 | id = orm.PrimaryKey(int, size=64)
48 | name = orm.Required(str)
49 | username = orm.Optional(str, nullable=True)
50 | started_at = orm.Required(datetime, default=datetime.utcnow)
51 |
52 |
53 | class LocalHistory(local_db.Entity):
54 | __table__ = "history"
55 | chat_id = orm.Required(int, size=64)
56 | author = orm.Required(str)
57 | name = orm.Required(str)
58 | message = orm.Required(str)
59 |
60 | local_db.bind(provider="sqlite", filename="local.db", create_db=True)
61 | logger.info("Local Database configured.")
62 |
63 | with orm.db_session:
64 | for user in RemoteUser.select():
65 | if local_user := LocalUser.get(id=user.id):
66 | local_user.name = user.name
67 | local_user.username = user.username or None
68 | local_user.started_at = user.started_at
69 | else:
70 | LocalUser(id=user.id, name=user.name, username=user.username, started_at=user.started_at)
71 | orm.commit()
72 |
73 | logger.info("Users transferred.")
74 |
75 | with orm.db_session:
76 | for history in LocalHistory.select():
77 | local_history = LocalHistory.get(chat_id=history.chat_id, author=history.author, message=history.message)
78 | if not local_history:
79 | LocalHistory(chat_id=history.chat_id, author=history.author, name=history.name, message=history.message)
80 | orm.commit()
81 |
82 | logger.info("History transferred.")
83 | logger.info("Transfer complete.")
--------------------------------------------------------------------------------
/bot/database.py:
--------------------------------------------------------------------------------
1 | import logging
2 | from enum import Enum
3 | from typing import Optional
4 |
5 | from pony import orm
6 | from datetime import datetime
7 |
8 | from config import Config
9 |
10 | logger = logging.getLogger(__name__)
11 | db = orm.Database()
12 |
13 |
14 | class User(db.Entity):
15 | __table__ = "users"
16 | id = orm.PrimaryKey(int, size=64)
17 | name = orm.Required(str)
18 | username = orm.Optional(str, nullable=True)
19 | started_at = orm.Required(datetime, default=datetime.utcnow)
20 |
21 |
22 | class History(db.Entity):
23 | __table__ = "history"
24 | chat_id = orm.Required(int, size=64)
25 | author = orm.Required(str)
26 | name = orm.Required(str)
27 | message = orm.Required(str)
28 |
29 |
30 | if not any((Config.DB_HOST, Config.DB_USER, Config.DB_PASSWORD, Config.DB_NAME)):
31 | logger.warning("External Database not configured. Using SQLite instead.")
32 | db.bind(provider="sqlite", filename="data.db", create_db=True)
33 |
34 | else:
35 | db.bind(
36 | provider="postgres",
37 | host=Config.DB_HOST,
38 | user=Config.DB_USER,
39 | password=Config.DB_PASSWORD,
40 | database=Config.DB_NAME,
41 | sslmode="require",
42 | )
43 | logger.info("Postgres Database configured.")
44 |
45 | db.generate_mapping(create_tables=True)
46 |
47 | # ==============================================================================
48 | # https://docs.ponyorm.org/integration_with_fastapi.html#async-and-db-session ||
49 | # ==============================================================================
50 |
51 | async def add_user(id: int, name: str, username: Optional[str]):
52 | with orm.db_session:
53 | if not User.exists(id=id):
54 | try:
55 | User(id=id, name=name, username=username)
56 | except (orm.IntegrityError, orm.TransactionIntegrityError):
57 | return False
58 | return True
59 | return False
60 |
61 |
62 | async def set_user_history(user_id: int, name: str, message: str):
63 | with orm.db_session:
64 | History(chat_id=user_id, name=name, message=message, author="user")
65 | db.commit()
66 |
67 |
68 | async def set_response_history(user_id: int, name: str, message: str):
69 | with orm.db_session:
70 | History(chat_id=user_id, name=name, message=message, author="model")
71 | db.commit()
72 |
73 |
74 | async def get_history(user_id: int) -> list[History]:
75 | with orm.db_session:
76 | return list(
77 | reversed(
78 | list(
79 | History.select(lambda h: h.chat_id == user_id)
80 | .order_by(orm.desc(History.id))
81 | .limit(20)
82 | )
83 | )
84 | )
85 |
86 |
87 | async def clear_history(user_id: int):
88 | with orm.db_session:
89 | orm.select(h for h in History if h.chat_id == user_id).delete(bulk=True)
90 |
--------------------------------------------------------------------------------
/bot/PaLM.py:
--------------------------------------------------------------------------------
1 | from typing import Dict, List
2 |
3 | import logging
4 |
5 | import google.generativeai as palm
6 | from google.generativeai.types.discuss_types import ChatResponse
7 |
8 | from bot.database import get_history, set_response_history, set_user_history
9 | from bot.helpers import trim_response
10 |
11 |
12 | EXAMPLES = [
13 | {"author": "user", "content": "Hi"},
14 | {"author": "model", "content": "Oh hi! Are you new here?"},
15 | {"author": "user", "content": "Yeah"},
16 | {
17 | "author": "model",
18 | "content": "Oh alright cool, hope you'll have a great time chatting with others!",
19 | },
20 | {"author": "user", "content": "Yeah, Mind if I promote my channel here?"},
21 | {
22 | "author": "model",
23 | "content": "Oh no, I'm afraid you can't as it's against out community rules.",
24 | },
25 | {"author": "user", "content": "Oh, Alright. Where do you live?"},
26 | {"author": "model", "content": "I live in Tokyo."},
27 | {"author": "user", "content": "Who are you?"},
28 | {"author": "model", "content": "I'm Ai"},
29 | ]
30 |
31 | SYSTEM_PROMPT = """
32 | My name is Ai. I'm a 18 year old girl who loves watching anime, loves gaming.
33 | I'm talking to a person called "{}".
34 | It's best to respond to user's message in a short sentence or two.
35 | Try to act kawaii and casual rather than being very formal.
36 | """
37 |
38 |
39 | class PaLMChat:
40 | """
41 | A chat class that uses PaLM to generate responses.
42 | """
43 |
44 | def __init__(self, api_key: str) -> None:
45 | palm.configure(api_key=api_key)
46 | self.logger = logging.getLogger(__name__)
47 | self.logger.info("Initialized PaLM Client")
48 |
49 | async def _generate_history(
50 | self, user_id: int, message: str
51 | ) -> List[Dict[str, str]]:
52 | """
53 | Helper function to generate history for PaLM.
54 | """
55 | histories = await get_history(user_id)
56 | self.logger.debug(
57 | "Found total %d histories for user[%d]", len(list(histories)), user_id
58 | )
59 | messages = [
60 | {
61 | "author": history.author,
62 | "content": history.message,
63 | }
64 | for history in histories
65 | ]
66 | messages.append(
67 | {
68 | "author": "user",
69 | "content": message,
70 | }
71 | )
72 | return messages
73 |
74 |
75 | async def get_reponse(self, user_id: int, name: str, message: str) -> ChatResponse:
76 | """
77 | Get a response from PaLM and save the history.
78 | """
79 | # Not using async here is bacause it can cause event loop eror when being hosted on vercel.
80 | response: ChatResponse = await palm.chat_async(
81 | model="models/chat-bison-001",
82 | messages=await self._generate_history(user_id, message),
83 | examples=EXAMPLES,
84 | context=SYSTEM_PROMPT.format(name),
85 | )
86 | if not response:
87 | return
88 | await set_user_history(user_id, name, message)
89 | await set_response_history(user_id, name, response.last or "*ignores you*")
90 | if not response.last:
91 | return response
92 | response.last = trim_response(response.last)
93 | self.logger.debug("Generated response for user[%d]", user_id)
94 | self.logger.debug(response)
95 | return response
96 |
--------------------------------------------------------------------------------
/.gitignore:
--------------------------------------------------------------------------------
1 | # Byte-compiled / optimized / DLL files
2 | __pycache__/
3 | *.py[cod]
4 | *$py.class
5 |
6 | # C extensions
7 | *.so
8 |
9 | # Distribution / packaging
10 | .Python
11 | build/
12 | develop-eggs/
13 | dist/
14 | downloads/
15 | eggs/
16 | .eggs/
17 | lib/
18 | lib64/
19 | parts/
20 | sdist/
21 | var/
22 | wheels/
23 | share/python-wheels/
24 | *.egg-info/
25 | .installed.cfg
26 | *.egg
27 | MANIFEST
28 |
29 | # PyInstaller
30 | # Usually these files are written by a python script from a template
31 | # before PyInstaller builds the exe, so as to inject date/other infos into it.
32 | *.manifest
33 | *.spec
34 |
35 | # Installer logs
36 | pip-log.txt
37 | pip-delete-this-directory.txt
38 |
39 | # Unit test / coverage reports
40 | htmlcov/
41 | .tox/
42 | .nox/
43 | .coverage
44 | .coverage.*
45 | .cache
46 | nosetests.xml
47 | coverage.xml
48 | *.cover
49 | *.py,cover
50 | .hypothesis/
51 | .pytest_cache/
52 | cover/
53 |
54 | # Translations
55 | *.mo
56 | *.pot
57 |
58 | # Django stuff:
59 | *.log
60 | local_settings.py
61 | db.sqlite3
62 | db.sqlite3-journal
63 |
64 | # Flask stuff:
65 | instance/
66 | .webassets-cache
67 |
68 | # Scrapy stuff:
69 | .scrapy
70 |
71 | # Sphinx documentation
72 | docs/_build/
73 |
74 | # PyBuilder
75 | .pybuilder/
76 | target/
77 |
78 | # Jupyter Notebook
79 | .ipynb_checkpoints
80 |
81 | # IPython
82 | profile_default/
83 | ipython_config.py
84 |
85 | # pyenv
86 | # For a library or package, you might want to ignore these files since the code is
87 | # intended to run in multiple environments; otherwise, check them in:
88 | # .python-version
89 |
90 | # pipenv
91 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
92 | # However, in case of collaboration, if having platform-specific dependencies or dependencies
93 | # having no cross-platform support, pipenv may install dependencies that don't work, or not
94 | # install all needed dependencies.
95 | #Pipfile.lock
96 |
97 | # poetry
98 | # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
99 | # This is especially recommended for binary packages to ensure reproducibility, and is more
100 | # commonly ignored for libraries.
101 | # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
102 | #poetry.lock
103 |
104 | # pdm
105 | # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
106 | #pdm.lock
107 | # pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
108 | # in version control.
109 | # https://pdm.fming.dev/#use-with-ide
110 | .pdm.toml
111 |
112 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
113 | __pypackages__/
114 |
115 | # Celery stuff
116 | celerybeat-schedule
117 | celerybeat.pid
118 |
119 | # SageMath parsed files
120 | *.sage.py
121 |
122 | # Environments
123 | .env
124 | .venv
125 | env/
126 | venv/
127 | ENV/
128 | env.bak/
129 | venv.bak/
130 |
131 | # Spyder project settings
132 | .spyderproject
133 | .spyproject
134 |
135 | # Rope project settings
136 | .ropeproject
137 |
138 | # mkdocs documentation
139 | /site
140 |
141 | # mypy
142 | .mypy_cache/
143 | .dmypy.json
144 | dmypy.json
145 |
146 | # Pyre type checker
147 | .pyre/
148 |
149 | # pytype static type analyzer
150 | .pytype/
151 |
152 | # Cython debug symbols
153 | cython_debug/
154 |
155 | # PyCharm
156 | # JetBrains specific template is maintained in a separate JetBrains.gitignore that can
157 | # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
158 | # and can be added to the global gitignore or merged into this file. For a more nuclear
159 | # option (not recommended) you can uncomment the following to ignore the entire idea folder.
160 | #.idea/
161 |
162 | *.env
163 | venv/
164 | db.json
165 | unknown_errors.txt
166 | data.db
167 | *.session*
--------------------------------------------------------------------------------
/bot/helpers.py:
--------------------------------------------------------------------------------
1 | import logging
2 | from typing import Callable
3 | from cachetools import TTLCache
4 | from pyrogram import types, filters, enums
5 |
6 | logger = logging.getLogger(__name__)
7 |
8 |
9 | # The PaLM API acts weird sometimes so q quick lazy fix for it
10 | def trim_response(text: str) -> str:
11 | return text[:text.find("USER:") or len(text)].strip()
12 |
13 | # This filter is used to check if the message has bot mentioned in it.
14 | async def mentioned_filter(_, __, message: types.Message):
15 | if message.mentioned:
16 | return True
17 | if message.entities:
18 | for ent in message.entities:
19 | if ent.type == enums.MessageEntityType.MENTION:
20 | return f"@{message._client.me.username}" in message.text
21 |
22 | mentioned = filters.create(mentioned_filter)
23 |
24 | def limiter(rate_limit_seconds: float) -> Callable:
25 | logger = logging.getLogger("limiter")
26 | # Hacky and efficient way for doing a time based cache.
27 | cache = TTLCache(maxsize=1_000, ttl=rate_limit_seconds)
28 | def decorator(func):
29 | async def wrapper(_, message: types.Message):
30 | logger.debug("Limiter was invoked for %s.", func.__name__)
31 | # Anon Admins
32 | user_id = message.from_user.id if message.from_user else message.sender_chat.id
33 | if user_id not in cache:
34 | cache[user_id] = False
35 | await func(_, message)
36 | # Added because the main message handler will pick this message update too.
37 | await message.stop_propagation()
38 | else:
39 | if cache[user_id] is False:
40 | await message.reply("You're sending messages too quickly. Please wait.")
41 | # Another hacky and minimal implementation to make the bot not spam this message.
42 | cache[user_id] = True
43 | return wrapper
44 | return decorator
45 |
46 |
47 |
48 |
49 |
50 |
51 |
52 |
53 |
54 |
55 |
56 |
57 |
58 |
59 |
60 |
61 |
62 |
63 |
64 |
65 |
66 |
67 | # ▄▄▄▄▄▄▄ ▄ ▄ ▄▄ ▄ ▄ ▄ ▄▄ ▄ ▄▄▄▄ ▄▄ ▄▄▄▄▄▄▄
68 | # █ ▄▄▄ █ ▀███▀▀ ██▄ █ █▀▀▄█▀▀█▄▄█▄▀█▄▀▄▀▄█▄▀█ ▄▄█ █ ▄▄▄ █
69 | # █ ███ █ ▀▀█ ▀ ▀▄▀▀█ ▄▄█▄ ▄▄▄█▄▄▄▄▀▀█▄ ▄ ▄█▀ ▄ █ █ ███ █
70 | # █▄▄▄▄▄█ █▀▄ ▄ █▀█▀█ ▄▀▄▀▄▀█ ▄ █ ▄ █ █ ▄▀█▀█▀█ █ ▄ █▄▄▄▄▄█
71 | # ▄▄▄▄▄ ▄▄▄▀ ██▄▀▄█████▄▀ █ █▄▄▄█▄▀▄ █▄▄█▀ ▄▀ ▄▄▀▄ ▄ ▄ ▄
72 | # █ █ ▀▄█▀█▀▄▀ ▀ █▀ ▄██▀█▀ ▀▀██▀▄ █▄█ ▀ ▀▀▀█ ▄▄▀▀ ▄█▀██▀
73 | # █ █ ▀▄██▀▀ ▀ ▀▄█▄ ▀█▀█▄█▄▀▀█▀█ ▄▄ ▄█ ▄▀█▀▄ ▀█ ▀ ▄ ▀██▄
74 | # ▄█▀▄ ▄▄█ █ ▄█ ▄▀▄▀ ▄▀▀▀▀ █ █▀▄▄█▄▀▄▀ ▀▀█▄ ▀▄▄▀█▄▀▄▄▀█▀
75 | # ▄▄▄▀▄▄ ▀██ ▀█ ▄▄▀▄█▄▀ ▄█ ▄▀▄█▄ ▄█▀▄▄█▀▀▄▀▄▀▀ ▀▄▄▀ ▀▀█▀
76 | # ▀ ██▄▀▄ ▀▀▄▄█▄▄ ▀ ▀▄ ▀██▀ ▀▀▀▀▀█▄▀▄▀ █ ▄▀▄▄▄▄▀█▄▄ ▀██▀
77 | # ▀▀██▀▀▄ ▄▄▀▄▄▄██ ▄▄▀ ▀██▀ ▀█▀▄ ▄ ▀▄▀▀▀▀▄▀▄ ▄▄███▄▀▀▀█
78 | # ▀█▀ ▀█▄ ▀▀ ▄▀▄▄█▄ ▄▀ ▄▀▀█▀▀▄▀█▀█ ▄ ▀▄█ █ █▀▄ ▀▄▀█▀ ▀ ██▀▀
79 | # █ ▀▄▄▄▄ ▀ █▄█ ▀▄▀▄▀▄▀ ▀▄█ ▀▀▀ ▀▀▄▀▄▄▀ ▀█▀▄ ▄ ▀▄▄▄█▀▀█▄▄
80 | # █▀ ███▄██ ██ █▀▄▀ █ ▀██▀▄██▄█▀█▄▄▄▄▀▄ ▀█▀▄ ▄ ██▄▄██▀▄▄▀
81 | # █ ▄█ ▄ █▄▀▀▀█▀█▄ ▄▄▄▀ ▄▀█ ▄ █▄ ▀███▄▄▀ ▀ ▀▀▄ ▄█ ▄ ███
82 | # ▄▀▀▄█▄▄▄█▄██ ▀ ▄▀▄█▄█ ▀█▀█▄▄▄█▄ ▄ ▄▀▄▄ ▀▀█▄█▄ ▄█▄▄▄███▄▀
83 | # ▀ ▀ ▄█▄ █▄ ██▀▀▄▄██▄▀▀▀▀▀ ▀▄█▄█ ▀█▄█▄▄▄▀▄█▄▀██ ▄█▀▄ ▀█ █
84 | # ▀████▄▄▀▀▄▀ ▀▄▀▀▄█▄▀ ▀█▀▄█▀▄▄█ █▄ ██▄▄▄▀▀▀ ▄▄▄▄█▀▄ ▀▄▀▀
85 | # ▄█ █▄█ █ █ █▄█ █▄▄▄▀ ▀█▀██▄▀▀█▀▄▀▄█▄▄▄▀█▀▀▀▄ ▀ ▄██▀▀
86 | # ▀ ▀▄█▄▀▀▄ ▀ ▀ ▄▀▄ ▄ ▀█ ▄ ▄▄▄█▄ ▀ ▀▄▄▀▄▀ ▄▄ ▀██ ▄▄▄▀
87 | # █▀▄█▄▄▄█▀█ █▄▀█▄ ▄███▀ ▀█▄▀▀ ▄▄▄ ▄▄█▄█▄▀▀▀██▀▄▀ ▄▀█ ▄▄▀▀
88 | # ▀ ▀█▀▀▄▀▄▀ ▄▀█▄ ▄▄▄▄▀█▀▀ █▄ ▄█▄ ▄▄▀▀█▄▄▀███▀▄▄ ▄ ▄▄▀▀▀
89 | # █▄▀ █ ▄▄ ██▀ ▄▄█ ▄▀ █ ▄ █▄█▀▀▀▄▄█ ▄▄▄ ▄▀▀▄▀▀ ▄██▀▀
90 | # █▀█ ▀█▄▀▀█▀ █▄▄▀ █▀ █▀▄▀ ▄▄█ ▄▀▄▄█▀▄▀▄▀██▀ ▄▄▄ ▄ ▄▀▄▀
91 | # ▀▀▀▀▀ ▄ ▄▀ ▀▀ ▀█▀█▄▄█ ▀▀▄███▄█▄▀▀▄▄▄▄ ▀ ▄▄▀▀ ▄▄██▄▄▀▀█
92 | # ▄▄▄▄▄▄▄ █ ▄ ▄ ▄▄▄▄█▄██▀█ ▄ █▀ ▀▄ ▀▄ █▀▀▀▄ █ ▄ █ ▄▄
93 | # █ ▄▄▄ █ ▄ ▄ ▀▄▀ ▄█ ▄▀▄█▄▄▄█▄▀▀▄ ▄ ▀▀ ▀█▄▀▄▀▀█▄▄▄█▀▀█
94 | # █ ███ █ █▄▄▄ ▀▄▄ █▀█▄▀ ▀▄▀▀ ▄▄ ▄▄ ▄▄ █▀█▀▀ ███▄▄▀ █
95 | # █▄▄▄▄▄█ █▄▄▀ ███▀ ▄█▀ ▀▀█▀ ▄██▀ ▄▀▄▄▄▀██▄▄▄▄▀▀█▀▄
96 | # Google Tranlate is trash for this job btw
--------------------------------------------------------------------------------