Поддержка множества аккаунтов ботов.
This commit is contained in:
parent
6088606810
commit
b4a60b59e8
19 changed files with 427 additions and 310 deletions
40
ai_agent.py
40
ai_agent.py
|
|
@ -40,7 +40,7 @@ class AiAgent:
|
|||
self.model_temp = model_temp
|
||||
self.client = OpenRouter(api_key=api_token, retry_config=retry_config)
|
||||
|
||||
async def get_group_chat_reply(self, chat_id: int,
|
||||
async def get_group_chat_reply(self, bot_id: int, chat_id: int,
|
||||
message: Message, forwarded_messages: List[Message]) -> Tuple[str, bool]:
|
||||
message_text = f"[{message.user_name}]: {message.text}"
|
||||
for fwd_message in forwarded_messages:
|
||||
|
|
@ -48,7 +48,7 @@ class AiAgent:
|
|||
message_text += fwd_message.text + '\n'
|
||||
message_text += '<Конец цитаты>'
|
||||
|
||||
context = self._get_chat_context(is_group_chat=True, chat_id=chat_id)
|
||||
context = self._get_chat_context(is_group_chat=True, bot_id=bot_id, chat_id=chat_id)
|
||||
context.append({"role": "user", "content": message_text})
|
||||
|
||||
try:
|
||||
|
|
@ -64,9 +64,9 @@ class AiAgent:
|
|||
ai_response = response.choices[0].message.content
|
||||
|
||||
# Add message and AI response to context
|
||||
self.db.context_add_message(chat_id=chat_id, role="user", content=message_text,
|
||||
self.db.context_add_message(bot_id, chat_id, role="user", content=message_text,
|
||||
message_id=message.message_id, max_messages=GROUP_CHAT_MAX_MESSAGES)
|
||||
self.db.context_add_message(chat_id=chat_id, role="assistant", content=ai_response,
|
||||
self.db.context_add_message(bot_id, chat_id, role="assistant", content=ai_response,
|
||||
message_id=None, max_messages=GROUP_CHAT_MAX_MESSAGES)
|
||||
|
||||
return ai_response, True
|
||||
|
|
@ -78,8 +78,8 @@ class AiAgent:
|
|||
print(f"Ошибка выполнения запроса к ИИ: {e}")
|
||||
return f"Извините, при обработке запроса произошла ошибка.", False
|
||||
|
||||
async def get_private_chat_reply(self, chat_id: int, message: str, message_id: int) -> Tuple[str, bool]:
|
||||
context = self._get_chat_context(is_group_chat=False, chat_id=chat_id)
|
||||
async def get_private_chat_reply(self, bot_id: int, chat_id: int, message: str, message_id: int) -> Tuple[str, bool]:
|
||||
context = self._get_chat_context(is_group_chat=False, bot_id=bot_id, chat_id=chat_id)
|
||||
context.append({"role": "user", "content": message})
|
||||
|
||||
try:
|
||||
|
|
@ -88,16 +88,16 @@ class AiAgent:
|
|||
model=self.model,
|
||||
messages=context,
|
||||
max_tokens=500,
|
||||
temperature=0.5
|
||||
temperature=self.model_temp
|
||||
)
|
||||
|
||||
# Extract AI response
|
||||
ai_response = response.choices[0].message.content
|
||||
|
||||
# Add message and AI response to context
|
||||
self.db.context_add_message(chat_id=chat_id, role="user", content=message,
|
||||
self.db.context_add_message(bot_id, chat_id, role="user", content=message,
|
||||
message_id=message_id, max_messages=PRIVATE_CHAT_MAX_MESSAGES)
|
||||
self.db.context_add_message(chat_id=chat_id, role="assistant", content=ai_response,
|
||||
self.db.context_add_message(bot_id, chat_id, role="assistant", content=ai_response,
|
||||
message_id=None, max_messages=PRIVATE_CHAT_MAX_MESSAGES)
|
||||
|
||||
return ai_response, True
|
||||
|
|
@ -109,23 +109,27 @@ class AiAgent:
|
|||
print(f"Ошибка выполнения запроса к ИИ: {e}")
|
||||
return f"Извините, при обработке запроса произошла ошибка.", False
|
||||
|
||||
def get_last_assistant_message_id(self, chat_id: int):
|
||||
return self.db.context_get_last_assistant_message_id(chat_id)
|
||||
def get_last_assistant_message_id(self, bot_id: int, chat_id: int):
|
||||
return self.db.context_get_last_assistant_message_id(bot_id, chat_id)
|
||||
|
||||
def set_last_response_id(self, chat_id: int, message_id: int):
|
||||
self.db.context_set_last_message_id(chat_id, message_id)
|
||||
def set_last_response_id(self, bot_id: int, chat_id: int, message_id: int):
|
||||
self.db.context_set_last_message_id(bot_id, chat_id, message_id)
|
||||
|
||||
def clear_chat_context(self, chat_id: int):
|
||||
self.db.context_clear(chat_id)
|
||||
def clear_chat_context(self, bot_id: int, chat_id: int):
|
||||
self.db.context_clear(bot_id, chat_id)
|
||||
|
||||
def _get_chat_context(self, is_group_chat: bool, chat_id: int) -> list[dict]:
|
||||
def _get_chat_context(self, is_group_chat: bool, bot_id: int, chat_id: int) -> list[dict]:
|
||||
prompt = GROUP_CHAT_SYSTEM_PROMPT if is_group_chat else PRIVATE_CHAT_SYSTEM_PROMPT
|
||||
|
||||
chat = self.db.create_chat_if_not_exists(chat_id)
|
||||
bot = self.db.get_bot(bot_id)
|
||||
if bot['ai_prompt'] is not None:
|
||||
prompt += '\n\n' + bot['ai_prompt']
|
||||
|
||||
chat = self.db.create_chat_if_not_exists(bot_id, chat_id)
|
||||
if chat['ai_prompt'] is not None:
|
||||
prompt += '\n\n' + chat['ai_prompt']
|
||||
|
||||
messages = self.db.context_get_messages(chat_id)
|
||||
messages = self.db.context_get_messages(bot_id, chat_id)
|
||||
return [{"role": "system", "content": prompt}] + messages
|
||||
|
||||
|
||||
|
|
|
|||
167
database.py
167
database.py
|
|
@ -12,166 +12,179 @@ class BasicDatabase:
|
|||
self.conn.setencoding(encoding='utf-8')
|
||||
self.cursor = self.conn.cursor()
|
||||
|
||||
def get_chats(self):
|
||||
self.cursor.execute("SELECT * FROM chats")
|
||||
def get_bots(self):
|
||||
self.cursor.execute("SELECT * FROM bots")
|
||||
return self._to_dict(self.cursor.fetchall())
|
||||
|
||||
def get_chat(self, chat_id: int):
|
||||
self.cursor.execute("SELECT * FROM chats WHERE id = ?", chat_id)
|
||||
def get_bot(self, bot_id: int):
|
||||
self.cursor.execute("SELECT * FROM bots WHERE id = ?", bot_id)
|
||||
return self._to_dict(self.cursor.fetchone())
|
||||
|
||||
def add_chat(self, chat_id: int):
|
||||
self.cursor.execute("INSERT INTO chats (id) VALUES (?)", chat_id)
|
||||
def get_chats(self, bot_id: int):
|
||||
self.cursor.execute("SELECT * FROM chats WHERE bot_id = ?", bot_id)
|
||||
return self._to_dict(self.cursor.fetchall())
|
||||
|
||||
def chat_update(self, chat_id: int, **kwargs):
|
||||
def get_chat(self, bot_id: int, chat_id: int):
|
||||
self.cursor.execute("SELECT * FROM chats WHERE bot_id = ? AND chat_id = ?", bot_id, chat_id)
|
||||
return self._to_dict(self.cursor.fetchone())
|
||||
|
||||
def add_chat(self, bot_id: int, chat_id: int):
|
||||
self.cursor.execute("INSERT INTO chats (bot_id, chat_id) VALUES (?, ?)", bot_id, chat_id)
|
||||
|
||||
def chat_update(self, bot_id: int, chat_id: int, **kwargs):
|
||||
self.cursor.execute("UPDATE chats SET " + ", ".join(f + " = ?" for f in kwargs) +
|
||||
" WHERE id = ?", list(kwargs.values()) + [chat_id])
|
||||
" WHERE bot_id = ? AND chat_id = ?", list(kwargs.values()) + [bot_id, chat_id])
|
||||
|
||||
def chat_delete(self, chat_id: int):
|
||||
self.cursor.execute("DELETE FROM chats WHERE id = ?", chat_id)
|
||||
def chat_delete(self, bot_id: int, chat_id: int):
|
||||
self.cursor.execute("DELETE FROM chats WHERE bot_id = ? AND chat_id = ?", bot_id, chat_id)
|
||||
|
||||
def get_user(self, chat_id: int, user_id: int):
|
||||
self.cursor.execute("SELECT * FROM users WHERE chat_id = ? AND user_id = ?", chat_id, user_id)
|
||||
def get_user(self, bot_id: int, chat_id: int, user_id: int):
|
||||
self.cursor.execute("SELECT * FROM users WHERE bot_id = ? AND chat_id = ? AND user_id = ?",
|
||||
bot_id, chat_id, user_id)
|
||||
return self._to_dict(self.cursor.fetchone())
|
||||
|
||||
def get_users(self, chat_id: int):
|
||||
self.cursor.execute("SELECT * FROM users WHERE chat_id = ?", chat_id)
|
||||
def get_users(self, bot_id: int, chat_id: int):
|
||||
self.cursor.execute("SELECT * FROM users WHERE bot_id = ? AND chat_id = ?", bot_id, chat_id)
|
||||
return self._to_dict(self.cursor.fetchall())
|
||||
|
||||
def add_user(self, chat_id: int, user_id: int):
|
||||
self.cursor.execute("INSERT INTO users (chat_id, user_id) VALUES (?, ?)", chat_id, user_id)
|
||||
def add_user(self, bot_id: int, chat_id: int, user_id: int):
|
||||
self.cursor.execute("INSERT INTO users (bot_id, chat_id, user_id) VALUES (?, ?, ?)",
|
||||
bot_id, chat_id, user_id)
|
||||
|
||||
def user_set_last_message(self, chat_id: int, user_id: int, last_message: int):
|
||||
self.user_update(chat_id, user_id, last_message=last_message)
|
||||
def user_set_last_message(self, bot_id: int, chat_id: int, user_id: int, last_message: int):
|
||||
self.user_update(bot_id, chat_id, user_id, last_message=last_message)
|
||||
|
||||
def user_increment_messages(self, chat_id: int, user_id: int):
|
||||
self.user_increment(chat_id, user_id, ['messages_today', 'messages_month'])
|
||||
def user_increment_messages(self, bot_id: int, chat_id: int, user_id: int):
|
||||
self.user_increment(bot_id, chat_id, user_id, ['messages_today', 'messages_month'])
|
||||
|
||||
def user_increment_warnings(self, chat_id: int, user_id: int):
|
||||
self.user_increment(chat_id, user_id, ['warnings'])
|
||||
def user_increment_warnings(self, bot_id: int, chat_id: int, user_id: int):
|
||||
self.user_increment(bot_id, chat_id, user_id, ['warnings'])
|
||||
|
||||
def user_increment(self, chat_id: int, user_id: int, fields: List[str]):
|
||||
def user_increment(self, bot_id: int, chat_id: int, user_id: int, fields: List[str]):
|
||||
self.cursor.execute("UPDATE users SET " + ", ".join(f + " = " + f + " + 1" for f in fields) +
|
||||
" WHERE chat_id = ? AND user_id = ?", chat_id, user_id)
|
||||
" WHERE bot_id = ? AND chat_id = ? AND user_id = ?", bot_id, chat_id, user_id)
|
||||
|
||||
def user_update(self, chat_id: int, user_id: int, **kwargs):
|
||||
def user_update(self, bot_id, chat_id: int, user_id: int, **kwargs):
|
||||
self.cursor.execute("UPDATE users SET " + ", ".join(f + " = ?" for f in kwargs) +
|
||||
" WHERE chat_id = ? AND user_id = ?", list(kwargs.values()) + [chat_id, user_id])
|
||||
" WHERE bot_id = ? AND chat_id = ? AND user_id = ?",
|
||||
list(kwargs.values()) + [bot_id, chat_id, user_id])
|
||||
|
||||
def delete_user(self, chat_id: int, user_id: int):
|
||||
self.cursor.execute("DELETE FROM users WHERE chat_id = ? AND user_id = ?", chat_id, user_id)
|
||||
def delete_user(self, bot_id: int, chat_id: int, user_id: int):
|
||||
self.cursor.execute("DELETE FROM users WHERE bot_id = ? AND chat_id = ? AND user_id = ?",
|
||||
bot_id, chat_id, user_id)
|
||||
|
||||
def get_top_messages_today(self, chat_id: int):
|
||||
def get_top_messages_today(self, bot_id: int, chat_id: int):
|
||||
self.cursor.execute("""
|
||||
SELECT user_id, messages_today AS value FROM users
|
||||
WHERE chat_id = ? AND messages_today > 0
|
||||
WHERE bot_id = ? AND chat_id = ? AND messages_today > 0
|
||||
ORDER BY messages_today DESC
|
||||
""", chat_id)
|
||||
""", bot_id, chat_id)
|
||||
return self._to_dict(self.cursor.fetchall())
|
||||
|
||||
def get_top_messages_month(self, chat_id: int):
|
||||
def get_top_messages_month(self, bot_id: int, chat_id: int):
|
||||
self.cursor.execute("""
|
||||
SELECT user_id, messages_month AS value FROM users
|
||||
WHERE chat_id = ? AND messages_month > 0
|
||||
WHERE bot_id = ? AND chat_id = ? AND messages_month > 0
|
||||
ORDER BY messages_month DESC
|
||||
""", chat_id)
|
||||
""", bot_id, chat_id)
|
||||
return self._to_dict(self.cursor.fetchall())
|
||||
|
||||
def get_top_silent(self, chat_id: int, threshold_days: int):
|
||||
def get_top_silent(self, bot_id: int, chat_id: int, threshold_days: int):
|
||||
current_time = int(datetime.now().timestamp())
|
||||
threshold = current_time - threshold_days * 86400
|
||||
self.cursor.execute("""
|
||||
SELECT user_id, (? - last_message) DIV 86400 as value
|
||||
FROM users
|
||||
WHERE chat_id = ? AND last_message <= ?
|
||||
WHERE bot_id = ? AND chat_id = ? AND last_message <= ?
|
||||
ORDER BY last_message ASC
|
||||
""", current_time, chat_id, threshold)
|
||||
""", current_time, bot_id, chat_id, threshold)
|
||||
result = self._to_dict(self.cursor.fetchall())
|
||||
for row in result:
|
||||
if row['value'] > 3650:
|
||||
row['value'] = 'никогда'
|
||||
return result
|
||||
|
||||
def get_top_warnings(self, chat_id: int):
|
||||
def get_top_warnings(self, bot_id: int, chat_id: int):
|
||||
self.cursor.execute("""
|
||||
SELECT user_id, warnings AS value FROM users
|
||||
WHERE chat_id = ? AND warnings > 0
|
||||
WHERE bot_id = ? AND chat_id = ? AND warnings > 0
|
||||
ORDER BY warnings DESC
|
||||
""", chat_id)
|
||||
""", bot_id, chat_id)
|
||||
return self._to_dict(self.cursor.fetchall())
|
||||
|
||||
def reset_messages_today(self):
|
||||
self.cursor.execute("UPDATE users SET messages_today = 0")
|
||||
def reset_messages_today(self, bot_id: int):
|
||||
self.cursor.execute("UPDATE users SET messages_today = 0 WHERE bot_id = ?", bot_id)
|
||||
|
||||
def reset_messages_month(self):
|
||||
self.cursor.execute("UPDATE users SET messages_month = 0")
|
||||
def reset_messages_month(self, bot_id: int):
|
||||
self.cursor.execute("UPDATE users SET messages_month = 0 WHERE bot_id = ?", bot_id)
|
||||
|
||||
def context_get_messages(self, chat_id: int) -> list[dict]:
|
||||
def context_get_messages(self, bot_id: int, chat_id: int) -> list[dict]:
|
||||
self.cursor.execute("""
|
||||
SELECT role, content FROM contexts
|
||||
WHERE chat_id = ? AND message_id IS NOT NULL
|
||||
WHERE bot_id = ? AND chat_id = ? AND message_id IS NOT NULL
|
||||
ORDER BY message_id
|
||||
""", chat_id)
|
||||
""", bot_id, chat_id)
|
||||
return self._to_dict(self.cursor.fetchall())
|
||||
|
||||
def context_get_count(self, chat_id: int) -> int:
|
||||
self.cursor.execute("SELECT COUNT(*) FROM contexts WHERE chat_id = ?", chat_id)
|
||||
def context_get_count(self, bot_id: int, chat_id: int) -> int:
|
||||
self.cursor.execute("SELECT COUNT(*) FROM contexts WHERE bot_id = ? AND chat_id = ?", bot_id, chat_id)
|
||||
return self.cursor.fetchval()
|
||||
|
||||
def context_get_last_assistant_message_id(self, chat_id: int) -> Optional[int]:
|
||||
def context_get_last_assistant_message_id(self, bot_id: int, chat_id: int) -> Optional[int]:
|
||||
return self.cursor.execute("""
|
||||
SELECT message_id FROM contexts
|
||||
WHERE chat_id = ? AND role = 'assistant' AND message_id IS NOT NULL
|
||||
WHERE bot_id = ? AND chat_id = ? AND role = 'assistant' AND message_id IS NOT NULL
|
||||
ORDER BY message_id DESC
|
||||
LIMIT 1
|
||||
""", chat_id).fetchval()
|
||||
""", bot_id, chat_id).fetchval()
|
||||
|
||||
def context_add_message(self, chat_id: int, role: str, content: str, message_id: Optional[int], max_messages: int):
|
||||
self._context_trim(chat_id, max_messages)
|
||||
def context_add_message(self, bot_id: int, chat_id: int, role: str, content: str, message_id: Optional[int], max_messages: int):
|
||||
self._context_trim(bot_id, chat_id, max_messages)
|
||||
|
||||
if message_id is not None:
|
||||
self.cursor.execute("INSERT INTO contexts (chat_id, message_id, role, content) VALUES (?, ?, ?, ?)",
|
||||
chat_id, message_id, role, content)
|
||||
self.cursor.execute(
|
||||
"INSERT INTO contexts (bot_id, chat_id, message_id, role, content) VALUES (?, ?, ?, ?, ?)",
|
||||
bot_id, chat_id, message_id, role, content)
|
||||
else:
|
||||
self.cursor.execute("INSERT INTO contexts (chat_id, role, content) VALUES (?, ?, ?)",
|
||||
chat_id, role, content)
|
||||
self.cursor.execute("INSERT INTO contexts (bot_id, chat_id, role, content) VALUES (?, ?, ?, ?)",
|
||||
bot_id, chat_id, role, content)
|
||||
|
||||
def context_set_last_message_id(self, chat_id: int, message_id: int):
|
||||
self.cursor.execute("UPDATE contexts SET message_id = ? WHERE chat_id = ? AND message_id IS NULL",
|
||||
message_id, chat_id)
|
||||
def context_set_last_message_id(self, bot_id: int, chat_id: int, message_id: int):
|
||||
self.cursor.execute("UPDATE contexts SET message_id = ? WHERE bot_id = ? AND chat_id = ? AND message_id IS NULL",
|
||||
message_id, bot_id, chat_id)
|
||||
|
||||
def _context_trim(self, chat_id: int, max_messages: int):
|
||||
current_count = self.context_get_count(chat_id)
|
||||
def _context_trim(self, bot_id: int, chat_id: int, max_messages: int):
|
||||
current_count = self.context_get_count(bot_id, chat_id)
|
||||
while current_count >= max_messages:
|
||||
oldest_message_id = self.cursor.execute("""
|
||||
SELECT message_id FROM contexts
|
||||
WHERE chat_id = ? AND message_id IS NOT NULL
|
||||
WHERE bot_id = ? AND chat_id = ? AND message_id IS NOT NULL
|
||||
ORDER BY message_id ASC
|
||||
LIMIT 1
|
||||
""", chat_id).fetchval()
|
||||
""", bot_id, chat_id).fetchval()
|
||||
|
||||
if oldest_message_id:
|
||||
self.cursor.execute("DELETE FROM contexts WHERE chat_id = ? AND message_id = ?",
|
||||
chat_id, oldest_message_id)
|
||||
self.cursor.execute("DELETE FROM contexts WHERE bot_id = ? AND chat_id = ? AND message_id = ?",
|
||||
bot_id, chat_id, oldest_message_id)
|
||||
current_count -= 1
|
||||
else:
|
||||
break
|
||||
|
||||
def context_clear(self, chat_id: int):
|
||||
self.cursor.execute("DELETE FROM contexts WHERE chat_id = ?", chat_id)
|
||||
def context_clear(self, bot_id: int, chat_id: int):
|
||||
self.cursor.execute("DELETE FROM contexts WHERE bot_id = ? AND chat_id = ?", bot_id, chat_id)
|
||||
|
||||
def create_chat_if_not_exists(self, chat_id: int):
|
||||
chat = self.get_chat(chat_id)
|
||||
def create_chat_if_not_exists(self, bot_id: int, chat_id: int):
|
||||
chat = self.get_chat(bot_id, chat_id)
|
||||
if chat is None:
|
||||
self.add_chat(chat_id)
|
||||
chat = self.get_chat(chat_id)
|
||||
self.add_chat(bot_id, chat_id)
|
||||
chat = self.get_chat(bot_id, chat_id)
|
||||
return chat
|
||||
|
||||
def create_user_if_not_exists(self, chat_id: int, user_id: int):
|
||||
user = self.get_user(chat_id, user_id)
|
||||
def create_user_if_not_exists(self, bot_id: int, chat_id: int, user_id: int):
|
||||
user = self.get_user(bot_id, chat_id, user_id)
|
||||
if user is None:
|
||||
self.add_user(chat_id, user_id)
|
||||
user = self.get_user(chat_id, user_id)
|
||||
self.add_user(bot_id, chat_id, user_id)
|
||||
user = self.get_user(bot_id, chat_id, user_id)
|
||||
return user
|
||||
|
||||
def _to_dict(self, args: Union[Row, List[Row], None]):
|
||||
|
|
|
|||
|
|
@ -16,18 +16,24 @@ async def main() -> None:
|
|||
config = json.load(file)
|
||||
print('Конфигурация загружена.')
|
||||
|
||||
bot = Bot(token=config['api_token'])
|
||||
database.create_database(config['db_connection_string'])
|
||||
|
||||
create_ai_agent(config['openrouter_token'],
|
||||
config['openrouter_model'],
|
||||
config['openrouter_model_temp'],
|
||||
database.DB)
|
||||
|
||||
bots: list[Bot] = []
|
||||
for item in database.DB.get_bots():
|
||||
bot = Bot(token=item['api_token'])
|
||||
asyncio.create_task(tasks.startup_task(bot))
|
||||
asyncio.create_task(tasks.daily_maintenance_task(bot))
|
||||
bots.append(bot)
|
||||
|
||||
dp = Dispatcher()
|
||||
dp.include_router(handlers.router)
|
||||
dp.startup.register(tasks.startup_task)
|
||||
asyncio.create_task(tasks.daily_maintenance_task(bot))
|
||||
await dp.start_polling(bot)
|
||||
await dp.start_polling(*bots)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
asyncio.run(main())
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
from aiogram import Router, F
|
||||
from aiogram import Router, F, Bot
|
||||
from aiogram.enums import ContentType, ParseMode
|
||||
from aiogram.types import Message
|
||||
|
||||
|
|
@ -9,9 +9,9 @@ router = Router()
|
|||
|
||||
|
||||
@router.message(F.content_type == ContentType.NEW_CHAT_MEMBERS)
|
||||
async def user_join_handler(message: Message):
|
||||
async def user_join_handler(message: Message, bot: Bot):
|
||||
chat_id = message.chat.id
|
||||
chat = database.DB.create_chat_if_not_exists(chat_id)
|
||||
chat = database.DB.create_chat_if_not_exists(bot.id, chat_id)
|
||||
if chat['active'] == 0:
|
||||
return
|
||||
|
||||
|
|
@ -25,9 +25,9 @@ async def user_join_handler(message: Message):
|
|||
|
||||
|
||||
@router.message(F.content_type == ContentType.LEFT_CHAT_MEMBER)
|
||||
async def user_join_handler(message: Message):
|
||||
async def user_join_handler(message: Message, bot: Bot):
|
||||
chat_id = message.chat.id
|
||||
chat = database.DB.create_chat_if_not_exists(chat_id)
|
||||
chat = database.DB.create_chat_if_not_exists(bot.id, chat_id)
|
||||
if chat['active'] == 0:
|
||||
return
|
||||
|
||||
|
|
@ -35,11 +35,11 @@ async def user_join_handler(message: Message):
|
|||
if member.is_bot:
|
||||
return
|
||||
|
||||
database.DB.delete_user(chat_id, member.id)
|
||||
database.DB.delete_user(bot.id, chat_id, member.id)
|
||||
|
||||
|
||||
@router.message(F.content_type == ContentType.MIGRATE_TO_CHAT_ID)
|
||||
async def migration_handler(message: Message):
|
||||
async def migration_handler(message: Message, bot: Bot):
|
||||
old_id, new_id = message.chat.id, message.migrate_to_chat_id
|
||||
database.DB.chat_delete(new_id)
|
||||
database.DB.chat_update(old_id, id=new_id)
|
||||
database.DB.chat_delete(bot.id, new_id)
|
||||
database.DB.chat_update(bot.id, old_id, id=new_id)
|
||||
|
|
|
|||
|
|
@ -21,20 +21,20 @@ async def tg_user_is_admin(chat_id: int, user_id: int, bot: Bot):
|
|||
@router.message(F.text == "!старт")
|
||||
async def start_handler(message: Message, bot: Bot):
|
||||
chat_id = message.chat.id
|
||||
database.DB.create_chat_if_not_exists(chat_id)
|
||||
database.DB.create_chat_if_not_exists(bot.id, chat_id)
|
||||
|
||||
if not await tg_user_is_admin(chat_id, message.from_user.id, bot):
|
||||
await message.answer(MESSAGE_PERMISSION_DENIED)
|
||||
return
|
||||
|
||||
database.DB.chat_update(chat_id, active=1)
|
||||
database.DB.chat_update(bot.id, chat_id, active=1)
|
||||
await message.answer('Готова к работе!')
|
||||
|
||||
|
||||
@router.message(F.text == "!правила")
|
||||
async def rules_handler(message: Message, bot: Bot):
|
||||
chat_id = message.chat.id
|
||||
chat = database.DB.create_chat_if_not_exists(chat_id)
|
||||
chat = database.DB.create_chat_if_not_exists(bot.id, chat_id)
|
||||
if chat['active'] == 0:
|
||||
await message.answer(MESSAGE_CHAT_NOT_ACTIVE)
|
||||
return
|
||||
|
|
@ -50,14 +50,14 @@ async def rules_handler(message: Message, bot: Bot):
|
|||
await message.answer(MESSAGE_PERMISSION_DENIED)
|
||||
return
|
||||
|
||||
database.DB.chat_update(chat_id, rules=message.reply_to_message.text)
|
||||
database.DB.chat_update(bot.id, chat_id, rules=message.reply_to_message.text)
|
||||
await message.answer('Правила чата изменены.')
|
||||
|
||||
|
||||
@router.message(F.text == "!приветствие")
|
||||
async def set_greeting_handler(message: Message, bot: Bot):
|
||||
chat_id = message.chat.id
|
||||
chat = database.DB.create_chat_if_not_exists(chat_id)
|
||||
chat = database.DB.create_chat_if_not_exists(bot.id, chat_id)
|
||||
if chat['active'] == 0:
|
||||
await message.answer(MESSAGE_CHAT_NOT_ACTIVE)
|
||||
return
|
||||
|
|
@ -70,14 +70,14 @@ async def set_greeting_handler(message: Message, bot: Bot):
|
|||
await message.answer(MESSAGE_NEED_REPLY)
|
||||
return
|
||||
|
||||
database.DB.chat_update(chat_id, greeting_join=message.reply_to_message.text)
|
||||
database.DB.chat_update(bot.id, chat_id, greeting_join=message.reply_to_message.text)
|
||||
await message.answer('Приветствие изменено.')
|
||||
|
||||
|
||||
@router.message(F.text == "!личность")
|
||||
async def set_ai_prompt_handler(message: Message, bot: Bot):
|
||||
chat_id = message.chat.id
|
||||
chat = database.DB.create_chat_if_not_exists(chat_id)
|
||||
chat = database.DB.create_chat_if_not_exists(bot.id, chat_id)
|
||||
if chat['active'] == 0:
|
||||
await message.answer(MESSAGE_CHAT_NOT_ACTIVE)
|
||||
return
|
||||
|
|
@ -90,15 +90,15 @@ async def set_ai_prompt_handler(message: Message, bot: Bot):
|
|||
await message.answer(MESSAGE_NEED_REPLY)
|
||||
return
|
||||
|
||||
database.DB.chat_update(chat_id, ai_prompt=message.reply_to_message.text)
|
||||
ai_agent.agent.clear_chat_context(chat_id)
|
||||
database.DB.chat_update(bot.id, chat_id, ai_prompt=message.reply_to_message.text)
|
||||
ai_agent.agent.clear_chat_context(bot.id, chat_id)
|
||||
await message.answer('Личность ИИ изменена.')
|
||||
|
||||
|
||||
@router.message(F.text == "!предупреждение")
|
||||
async def warning_handler(message: Message, bot: Bot):
|
||||
chat_id = message.chat.id
|
||||
chat = database.DB.create_chat_if_not_exists(chat_id)
|
||||
chat = database.DB.create_chat_if_not_exists(bot.id, chat_id)
|
||||
if chat['active'] == 0:
|
||||
await message.answer(MESSAGE_CHAT_NOT_ACTIVE)
|
||||
return
|
||||
|
|
@ -112,10 +112,10 @@ async def warning_handler(message: Message, bot: Bot):
|
|||
return
|
||||
|
||||
user_id = message.reply_to_message.from_user.id
|
||||
database.DB.create_user_if_not_exists(chat_id, user_id)
|
||||
database.DB.create_user_if_not_exists(bot.id, chat_id, user_id)
|
||||
|
||||
database.DB.user_increment_warnings(chat_id, user_id)
|
||||
user = database.DB.get_user(chat_id, user_id)
|
||||
database.DB.user_increment_warnings(bot.id, chat_id, user_id)
|
||||
user = database.DB.get_user(bot.id, chat_id, user_id)
|
||||
|
||||
user_info = message.reply_to_message.from_user
|
||||
# TODO: родительный падеж имени и фамилии, если возможно
|
||||
|
|
|
|||
|
|
@ -1,9 +1,7 @@
|
|||
from functools import partial
|
||||
from typing import Optional
|
||||
|
||||
from aiogram import Router, F
|
||||
from aiogram import Router, F, Bot
|
||||
from aiogram.types import Message
|
||||
from aiogram.types.user import User
|
||||
from aiogram.enums.content_type import ContentType
|
||||
|
||||
import ai_agent
|
||||
|
|
@ -15,8 +13,6 @@ from tg.utils import get_user_name_for_ai
|
|||
|
||||
router = Router()
|
||||
|
||||
bot_user: Optional[User] = None
|
||||
|
||||
ACCEPTED_CONTENT_TYPES: list[ContentType] = [
|
||||
ContentType.TEXT,
|
||||
ContentType.ANIMATION,
|
||||
|
|
@ -38,9 +34,9 @@ ACCEPTED_CONTENT_TYPES: list[ContentType] = [
|
|||
|
||||
|
||||
@router.message(F.content_type.in_(ACCEPTED_CONTENT_TYPES))
|
||||
async def any_message_handler(message: Message):
|
||||
async def any_message_handler(message: Message, bot: Bot):
|
||||
chat_id = message.chat.id
|
||||
chat = database.DB.create_chat_if_not_exists(chat_id)
|
||||
chat = database.DB.create_chat_if_not_exists(bot.id, chat_id)
|
||||
if chat['active'] == 0:
|
||||
return
|
||||
|
||||
|
|
@ -49,13 +45,11 @@ async def any_message_handler(message: Message):
|
|||
return
|
||||
|
||||
user_id = message.from_user.id
|
||||
database.DB.create_user_if_not_exists(chat_id, user_id)
|
||||
database.DB.user_set_last_message(chat_id, user_id, utils.posix_time())
|
||||
database.DB.user_increment_messages(chat_id, user_id)
|
||||
database.DB.create_user_if_not_exists(bot.id, chat_id, user_id)
|
||||
database.DB.user_set_last_message(bot.id, chat_id, user_id, utils.posix_time())
|
||||
database.DB.user_increment_messages(bot.id, chat_id, user_id)
|
||||
|
||||
global bot_user
|
||||
if bot_user is None:
|
||||
bot_user = await message.bot.get_me()
|
||||
bot_user = await bot.me()
|
||||
|
||||
ai_message = ai_agent.Message()
|
||||
ai_fwd_messages: list[ai_agent.Message] = []
|
||||
|
|
@ -82,7 +76,7 @@ async def any_message_handler(message: Message):
|
|||
await message.reply(MESSAGE_NOT_TEXT)
|
||||
return
|
||||
|
||||
last_id = ai_agent.agent.get_last_assistant_message_id(chat_id)
|
||||
last_id = ai_agent.agent.get_last_assistant_message_id(bot.id, chat_id)
|
||||
if message.reply_to_message.message_id != last_id:
|
||||
# Оригинального сообщения нет в контексте, или оно не последнее
|
||||
if message.content_type == ContentType.TEXT:
|
||||
|
|
@ -99,10 +93,10 @@ async def any_message_handler(message: Message):
|
|||
ai_message.message_id = message.message_id
|
||||
|
||||
answer, success = await utils.run_with_progress(
|
||||
partial(ai_agent.agent.get_group_chat_reply, chat_id, ai_message, ai_fwd_messages),
|
||||
partial(ai_agent.agent.get_group_chat_reply, bot.id, chat_id, ai_message, ai_fwd_messages),
|
||||
partial(message.bot.send_chat_action, chat_id, 'typing'),
|
||||
interval=4)
|
||||
|
||||
answer_id = (await message.reply(answer)).message_id
|
||||
if success:
|
||||
ai_agent.agent.set_last_response_id(chat_id, answer_id)
|
||||
ai_agent.agent.set_last_response_id(bot.id, chat_id, answer_id)
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
from functools import partial
|
||||
|
||||
from aiogram import Router, F
|
||||
from aiogram import Router, F, Bot
|
||||
from aiogram.enums import ChatType, ContentType
|
||||
from aiogram.filters import Command, CommandObject, CommandStart
|
||||
from aiogram.types import Message
|
||||
|
|
@ -16,33 +16,33 @@ router = Router()
|
|||
|
||||
|
||||
@router.message(CommandStart(), F.chat.type == ChatType.PRIVATE)
|
||||
async def start_handler(message: Message):
|
||||
async def start_handler(message: Message, bot: Bot):
|
||||
chat_id = message.chat.id
|
||||
database.DB.create_chat_if_not_exists(chat_id)
|
||||
database.DB.chat_update(chat_id, active=1)
|
||||
database.DB.create_chat_if_not_exists(bot.id, chat_id)
|
||||
database.DB.chat_update(bot.id, chat_id, active=1)
|
||||
await message.answer("Привет!")
|
||||
|
||||
|
||||
@router.message(Command("личность", prefix="!"), F.chat.type == ChatType.PRIVATE)
|
||||
async def set_prompt_handler(message: Message, command: CommandObject):
|
||||
async def set_prompt_handler(message: Message, command: CommandObject, bot: Bot):
|
||||
chat_id = message.chat.id
|
||||
database.DB.create_chat_if_not_exists(chat_id)
|
||||
database.DB.create_chat_if_not_exists(bot.id, chat_id)
|
||||
|
||||
database.DB.chat_update(chat_id, ai_prompt=command.args)
|
||||
database.DB.chat_update(bot.id, chat_id, ai_prompt=command.args)
|
||||
await message.answer("Личность ИИ изменена.")
|
||||
|
||||
|
||||
@router.message(Command("сброс", prefix="!"), F.chat.type == ChatType.PRIVATE)
|
||||
async def reset_context_handler(message: Message):
|
||||
async def reset_context_handler(message: Message, bot: Bot):
|
||||
chat_id = message.chat.id
|
||||
database.DB.create_chat_if_not_exists(chat_id)
|
||||
database.DB.create_chat_if_not_exists(bot.id, chat_id)
|
||||
|
||||
ai_agent.agent.clear_chat_context(chat_id)
|
||||
ai_agent.agent.clear_chat_context(bot.id, chat_id)
|
||||
await message.answer("Контекст очищен.")
|
||||
|
||||
|
||||
@router.message(F.content_type.in_(ACCEPTED_CONTENT_TYPES), F.chat.type == ChatType.PRIVATE)
|
||||
async def any_message_handler(message: Message):
|
||||
async def any_message_handler(message: Message, bot: Bot):
|
||||
chat_id = message.chat.id
|
||||
|
||||
if message.content_type != ContentType.TEXT:
|
||||
|
|
@ -50,10 +50,10 @@ async def any_message_handler(message: Message):
|
|||
return
|
||||
|
||||
answer, success = await utils.run_with_progress(
|
||||
partial(ai_agent.agent.get_private_chat_reply, chat_id, message.text, message.message_id),
|
||||
partial(ai_agent.agent.get_private_chat_reply, bot.id, chat_id, message.text, message.message_id),
|
||||
partial(message.bot.send_chat_action, chat_id, 'typing'),
|
||||
interval=4)
|
||||
|
||||
answer_id = (await message.answer(answer)).message_id
|
||||
if success:
|
||||
ai_agent.agent.set_last_response_id(chat_id, answer_id)
|
||||
ai_agent.agent.set_last_response_id(bot.id, chat_id, answer_id)
|
||||
|
|
|
|||
|
|
@ -36,9 +36,9 @@ def calculate_total_messages(top_users: List[Any]) -> int:
|
|||
|
||||
|
||||
@router.message(F.text == '!помощь')
|
||||
async def help_handler(message: Message):
|
||||
async def help_handler(message: Message, bot: Bot):
|
||||
chat_id = message.chat.id
|
||||
chat = database.DB.create_chat_if_not_exists(chat_id)
|
||||
chat = database.DB.create_chat_if_not_exists(bot.id, chat_id)
|
||||
if chat['active'] == 0:
|
||||
await message.answer(MESSAGE_CHAT_NOT_ACTIVE)
|
||||
return
|
||||
|
|
@ -67,15 +67,15 @@ async def help_handler(message: Message):
|
|||
|
||||
|
||||
@router.message(F.text == "!я")
|
||||
async def about_handler(message: Message):
|
||||
async def about_handler(message: Message, bot: Bot):
|
||||
chat_id = message.chat.id
|
||||
chat = database.DB.create_chat_if_not_exists(chat_id)
|
||||
chat = database.DB.create_chat_if_not_exists(bot.id, chat_id)
|
||||
if chat['active'] == 0:
|
||||
await message.answer(MESSAGE_CHAT_NOT_ACTIVE)
|
||||
return
|
||||
|
||||
target_user = message.from_user
|
||||
user = database.DB.create_user_if_not_exists(chat_id, target_user.id)
|
||||
user = database.DB.create_user_if_not_exists(bot.id, chat_id, target_user.id)
|
||||
if message.reply_to_message is None:
|
||||
if user['about'] is not None:
|
||||
response = Italic(utils.full_name(target_user.first_name, target_user.last_name)) + '\n' + user['about']
|
||||
|
|
@ -84,16 +84,16 @@ async def about_handler(message: Message):
|
|||
await message.answer('Вы не установили визитку.')
|
||||
else:
|
||||
if message.reply_to_message.content_type == ContentType.TEXT:
|
||||
database.DB.user_update(chat_id, target_user.id, about=message.reply_to_message.text)
|
||||
database.DB.user_update(bot.id, chat_id, target_user.id, about=message.reply_to_message.text)
|
||||
await message.answer('Визитка изменена.')
|
||||
else:
|
||||
await message.answer('Извините, но я понимаю только текст.')
|
||||
|
||||
|
||||
@router.message(F.text == "!кто")
|
||||
async def whois_handler(message: Message):
|
||||
async def whois_handler(message: Message, bot: Bot):
|
||||
chat_id = message.chat.id
|
||||
chat = database.DB.create_chat_if_not_exists(chat_id)
|
||||
chat = database.DB.create_chat_if_not_exists(bot.id, chat_id)
|
||||
if chat['active'] == 0:
|
||||
await message.answer(MESSAGE_CHAT_NOT_ACTIVE)
|
||||
return
|
||||
|
|
@ -104,7 +104,7 @@ async def whois_handler(message: Message):
|
|||
|
||||
target_user = message.reply_to_message.from_user
|
||||
full_name = utils.full_name(target_user.first_name, target_user.last_name)
|
||||
user = database.DB.get_user(chat_id, target_user.id)
|
||||
user = database.DB.get_user(bot.id, chat_id, target_user.id)
|
||||
if user is not None and user['about'] is not None:
|
||||
response = Italic(full_name) + '\n' + user['about']
|
||||
await message.answer(**response.as_kwargs())
|
||||
|
|
@ -115,12 +115,12 @@ async def whois_handler(message: Message):
|
|||
@router.message(F.text == "!сегодня")
|
||||
async def stats_today_handler(message: Message, bot: Bot):
|
||||
chat_id = message.chat.id
|
||||
chat = database.DB.create_chat_if_not_exists(chat_id)
|
||||
chat = database.DB.create_chat_if_not_exists(bot.id, chat_id)
|
||||
if chat['active'] == 0:
|
||||
await message.answer(MESSAGE_CHAT_NOT_ACTIVE)
|
||||
return
|
||||
|
||||
top_users = database.DB.get_top_messages_today(chat_id)
|
||||
top_users = database.DB.get_top_messages_today(bot.id, chat_id)
|
||||
if len(top_users) == 0:
|
||||
await message.answer('Сегодня еще никто не писал.')
|
||||
return
|
||||
|
|
@ -134,12 +134,12 @@ async def stats_today_handler(message: Message, bot: Bot):
|
|||
@router.message(F.text == "!месяц")
|
||||
async def stats_month_handler(message: Message, bot: Bot):
|
||||
chat_id = message.chat.id
|
||||
chat = database.DB.create_chat_if_not_exists(chat_id)
|
||||
chat = database.DB.create_chat_if_not_exists(bot.id, chat_id)
|
||||
if chat['active'] == 0:
|
||||
await message.answer(MESSAGE_CHAT_NOT_ACTIVE)
|
||||
return
|
||||
|
||||
top_users = database.DB.get_top_messages_month(chat_id)
|
||||
top_users = database.DB.get_top_messages_month(bot.id, chat_id)
|
||||
if len(top_users) == 0:
|
||||
await message.answer('В этом месяце еще никто не писал.')
|
||||
return
|
||||
|
|
@ -153,12 +153,12 @@ async def stats_month_handler(message: Message, bot: Bot):
|
|||
@router.message(F.text == "!молчуны")
|
||||
async def silent_handler(message: Message, bot: Bot):
|
||||
chat_id = message.chat.id
|
||||
chat = database.DB.create_chat_if_not_exists(chat_id)
|
||||
chat = database.DB.create_chat_if_not_exists(bot.id, chat_id)
|
||||
if chat['active'] == 0:
|
||||
await message.answer(MESSAGE_CHAT_NOT_ACTIVE)
|
||||
return
|
||||
|
||||
top_users = database.DB.get_top_silent(chat_id, 14)
|
||||
top_users = database.DB.get_top_silent(bot.id, chat_id, 14)
|
||||
if len(top_users) == 0:
|
||||
await message.answer('Молчунов нет. Все молодцы!')
|
||||
return
|
||||
|
|
@ -171,12 +171,12 @@ async def silent_handler(message: Message, bot: Bot):
|
|||
@router.message(F.text == "!предупреждения")
|
||||
async def warnings_handler(message: Message, bot: Bot):
|
||||
chat_id = message.chat.id
|
||||
chat = database.DB.create_chat_if_not_exists(chat_id)
|
||||
chat = database.DB.create_chat_if_not_exists(bot.id, chat_id)
|
||||
if chat['active'] == 0:
|
||||
await message.answer(MESSAGE_CHAT_NOT_ACTIVE)
|
||||
return
|
||||
|
||||
top_users = database.DB.get_top_warnings(chat_id)
|
||||
top_users = database.DB.get_top_warnings(bot.id, chat_id)
|
||||
if len(top_users) == 0:
|
||||
await message.answer('Пока все спокойно. Продолжайте в том же духе.')
|
||||
return
|
||||
|
|
@ -189,7 +189,7 @@ async def warnings_handler(message: Message, bot: Bot):
|
|||
@router.message(F.text == "!проверка")
|
||||
async def check_rules_violation_handler(message: Message, bot: Bot):
|
||||
chat_id = message.chat.id
|
||||
chat = database.DB.create_chat_if_not_exists(chat_id)
|
||||
chat = database.DB.create_chat_if_not_exists(bot.id, chat_id)
|
||||
if chat['active'] == 0:
|
||||
await message.answer(MESSAGE_CHAT_NOT_ACTIVE)
|
||||
return
|
||||
|
|
@ -207,14 +207,12 @@ async def check_rules_violation_handler(message: Message, bot: Bot):
|
|||
prompt += chat_rules + '\n\n'
|
||||
prompt += 'Проверь, не нарушают ли правила следующие сообщения (если нарушают, то укажи пункты правил):'
|
||||
|
||||
chat_prompt = chat['ai_prompt']
|
||||
|
||||
ai_message = ai_agent.Message(user_name=await get_user_name_for_ai(message.from_user), text=prompt)
|
||||
ai_fwd_messages = [ai_agent.Message(user_name=await get_user_name_for_ai(message.reply_to_message.from_user),
|
||||
text=message.reply_to_message.text)]
|
||||
|
||||
await message.answer(
|
||||
await utils.run_with_progress(
|
||||
partial(ai_agent.agent.get_group_chat_reply, chat_id, chat_prompt, ai_message, ai_fwd_messages),
|
||||
partial(ai_agent.agent.get_group_chat_reply, bot.id, chat_id, ai_message, ai_fwd_messages),
|
||||
partial(bot.send_chat_action, chat_id, 'typing'),
|
||||
interval=4))
|
||||
|
|
|
|||
54
tg/tasks.py
54
tg/tasks.py
|
|
@ -2,8 +2,9 @@ import datetime
|
|||
import traceback
|
||||
|
||||
from asyncio import sleep
|
||||
|
||||
from aiogram import Bot
|
||||
from aiogram.exceptions import TelegramBadRequest, TelegramMigrateToChat
|
||||
from aiogram.exceptions import TelegramBadRequest
|
||||
from aiogram.types import ChatMemberBanned, ChatMemberLeft
|
||||
from aiogram.utils.formatting import Bold
|
||||
|
||||
|
|
@ -25,51 +26,50 @@ async def is_user_in_chat(bot: Bot, chat_id: int, user_id: int) -> bool:
|
|||
|
||||
|
||||
async def cleanup_chats(bot: Bot):
|
||||
me = await bot.get_me()
|
||||
|
||||
for chat in database.DB.get_chats():
|
||||
chat_id = chat['id']
|
||||
for chat in database.DB.get_chats(bot.id):
|
||||
chat_id = chat['chat_id']
|
||||
if chat_id > 0:
|
||||
# TODO
|
||||
continue
|
||||
|
||||
if not await is_user_in_chat(bot, chat_id, me.id):
|
||||
database.DB.chat_delete(chat_id)
|
||||
if not await is_user_in_chat(bot, chat_id, bot.id):
|
||||
database.DB.chat_delete(bot.id, chat_id)
|
||||
print(f"Чат (bot_id={bot.id}, chat_id={chat_id}) удален.")
|
||||
|
||||
|
||||
async def cleanup_users(bot: Bot):
|
||||
for chat in database.DB.get_chats():
|
||||
for chat in database.DB.get_chats(bot.id):
|
||||
if chat['active'] == 0:
|
||||
continue
|
||||
|
||||
chat_id = chat['id']
|
||||
|
||||
for user in database.DB.get_top_silent(chat_id=chat_id, threshold_days=14):
|
||||
chat_id = chat['chat_id']
|
||||
for user in database.DB.get_top_silent(bot_id=bot.id, chat_id=chat_id, threshold_days=14):
|
||||
user_id = user['user_id']
|
||||
if not await is_user_in_chat(bot, chat_id, user_id):
|
||||
database.DB.delete_user(chat_id, user_id)
|
||||
print(f'Из чата (id={chat_id}) удален пользователь (id={user_id})')
|
||||
database.DB.delete_user(bot.id, chat_id, user_id)
|
||||
print(f'Из чата (bot_id={bot.id}, chat_id={chat_id}) удален пользователь (id={user_id})')
|
||||
|
||||
|
||||
async def reset_counters(reset_month: bool, bot: Bot):
|
||||
database.DB.reset_messages_today()
|
||||
print('Дневные счетчики сброшены.')
|
||||
database.DB.reset_messages_today(bot.id)
|
||||
|
||||
if reset_month:
|
||||
for chat in database.DB.get_chats():
|
||||
if chat['active'] == 0:
|
||||
continue
|
||||
if not reset_month:
|
||||
return
|
||||
|
||||
top_users = database.DB.get_top_messages_month(chat['id'])
|
||||
if len(top_users) == 0:
|
||||
continue
|
||||
for chat in database.DB.get_chats(bot.id):
|
||||
if chat['active'] == 0:
|
||||
continue
|
||||
|
||||
message = Bold('Итоговая статистика за прошедший месяц') + '\n'
|
||||
message += await format_rating(chat['id'], top_users, bot)
|
||||
await bot.send_message(chat_id=chat['id'], **message.as_kwargs())
|
||||
chat_id = chat['chat_id']
|
||||
top_users = database.DB.get_top_messages_month(bot.id, chat['chat_id'])
|
||||
if len(top_users) == 0:
|
||||
continue
|
||||
|
||||
database.DB.reset_messages_month()
|
||||
print('Месячные счетчики сброшены.')
|
||||
message = Bold('Итоговая статистика за прошедший месяц') + '\n'
|
||||
message += await format_rating(chat_id, top_users, bot)
|
||||
await bot.send_message(chat_id=chat_id, **message.as_kwargs())
|
||||
|
||||
database.DB.reset_messages_month(bot.id)
|
||||
|
||||
|
||||
async def wait_until(target_time: datetime.datetime):
|
||||
|
|
|
|||
|
|
@ -6,17 +6,31 @@ class TgDatabase(database.BasicDatabase):
|
|||
super().__init__(connection_string)
|
||||
|
||||
self.cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS chats (
|
||||
CREATE TABLE IF NOT EXISTS bots (
|
||||
id BIGINT NOT NULL,
|
||||
active TINYINT NOT NULL DEFAULT 0,
|
||||
rules VARCHAR(4000),
|
||||
greeting_join VARCHAR(2000),
|
||||
ai_prompt VARCHAR(4000),
|
||||
owner_id BIGINT NOT NULL,
|
||||
api_token VARCHAR(64) NOT NULL,
|
||||
ai_prompt VARCHAR(4000) DEFAULT NULL,
|
||||
group_chats_allowed TINYINT NOT NULL DEFAULT 1,
|
||||
private_chats_allowed TINYINT NOT NULL DEFAULT 1,
|
||||
PRIMARY KEY (id))
|
||||
""")
|
||||
|
||||
self.cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS chats (
|
||||
bot_id BIGINT NOT NULL,
|
||||
chat_id BIGINT NOT NULL,
|
||||
active TINYINT NOT NULL DEFAULT 0,
|
||||
rules VARCHAR(4000),
|
||||
greeting_join VARCHAR(2000),
|
||||
ai_prompt VARCHAR(4000),
|
||||
PRIMARY KEY (bot_id, chat_id),
|
||||
CONSTRAINT fk_chats_bots FOREIGN KEY (bot_id) REFERENCES bots (id) ON DELETE CASCADE ON UPDATE CASCADE)
|
||||
""")
|
||||
|
||||
self.cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS users (
|
||||
bot_id BIGINT NOT NULL,
|
||||
chat_id BIGINT NOT NULL,
|
||||
user_id BIGINT NOT NULL,
|
||||
last_message BIGINT NOT NULL DEFAULT 0,
|
||||
|
|
@ -24,18 +38,19 @@ class TgDatabase(database.BasicDatabase):
|
|||
messages_month SMALLINT NOT NULL DEFAULT 0,
|
||||
warnings TINYINT NOT NULL DEFAULT 0,
|
||||
about VARCHAR(1000),
|
||||
PRIMARY KEY (chat_id, user_id),
|
||||
CONSTRAINT fk_users_chats FOREIGN KEY (chat_id) REFERENCES chats (id) ON UPDATE CASCADE ON DELETE CASCADE)
|
||||
PRIMARY KEY (bot_id, chat_id, user_id),
|
||||
CONSTRAINT fk_users_chats FOREIGN KEY (bot_id, chat_id) REFERENCES chats (bot_id, chat_id) ON UPDATE CASCADE ON DELETE CASCADE)
|
||||
""")
|
||||
|
||||
self.cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS contexts (
|
||||
bot_id BIGINT NOT NULL,
|
||||
chat_id BIGINT NOT NULL,
|
||||
message_id BIGINT,
|
||||
role VARCHAR(16) NOT NULL,
|
||||
content VARCHAR(2000) NOT NULL,
|
||||
UNIQUE KEY contexts_unique (chat_id, message_id),
|
||||
CONSTRAINT fk_contexts_chats FOREIGN KEY (chat_id) REFERENCES chats (id) ON UPDATE CASCADE ON DELETE CASCADE)
|
||||
UNIQUE KEY contexts_unique (bot_id, chat_id, message_id),
|
||||
CONSTRAINT fk_contexts_chats FOREIGN KEY (bot_id, chat_id) REFERENCES chats (bot_id, chat_id) ON UPDATE CASCADE ON DELETE CASCADE)
|
||||
""")
|
||||
|
||||
self.conn.commit()
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
import json
|
||||
|
||||
from vkbottle.bot import Bot as VkBot
|
||||
from vkbottle.bot import Bot, run_multibot
|
||||
|
||||
from ai_agent import create_ai_agent
|
||||
|
||||
|
|
@ -8,6 +8,7 @@ import vk.vk_database as database
|
|||
|
||||
from . import handlers
|
||||
from . import tasks
|
||||
from .utils import MyAPI
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
|
|
@ -15,13 +16,20 @@ if __name__ == '__main__':
|
|||
config = json.load(file)
|
||||
print('Конфигурация загружена.')
|
||||
|
||||
bot = VkBot(config['api_token'], labeler=handlers.labeler)
|
||||
database.create_database(config['db_connection_string'])
|
||||
|
||||
create_ai_agent(config['openrouter_token'],
|
||||
config['openrouter_model'],
|
||||
config['openrouter_model_temp'],
|
||||
database.DB)
|
||||
|
||||
bot.loop_wrapper.on_startup.append(tasks.startup_task(bot.api))
|
||||
bot.loop_wrapper.add_task(tasks.daily_maintenance_task(bot.api))
|
||||
bot.run_forever()
|
||||
bot = Bot(labeler=handlers.labeler)
|
||||
|
||||
apis: list[MyAPI] = []
|
||||
for item in database.DB.get_bots():
|
||||
api = MyAPI(item['id'], item['api_token'])
|
||||
bot.loop_wrapper.on_startup.append(tasks.startup_task(api))
|
||||
bot.loop_wrapper.add_task(tasks.daily_maintenance_task(api))
|
||||
apis.append(api)
|
||||
|
||||
run_multibot(bot, apis)
|
||||
|
|
|
|||
|
|
@ -3,14 +3,16 @@ from vkbottle.framework.labeler import BotLabeler
|
|||
|
||||
from messages import *
|
||||
import vk.vk_database as database
|
||||
from vk.utils import *
|
||||
|
||||
labeler = BotLabeler()
|
||||
|
||||
|
||||
@labeler.chat_message(action=['chat_invite_user', 'chat_invite_user_by_link'])
|
||||
async def user_join_handler(message: Message):
|
||||
bot_id = get_bot_id(message.ctx_api)
|
||||
chat_id = message.peer_id
|
||||
chat = database.DB.create_chat_if_not_exists(chat_id)
|
||||
chat = database.DB.create_chat_if_not_exists(bot_id, chat_id)
|
||||
if chat['active'] == 0:
|
||||
return
|
||||
|
||||
|
|
@ -37,11 +39,12 @@ async def user_join_handler(message: Message):
|
|||
|
||||
@labeler.chat_message(action=['chat_kick_user'])
|
||||
async def user_leave_handler(message: Message):
|
||||
bot_id = get_bot_id(message.ctx_api)
|
||||
chat_id = message.peer_id
|
||||
chat = database.DB.create_chat_if_not_exists(chat_id)
|
||||
chat = database.DB.create_chat_if_not_exists(bot_id, chat_id)
|
||||
if chat['active'] == 0:
|
||||
return
|
||||
|
||||
user_id = message.action.member_id
|
||||
if user_id != message.from_id:
|
||||
database.DB.delete_user(chat_id, user_id)
|
||||
database.DB.delete_user(bot_id, chat_id, user_id)
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ import utils
|
|||
from messages import *
|
||||
|
||||
import vk.vk_database as database
|
||||
from vk.utils import *
|
||||
|
||||
labeler = BotLabeler()
|
||||
|
||||
|
|
@ -22,8 +23,9 @@ def vk_user_is_admin(user_id: int, chat_members: MessagesGetConversationMembers)
|
|||
|
||||
@labeler.chat_message(text="!старт")
|
||||
async def start_handler(message: Message):
|
||||
bot_id = get_bot_id(message.ctx_api)
|
||||
chat_id = message.peer_id
|
||||
database.DB.create_chat_if_not_exists(chat_id)
|
||||
database.DB.create_chat_if_not_exists(bot_id, chat_id)
|
||||
|
||||
chat_members = await message.ctx_api.messages.get_conversation_members(peer_id=chat_id, extended=False)
|
||||
|
||||
|
|
@ -31,21 +33,22 @@ async def start_handler(message: Message):
|
|||
await message.answer(MESSAGE_PERMISSION_DENIED)
|
||||
return
|
||||
|
||||
database.DB.chat_update(chat_id, active=1)
|
||||
database.DB.chat_update(bot_id, chat_id, active=1)
|
||||
|
||||
for member in chat_members.items:
|
||||
# Пропустить ботов
|
||||
if member.member_id < 0:
|
||||
continue
|
||||
database.DB.create_user_if_not_exists(chat_id, member.member_id)
|
||||
database.DB.create_user_if_not_exists(bot_id, chat_id, member.member_id)
|
||||
|
||||
await message.answer('Готова к работе!')
|
||||
|
||||
|
||||
@labeler.chat_message(text="!правила")
|
||||
async def rules_handler(message: Message):
|
||||
bot_id = get_bot_id(message.ctx_api)
|
||||
chat_id = message.peer_id
|
||||
chat = database.DB.create_chat_if_not_exists(chat_id)
|
||||
chat = database.DB.create_chat_if_not_exists(bot_id, chat_id)
|
||||
if chat['active'] == 0:
|
||||
await message.answer(MESSAGE_CHAT_NOT_ACTIVE)
|
||||
return
|
||||
|
|
@ -61,14 +64,15 @@ async def rules_handler(message: Message):
|
|||
await message.answer(MESSAGE_PERMISSION_DENIED)
|
||||
return
|
||||
|
||||
database.DB.chat_update(chat_id, rules=message.reply_message.text)
|
||||
database.DB.chat_update(bot_id, chat_id, rules=message.reply_message.text)
|
||||
await message.answer('Правила чата изменены.')
|
||||
|
||||
|
||||
@labeler.chat_message(text="!приветствие")
|
||||
async def set_greeting_join_handler(message: Message):
|
||||
bot_id = get_bot_id(message.ctx_api)
|
||||
chat_id = message.peer_id
|
||||
chat = database.DB.create_chat_if_not_exists(chat_id)
|
||||
chat = database.DB.create_chat_if_not_exists(bot_id, chat_id)
|
||||
if chat['active'] == 0:
|
||||
await message.answer(MESSAGE_CHAT_NOT_ACTIVE)
|
||||
return
|
||||
|
|
@ -82,14 +86,15 @@ async def set_greeting_join_handler(message: Message):
|
|||
await message.answer(MESSAGE_NEED_REPLY)
|
||||
return
|
||||
|
||||
database.DB.chat_update(chat_id, greeting_join=message.reply_message.text)
|
||||
database.DB.chat_update(bot_id, chat_id, greeting_join=message.reply_message.text)
|
||||
await message.answer('Приветствие изменено.')
|
||||
|
||||
|
||||
@labeler.chat_message(text="!возвращение")
|
||||
async def set_greeting_rejoin_handler(message: Message):
|
||||
bot_id = get_bot_id(message.ctx_api)
|
||||
chat_id = message.peer_id
|
||||
chat = database.DB.create_chat_if_not_exists(chat_id)
|
||||
chat = database.DB.create_chat_if_not_exists(bot_id, chat_id)
|
||||
if chat['active'] == 0:
|
||||
await message.answer(MESSAGE_CHAT_NOT_ACTIVE)
|
||||
return
|
||||
|
|
@ -103,14 +108,15 @@ async def set_greeting_rejoin_handler(message: Message):
|
|||
await message.answer(MESSAGE_NEED_REPLY)
|
||||
return
|
||||
|
||||
database.DB.chat_update(chat_id, greeting_rejoin=message.reply_message.text)
|
||||
database.DB.chat_update(bot_id, chat_id, greeting_rejoin=message.reply_message.text)
|
||||
await message.answer('Приветствие при возвращении изменено.')
|
||||
|
||||
|
||||
@labeler.chat_message(text="!деньрождения")
|
||||
async def set_birthday_handler(message: Message):
|
||||
bot_id = get_bot_id(message.ctx_api)
|
||||
chat_id = message.peer_id
|
||||
chat = database.DB.create_chat_if_not_exists(chat_id)
|
||||
chat = database.DB.create_chat_if_not_exists(bot_id, chat_id)
|
||||
if chat['active'] == 0:
|
||||
await message.answer(MESSAGE_CHAT_NOT_ACTIVE)
|
||||
return
|
||||
|
|
@ -124,14 +130,15 @@ async def set_birthday_handler(message: Message):
|
|||
await message.answer(MESSAGE_NEED_REPLY)
|
||||
return
|
||||
|
||||
database.DB.chat_update(chat_id, birthday_message=message.reply_message.text)
|
||||
database.DB.chat_update(bot_id, chat_id, birthday_message=message.reply_message.text)
|
||||
await message.answer('Уведомление о дне рождения изменено.')
|
||||
|
||||
|
||||
@labeler.chat_message(text="!личность")
|
||||
async def set_ai_prompt_handler(message: Message):
|
||||
bot_id = get_bot_id(message.ctx_api)
|
||||
chat_id = message.peer_id
|
||||
chat = database.DB.create_chat_if_not_exists(chat_id)
|
||||
chat = database.DB.create_chat_if_not_exists(bot_id, chat_id)
|
||||
if chat['active'] == 0:
|
||||
await message.answer(MESSAGE_CHAT_NOT_ACTIVE)
|
||||
return
|
||||
|
|
@ -145,15 +152,16 @@ async def set_ai_prompt_handler(message: Message):
|
|||
await message.answer(MESSAGE_NEED_REPLY)
|
||||
return
|
||||
|
||||
database.DB.chat_update(chat_id, ai_prompt=message.reply_message.text)
|
||||
ai_agent.agent.clear_chat_context(chat_id)
|
||||
database.DB.chat_update(bot_id, chat_id, ai_prompt=message.reply_message.text)
|
||||
ai_agent.agent.clear_chat_context(bot_id, chat_id)
|
||||
await message.answer('Личность ИИ изменена.')
|
||||
|
||||
|
||||
@labeler.chat_message(text="!предупреждение")
|
||||
async def warning_handler(message: Message):
|
||||
bot_id = get_bot_id(message.ctx_api)
|
||||
chat_id = message.peer_id
|
||||
chat = database.DB.create_chat_if_not_exists(chat_id)
|
||||
chat = database.DB.create_chat_if_not_exists(bot_id, chat_id)
|
||||
if chat['active'] == 0:
|
||||
await message.answer(MESSAGE_CHAT_NOT_ACTIVE)
|
||||
return
|
||||
|
|
@ -169,10 +177,10 @@ async def warning_handler(message: Message):
|
|||
return
|
||||
|
||||
user_id = message.reply_message.from_id
|
||||
database.DB.create_user_if_not_exists(chat_id, user_id)
|
||||
database.DB.create_user_if_not_exists(bot_id, chat_id, user_id)
|
||||
|
||||
database.DB.user_increment_warnings(chat_id, user_id)
|
||||
user = database.DB.get_user(chat_id, user_id)
|
||||
database.DB.user_increment_warnings(bot_id, chat_id, user_id)
|
||||
user = database.DB.get_user(bot_id, chat_id, user_id)
|
||||
|
||||
user_info = await message.ctx_api.users.get(user_ids=[user_id], name_case='gen')
|
||||
if len(user_info) == 1:
|
||||
|
|
@ -186,8 +194,9 @@ async def warning_handler(message: Message):
|
|||
|
||||
@labeler.chat_message(text="!исключить")
|
||||
async def ban_handler(message: Message):
|
||||
bot_id = get_bot_id(message.ctx_api)
|
||||
chat_id = message.peer_id
|
||||
chat = database.DB.create_chat_if_not_exists(chat_id)
|
||||
chat = database.DB.create_chat_if_not_exists(bot_id, chat_id)
|
||||
if chat['active'] == 0:
|
||||
await message.answer(MESSAGE_CHAT_NOT_ACTIVE)
|
||||
return
|
||||
|
|
@ -205,10 +214,10 @@ async def ban_handler(message: Message):
|
|||
user_id = message.reply_message.from_id
|
||||
|
||||
try:
|
||||
await message.ctx_api.messages.remove_chat_user(chat_id - 2000000000, member_id=user_id)
|
||||
await message.ctx_api.messages.remove_chat_user(bot_id, chat_id - 2000000000, member_id=user_id)
|
||||
|
||||
if user_id > 0:
|
||||
database.DB.delete_user(chat_id, user_id)
|
||||
database.DB.delete_user(bot_id, chat_id, user_id)
|
||||
|
||||
user_info = await message.ctx_api.users.get(user_ids=[user_id])
|
||||
await message.answer('Участник {} {} исключен.'.format(user_info[0].first_name, user_info[0].last_name))
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@ import utils
|
|||
from messages import *
|
||||
|
||||
import vk.vk_database as database
|
||||
from vk.utils import get_user_name_for_ai
|
||||
from vk.utils import *
|
||||
|
||||
labeler = BotLabeler()
|
||||
|
||||
|
|
@ -21,8 +21,9 @@ bot_user: Optional[GroupsGroup] = None
|
|||
# Обычные сообщения (не команды и не действия)
|
||||
@labeler.chat_message()
|
||||
async def any_message_handler(message: Message):
|
||||
bot_id = get_bot_id(message.ctx_api)
|
||||
chat_id = message.peer_id
|
||||
chat = database.DB.create_chat_if_not_exists(chat_id)
|
||||
chat = database.DB.create_chat_if_not_exists(bot_id, chat_id)
|
||||
if chat['active'] == 0:
|
||||
return
|
||||
|
||||
|
|
@ -35,9 +36,9 @@ async def any_message_handler(message: Message):
|
|||
return
|
||||
|
||||
user_id = message.from_id
|
||||
database.DB.create_user_if_not_exists(chat_id, user_id)
|
||||
database.DB.user_set_last_message(chat_id, user_id, utils.posix_time())
|
||||
database.DB.user_increment_messages(chat_id, user_id)
|
||||
database.DB.create_user_if_not_exists(bot_id, chat_id, user_id)
|
||||
database.DB.user_set_last_message(bot_id, chat_id, user_id, utils.posix_time())
|
||||
database.DB.user_increment_messages(bot_id, chat_id, user_id)
|
||||
|
||||
global bot_user
|
||||
if bot_user is None:
|
||||
|
|
@ -65,7 +66,7 @@ async def any_message_handler(message: Message):
|
|||
await message.reply(MESSAGE_NOT_TEXT)
|
||||
return
|
||||
|
||||
last_id = ai_agent.agent.get_last_assistant_message_id(chat_id)
|
||||
last_id = ai_agent.agent.get_last_assistant_message_id(bot_id, chat_id)
|
||||
if message.reply_message.message_id != last_id:
|
||||
# Оригинального сообщения нет в контексте, или оно не последнее
|
||||
if len(message.reply_message.text) > 0:
|
||||
|
|
@ -83,8 +84,9 @@ async def any_message_handler(message: Message):
|
|||
if message.reply_message:
|
||||
if len(message.reply_message.text) > 0:
|
||||
ai_fwd_messages.append(
|
||||
ai_agent.Message(user_name=await get_user_name_for_ai(message.ctx_api, message.reply_message.from_id),
|
||||
text=message.reply_message.text))
|
||||
ai_agent.Message(
|
||||
user_name=await get_user_name_for_ai(message.ctx_api, message.reply_message.from_id),
|
||||
text=message.reply_message.text))
|
||||
else:
|
||||
await message.reply(MESSAGE_NOT_TEXT)
|
||||
return
|
||||
|
|
@ -102,10 +104,10 @@ async def any_message_handler(message: Message):
|
|||
ai_message.message_id = message.conversation_message_id
|
||||
|
||||
answer, success = await utils.run_with_progress(
|
||||
partial(ai_agent.agent.get_group_chat_reply, chat_id, ai_message, ai_fwd_messages),
|
||||
partial(ai_agent.agent.get_group_chat_reply, bot_id, chat_id, ai_message, ai_fwd_messages),
|
||||
partial(message.ctx_api.messages.set_activity, peer_id=chat_id, type='typing'),
|
||||
interval=4)
|
||||
|
||||
answer_id = (await message.reply(answer)).conversation_message_id
|
||||
if success:
|
||||
ai_agent.agent.set_last_response_id(chat_id, answer_id)
|
||||
ai_agent.agent.set_last_response_id(bot_id, chat_id, answer_id)
|
||||
|
|
|
|||
|
|
@ -9,38 +9,43 @@ import utils
|
|||
from messages import *
|
||||
|
||||
import vk.vk_database as database
|
||||
from vk.utils import *
|
||||
|
||||
labeler = BotLabeler()
|
||||
|
||||
|
||||
@labeler.private_message(text="!старт")
|
||||
async def start_handler(message: Message):
|
||||
bot_id = get_bot_id(message.ctx_api)
|
||||
chat_id = message.peer_id
|
||||
database.DB.create_chat_if_not_exists(chat_id)
|
||||
database.DB.chat_update(chat_id, active=1)
|
||||
database.DB.create_chat_if_not_exists(bot_id, chat_id)
|
||||
database.DB.chat_update(bot_id, chat_id, active=1)
|
||||
await message.answer("Привет!")
|
||||
|
||||
|
||||
@labeler.private_message(RegexRule(r"^!личность ((?:.|\n)+)"))
|
||||
async def set_prompt_handler(message: Message, match):
|
||||
bot_id = get_bot_id(message.ctx_api)
|
||||
chat_id = message.peer_id
|
||||
database.DB.create_chat_if_not_exists(chat_id)
|
||||
database.DB.create_chat_if_not_exists(bot_id, chat_id)
|
||||
|
||||
database.DB.chat_update(chat_id, ai_prompt=match[0])
|
||||
database.DB.chat_update(bot_id, chat_id, ai_prompt=match[0])
|
||||
await message.answer("Личность ИИ изменена.")
|
||||
|
||||
|
||||
@labeler.private_message(text="!сброс")
|
||||
async def reset_context_handler(message: Message):
|
||||
bot_id = get_bot_id(message.ctx_api)
|
||||
chat_id = message.peer_id
|
||||
database.DB.create_chat_if_not_exists(chat_id)
|
||||
database.DB.create_chat_if_not_exists(bot_id, chat_id)
|
||||
|
||||
ai_agent.agent.clear_chat_context(chat_id)
|
||||
ai_agent.agent.clear_chat_context(bot_id, chat_id)
|
||||
await message.answer("Контекст очищен.")
|
||||
|
||||
|
||||
@labeler.private_message()
|
||||
async def any_message_handler(message: Message):
|
||||
bot_id = get_bot_id(message.ctx_api)
|
||||
chat_id = message.peer_id
|
||||
|
||||
if len(message.text) == 0:
|
||||
|
|
@ -48,10 +53,10 @@ async def any_message_handler(message: Message):
|
|||
return
|
||||
|
||||
answer, success = await utils.run_with_progress(
|
||||
partial(ai_agent.agent.get_private_chat_reply, chat_id, message.text, message.message_id),
|
||||
partial(ai_agent.agent.get_private_chat_reply, bot_id, chat_id, message.text, message.message_id),
|
||||
partial(message.ctx_api.messages.set_activity, peer_id=chat_id, type='typing'),
|
||||
interval=4)
|
||||
|
||||
answer_id = (await message.answer(answer)).message_id
|
||||
if success:
|
||||
ai_agent.agent.set_last_response_id(chat_id, answer_id)
|
||||
ai_agent.agent.set_last_response_id(bot_id, chat_id, answer_id)
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
from functools import partial
|
||||
from typing import List, Any
|
||||
|
||||
from vkbottle import bold, italic, API
|
||||
from vkbottle import bold, italic
|
||||
from vkbottle.bot import Message
|
||||
from vkbottle.framework.labeler import BotLabeler
|
||||
|
||||
|
|
@ -10,7 +10,7 @@ import utils
|
|||
from messages import *
|
||||
|
||||
import vk.vk_database as database
|
||||
from vk.utils import get_user_name_for_ai
|
||||
from vk.utils import *
|
||||
|
||||
labeler = BotLabeler()
|
||||
|
||||
|
|
@ -42,8 +42,9 @@ def calculate_total_messages(top_users: List[Any]) -> int:
|
|||
# noinspection SpellCheckingInspection
|
||||
@labeler.chat_message(text="!помощь")
|
||||
async def rules_handler(message: Message):
|
||||
bot_id = get_bot_id(message.ctx_api)
|
||||
chat_id = message.peer_id
|
||||
chat = database.DB.create_chat_if_not_exists(chat_id)
|
||||
chat = database.DB.create_chat_if_not_exists(bot_id, chat_id)
|
||||
if chat['active'] == 0:
|
||||
await message.answer(MESSAGE_CHAT_NOT_ACTIVE)
|
||||
return
|
||||
|
|
@ -77,14 +78,15 @@ async def rules_handler(message: Message):
|
|||
|
||||
@labeler.chat_message(text="!я")
|
||||
async def about_handler(message: Message):
|
||||
bot_id = get_bot_id(message.ctx_api)
|
||||
chat_id = message.peer_id
|
||||
chat = database.DB.create_chat_if_not_exists(chat_id)
|
||||
chat = database.DB.create_chat_if_not_exists(bot_id, chat_id)
|
||||
if chat['active'] == 0:
|
||||
await message.answer(MESSAGE_CHAT_NOT_ACTIVE)
|
||||
return
|
||||
|
||||
user_id = message.from_id
|
||||
user = database.DB.create_user_if_not_exists(chat_id, user_id)
|
||||
user = database.DB.create_user_if_not_exists(bot_id, chat_id, user_id)
|
||||
if message.reply_message is None:
|
||||
info = await message.ctx_api.users.get(user_ids=[user_id])
|
||||
if len(info) == 1:
|
||||
|
|
@ -96,7 +98,7 @@ async def about_handler(message: Message):
|
|||
await message.answer('Не могу получить информацию об участнике.')
|
||||
else:
|
||||
if len(message.reply_message.text) > 0:
|
||||
database.DB.user_update(chat_id, user_id, about=message.reply_message.text)
|
||||
database.DB.user_update(bot_id, chat_id, user_id, about=message.reply_message.text)
|
||||
await message.answer('Визитка изменена.')
|
||||
else:
|
||||
await message.answer('Извините, но я понимаю только текст.')
|
||||
|
|
@ -104,8 +106,9 @@ async def about_handler(message: Message):
|
|||
|
||||
@labeler.chat_message(text="!кто")
|
||||
async def whois_handler(message: Message):
|
||||
bot_id = get_bot_id(message.ctx_api)
|
||||
chat_id = message.peer_id
|
||||
chat = database.DB.create_chat_if_not_exists(chat_id)
|
||||
chat = database.DB.create_chat_if_not_exists(bot_id, chat_id)
|
||||
if chat['active'] == 0:
|
||||
await message.answer(MESSAGE_CHAT_NOT_ACTIVE)
|
||||
return
|
||||
|
|
@ -117,7 +120,7 @@ async def whois_handler(message: Message):
|
|||
target_user_id = message.reply_message.from_id
|
||||
info = await message.ctx_api.users.get(user_ids=[target_user_id])
|
||||
if len(info) == 1:
|
||||
user = database.DB.get_user(chat_id, target_user_id)
|
||||
user = database.DB.get_user(bot_id, chat_id, target_user_id)
|
||||
if user is not None and user['about'] is not None:
|
||||
await message.answer(italic(f'{info[0].first_name} {info[0].last_name}') + '\n' + user['about'])
|
||||
else:
|
||||
|
|
@ -128,13 +131,14 @@ async def whois_handler(message: Message):
|
|||
|
||||
@labeler.chat_message(text="!сегодня")
|
||||
async def stats_today_handler(message: Message):
|
||||
bot_id = get_bot_id(message.ctx_api)
|
||||
chat_id = message.peer_id
|
||||
chat = database.DB.create_chat_if_not_exists(chat_id)
|
||||
chat = database.DB.create_chat_if_not_exists(bot_id, chat_id)
|
||||
if chat['active'] == 0:
|
||||
await message.answer(MESSAGE_CHAT_NOT_ACTIVE)
|
||||
return
|
||||
|
||||
top_users = database.DB.get_top_messages_today(chat_id)
|
||||
top_users = database.DB.get_top_messages_today(bot_id, chat_id)
|
||||
if len(top_users) == 0:
|
||||
await message.answer('Сегодня еще никто не писал.')
|
||||
return
|
||||
|
|
@ -147,13 +151,14 @@ async def stats_today_handler(message: Message):
|
|||
|
||||
@labeler.chat_message(text="!месяц")
|
||||
async def stats_month_handler(message: Message):
|
||||
bot_id = get_bot_id(message.ctx_api)
|
||||
chat_id = message.peer_id
|
||||
chat = database.DB.create_chat_if_not_exists(chat_id)
|
||||
chat = database.DB.create_chat_if_not_exists(bot_id, chat_id)
|
||||
if chat['active'] == 0:
|
||||
await message.answer(MESSAGE_CHAT_NOT_ACTIVE)
|
||||
return
|
||||
|
||||
top_users = database.DB.get_top_messages_month(chat_id)
|
||||
top_users = database.DB.get_top_messages_month(bot_id, chat_id)
|
||||
if len(top_users) == 0:
|
||||
await message.answer('В этом месяце еще никто не писал.')
|
||||
return
|
||||
|
|
@ -166,13 +171,14 @@ async def stats_month_handler(message: Message):
|
|||
|
||||
@labeler.chat_message(text="!молчуны")
|
||||
async def silent_handler(message: Message):
|
||||
bot_id = get_bot_id(message.ctx_api)
|
||||
chat_id = message.peer_id
|
||||
chat = database.DB.create_chat_if_not_exists(chat_id)
|
||||
chat = database.DB.create_chat_if_not_exists(bot_id, chat_id)
|
||||
if chat['active'] == 0:
|
||||
await message.answer(MESSAGE_CHAT_NOT_ACTIVE)
|
||||
return
|
||||
|
||||
top_users = database.DB.get_top_silent(chat_id, 14)
|
||||
top_users = database.DB.get_top_silent(bot_id, chat_id, 14)
|
||||
if len(top_users) == 0:
|
||||
await message.answer('Молчунов нет. Все молодцы!')
|
||||
return
|
||||
|
|
@ -184,13 +190,14 @@ async def silent_handler(message: Message):
|
|||
|
||||
@labeler.chat_message(text="!предупреждения")
|
||||
async def warnings_handler(message: Message):
|
||||
bot_id = get_bot_id(message.ctx_api)
|
||||
chat_id = message.peer_id
|
||||
chat = database.DB.create_chat_if_not_exists(chat_id)
|
||||
chat = database.DB.create_chat_if_not_exists(bot_id, chat_id)
|
||||
if chat['active'] == 0:
|
||||
await message.answer(MESSAGE_CHAT_NOT_ACTIVE)
|
||||
return
|
||||
|
||||
top_users = database.DB.get_top_warnings(chat_id)
|
||||
top_users = database.DB.get_top_warnings(bot_id, chat_id)
|
||||
if len(top_users) == 0:
|
||||
await message.answer('Пока все спокойно. Продолжайте в том же духе.')
|
||||
return
|
||||
|
|
@ -202,16 +209,17 @@ async def warnings_handler(message: Message):
|
|||
|
||||
@labeler.chat_message(text="!поздравление")
|
||||
async def no_birthday_handler(message: Message):
|
||||
bot_id = get_bot_id(message.ctx_api)
|
||||
chat_id = message.peer_id
|
||||
chat = database.DB.create_chat_if_not_exists(chat_id)
|
||||
chat = database.DB.create_chat_if_not_exists(bot_id, chat_id)
|
||||
if chat['active'] == 0:
|
||||
await message.answer(MESSAGE_CHAT_NOT_ACTIVE)
|
||||
return
|
||||
|
||||
user_id = message.from_id
|
||||
user = database.DB.create_user_if_not_exists(chat_id, user_id)
|
||||
user = database.DB.create_user_if_not_exists(bot_id, chat_id, user_id)
|
||||
happy_birthday = 1 if user['happy_birthday'] == 0 else 0
|
||||
database.DB.user_toggle_happy_birthday(chat_id, user_id, happy_birthday)
|
||||
database.DB.user_toggle_happy_birthday(bot_id, chat_id, user_id, happy_birthday)
|
||||
|
||||
if happy_birthday == 1:
|
||||
await message.answer('Хорошо, я буду поздравлять тебя с днем рождения, если его дата не скрыта.')
|
||||
|
|
@ -221,8 +229,9 @@ async def no_birthday_handler(message: Message):
|
|||
|
||||
@labeler.chat_message(text="!проверка")
|
||||
async def check_rules_violation_handler(message: Message):
|
||||
bot_id = get_bot_id(message.ctx_api)
|
||||
chat_id = message.peer_id
|
||||
chat = database.DB.create_chat_if_not_exists(chat_id)
|
||||
chat = database.DB.create_chat_if_not_exists(bot_id, chat_id)
|
||||
if chat['active'] == 0:
|
||||
await message.answer(MESSAGE_CHAT_NOT_ACTIVE)
|
||||
return
|
||||
|
|
@ -236,8 +245,6 @@ async def check_rules_violation_handler(message: Message):
|
|||
prompt += chat_rules + '\n\n'
|
||||
prompt += 'Проверь, не нарушают ли правила следующие сообщения (если нарушают, то укажи пункты правил):'
|
||||
|
||||
chat_prompt = chat['ai_prompt']
|
||||
|
||||
ai_message = ai_agent.Message(user_name=await get_user_name_for_ai(message.ctx_api, message.from_id), text=prompt)
|
||||
ai_fwd_messages: list[ai_agent.Message] = []
|
||||
if message.reply_message is not None and len(message.reply_message.text) > 0:
|
||||
|
|
@ -257,6 +264,6 @@ async def check_rules_violation_handler(message: Message):
|
|||
|
||||
await message.answer(
|
||||
await utils.run_with_progress(
|
||||
partial(ai_agent.agent.get_group_chat_reply, chat_id, chat_prompt, ai_message, ai_fwd_messages),
|
||||
partial(ai_agent.agent.get_group_chat_reply, bot_id, chat_id, ai_message, ai_fwd_messages),
|
||||
partial(message.ctx_api.messages.set_activity, peer_id=chat_id, type='typing'),
|
||||
interval=4))
|
||||
|
|
|
|||
77
vk/tasks.py
77
vk/tasks.py
|
|
@ -2,22 +2,32 @@ import datetime
|
|||
import traceback
|
||||
|
||||
from asyncio import sleep
|
||||
|
||||
from vkbottle import API, bold
|
||||
from vkbottle_types.codegen.objects import UsersFields
|
||||
|
||||
from messages import *
|
||||
|
||||
import vk.vk_database as database
|
||||
from vk.handlers.user import format_rating
|
||||
from vk.utils import *
|
||||
|
||||
|
||||
async def cleanup_chats(api: API):
|
||||
# TODO
|
||||
pass
|
||||
|
||||
|
||||
async def cleanup_users(api: API):
|
||||
for chat in database.DB.get_chats():
|
||||
bot_id = get_bot_id(api)
|
||||
for chat in database.DB.get_chats(bot_id):
|
||||
if chat['active'] == 0:
|
||||
continue
|
||||
|
||||
chat_id = chat['id']
|
||||
chat_id = chat['chat_id']
|
||||
members = await api.messages.get_conversation_members(peer_id=chat_id, extended=False)
|
||||
|
||||
for user in database.DB.get_top_silent(chat_id=chat_id, threshold_days=14):
|
||||
for user in database.DB.get_top_silent(bot_id, chat_id, threshold_days=14):
|
||||
user_id = user['user_id']
|
||||
found = False
|
||||
for item in members.items:
|
||||
|
|
@ -26,48 +36,52 @@ async def cleanup_users(api: API):
|
|||
break
|
||||
|
||||
if not found:
|
||||
database.DB.delete_user(chat_id, user_id)
|
||||
print(f'Из чата (id={chat_id}) удален пользователь (id={user_id})')
|
||||
database.DB.delete_user(bot_id, chat_id, user_id)
|
||||
print(f'Из чата (bot_id={bot_id}, chat_id={chat_id}) удален пользователь (id={user_id})')
|
||||
|
||||
|
||||
async def reset_counters(reset_month: bool, api: API):
|
||||
database.DB.reset_messages_today()
|
||||
print('Дневные счетчики сброшены.')
|
||||
bot_id = get_bot_id(api)
|
||||
database.DB.reset_messages_today(bot_id)
|
||||
|
||||
if reset_month:
|
||||
for chat in database.DB.get_chats():
|
||||
if chat['active'] == 0:
|
||||
continue
|
||||
if not reset_month:
|
||||
return
|
||||
|
||||
top_users = database.DB.get_top_messages_month(chat['id'])
|
||||
if len(top_users) == 0:
|
||||
continue
|
||||
for chat in database.DB.get_chats(bot_id):
|
||||
if chat['active'] == 0:
|
||||
continue
|
||||
|
||||
message = bold('Итоговая статистика за прошедший месяц') + '\n'
|
||||
message += await format_rating(top_users, api)
|
||||
# noinspection PyArgumentList
|
||||
await api.messages.send(random_id=0, peer_id=chat['id'],
|
||||
message=str(message), format_data=message.as_raw_data())
|
||||
chat_id = chat['chat_id']
|
||||
top_users = database.DB.get_top_messages_month(bot_id, chat_id)
|
||||
if len(top_users) == 0:
|
||||
continue
|
||||
|
||||
database.DB.reset_messages_month()
|
||||
print('Месячные счетчики сброшены.')
|
||||
message = bold('Итоговая статистика за прошедший месяц') + '\n'
|
||||
message += await format_rating(top_users, api)
|
||||
# noinspection PyArgumentList
|
||||
await api.messages.send(random_id=0, peer_id=chat_id,
|
||||
message=str(message), format_data=message.as_raw_data())
|
||||
|
||||
database.DB.reset_messages_month(bot_id)
|
||||
|
||||
|
||||
async def check_birthdays(api: API):
|
||||
chats = database.DB.get_chats()
|
||||
bot_id = get_bot_id(api)
|
||||
chats = database.DB.get_chats(bot_id)
|
||||
for chat in chats:
|
||||
if chat['active'] == 0:
|
||||
continue
|
||||
|
||||
chat_id = chat['id']
|
||||
members = await api.messages.get_conversation_members(peer_id=chat_id, extended=False, fields=['bdate'])
|
||||
chat_id = chat['chat_id']
|
||||
members = await api.messages.get_conversation_members(peer_id=chat_id, extended=False,
|
||||
fields=[UsersFields.BDATE])
|
||||
today = datetime.datetime.today()
|
||||
|
||||
for item in members.items:
|
||||
user_id = item.member_id
|
||||
if user_id < 0:
|
||||
continue
|
||||
user = database.DB.create_user_if_not_exists(chat_id, user_id)
|
||||
user = database.DB.create_user_if_not_exists(bot_id, chat_id, user_id)
|
||||
|
||||
for profile in members.profiles:
|
||||
if profile.id == user_id:
|
||||
|
|
@ -111,9 +125,22 @@ async def daily_maintenance_task(api: API):
|
|||
while True:
|
||||
await wait_until(target_datetime)
|
||||
|
||||
try:
|
||||
await cleanup_chats(api)
|
||||
except Exception:
|
||||
print(traceback.format_exc())
|
||||
|
||||
try:
|
||||
await cleanup_users(api)
|
||||
except Exception:
|
||||
print(traceback.format_exc())
|
||||
|
||||
try:
|
||||
await reset_counters(target_datetime.day == 1, api)
|
||||
except Exception:
|
||||
print(traceback.format_exc())
|
||||
|
||||
try:
|
||||
await check_birthdays(api)
|
||||
except Exception:
|
||||
print(traceback.format_exc())
|
||||
|
|
|
|||
11
vk/utils.py
11
vk/utils.py
|
|
@ -1,6 +1,17 @@
|
|||
from vkbottle import API
|
||||
|
||||
|
||||
class MyAPI(API):
|
||||
def __init__(self, bot_id: int, api_token: str):
|
||||
super().__init__(api_token)
|
||||
self.bot_id = bot_id
|
||||
|
||||
|
||||
def get_bot_id(api: API) -> int:
|
||||
my_api: MyAPI = api
|
||||
return my_api.bot_id
|
||||
|
||||
|
||||
async def get_user_name_for_ai(api: API, user_id: int):
|
||||
user = await api.users.get(user_ids=[user_id])
|
||||
if len(user) == 1:
|
||||
|
|
|
|||
|
|
@ -6,19 +6,33 @@ class VkDatabase(database.BasicDatabase):
|
|||
super().__init__(connection_string)
|
||||
|
||||
self.cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS chats (
|
||||
CREATE TABLE IF NOT EXISTS bots (
|
||||
id BIGINT NOT NULL,
|
||||
owner_id BIGINT NOT NULL,
|
||||
api_token VARCHAR(256) NOT NULL,
|
||||
ai_prompt VARCHAR(4000) DEFAULT NULL,
|
||||
group_chats_allowed TINYINT NOT NULL DEFAULT 1,
|
||||
private_chats_allowed TINYINT NOT NULL DEFAULT 1,
|
||||
PRIMARY KEY (id))
|
||||
""")
|
||||
|
||||
self.cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS chats (
|
||||
bot_id BIGINT NOT NULL,
|
||||
chat_id BIGINT NOT NULL,
|
||||
active TINYINT NOT NULL DEFAULT 0,
|
||||
rules VARCHAR(4000),
|
||||
greeting_join VARCHAR(2000),
|
||||
greeting_rejoin VARCHAR(2000),
|
||||
birthday_message VARCHAR(2000),
|
||||
ai_prompt VARCHAR(4000),
|
||||
PRIMARY KEY (id))
|
||||
PRIMARY KEY (bot_id, chat_id),
|
||||
CONSTRAINT fk_chats_bots FOREIGN KEY (bot_id) REFERENCES bots (id) ON DELETE CASCADE ON UPDATE CASCADE)
|
||||
""")
|
||||
|
||||
self.cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS users (
|
||||
bot_id BIGINT NOT NULL,
|
||||
chat_id BIGINT NOT NULL,
|
||||
user_id BIGINT NOT NULL,
|
||||
last_message BIGINT NOT NULL DEFAULT 0,
|
||||
|
|
@ -27,24 +41,25 @@ class VkDatabase(database.BasicDatabase):
|
|||
warnings TINYINT NOT NULL DEFAULT 0,
|
||||
happy_birthday TINYINT NOT NULL DEFAULT 1,
|
||||
about VARCHAR(1000),
|
||||
PRIMARY KEY (chat_id, user_id),
|
||||
CONSTRAINT fk_users_chats FOREIGN KEY (chat_id) REFERENCES chats (id) ON UPDATE CASCADE ON DELETE CASCADE)
|
||||
PRIMARY KEY (bot_id, chat_id, user_id),
|
||||
CONSTRAINT fk_users_chats FOREIGN KEY (bot_id, chat_id) REFERENCES chats (bot_id, chat_id) ON UPDATE CASCADE ON DELETE CASCADE)
|
||||
""")
|
||||
|
||||
self.cursor.execute("""
|
||||
CREATE TABLE IF NOT EXISTS contexts (
|
||||
bot_id BIGINT NOT NULL,
|
||||
chat_id BIGINT NOT NULL,
|
||||
message_id BIGINT,
|
||||
role VARCHAR(16) NOT NULL,
|
||||
content VARCHAR(2000) NOT NULL,
|
||||
UNIQUE KEY contexts_unique (chat_id, message_id),
|
||||
CONSTRAINT fk_contexts_chats FOREIGN KEY (chat_id) REFERENCES chats (id) ON UPDATE CASCADE ON DELETE CASCADE)
|
||||
UNIQUE KEY contexts_unique (bot_id, chat_id, message_id),
|
||||
CONSTRAINT fk_contexts_chats FOREIGN KEY (bot_id, chat_id) REFERENCES chats (bot_id, chat_id) ON UPDATE CASCADE ON DELETE CASCADE)
|
||||
""")
|
||||
|
||||
self.conn.commit()
|
||||
|
||||
def user_toggle_happy_birthday(self, chat_id: int, user_id: int, happy_birthday: int):
|
||||
self.user_update(chat_id, user_id, happy_birthday=happy_birthday)
|
||||
def user_toggle_happy_birthday(self, bot_id: int, chat_id: int, user_id: int, happy_birthday: int):
|
||||
self.user_update(bot_id, chat_id, user_id, happy_birthday=happy_birthday)
|
||||
|
||||
|
||||
DB: VkDatabase
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue