Hidosi/umb
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor: review fixes and code improvements

Browse Source 
- 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
This commit is contained in:
Галингер Р.С.
2026-07-07 18:42:10 +07:00
parent 76e5701eba
commit 0f674f8832
23 changed files with 894 additions and 578 deletions
Download Patch File Download Diff File Expand all files Collapse all files
bot/utils/database.py
+132 -103
View File
@@ -1,7 +1,18 @@
import time
from datetime import datetime
from sqlalchemy import Column, Integer, String, Float, BigInteger, Text
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
@@ -20,6 +31,10 @@ class UserContext(Base):
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"
@@ -31,6 +46,10 @@ class AiBlockedUser(Base):
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"
@@ -50,6 +69,10 @@ class StickerBan(Base):
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"
@@ -61,6 +84,10 @@ class ChatUser(Base):
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"
@@ -70,10 +97,14 @@ class DialogueSession(Base):
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)
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"
@@ -85,6 +116,10 @@ class ConversationSummary(Base):
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)
@@ -96,13 +131,11 @@ async def init_db():
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,
text=text or "",
role=role,
timestamp=time.time(),
)
@@ -112,8 +145,6 @@ async def add_context_message(user_id: int, chat_id: int, text: str, role: str)
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)
@@ -125,10 +156,28 @@ async def get_user_context(user_id: int, chat_id: int, limit: int = 10) -> list[
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:
from sqlalchemy import select
stmt = (
select(AiBlockedUser)
.where(AiBlockedUser.user_id == user_id, AiBlockedUser.chat_id == chat_id)
@@ -138,15 +187,16 @@ async def is_ai_blocked(user_id: int, chat_id: int) -> bool:
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 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(
delete(AiBlockedUser).where(
AiBlockedUser.user_id == user_id,
AiBlockedUser.chat_id == chat_id,
)
)
await session.execute(stmt)
now = time.time()
entry = AiBlockedUser(
@@ -162,8 +212,6 @@ async def block_user_from_ai(user_id: int, chat_id: int, blocked_by: int, durati
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,
@@ -175,8 +223,6 @@ async def unblock_user_from_ai(user_id: int, chat_id: int) -> bool:
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,
@@ -195,8 +241,6 @@ async def get_sticker_ban(user_id: int, chat_id: int) -> dict | None:
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,
@@ -229,8 +273,6 @@ async def add_sticker_message(user_id: int, chat_id: int) -> int:
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,
@@ -246,8 +288,6 @@ async def ban_user_stickers(user_id: int, chat_id: int, duration: int = 300) ->
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,
@@ -268,10 +308,26 @@ async def is_user_sticker_banned(user_id: int, chat_id: int) -> bool:
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:
from sqlalchemy import select
stmt = select(FileIdCache).where(FileIdCache.file_key == file_key)
result = await session.execute(stmt)
row = result.scalar_one_or_none()
@@ -280,18 +336,22 @@ async def get_file_id(file_key: str) -> str | 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))
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 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,
@@ -306,20 +366,20 @@ async def save_chat_user(user_id: int, chat_id: int, username: str | None, full_
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,
))
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)
@@ -334,8 +394,6 @@ async def get_chat_users(chat_id: int) -> list[dict]:
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,
@@ -351,7 +409,7 @@ async def get_or_create_dialogue(user_id: int, chat_id: int) -> dict:
chat_id=chat_id,
phase=1,
msg_count=0,
is_active=1,
is_active=True,
last_activity=now,
)
session.add(row)
@@ -366,7 +424,7 @@ async def get_or_create_dialogue(user_id: int, chat_id: int) -> dict:
"blocked_until": row.blocked_until,
}
row.is_active = 1
row.is_active = True
row.last_activity = now
await session.commit()
@@ -380,8 +438,6 @@ async def get_or_create_dialogue(user_id: int, chat_id: int) -> dict:
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,
@@ -401,8 +457,6 @@ async def increment_dialogue_count(user_id: int, chat_id: int) -> dict:
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,
@@ -419,8 +473,6 @@ async def reset_dialogue_to_phase2(user_id: int, chat_id: int) -> None:
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,
@@ -430,14 +482,12 @@ async def block_dialogue(user_id: int, chat_id: int, duration: int = 3600) -> No
if row:
row.blocked_until = time.time() + duration
row.is_active = 0
row.is_active = False
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,
@@ -448,67 +498,46 @@ async def clear_dialogue(user_id: int, chat_id: int) -> None:
if row:
row.phase = 1
row.msg_count = 0
row.is_active = 0
row.is_active = False
row.blocked_until = None
row.last_activity = time.time()
await session.commit()
async def clear_user_context(user_id: int, chat_id: int) -> None:
async def save_conversation_summary(
user_id: int, chat_id: int, summary: str, embedding: str | None = None
) -> 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,
session.add(
ConversationSummary(
user_id=user_id,
chat_id=chat_id,
summary=summary,
embedding=embedding,
created_at=time.time(),
)
)
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)
.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}
{
"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