kiriru/vk_chat_bot
1
0
Fork 0
Code Issues 4 Pull requests Projects Releases Packages Activity Actions

Переход с pyodbc на mariadb.

Browse source
This commit is contained in:
Kirill Kirilenko 2026-06-14 00:02:29 +03:00
parent 56ab157a96
commit 5b9f5cd1d6
6 changed files with 284 additions and 242 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

296
database.py
View file

@ -1,55 +1,76 @@
from datetime import datetime from datetime import datetime
from typing import Dict, List, Optional from typing import Dict, List, Optional, Any, Type
from pyodbc import connect, SQL_CHAR, SQL_WCHAR, Row from mariadb import create_pool
class BasicDatabase: class BasicDatabase:
def __init__(self, connection_string: str): def __init__(self, hostname: str, user: str, password: str, dbname: str):
self.conn = connect(connection_string, autocommit=True) self.hostname = hostname
self.conn.setdecoding(SQL_CHAR, encoding='utf-8') self.user = user
self.conn.setdecoding(SQL_WCHAR, encoding='utf-8') self.password = password
self.conn.setencoding(encoding='utf-8') self.dbname = dbname
self.cursor = self.conn.cursor() self.pool = create_pool(host=hostname, user=user, password=password, database=dbname,
autocommit=True, dictionary=True, min_size=1, max_size=5)
def get_bots(self): def get_bots(self) -> List[Dict[str, Any]]:
self.cursor.execute("SELECT * FROM bots") with self.pool.acquire() as conn:
return self._to_list(self.cursor.fetchall()) with conn.cursor() as cursor:
cursor.execute("SELECT * FROM bots")
return cursor.fetchall()
def get_bot(self, bot_id: int): def get_bot(self, bot_id: int) -> Optional[Dict[str, Any]]:
self.cursor.execute("SELECT * FROM bots WHERE id = ?", bot_id) with self.pool.acquire() as conn:
return self._to_dict(self.cursor.fetchone()) with conn.cursor() as cursor:
cursor.execute("SELECT * FROM bots WHERE id = ?", (bot_id,))
return cursor.fetchone()
def get_chats(self, bot_id: int): def get_chats(self, bot_id: int) -> List[Dict[str, Any]]:
self.cursor.execute("SELECT * FROM chats WHERE bot_id = ?", bot_id) with self.pool.acquire() as conn:
return self._to_list(self.cursor.fetchall()) with conn.cursor() as cursor:
cursor.execute("SELECT * FROM chats WHERE bot_id = ?", (bot_id,))
return cursor.fetchall()
def get_chat(self, bot_id: int, chat_id: int): def get_chat(self, bot_id: int, chat_id: int) -> Optional[Dict[str, Any]]:
self.cursor.execute("SELECT * FROM chats WHERE bot_id = ? AND chat_id = ?", bot_id, chat_id) with self.pool.acquire() as conn:
return self._to_dict(self.cursor.fetchone()) with conn.cursor() as cursor:
cursor.execute("SELECT * FROM chats WHERE bot_id = ? AND chat_id = ?", (bot_id, chat_id))
return cursor.fetchone()
def add_chat(self, bot_id: int, chat_id: int): 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) with self.pool.acquire() as conn:
with conn.cursor() as cursor:
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): def chat_update(self, bot_id: int, chat_id: int, **kwargs):
self.cursor.execute("UPDATE chats SET " + ", ".join(f + " = ?" for f in kwargs) + with self.pool.acquire() as conn:
" WHERE bot_id = ? AND chat_id = ?", list(kwargs.values()) + [bot_id, chat_id]) with conn.cursor() as cursor:
cursor.execute("UPDATE chats SET " + ", ".join(f + " = ?" for f in kwargs) +
" WHERE bot_id = ? AND chat_id = ?", list(kwargs.values()) + [bot_id, chat_id])
def chat_delete(self, bot_id: int, chat_id: int): 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) with self.pool.acquire() as conn:
with conn.cursor() as cursor:
cursor.execute("DELETE FROM chats WHERE bot_id = ? AND chat_id = ?", (bot_id, chat_id))
def get_user(self, bot_id: int, chat_id: int, user_id: int): def get_user(self, bot_id: int, chat_id: int, user_id: int) -> Optional[Dict[str, Any]]:
self.cursor.execute("SELECT * FROM users WHERE bot_id = ? AND chat_id = ? AND user_id = ?", with self.pool.acquire() as conn:
bot_id, chat_id, user_id) with conn.cursor() as cursor:
return self._to_dict(self.cursor.fetchone()) cursor.execute("SELECT * FROM users WHERE bot_id = ? AND chat_id = ? AND user_id = ?",
(bot_id, chat_id, user_id))
return cursor.fetchone()
def get_users(self, bot_id: int, chat_id: int): def get_users(self, bot_id: int, chat_id: int) -> List[Dict[str, Any]]:
self.cursor.execute("SELECT * FROM users WHERE bot_id = ? AND chat_id = ?", bot_id, chat_id) with self.pool.acquire() as conn:
return self._to_list(self.cursor.fetchall()) with conn.cursor() as cursor:
cursor.execute("SELECT * FROM users WHERE bot_id = ? AND chat_id = ?", (bot_id, chat_id))
return cursor.fetchall()
def add_user(self, bot_id: int, chat_id: int, user_id: int): 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 (?, ?, ?)", with self.pool.acquire() as conn:
bot_id, chat_id, user_id) with conn.cursor() as cursor:
cursor.execute("INSERT INTO users (bot_id, chat_id, user_id) VALUES (?, ?, ?)",
(bot_id, chat_id, user_id))
def user_set_last_message(self, bot_id: int, chat_id: int, user_id: int, last_message: int): 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) self.user_update(bot_id, chat_id, user_id, last_message=last_message)
@ -61,83 +82,107 @@ class BasicDatabase:
self.user_increment(bot_id, chat_id, user_id, ['warnings']) self.user_increment(bot_id, chat_id, user_id, ['warnings'])
def user_increment(self, bot_id: int, 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) + with self.pool.acquire() as conn:
" WHERE bot_id = ? AND chat_id = ? AND user_id = ?", bot_id, chat_id, user_id) with conn.cursor() as cursor:
cursor.execute("UPDATE users SET " + ", ".join(f + " = " + f + " + 1" for f in fields) +
" WHERE bot_id = ? AND chat_id = ? AND user_id = ?", (bot_id, chat_id, user_id))
def user_update(self, bot_id, chat_id: int, user_id: int, **kwargs): def user_update(self, bot_id: int, chat_id: int, user_id: int, **kwargs):
self.cursor.execute("UPDATE users SET " + ", ".join(f + " = ?" for f in kwargs) + with self.pool.acquire() as conn:
" WHERE bot_id = ? AND chat_id = ? AND user_id = ?", with conn.cursor() as cursor:
list(kwargs.values()) + [bot_id, chat_id, user_id]) cursor.execute("UPDATE users SET " + ", ".join(f + " = ?" for f in kwargs) +
" WHERE bot_id = ? AND chat_id = ? AND user_id = ?",
list(kwargs.values()) + [bot_id, chat_id, user_id])
def delete_user(self, bot_id: int, chat_id: int, user_id: int): 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 = ?", with self.pool.acquire() as conn:
bot_id, chat_id, user_id) with conn.cursor() as cursor:
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, bot_id: int, chat_id: int): def get_top_messages_today(self, bot_id: int, chat_id: int) -> List[Dict[str, Any]]:
self.cursor.execute(""" with self.pool.acquire() as conn:
SELECT user_id, messages_today AS value FROM users with conn.cursor() as cursor:
WHERE bot_id = ? AND chat_id = ? AND messages_today > 0 cursor.execute("""
ORDER BY messages_today DESC SELECT user_id, messages_today AS value FROM users
""", bot_id, chat_id) WHERE bot_id = ? AND chat_id = ? AND messages_today > 0
return self._to_list(self.cursor.fetchall()) ORDER BY messages_today DESC
""", (bot_id, chat_id))
return cursor.fetchall()
def get_top_messages_month(self, bot_id: int, chat_id: int): def get_top_messages_month(self, bot_id: int, chat_id: int) -> List[Dict[str, Any]]:
self.cursor.execute(""" with self.pool.acquire() as conn:
SELECT user_id, messages_month AS value FROM users with conn.cursor() as cursor:
WHERE bot_id = ? AND chat_id = ? AND messages_month > 0 cursor.execute("""
ORDER BY messages_month DESC SELECT user_id, messages_month AS value FROM users
""", bot_id, chat_id) WHERE bot_id = ? AND chat_id = ? AND messages_month > 0
return self._to_list(self.cursor.fetchall()) ORDER BY messages_month DESC
""", (bot_id, chat_id))
return cursor.fetchall()
def get_top_silent(self, bot_id: int, chat_id: int, threshold_days: int): def get_top_silent(self, bot_id: int, chat_id: int, threshold_days: int) -> List[Dict[str, Any]]:
current_time = int(datetime.now().timestamp()) with self.pool.acquire() as conn:
threshold = current_time - threshold_days * 86400 with conn.cursor() as cursor:
self.cursor.execute(""" current_time = int(datetime.now().timestamp())
SELECT user_id, (? - last_message) DIV 86400 as value threshold = current_time - threshold_days * 86400
FROM users cursor.execute("""
WHERE bot_id = ? AND chat_id = ? AND last_message <= ? SELECT user_id, (? - last_message) DIV 86400 as value
ORDER BY last_message ASC FROM users
""", current_time, bot_id, chat_id, threshold) WHERE bot_id = ? AND chat_id = ? AND last_message <= ?
result = self._to_list(self.cursor.fetchall()) ORDER BY last_message ASC
for row in result: """, (current_time, bot_id, chat_id, threshold))
if row['value'] > 3650: result = cursor.fetchall()
row['value'] = 'никогда' for row in result:
return result if row['value'] > 3650:
row['value'] = 'никогда'
return result
def get_top_warnings(self, bot_id: int, chat_id: int): def get_top_warnings(self, bot_id: int, chat_id: int) -> List[Dict[str, Any]]:
self.cursor.execute(""" with self.pool.acquire() as conn:
SELECT user_id, warnings AS value FROM users with conn.cursor() as cursor:
WHERE bot_id = ? AND chat_id = ? AND warnings > 0 cursor.execute("""
ORDER BY warnings DESC SELECT user_id, warnings AS value FROM users
""", bot_id, chat_id) WHERE bot_id = ? AND chat_id = ? AND warnings > 0
return self._to_list(self.cursor.fetchall()) ORDER BY warnings DESC
""", (bot_id, chat_id))
return cursor.fetchall()
def reset_messages_today(self, bot_id: int): def reset_messages_today(self, bot_id: int):
self.cursor.execute("UPDATE users SET messages_today = 0 WHERE bot_id = ?", bot_id) with self.pool.acquire() as conn:
with conn.cursor() as cursor:
cursor.execute("UPDATE users SET messages_today = 0 WHERE bot_id = ?", (bot_id,))
def reset_messages_month(self, bot_id: int): def reset_messages_month(self, bot_id: int):
self.cursor.execute("UPDATE users SET messages_month = 0 WHERE bot_id = ?", bot_id) with self.pool.acquire() as conn:
with conn.cursor() as cursor:
cursor.execute("UPDATE users SET messages_month = 0 WHERE bot_id = ?", (bot_id,))
def context_get_messages(self, bot_id: int, chat_id: int) -> List[Dict]: def context_get_messages(self, bot_id: int, chat_id: int) -> List[Dict[str, Any]]:
self.cursor.execute(""" with self.pool.acquire() as conn:
SELECT role, text, image FROM contexts with conn.cursor() as cursor:
WHERE bot_id = ? AND chat_id = ? cursor.execute("""
ORDER BY id SELECT role, text, image FROM contexts
""", bot_id, chat_id) WHERE bot_id = ? AND chat_id = ?
result = self._to_list(self.cursor.fetchall()) ORDER BY id
return result """, (bot_id, chat_id))
return cursor.fetchall()
def context_get_count(self, bot_id: int, chat_id: int) -> int: 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) with self.pool.acquire() as conn:
return self.cursor.fetchval() with conn.cursor() as cursor:
cursor.execute("SELECT COUNT(*) FROM contexts WHERE bot_id = ? AND chat_id = ?", (bot_id, chat_id))
return _to_val(int, cursor.fetchone())
def context_get_last_assistant_message_id(self, bot_id: int, 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(""" with self.pool.acquire() as conn:
SELECT message_id FROM contexts with conn.cursor() as cursor:
WHERE bot_id = ? AND chat_id = ? AND role = 'assistant' cursor.execute("""
ORDER BY id DESC SELECT message_id FROM contexts
LIMIT 1 WHERE bot_id = ? AND chat_id = ? AND role = 'assistant'
""", bot_id, chat_id).fetchval() ORDER BY id DESC
LIMIT 1
""", (bot_id, chat_id))
return _to_val(int, cursor.fetchone())
def context_add_message(self, bot_id: int, chat_id: int, role: str, def context_add_message(self, bot_id: int, chat_id: int, role: str,
text: Optional[str], image: Optional[bytes], text: Optional[str], image: Optional[bytes],
@ -158,60 +203,55 @@ class BasicDatabase:
placeholders = ', '.join(['?' for _ in columns]) placeholders = ', '.join(['?' for _ in columns])
values = tuple(data[k] for k in columns) values = tuple(data[k] for k in columns)
query = f"INSERT INTO contexts ({', '.join(columns)}) VALUES ({placeholders})" with self.pool.acquire() as conn:
self.cursor.execute(query, values) with conn.cursor() as cursor:
query = f"INSERT INTO contexts ({', '.join(columns)}) VALUES ({placeholders})"
cursor.execute(query, values)
def context_set_last_message_id(self, bot_id: int, chat_id: int, message_id: int): def context_set_last_message_id(self, bot_id: int, chat_id: int, message_id: int):
self.cursor.execute( with self.pool.acquire() as conn:
"UPDATE contexts SET message_id = ? WHERE bot_id = ? AND chat_id = ? AND message_id IS NULL", with conn.cursor() as cursor:
message_id, bot_id, chat_id) 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, bot_id: int, chat_id: int, max_messages: int): def _context_trim(self, bot_id: int, chat_id: int, max_messages: int):
current_count = self.context_get_count(bot_id, chat_id) current_count = self.context_get_count(bot_id, chat_id)
while current_count >= max_messages: while current_count >= max_messages:
oldest_message_id = self.cursor.execute( with self.pool.acquire() as conn:
"SELECT id FROM contexts WHERE bot_id = ? AND chat_id = ? ORDER BY id LIMIT 1", with conn.cursor() as cursor:
bot_id, chat_id).fetchval() cursor.execute(
"SELECT id FROM contexts WHERE bot_id = ? AND chat_id = ? ORDER BY id LIMIT 1",
(bot_id, chat_id))
oldest_message_id = _to_val(int, cursor.fetchone())
if oldest_message_id: if oldest_message_id:
self.cursor.execute("DELETE FROM contexts WHERE id = ?", oldest_message_id) cursor.execute("DELETE FROM contexts WHERE id = ?", oldest_message_id)
current_count -= 1 current_count -= 1
else: else:
break break
def context_clear(self, bot_id: int, chat_id: int): 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) with self.pool.acquire() as conn:
with conn.cursor() as cursor:
cursor.execute("DELETE FROM contexts WHERE bot_id = ? AND chat_id = ?", (bot_id, chat_id))
def create_chat_if_not_exists(self, bot_id: int, chat_id: int): def create_chat_if_not_exists(self, bot_id: int, chat_id: int) -> Dict[str, Any]:
chat = self.get_chat(bot_id, chat_id) chat = self.get_chat(bot_id, chat_id)
if chat is None: if chat is None:
self.add_chat(bot_id, chat_id) self.add_chat(bot_id, chat_id)
chat = self.get_chat(bot_id, chat_id) chat = self.get_chat(bot_id, chat_id)
assert (chat is not None)
return chat return chat
def create_user_if_not_exists(self, bot_id: int, chat_id: int, user_id: int): def create_user_if_not_exists(self, bot_id: int, chat_id: int, user_id: int) -> Dict[str, Any]:
user = self.get_user(bot_id, chat_id, user_id) user = self.get_user(bot_id, chat_id, user_id)
if user is None: if user is None:
self.add_user(bot_id, chat_id, user_id) self.add_user(bot_id, chat_id, user_id)
user = self.get_user(bot_id, chat_id, user_id) user = self.get_user(bot_id, chat_id, user_id)
assert (user is not None)
return user return user
def _to_dict(self, args: Optional[Row]) -> Optional[Dict]:
columns = [column[0] for column in self.cursor.description]
if args is not None:
result = {}
for i, column in enumerate(columns):
result[column] = args[i]
return result
else:
return None
def _to_list(self, args: List[Row]) -> List[Dict]: def _to_val[T](_: Type[T], arg: Any) -> T:
columns = [column[0] for column in self.cursor.description] return next(iter(arg.values())) if arg is not None else None
results: list[dict] = []
for row in args:
row_dict = {}
for i, column in enumerate(columns):
row_dict[column] = row[i]
results.append(row_dict)
return results

2
requirements.txt
View file

@ -2,7 +2,7 @@ aiogram~=3.28.2
aiohttp~=3.13.5 aiohttp~=3.13.5
vkbottle~=4.8.2 vkbottle~=4.8.2
vkbottle-types~=5.199.99.20 vkbottle-types~=5.199.99.20
pyodbc~=5.3.0 mariadb[pool]~=2.0.0rc2
openrouter==0.9.1 openrouter==0.9.1
replicate~=1.0.7 replicate~=1.0.7
tavily~=1.1.0 tavily~=1.1.0

2
tg/__main__.py
View file

@ -22,7 +22,7 @@ async def main() -> None:
config = json.load(file) config = json.load(file)
print('Конфигурация загружена.') print('Конфигурация загружена.')
database.create_database(config['db_connection_string']) database.create_database(config['db_hostname'], config['db_user'], config['db_password'], config['db_database'])
create_ai_agent(config['openrouter_token'], config['openrouter_model'], create_ai_agent(config['openrouter_token'], config['openrouter_model'],
config['replicate_token'], config['tavily_token'], config['replicate_token'], config['tavily_token'],

109
tg/tg_database.py
View file

@ -2,67 +2,68 @@ import database
class TgDatabase(database.BasicDatabase): class TgDatabase(database.BasicDatabase):
def __init__(self, connection_string: str): def __init__(self, hostname: str, user: str, password: str, dbname: str):
super().__init__(connection_string) super().__init__(hostname, user, password, dbname)
self.cursor.execute(""" with self.pool.acquire() as conn:
CREATE TABLE IF NOT EXISTS bots ( with conn.cursor() as cursor:
id BIGINT NOT NULL, cursor.execute("""
owner_id BIGINT NOT NULL, CREATE TABLE IF NOT EXISTS bots (
api_token VARCHAR(64) NOT NULL, id BIGINT NOT NULL,
ai_prompt VARCHAR(2000) DEFAULT NULL, owner_id BIGINT NOT NULL,
group_chats_allowed TINYINT NOT NULL DEFAULT 1, api_token VARCHAR(64) NOT NULL,
private_chats_allowed TINYINT NOT NULL DEFAULT 1, ai_prompt VARCHAR(2000) DEFAULT NULL,
PRIMARY KEY (id)) group_chats_allowed TINYINT NOT NULL DEFAULT 1,
""") private_chats_allowed TINYINT NOT NULL DEFAULT 1,
PRIMARY KEY (id)
)""")
self.cursor.execute(""" cursor.execute("""
CREATE TABLE IF NOT EXISTS chats ( CREATE TABLE IF NOT EXISTS chats (
bot_id BIGINT NOT NULL, bot_id BIGINT NOT NULL,
chat_id BIGINT NOT NULL, chat_id BIGINT NOT NULL,
active TINYINT NOT NULL DEFAULT 0, active TINYINT NOT NULL DEFAULT 0,
rules VARCHAR(4000), rules VARCHAR(4000),
greeting_join VARCHAR(2000), greeting_join VARCHAR(2000),
ai_prompt VARCHAR(2000), ai_prompt VARCHAR(2000),
PRIMARY KEY (bot_id, chat_id), PRIMARY KEY (bot_id, chat_id),
CONSTRAINT fk_chats_bots FOREIGN KEY (bot_id) REFERENCES bots (id) ON DELETE CASCADE ON UPDATE CASCADE) CONSTRAINT fk_chats_bots FOREIGN KEY (bot_id) REFERENCES bots (id)
""") ON UPDATE CASCADE ON DELETE CASCADE
)""")
self.cursor.execute(""" cursor.execute("""
CREATE TABLE IF NOT EXISTS users ( CREATE TABLE IF NOT EXISTS users (
bot_id BIGINT NOT NULL, bot_id BIGINT NOT NULL,
chat_id BIGINT NOT NULL, chat_id BIGINT NOT NULL,
user_id BIGINT NOT NULL, user_id BIGINT NOT NULL,
last_message BIGINT NOT NULL DEFAULT 0, last_message BIGINT NOT NULL DEFAULT 0,
messages_today SMALLINT NOT NULL DEFAULT 0, messages_today SMALLINT NOT NULL DEFAULT 0,
messages_month SMALLINT NOT NULL DEFAULT 0, messages_month SMALLINT NOT NULL DEFAULT 0,
warnings TINYINT NOT NULL DEFAULT 0, warnings TINYINT NOT NULL DEFAULT 0,
about VARCHAR(1000), about VARCHAR(1000),
PRIMARY KEY (bot_id, chat_id, user_id), PRIMARY KEY (bot_id, chat_id, user_id),
CONSTRAINT fk_users_chats FOREIGN KEY (bot_id, chat_id) REFERENCES chats (bot_id, chat_id) CONSTRAINT fk_users_chats FOREIGN KEY (bot_id, chat_id) REFERENCES chats (bot_id, chat_id)
ON UPDATE CASCADE ON DELETE CASCADE ON UPDATE CASCADE ON DELETE CASCADE
)""") )""")
self.cursor.execute(""" cursor.execute("""
CREATE TABLE IF NOT EXISTS contexts ( CREATE TABLE IF NOT EXISTS contexts (
id BIGINT NOT NULL auto_increment, id BIGINT NOT NULL auto_increment,
bot_id BIGINT NOT NULL, bot_id BIGINT NOT NULL,
chat_id BIGINT NOT NULL, chat_id BIGINT NOT NULL,
message_id BIGINT, message_id BIGINT,
role VARCHAR(16) NOT NULL, role VARCHAR(16) NOT NULL,
text VARCHAR(4000), text VARCHAR(4000),
image MEDIUMBLOB, image MEDIUMBLOB,
PRIMARY KEY (id), PRIMARY KEY (id),
CONSTRAINT fk_contexts_chats FOREIGN KEY (bot_id, chat_id) REFERENCES chats (bot_id, chat_id) CONSTRAINT fk_contexts_chats FOREIGN KEY (bot_id, chat_id) REFERENCES chats (bot_id, chat_id)
ON UPDATE CASCADE ON DELETE CASCADE ON UPDATE CASCADE ON DELETE CASCADE
)""") )""")
self.conn.commit()
DB: TgDatabase DB: TgDatabase
def create_database(connection_string: str): def create_database(hostname: str, user: str, password: str, dbname: str):
global DB global DB
DB = TgDatabase(connection_string) DB = TgDatabase(hostname, user, password, dbname)

2
vk/__main__.py
View file

@ -22,7 +22,7 @@ if __name__ == '__main__':
config = json.load(file) config = json.load(file)
print('Конфигурация загружена.') print('Конфигурация загружена.')
database.create_database(config['db_connection_string']) database.create_database(config['db_hostname'], config['db_user'], config['db_password'], config['db_database'])
create_ai_agent(config['openrouter_token'], config['openrouter_model'], create_ai_agent(config['openrouter_token'], config['openrouter_model'],
config['replicate_token'], config['tavily_token'], config['replicate_token'], config['tavily_token'],

115
vk/vk_database.py
View file

@ -2,65 +2,66 @@ import database
class VkDatabase(database.BasicDatabase): class VkDatabase(database.BasicDatabase):
def __init__(self, connection_string: str): def __init__(self, hostname: str, user: str, password: str, dbname: str):
super().__init__(connection_string) super().__init__(hostname, user, password, dbname)
self.cursor.execute(""" with self.pool.acquire() as conn:
CREATE TABLE IF NOT EXISTS bots ( with conn.cursor() as cursor:
id BIGINT NOT NULL, cursor.execute("""
owner_id BIGINT NOT NULL, CREATE TABLE IF NOT EXISTS bots (
api_token VARCHAR(256) NOT NULL, id BIGINT NOT NULL,
ai_prompt VARCHAR(2000) DEFAULT NULL, owner_id BIGINT NOT NULL,
group_chats_allowed TINYINT NOT NULL DEFAULT 1, api_token VARCHAR(256) NOT NULL,
private_chats_allowed TINYINT NOT NULL DEFAULT 1, ai_prompt VARCHAR(2000) DEFAULT NULL,
PRIMARY KEY (id)) group_chats_allowed TINYINT NOT NULL DEFAULT 1,
""") private_chats_allowed TINYINT NOT NULL DEFAULT 1,
PRIMARY KEY (id)
)""")
self.cursor.execute(""" cursor.execute("""
CREATE TABLE IF NOT EXISTS chats ( CREATE TABLE IF NOT EXISTS chats (
bot_id BIGINT NOT NULL, bot_id BIGINT NOT NULL,
chat_id BIGINT NOT NULL, chat_id BIGINT NOT NULL,
active TINYINT NOT NULL DEFAULT 0, active TINYINT NOT NULL DEFAULT 0,
rules VARCHAR(4000), rules VARCHAR(4000),
greeting_join VARCHAR(2000), greeting_join VARCHAR(2000),
greeting_rejoin VARCHAR(2000), greeting_rejoin VARCHAR(2000),
birthday_message VARCHAR(2000), birthday_message VARCHAR(2000),
ai_prompt VARCHAR(2000), ai_prompt VARCHAR(2000),
PRIMARY KEY (bot_id, chat_id), PRIMARY KEY (bot_id, chat_id),
CONSTRAINT fk_chats_bots FOREIGN KEY (bot_id) REFERENCES bots (id) ON DELETE CASCADE ON UPDATE CASCADE) CONSTRAINT fk_chats_bots FOREIGN KEY (bot_id) REFERENCES bots (id)
""") ON UPDATE CASCADE ON DELETE CASCADE
)""")
self.cursor.execute(""" cursor.execute("""
CREATE TABLE IF NOT EXISTS users ( CREATE TABLE IF NOT EXISTS users (
bot_id BIGINT NOT NULL, bot_id BIGINT NOT NULL,
chat_id BIGINT NOT NULL, chat_id BIGINT NOT NULL,
user_id BIGINT NOT NULL, user_id BIGINT NOT NULL,
last_message BIGINT NOT NULL DEFAULT 0, last_message BIGINT NOT NULL DEFAULT 0,
messages_today SMALLINT NOT NULL DEFAULT 0, messages_today SMALLINT NOT NULL DEFAULT 0,
messages_month SMALLINT NOT NULL DEFAULT 0, messages_month SMALLINT NOT NULL DEFAULT 0,
warnings TINYINT NOT NULL DEFAULT 0, warnings TINYINT NOT NULL DEFAULT 0,
happy_birthday TINYINT NOT NULL DEFAULT 1, happy_birthday TINYINT NOT NULL DEFAULT 1,
about VARCHAR(1000), about VARCHAR(1000),
PRIMARY KEY (bot_id, chat_id, user_id), PRIMARY KEY (bot_id, chat_id, user_id),
CONSTRAINT fk_users_chats FOREIGN KEY (bot_id, chat_id) REFERENCES chats (bot_id, chat_id) CONSTRAINT fk_users_chats FOREIGN KEY (bot_id, chat_id) REFERENCES chats (bot_id, chat_id)
ON UPDATE CASCADE ON DELETE CASCADE ON UPDATE CASCADE ON DELETE CASCADE
)""") )""")
self.cursor.execute(""" cursor.execute("""
CREATE TABLE IF NOT EXISTS contexts ( CREATE TABLE IF NOT EXISTS contexts (
id BIGINT NOT NULL auto_increment, id BIGINT NOT NULL auto_increment,
bot_id BIGINT NOT NULL, bot_id BIGINT NOT NULL,
chat_id BIGINT NOT NULL, chat_id BIGINT NOT NULL,
message_id BIGINT, message_id BIGINT,
role VARCHAR(16) NOT NULL, role VARCHAR(16) NOT NULL,
text VARCHAR(4000), text VARCHAR(4000),
image MEDIUMBLOB, image MEDIUMBLOB,
PRIMARY KEY (id), PRIMARY KEY (id),
CONSTRAINT fk_contexts_chats FOREIGN KEY (bot_id, chat_id) REFERENCES chats (bot_id, chat_id) CONSTRAINT fk_contexts_chats FOREIGN KEY (bot_id, chat_id) REFERENCES chats (bot_id, chat_id)
ON UPDATE CASCADE ON DELETE CASCADE ON UPDATE CASCADE ON DELETE CASCADE
)""") )""")
self.conn.commit()
def user_toggle_happy_birthday(self, bot_id: int, chat_id: int, user_id: int, happy_birthday: int): 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) self.user_update(bot_id, chat_id, user_id, happy_birthday=happy_birthday)
@ -69,6 +70,6 @@ class VkDatabase(database.BasicDatabase):
DB: VkDatabase DB: VkDatabase
def create_database(connection_string: str): def create_database(hostname: str, user: str, password: str, dbname: str):
global DB global DB
DB = VkDatabase(connection_string) DB = VkDatabase(hostname, user, password, dbname)