Initial commit: UMB Telegram Bot
This commit is contained in:
Галингер Р.С.
@@ -0,0 +1,514 @@
|
||||
import time
|
||||
from datetime import datetime
|
||||
|
||||
from sqlalchemy import Column, Integer, String, Float, BigInteger, Text
|
||||
from sqlalchemy.orm import declarative_base
|
||||
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
|
||||
|
||||
from config import DATABASE_URL
|
||||
|
||||
Base = declarative_base()
|
||||
|
||||
|
||||
class UserContext(Base):
|
||||
__tablename__ = "user_context"
|
||||
|
||||
id = Column(Integer, primary_key=True, autoincrement=True)
|
||||
user_id = Column(BigInteger, nullable=False, index=True)
|
||||
chat_id = Column(BigInteger, nullable=False, index=True)
|
||||
text = Column(Text, nullable=False)
|
||||
role = Column(String(16), nullable=False)
|
||||
timestamp = Column(Float, nullable=False)
|
||||
|
||||
|
||||
class AiBlockedUser(Base):
|
||||
__tablename__ = "ai_blocked_users"
|
||||
|
||||
id = Column(Integer, primary_key=True, autoincrement=True)
|
||||
user_id = Column(BigInteger, nullable=False, index=True)
|
||||
chat_id = Column(BigInteger, nullable=False, index=True)
|
||||
blocked_by = Column(BigInteger, nullable=False)
|
||||
blocked_at = Column(Float, nullable=False)
|
||||
expires_at = Column(Float, nullable=False)
|
||||
|
||||
|
||||
class FileIdCache(Base):
|
||||
__tablename__ = "file_ids"
|
||||
|
||||
file_key = Column(String, primary_key=True)
|
||||
file_id = Column(String, nullable=False)
|
||||
|
||||
|
||||
class StickerBan(Base):
|
||||
__tablename__ = "sticker_bans"
|
||||
|
||||
id = Column(Integer, primary_key=True, autoincrement=True)
|
||||
user_id = Column(BigInteger, nullable=False, index=True)
|
||||
chat_id = Column(BigInteger, nullable=False, index=True)
|
||||
count = Column(Integer, default=0)
|
||||
start_time = Column(Float, nullable=False)
|
||||
ban_until = Column(Float, nullable=True)
|
||||
ban_trigger = Column(Integer, default=0)
|
||||
|
||||
|
||||
class ChatUser(Base):
|
||||
__tablename__ = "chat_users"
|
||||
|
||||
id = Column(Integer, primary_key=True, autoincrement=True)
|
||||
chat_id = Column(BigInteger, nullable=False, index=True)
|
||||
user_id = Column(BigInteger, nullable=False, index=True)
|
||||
username = Column(String, nullable=True)
|
||||
full_name = Column(String, nullable=False)
|
||||
last_seen = Column(Float, nullable=False)
|
||||
|
||||
|
||||
class DialogueSession(Base):
|
||||
__tablename__ = "dialogue_sessions"
|
||||
|
||||
id = Column(Integer, primary_key=True, autoincrement=True)
|
||||
user_id = Column(BigInteger, nullable=False, index=True)
|
||||
chat_id = Column(BigInteger, nullable=False, index=True)
|
||||
phase = Column(Integer, default=1)
|
||||
msg_count = Column(Integer, default=0)
|
||||
is_active = Column(Integer, default=0)
|
||||
blocked_until = Column(Float, nullable=True)
|
||||
last_activity = Column(Float, nullable=False)
|
||||
|
||||
|
||||
class ConversationSummary(Base):
|
||||
__tablename__ = "conversation_summaries"
|
||||
|
||||
id = Column(Integer, primary_key=True, autoincrement=True)
|
||||
user_id = Column(BigInteger, nullable=False, index=True)
|
||||
chat_id = Column(BigInteger, nullable=False, index=True)
|
||||
summary = Column(Text, nullable=False)
|
||||
embedding = Column(Text, nullable=True)
|
||||
created_at = Column(Float, nullable=False)
|
||||
|
||||
|
||||
engine = create_async_engine(DATABASE_URL, echo=False)
|
||||
async_session = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
|
||||
|
||||
|
||||
async def init_db():
|
||||
async with engine.begin() as conn:
|
||||
await conn.run_sync(Base.metadata.create_all)
|
||||
|
||||
|
||||
async def add_context_message(user_id: int, chat_id: int, text: str, role: str) -> None:
|
||||
if text is None:
|
||||
text = ""
|
||||
async with async_session() as session:
|
||||
msg = UserContext(
|
||||
user_id=user_id,
|
||||
chat_id=chat_id,
|
||||
text=text,
|
||||
role=role,
|
||||
timestamp=time.time(),
|
||||
)
|
||||
session.add(msg)
|
||||
await session.commit()
|
||||
|
||||
|
||||
async def get_user_context(user_id: int, chat_id: int, limit: int = 10) -> list[dict]:
|
||||
async with async_session() as session:
|
||||
from sqlalchemy import select
|
||||
|
||||
stmt = (
|
||||
select(UserContext)
|
||||
.where(UserContext.user_id == user_id, UserContext.chat_id == chat_id)
|
||||
.order_by(UserContext.timestamp.desc())
|
||||
.limit(limit)
|
||||
)
|
||||
result = await session.execute(stmt)
|
||||
messages = result.scalars().all()
|
||||
return [{"role": m.role, "content": m.text} for m in reversed(messages)]
|
||||
|
||||
|
||||
async def is_ai_blocked(user_id: int, chat_id: int) -> bool:
|
||||
async with async_session() as session:
|
||||
from sqlalchemy import select
|
||||
|
||||
stmt = (
|
||||
select(AiBlockedUser)
|
||||
.where(AiBlockedUser.user_id == user_id, AiBlockedUser.chat_id == chat_id)
|
||||
.where(AiBlockedUser.expires_at > time.time())
|
||||
)
|
||||
result = await session.execute(stmt)
|
||||
return result.scalar_one_or_none() is not None
|
||||
|
||||
|
||||
async def block_user_from_ai(user_id: int, chat_id: int, blocked_by: int, duration: int = 86400) -> None:
|
||||
async with async_session() as session:
|
||||
from sqlalchemy import select, delete
|
||||
|
||||
stmt = delete(AiBlockedUser).where(
|
||||
AiBlockedUser.user_id == user_id,
|
||||
AiBlockedUser.chat_id == chat_id,
|
||||
)
|
||||
await session.execute(stmt)
|
||||
|
||||
now = time.time()
|
||||
entry = AiBlockedUser(
|
||||
user_id=user_id,
|
||||
chat_id=chat_id,
|
||||
blocked_by=blocked_by,
|
||||
blocked_at=now,
|
||||
expires_at=now + duration,
|
||||
)
|
||||
session.add(entry)
|
||||
await session.commit()
|
||||
|
||||
|
||||
async def unblock_user_from_ai(user_id: int, chat_id: int) -> bool:
|
||||
async with async_session() as session:
|
||||
from sqlalchemy import select, delete
|
||||
|
||||
stmt = delete(AiBlockedUser).where(
|
||||
AiBlockedUser.user_id == user_id,
|
||||
AiBlockedUser.chat_id == chat_id,
|
||||
)
|
||||
result = await session.execute(stmt)
|
||||
await session.commit()
|
||||
return result.rowcount > 0
|
||||
|
||||
|
||||
async def get_sticker_ban(user_id: int, chat_id: int) -> dict | None:
|
||||
async with async_session() as session:
|
||||
from sqlalchemy import select
|
||||
|
||||
stmt = select(StickerBan).where(
|
||||
StickerBan.user_id == user_id,
|
||||
StickerBan.chat_id == chat_id,
|
||||
)
|
||||
result = await session.execute(stmt)
|
||||
row = result.scalar_one_or_none()
|
||||
if not row:
|
||||
return None
|
||||
return {
|
||||
"count": row.count,
|
||||
"start_time": row.start_time,
|
||||
"ban_until": row.ban_until,
|
||||
"ban_trigger": row.ban_trigger,
|
||||
}
|
||||
|
||||
|
||||
async def add_sticker_message(user_id: int, chat_id: int) -> int:
|
||||
async with async_session() as session:
|
||||
from sqlalchemy import select
|
||||
|
||||
stmt = select(StickerBan).where(
|
||||
StickerBan.user_id == user_id,
|
||||
StickerBan.chat_id == chat_id,
|
||||
)
|
||||
result = await session.execute(stmt)
|
||||
row = result.scalar_one_or_none()
|
||||
|
||||
now = time.time()
|
||||
|
||||
if not row:
|
||||
row = StickerBan(
|
||||
user_id=user_id,
|
||||
chat_id=chat_id,
|
||||
count=1,
|
||||
start_time=now,
|
||||
)
|
||||
session.add(row)
|
||||
else:
|
||||
if now - row.start_time > 60:
|
||||
row.count = 1
|
||||
row.start_time = now
|
||||
row.ban_until = None
|
||||
row.ban_trigger = 0
|
||||
else:
|
||||
row.count += 1
|
||||
|
||||
await session.commit()
|
||||
return row.count
|
||||
|
||||
|
||||
async def ban_user_stickers(user_id: int, chat_id: int, duration: int = 300) -> None:
|
||||
async with async_session() as session:
|
||||
from sqlalchemy import select
|
||||
|
||||
stmt = select(StickerBan).where(
|
||||
StickerBan.user_id == user_id,
|
||||
StickerBan.chat_id == chat_id,
|
||||
)
|
||||
result = await session.execute(stmt)
|
||||
row = result.scalar_one_or_none()
|
||||
|
||||
if row:
|
||||
row.ban_until = time.time() + duration
|
||||
row.ban_trigger = row.count
|
||||
await session.commit()
|
||||
|
||||
|
||||
async def is_user_sticker_banned(user_id: int, chat_id: int) -> bool:
|
||||
async with async_session() as session:
|
||||
from sqlalchemy import select
|
||||
|
||||
stmt = select(StickerBan).where(
|
||||
StickerBan.user_id == user_id,
|
||||
StickerBan.chat_id == chat_id,
|
||||
)
|
||||
result = await session.execute(stmt)
|
||||
row = result.scalar_one_or_none()
|
||||
|
||||
if not row or not row.ban_until:
|
||||
return False
|
||||
|
||||
if time.time() > row.ban_until:
|
||||
row.ban_until = None
|
||||
row.count = 0
|
||||
row.ban_trigger = 0
|
||||
await session.commit()
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
async def get_file_id(file_key: str) -> str | None:
|
||||
async with async_session() as session:
|
||||
from sqlalchemy import select
|
||||
|
||||
stmt = select(FileIdCache).where(FileIdCache.file_key == file_key)
|
||||
result = await session.execute(stmt)
|
||||
row = result.scalar_one_or_none()
|
||||
return row.file_id if row else None
|
||||
|
||||
|
||||
async def save_file_id(file_key: str, file_id: str | None) -> None:
|
||||
async with async_session() as session:
|
||||
from sqlalchemy import delete
|
||||
|
||||
await session.execute(delete(FileIdCache).where(FileIdCache.file_key == file_key))
|
||||
if file_id is not None:
|
||||
session.add(FileIdCache(file_key=file_key, file_id=file_id))
|
||||
await session.commit()
|
||||
|
||||
|
||||
async def save_chat_user(user_id: int, chat_id: int, username: str | None, full_name: str) -> None:
|
||||
async with async_session() as session:
|
||||
from sqlalchemy import select
|
||||
|
||||
stmt = select(ChatUser).where(
|
||||
ChatUser.user_id == user_id,
|
||||
ChatUser.chat_id == chat_id,
|
||||
)
|
||||
result = await session.execute(stmt)
|
||||
row = result.scalar_one_or_none()
|
||||
|
||||
now = time.time()
|
||||
|
||||
if row:
|
||||
row.username = username
|
||||
row.full_name = full_name
|
||||
row.last_seen = now
|
||||
else:
|
||||
session.add(ChatUser(
|
||||
user_id=user_id,
|
||||
chat_id=chat_id,
|
||||
username=username,
|
||||
full_name=full_name,
|
||||
last_seen=now,
|
||||
))
|
||||
await session.commit()
|
||||
|
||||
|
||||
async def get_chat_users(chat_id: int) -> list[dict]:
|
||||
async with async_session() as session:
|
||||
from sqlalchemy import select
|
||||
|
||||
stmt = (
|
||||
select(ChatUser)
|
||||
.where(ChatUser.chat_id == chat_id)
|
||||
.order_by(ChatUser.full_name)
|
||||
)
|
||||
result = await session.execute(stmt)
|
||||
return [
|
||||
{"user_id": u.user_id, "username": u.username, "full_name": u.full_name}
|
||||
for u in result.scalars().all()
|
||||
]
|
||||
|
||||
|
||||
async def get_or_create_dialogue(user_id: int, chat_id: int) -> dict:
|
||||
async with async_session() as session:
|
||||
from sqlalchemy import select
|
||||
|
||||
stmt = select(DialogueSession).where(
|
||||
DialogueSession.user_id == user_id,
|
||||
DialogueSession.chat_id == chat_id,
|
||||
)
|
||||
result = await session.execute(stmt)
|
||||
row = result.scalar_one_or_none()
|
||||
|
||||
now = time.time()
|
||||
|
||||
if not row:
|
||||
row = DialogueSession(
|
||||
user_id=user_id,
|
||||
chat_id=chat_id,
|
||||
phase=1,
|
||||
msg_count=0,
|
||||
is_active=1,
|
||||
last_activity=now,
|
||||
)
|
||||
session.add(row)
|
||||
await session.commit()
|
||||
return {"phase": 1, "msg_count": 0, "is_active": True, "blocked_until": None}
|
||||
|
||||
if row.blocked_until and now < row.blocked_until:
|
||||
return {
|
||||
"phase": row.phase,
|
||||
"msg_count": row.msg_count,
|
||||
"is_active": False,
|
||||
"blocked_until": row.blocked_until,
|
||||
}
|
||||
|
||||
row.is_active = 1
|
||||
row.last_activity = now
|
||||
await session.commit()
|
||||
|
||||
return {
|
||||
"phase": row.phase,
|
||||
"msg_count": row.msg_count,
|
||||
"is_active": True,
|
||||
"blocked_until": None,
|
||||
}
|
||||
|
||||
|
||||
async def increment_dialogue_count(user_id: int, chat_id: int) -> dict:
|
||||
async with async_session() as session:
|
||||
from sqlalchemy import select
|
||||
|
||||
stmt = select(DialogueSession).where(
|
||||
DialogueSession.user_id == user_id,
|
||||
DialogueSession.chat_id == chat_id,
|
||||
)
|
||||
result = await session.execute(stmt)
|
||||
row = result.scalar_one_or_none()
|
||||
|
||||
if not row:
|
||||
return {"phase": 1, "msg_count": 0, "limit_reached": False}
|
||||
|
||||
row.msg_count += 1
|
||||
row.last_activity = time.time()
|
||||
await session.commit()
|
||||
|
||||
return {"phase": row.phase, "msg_count": row.msg_count, "limit_reached": False}
|
||||
|
||||
|
||||
async def reset_dialogue_to_phase2(user_id: int, chat_id: int) -> None:
|
||||
async with async_session() as session:
|
||||
from sqlalchemy import select
|
||||
|
||||
stmt = select(DialogueSession).where(
|
||||
DialogueSession.user_id == user_id,
|
||||
DialogueSession.chat_id == chat_id,
|
||||
)
|
||||
result = await session.execute(stmt)
|
||||
row = result.scalar_one_or_none()
|
||||
|
||||
if row:
|
||||
row.phase = 2
|
||||
row.msg_count = 0
|
||||
row.last_activity = time.time()
|
||||
await session.commit()
|
||||
|
||||
|
||||
async def block_dialogue(user_id: int, chat_id: int, duration: int = 3600) -> None:
|
||||
async with async_session() as session:
|
||||
from sqlalchemy import select
|
||||
|
||||
stmt = select(DialogueSession).where(
|
||||
DialogueSession.user_id == user_id,
|
||||
DialogueSession.chat_id == chat_id,
|
||||
)
|
||||
result = await session.execute(stmt)
|
||||
row = result.scalar_one_or_none()
|
||||
|
||||
if row:
|
||||
row.blocked_until = time.time() + duration
|
||||
row.is_active = 0
|
||||
await session.commit()
|
||||
|
||||
|
||||
async def clear_dialogue(user_id: int, chat_id: int) -> None:
|
||||
async with async_session() as session:
|
||||
from sqlalchemy import select
|
||||
|
||||
stmt = select(DialogueSession).where(
|
||||
DialogueSession.user_id == user_id,
|
||||
DialogueSession.chat_id == chat_id,
|
||||
)
|
||||
result = await session.execute(stmt)
|
||||
row = result.scalar_one_or_none()
|
||||
|
||||
if row:
|
||||
row.phase = 1
|
||||
row.msg_count = 0
|
||||
row.is_active = 0
|
||||
row.last_activity = time.time()
|
||||
await session.commit()
|
||||
|
||||
|
||||
async def clear_user_context(user_id: int, chat_id: int) -> None:
|
||||
async with async_session() as session:
|
||||
from sqlalchemy import delete
|
||||
|
||||
stmt = delete(UserContext).where(
|
||||
UserContext.user_id == user_id,
|
||||
UserContext.chat_id == chat_id,
|
||||
)
|
||||
await session.execute(stmt)
|
||||
await session.commit()
|
||||
|
||||
|
||||
async def save_conversation_summary(user_id: int, chat_id: int, summary: str, embedding: str | None = None) -> None:
|
||||
async with async_session() as session:
|
||||
session.add(ConversationSummary(
|
||||
user_id=user_id,
|
||||
chat_id=chat_id,
|
||||
summary=summary,
|
||||
embedding=embedding,
|
||||
created_at=time.time(),
|
||||
))
|
||||
await session.commit()
|
||||
|
||||
|
||||
async def get_summaries(user_id: int, chat_id: int, limit: int = 5) -> list[dict]:
|
||||
async with async_session() as session:
|
||||
from sqlalchemy import select
|
||||
|
||||
stmt = (
|
||||
select(ConversationSummary)
|
||||
.where(ConversationSummary.user_id == user_id, ConversationSummary.chat_id == chat_id)
|
||||
.order_by(ConversationSummary.created_at.desc())
|
||||
.limit(limit)
|
||||
)
|
||||
result = await session.execute(stmt)
|
||||
return [
|
||||
{"id": s.id, "summary": s.summary, "embedding": s.embedding, "created_at": s.created_at}
|
||||
for s in result.scalars().all()
|
||||
]
|
||||
|
||||
|
||||
async def unban_user_stickers(user_id: int, chat_id: int) -> bool:
|
||||
async with async_session() as session:
|
||||
from sqlalchemy import select
|
||||
|
||||
stmt = select(StickerBan).where(
|
||||
StickerBan.user_id == user_id,
|
||||
StickerBan.chat_id == chat_id,
|
||||
)
|
||||
result = await session.execute(stmt)
|
||||
row = result.scalar_one_or_none()
|
||||
|
||||
if row:
|
||||
row.ban_until = None
|
||||
row.count = 0
|
||||
row.ban_trigger = 0
|
||||
await session.commit()
|
||||
return True
|
||||
return False
|
||||
Reference in New Issue
Block a user