Галингер Р.С.
0f674f8832
- Add shared proxy module (bot/utils/proxy.py) to eliminate duplicate SOCKS5 parsing - Fix AI client: escape HTML before Markdown→HTML conversion, unify timeouts, make health check optional and disabled by default, handle 429 retries - Fix ai.py: correct forwarded message handling (aiogram 3.x forward_origin), pass relevant summaries via extra_system_content - Fix dialogue.py: only respond to Astra's messages, use system context instead of prompt injection, answer on the limit message before phase transition - Fix voice.py: load Whisper model in thread pool, safe WAV path generation - Improve database.py: composite indexes, Boolean is_active, upsert file_id cache, add context cleanup helper - Update weather.py and yadisk_download.py to use shared proxy connector - Update yadisk.py: validate URL before cache clear, add download size limit, wrap sync file ops in to_thread - Reuse S3 client via lru_cache - Update setup_commands with /aiclear and /aiuser - Update README, Dockerfile (Python 3.11), docker-compose (mount models) - Pin dependency versions, remove unused httpx[socks] - Add basic pytest tests for layout converter and voice normalization
544 lines
16 KiB
Python
544 lines
16 KiB
Python
import time
|
|
|
|
from sqlalchemy import (
|
|
Column,
|
|
Integer,
|
|
String,
|
|
Float,
|
|
BigInteger,
|
|
Text,
|
|
Boolean,
|
|
Index,
|
|
select,
|
|
delete,
|
|
)
|
|
from sqlalchemy.dialects.sqlite import insert
|
|
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)
|
|
|
|
__table_args__ = (
|
|
Index("ix_user_context_user_chat", "user_id", "chat_id"),
|
|
)
|
|
|
|
|
|
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)
|
|
|
|
__table_args__ = (
|
|
Index("ix_ai_blocked_user_chat", "user_id", "chat_id"),
|
|
)
|
|
|
|
|
|
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)
|
|
|
|
__table_args__ = (
|
|
Index("ix_sticker_ban_user_chat", "user_id", "chat_id"),
|
|
)
|
|
|
|
|
|
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)
|
|
|
|
__table_args__ = (
|
|
Index("ix_chat_user_chat_user", "chat_id", "user_id"),
|
|
)
|
|
|
|
|
|
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(Boolean, default=False)
|
|
blocked_until = Column(Float, nullable=True)
|
|
last_activity = Column(Float, nullable=False)
|
|
|
|
__table_args__ = (
|
|
Index("ix_dialogue_session_user_chat", "user_id", "chat_id"),
|
|
)
|
|
|
|
|
|
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)
|
|
|
|
__table_args__ = (
|
|
Index("ix_conversation_summary_user_chat", "user_id", "chat_id"),
|
|
)
|
|
|
|
|
|
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:
|
|
async with async_session() as session:
|
|
msg = UserContext(
|
|
user_id=user_id,
|
|
chat_id=chat_id,
|
|
text=text or "",
|
|
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:
|
|
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 clear_user_context(user_id: int, chat_id: int) -> None:
|
|
async with async_session() as session:
|
|
stmt = delete(UserContext).where(
|
|
UserContext.user_id == user_id,
|
|
UserContext.chat_id == chat_id,
|
|
)
|
|
await session.execute(stmt)
|
|
await session.commit()
|
|
|
|
|
|
async def cleanup_old_context_messages(older_than_days: int = 30) -> int:
|
|
"""Remove user context messages older than the given number of days."""
|
|
cutoff = time.time() - (older_than_days * 86400)
|
|
async with async_session() as session:
|
|
stmt = delete(UserContext).where(UserContext.timestamp < cutoff)
|
|
result = await session.execute(stmt)
|
|
await session.commit()
|
|
return result.rowcount
|
|
|
|
|
|
async def is_ai_blocked(user_id: int, chat_id: int) -> bool:
|
|
async with async_session() as session:
|
|
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:
|
|
await session.execute(
|
|
delete(AiBlockedUser).where(
|
|
AiBlockedUser.user_id == user_id,
|
|
AiBlockedUser.chat_id == chat_id,
|
|
)
|
|
)
|
|
|
|
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:
|
|
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:
|
|
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:
|
|
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:
|
|
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:
|
|
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 unban_user_stickers(user_id: int, chat_id: int) -> bool:
|
|
async with async_session() as session:
|
|
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
|
|
|
|
|
|
async def get_file_id(file_key: str) -> str | None:
|
|
async with async_session() as session:
|
|
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:
|
|
if file_id is None:
|
|
await session.execute(delete(FileIdCache).where(FileIdCache.file_key == file_key))
|
|
else:
|
|
stmt = (
|
|
insert(FileIdCache)
|
|
.values(file_key=file_key, file_id=file_id)
|
|
.on_conflict_do_update(index_elements=["file_key"], set_={"file_id": file_id})
|
|
)
|
|
await session.execute(stmt)
|
|
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:
|
|
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:
|
|
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:
|
|
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=True,
|
|
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 = True
|
|
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:
|
|
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:
|
|
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:
|
|
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 = False
|
|
await session.commit()
|
|
|
|
|
|
async def clear_dialogue(user_id: int, chat_id: int) -> None:
|
|
async with async_session() as session:
|
|
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 = False
|
|
row.blocked_until = None
|
|
row.last_activity = time.time()
|
|
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:
|
|
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()
|
|
]
|