From 01de81cd3f950c07ec2b6fe842442ff6ce293f09 Mon Sep 17 00:00:00 2001 From: Kirill Kirilenko Date: Sat, 24 Jan 2026 01:40:30 +0300 Subject: [PATCH] =?UTF-8?q?=D0=A1=D0=BE=D1=85=D1=80=D0=B0=D0=BD=D0=B5?= =?UTF-8?q?=D0=BD=D0=B8=D0=B5=20=D0=BA=D0=BE=D0=BD=D1=82=D0=B5=D0=BA=D1=81?= =?UTF-8?q?=D1=82=D0=B0=20=D0=98=D0=98=20=D0=B2=20=D0=91=D0=94.=20=D0=98?= =?UTF-8?q?=D0=BC=D1=8F=20=D0=98=D0=98-=D0=BC=D0=BE=D0=B4=D0=B5=D0=BB?= =?UTF-8?q?=D0=B8=20=D0=B8=20=D1=82=D0=B5=D0=BC=D0=BF=D0=B5=D1=80=D0=B0?= =?UTF-8?q?=D1=82=D1=83=D1=80=D0=B0=20=D0=B7=D0=B0=D0=B4=D0=B0=D1=8E=D1=82?= =?UTF-8?q?=D1=81=D1=8F=20=D0=B2=20=D0=BA=D0=BE=D0=BD=D1=84=D0=B8=D0=B3?= =?UTF-8?q?=D0=B5.?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- ai_agent.py | 144 +++++++++++++++++++---------------------- database.py | 56 +++++++++++++++- tg/__main__.py | 10 ++- tg/handlers/default.py | 13 ++-- tg/tg_database.py | 16 ++++- vk/__main__.py | 10 ++- vk/handlers/default.py | 13 ++-- vk/vk_database.py | 16 ++++- 8 files changed, 176 insertions(+), 102 deletions(-) diff --git a/ai_agent.py b/ai_agent.py index 7328cbe..8cf8691 100644 --- a/ai_agent.py +++ b/ai_agent.py @@ -1,9 +1,10 @@ from dataclasses import dataclass -from typing import List, Dict, Optional +from typing import List, Tuple from openrouter import OpenRouter, RetryConfig from openrouter.utils import BackoffStrategy +from database import BasicDatabase GROUP_CHAT_SYSTEM_PROMPT = """ Ты - ИИ-помощник в групповом чате.\n @@ -12,63 +13,80 @@ GROUP_CHAT_SYSTEM_PROMPT = """ Сообщения пользователей будут приходить в следующем формате: '[Имя]: текст сообщения'\n При ответе НЕ нужно указывать пользователя, которому он предназначен. """ +GROUP_CHAT_MAX_MESSAGES = 20 PRIVATE_CHAT_SYSTEM_PROMPT = """ Ты - ИИ-помощник в чате c пользователем.\n Отвечай на вопросы и поддерживай контекст беседы. """ - - -class ChatContext: - def __init__(self, max_messages: int): - self.max_messages: int = max_messages - self.messages: List[Dict[str, str]] = [] - - def add_message(self, role: str, content: str): - if len(self.messages) == self.max_messages: - # Всегда сохраняем в контексте системное сообщение - self.messages.pop(1) - self.messages.append({"role": role, "content": content}) - - def get_messages_for_api(self) -> List[Dict[str, str]]: - return self.messages - - def remove_last_message(self): - self.messages.pop() +PRIVATE_CHAT_MAX_MESSAGES = 40 @dataclass() class Message: user_name: str = None text: str = None + message_id: int = None class AiAgent: - def __init__(self, api_token: str): + def __init__(self, api_token: str, model: str, model_temp: float, db: BasicDatabase): retry_config = RetryConfig(strategy="backoff", backoff=BackoffStrategy( initial_interval=2000, max_interval=8000, exponent=2, max_elapsed_time=14000), retry_connection_errors=True) - self.model = "meta-llama/llama-3.3-70b-instruct:free" + self.db = db + self.model = model + self.model_temp = model_temp self.client = OpenRouter(api_key=api_token, retry_config=retry_config) - self.chat_contexts: Dict[int, ChatContext] = {} - async def get_group_chat_reply(self, chat_id: int, chat_prompt: str, - message: Message, forwarded_messages: List[Message]) -> str: - message_text = message.text + async def get_group_chat_reply(self, 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: message_text += '\n<Цитируемое сообщение от {}>\n'.format(fwd_message.user_name) message_text += fwd_message.text + '\n' message_text += '<Конец цитаты>' - context = self._get_chat_context(is_group_chat=True, chat_id=chat_id, chat_prompt=chat_prompt) - context.add_message(role="user", content=f"[{message.user_name}]: {message_text}") + context = self._get_chat_context(is_group_chat=True, chat_id=chat_id) + context.append({"role": "user", "content": message_text}) try: # Get response from OpenRouter response = await self.client.chat.send_async( model=self.model, - messages=context.get_messages_for_api(), + messages=context, + max_tokens=500, + 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_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, + message_id=None, max_messages=GROUP_CHAT_MAX_MESSAGES) + + return ai_response, True + + except Exception as e: + if str(e).find("Rate limit exceeded") != -1: + return "Извините, достигнут дневной лимит запросов к ИИ (обновляется в 03:00 МСК).", False + else: + 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) + context.append({"role": "user", "content": message}) + + try: + # Get response from OpenRouter + response = await self.client.chat.send_async( + model=self.model, + messages=context, max_tokens=500, temperature=0.5 ) @@ -76,71 +94,41 @@ class AiAgent: # Extract AI response ai_response = response.choices[0].message.content - # Add AI response to context - context.add_message(role="assistant", content=ai_response) + # Add message and AI response to context + self.db.context_add_message(chat_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, + message_id=None, max_messages=PRIVATE_CHAT_MAX_MESSAGES) - return ai_response + return ai_response, True except Exception as e: - context.remove_last_message() if str(e).find("Rate limit exceeded") != -1: - return "Извините, достигнут дневной лимит запросов к ИИ (обновляется в 03:00 МСК)." + return "Извините, достигнут дневной лимит запросов к ИИ (обновляется в 03:00 МСК).", False else: print(f"Ошибка выполнения запроса к ИИ: {e}") - return f"Извините, при обработке запроса произошла ошибка." + return f"Извините, при обработке запроса произошла ошибка.", False - async def get_private_chat_reply(self, chat_id: int, chat_prompt: str, message: str) -> str: - context = self._get_chat_context(is_group_chat=False, chat_id=chat_id, chat_prompt=chat_prompt) - context.add_message(role="user", content=message) - - try: - # Get response from OpenRouter - response = await self.client.chat.send_async( - model=self.model, - messages=context.get_messages_for_api(), - max_tokens=500, - temperature=0.5 - ) - - # Extract AI response - ai_response = response.choices[0].message.content - - # Add AI response to context - context.add_message(role="assistant", content=ai_response) - - return ai_response - - except Exception as e: - context.remove_last_message() - if str(e).find("Rate limit exceeded") != -1: - return "Извините, достигнут дневной лимит запросов к ИИ (обновляется в 03:00 МСК)." - else: - print(f"Ошибка выполнения запроса к ИИ: {e}") - return f"Извините, при обработке запроса произошла ошибка." + def set_last_response_id(self, chat_id: int, message_id: int): + self.db.context_set_last_message_id(chat_id, message_id) def clear_chat_context(self, chat_id: int): - self.chat_contexts.pop(chat_id, None) + self.db.context_clear(chat_id) - def _get_chat_context(self, is_group_chat: bool, chat_id: int, chat_prompt: Optional[str]) -> ChatContext: - """Get or create chat context for a specific chat""" - if chat_id not in self.chat_contexts: - if is_group_chat: - self.chat_contexts[chat_id] = ChatContext(max_messages=20) - prompt = GROUP_CHAT_SYSTEM_PROMPT - else: - self.chat_contexts[chat_id] = ChatContext(max_messages=40) - prompt = PRIVATE_CHAT_SYSTEM_PROMPT + def _get_chat_context(self, is_group_chat: bool, chat_id: int) -> list[dict]: + prompt = GROUP_CHAT_SYSTEM_PROMPT if is_group_chat else PRIVATE_CHAT_SYSTEM_PROMPT - if chat_prompt is not None: - prompt += '\n\n' + chat_prompt + chat = self.db.create_chat_if_not_exists(chat_id) + if chat['ai_prompt'] is not None: + prompt += '\n\n' + chat['ai_prompt'] - self.chat_contexts[chat_id].add_message(role="system", content=prompt) - return self.chat_contexts[chat_id] + messages = self.db.context_get_messages(chat_id) + return [{"role": "system", "content": prompt}] + messages agent: AiAgent -def create_ai_agent(api_token: str): +def create_ai_agent(api_token: str, model: str, model_temp: float, db: BasicDatabase): global agent - agent = AiAgent(api_token) + agent = AiAgent(api_token, model, model_temp, db) diff --git a/database.py b/database.py index 3177ba4..7a6fcf0 100644 --- a/database.py +++ b/database.py @@ -1,5 +1,5 @@ from datetime import datetime -from typing import List, Union +from typing import List, Optional, Union from pyodbc import connect, SQL_CHAR, SQL_WCHAR, Row @@ -106,6 +106,58 @@ class BasicDatabase: def reset_messages_month(self): self.cursor.execute("UPDATE users SET messages_month = 0") + def context_get_messages(self, chat_id: int) -> list[dict]: + self.cursor.execute(""" + SELECT role, content FROM contexts + WHERE chat_id = ? AND message_id IS NOT NULL + ORDER BY message_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) + return self.cursor.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) + + 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) + else: + self.cursor.execute(""" + INSERT INTO contexts (chat_id, role, content) + VALUES (?, ?, ?) + """, 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_trim(self, chat_id: int, max_messages: int): + current_count = self.context_get_count(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 + ORDER BY message_id ASC + LIMIT 1 + """, chat_id).fetchval() + + if oldest_message_id: + self.cursor.execute("DELETE FROM contexts WHERE chat_id = ? AND message_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 create_chat_if_not_exists(self, chat_id: int): chat = self.get_chat(chat_id) if chat is None: @@ -130,7 +182,7 @@ class BasicDatabase: result[column] = args[i] return result elif isinstance(args, list) and all(isinstance(item, Row) for item in args): - results = [] + results: list[dict] = [] for row in args: row_dict = {} for i, column in enumerate(columns): diff --git a/tg/__main__.py b/tg/__main__.py index 5c96ff9..7311ecd 100644 --- a/tg/__main__.py +++ b/tg/__main__.py @@ -4,7 +4,8 @@ import json from aiogram import Bot, Dispatcher from ai_agent import create_ai_agent -from tg.tg_database import create_database + +import tg.tg_database as database from . import handlers from . import tasks @@ -16,8 +17,11 @@ async def main() -> None: print('Конфигурация загружена.') bot = Bot(token=config['api_token']) - create_database(config['db_connection_string']) - create_ai_agent(config['openrouter_token']) + database.create_database(config['db_connection_string']) + create_ai_agent(config['openrouter_token'], + config['openrouter_model'], + config['openrouter_model_temp'], + database.DB) dp = Dispatcher() dp.include_router(handlers.router) diff --git a/tg/handlers/default.py b/tg/handlers/default.py index 471261d..0b7f0b3 100644 --- a/tg/handlers/default.py +++ b/tg/handlers/default.py @@ -85,10 +85,13 @@ async def any_message_handler(message: Message): return ai_message.user_name = await get_user_name_for_ai(message.from_user) - chat_prompt = chat['ai_prompt'] + ai_message.message_id = message.message_id - await message.reply( - await utils.run_with_progress( - partial(ai_agent.agent.get_group_chat_reply, chat_id, chat_prompt, ai_message, ai_fwd_messages), + answer, success = await utils.run_with_progress( + partial(ai_agent.agent.get_group_chat_reply, chat_id, ai_message, ai_fwd_messages), partial(message.bot.send_chat_action, chat_id, 'typing'), - interval=4)) + interval=4) + + answer_id = (await message.reply(answer)).message_id + if success: + ai_agent.agent.set_last_response_id(chat_id, answer_id) diff --git a/tg/tg_database.py b/tg/tg_database.py index 2f3bc03..832d03f 100644 --- a/tg/tg_database.py +++ b/tg/tg_database.py @@ -7,7 +7,7 @@ class TgDatabase(database.BasicDatabase): self.cursor.execute(""" CREATE TABLE IF NOT EXISTS chats ( - id BIGINT, + id BIGINT NOT NULL, active TINYINT NOT NULL DEFAULT 0, rules VARCHAR(4000), greeting_join VARCHAR(2000), @@ -17,8 +17,8 @@ class TgDatabase(database.BasicDatabase): self.cursor.execute(""" CREATE TABLE IF NOT EXISTS users ( - chat_id BIGINT, - user_id BIGINT, + chat_id BIGINT NOT NULL, + user_id BIGINT NOT NULL, last_message BIGINT NOT NULL DEFAULT 0, messages_today SMALLINT NOT NULL DEFAULT 0, messages_month SMALLINT NOT NULL DEFAULT 0, @@ -28,6 +28,16 @@ class TgDatabase(database.BasicDatabase): CONSTRAINT fk_users_chats FOREIGN KEY (chat_id) REFERENCES chats (id) ON UPDATE CASCADE ON DELETE CASCADE) """) + self.cursor.execute(""" + CREATE TABLE IF NOT EXISTS contexts ( + 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) + """) + self.conn.commit() diff --git a/vk/__main__.py b/vk/__main__.py index dccf44e..a01373e 100644 --- a/vk/__main__.py +++ b/vk/__main__.py @@ -3,7 +3,8 @@ import json from vkbottle.bot import Bot as VkBot from ai_agent import create_ai_agent -from vk.vk_database import create_database + +import vk.vk_database as database from . import handlers from . import tasks @@ -15,8 +16,11 @@ if __name__ == '__main__': print('Конфигурация загружена.') bot = VkBot(config['api_token'], labeler=handlers.labeler) - create_database(config['db_connection_string']) - create_ai_agent(config["openrouter_token"]) + 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)) diff --git a/vk/handlers/default.py b/vk/handlers/default.py index 7f9a524..b126f06 100644 --- a/vk/handlers/default.py +++ b/vk/handlers/default.py @@ -87,10 +87,13 @@ async def any_message_handler(message: Message): return ai_message.user_name = await get_user_name_for_ai(message.ctx_api, message.from_id) - chat_prompt = chat['ai_prompt'] + ai_message.message_id = message.conversation_message_id - await message.reply( - await utils.run_with_progress( - partial(ai_agent.agent.get_group_chat_reply, chat_id, chat_prompt, ai_message, ai_fwd_messages), + answer, success = await utils.run_with_progress( + partial(ai_agent.agent.get_group_chat_reply, chat_id, ai_message, ai_fwd_messages), partial(message.ctx_api.messages.set_activity, peer_id=chat_id, type='typing'), - interval=4)) + interval=4) + + answer_id = (await message.reply(answer)).conversation_message_id + if success: + ai_agent.agent.set_last_response_id(chat_id, answer_id) diff --git a/vk/vk_database.py b/vk/vk_database.py index 2a41407..ec9ea59 100644 --- a/vk/vk_database.py +++ b/vk/vk_database.py @@ -7,7 +7,7 @@ class VkDatabase(database.BasicDatabase): self.cursor.execute(""" CREATE TABLE IF NOT EXISTS chats ( - id BIGINT, + id BIGINT NOT NULL, active TINYINT NOT NULL DEFAULT 0, rules VARCHAR(4000), greeting_join VARCHAR(2000), @@ -19,8 +19,8 @@ class VkDatabase(database.BasicDatabase): self.cursor.execute(""" CREATE TABLE IF NOT EXISTS users ( - chat_id BIGINT, - user_id BIGINT, + chat_id BIGINT NOT NULL, + user_id BIGINT NOT NULL, last_message BIGINT NOT NULL DEFAULT 0, messages_today SMALLINT NOT NULL DEFAULT 0, messages_month SMALLINT NOT NULL DEFAULT 0, @@ -31,6 +31,16 @@ class VkDatabase(database.BasicDatabase): CONSTRAINT fk_users_chats FOREIGN KEY (chat_id) REFERENCES chats (id) ON UPDATE CASCADE ON DELETE CASCADE) """) + self.cursor.execute(""" + CREATE TABLE IF NOT EXISTS contexts ( + 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) + """) + self.conn.commit() def user_toggle_happy_birthday(self, chat_id: int, user_id: int, happy_birthday: int):