515 lines
16 KiB
Python
515 lines
16 KiB
Python
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
|