diff --git a/botspot/components/data/user_data.py b/botspot/components/data/user_data.py index 246f37d..9eff2ee 100644 --- a/botspot/components/data/user_data.py +++ b/botspot/components/data/user_data.py @@ -113,8 +113,10 @@ async def has_user(self, user_id: int) -> bool: def users_collection(self) -> "AsyncIOMotorCollection": return self.db[self.collection] - async def get_users(self, query: dict = {}) -> list[UserT]: + async def get_users(self, query: Optional[dict] = None) -> list[UserT]: """Get all users matching the query""" + if query is None: + query = {} data = await self.users_collection.find(query).to_list() return [self.user_class(**item) for item in data] @@ -188,9 +190,27 @@ async def sync_user_types(self) -> None: """ # Update admins if self.settings.admins: + # Convert usernames to user IDs + admin_ids = [] + for admin in self.settings.admins: + if admin.startswith("@"): + username = admin[1:] # Remove the '@' prefix + user = await self.users_collection.find_one({"username": username}) + if user and "user_id" in user: + admin_ids.append(user["user_id"]) + logger.info(f"Mapped username {admin} to user_id {user['user_id']}") + else: + logger.warning(f"Could not find user_id for username {admin}") + else: + # Assume it's already a user_id + try: + admin_ids.append(int(admin)) + except ValueError: + logger.warning(f"Invalid user_id format in admin list: {admin}") + # Promote current admins result = await self.users_collection.update_many( - {"user_id": {"$in": list(self.settings.admins)}}, + {"user_id": {"$in": admin_ids}}, {"$set": {"user_type": UserType.ADMIN}}, ) if result.modified_count: @@ -199,7 +219,7 @@ async def sync_user_types(self) -> None: # Demote former admins result = await self.users_collection.update_many( { - "user_id": {"$nin": list(self.settings.admins)}, + "user_id": {"$nin": admin_ids}, "user_type": UserType.ADMIN, }, {"$set": {"user_type": UserType.REGULAR}}, @@ -372,7 +392,7 @@ def initialize(settings: "BotspotSettings", user_class=None) -> UserManager: ) -def get_user_manager(): +def get_user_manager() -> UserManager: """Get UserManager instance from dependency manager.""" from botspot.core.dependency_manager import get_dependency_manager diff --git a/botspot/components/new/llm_provider.py b/botspot/components/new/llm_provider.py index c7117be..910b5cd 100644 --- a/botspot/components/new/llm_provider.py +++ b/botspot/components/new/llm_provider.py @@ -383,9 +383,13 @@ async def _prepare_request( user_message = "You don't have access to AI features" if deps.botspot_settings.admins and len(deps.botspot_settings.admins) > 0: if len(deps.botspot_settings.admins) > 1: - admin_contact = "admins (@" + ", @".join(deps.botspot_settings.admins) + ")" + admin_contact = ( + "admins (@" + + ", @".join(a.lstrip("@") for a in deps.botspot_settings.admins) + + ")" + ) else: - admin_contact = "admin @" + deps.botspot_settings.admins[0] + admin_contact = "admin @" + deps.botspot_settings.admins[0].lstrip("@") user_message += f", please write to {admin_contact} to request access" raise LLMPermissionError( message=f"User {user} is not allowed to use LLM features", diff --git a/botspot/components/new/queue_manager.py b/botspot/components/new/queue_manager.py index 6b7b1a0..2310c2a 100644 --- a/botspot/components/new/queue_manager.py +++ b/botspot/components/new/queue_manager.py @@ -5,8 +5,9 @@ from datetime import datetime from typing import Any, Dict, Generic, List, Optional, Type, TypeVar -from pydantic import BaseModel +from pydantic import BaseModel, Field from pydantic_settings import BaseSettings +from bson import ObjectId # Placeholder for database connection (adjust as needed) @@ -23,11 +24,13 @@ class Config: class QueueItem(BaseModel): + id: Optional[ObjectId] = Field(default=None, alias="_id") data: str class Config: populate_by_name = True arbitrary_types_allowed = True + allow_population_by_field_name = True T = TypeVar("T", bound=QueueItem) @@ -86,7 +89,8 @@ async def add_item(self, item: T, user_id: Optional[int] = None): raise QueuePermissionError("user_id is required unless single_user_mode is enabled") doc = self.enrich_item(item, user_id) - await self.collection.insert_one(doc) + # Use by_alias for MongoDB compatibility + await self.collection.insert_one(item.model_dump(by_alias=True, exclude_none=True)) async def get_items( self, @@ -106,6 +110,7 @@ async def get_items( if limit is not None: cursor = cursor.limit(limit) items = await cursor.to_list(length=limit) + # Use model_validate to load MongoDB docs into Pydantic models return [self.item_model.model_validate(item) for item in items] async def get_records( @@ -123,6 +128,65 @@ async def get_records( cursor = cursor.limit(limit) return await cursor.to_list(length=limit) + async def update_item(self, item: T): + """Update an item in the queue by its id.""" + assert item.id is not None, "Item must have an id to update." + await self.collection.update_one( + {"_id": item.id}, {"$set": item.model_dump(by_alias=True, exclude_none=True)} + ) + + async def delete_item(self, item_id: ObjectId): + """Delete an item from the queue by its id.""" + await self.collection.delete_one({"_id": item_id}) + + async def mark_done(self, item_id: ObjectId): + """Mark an item as done. Only if use_done is enabled.""" + assert self.use_done, "mark_done requires use_done to be enabled." + await self.collection.update_one({"_id": item_id}, {"$set": {"done": True}}) + + async def set_priority(self, item_id: ObjectId, priority: int): + """Set the priority of an item. Only if use_priority is enabled.""" + assert self.use_priority, "set_priority requires use_priority to be enabled." + await self.collection.update_one({"_id": item_id}, {"$set": {"priority": priority}}) + + async def mark_undone(self, item_id: ObjectId): + """Mark an item as not done (undone). Only if use_done is enabled.""" + assert self.use_done, "mark_undone requires use_done to be enabled." + await self.collection.update_one({"_id": item_id}, {"$set": {"done": False}}) + + async def pop( + self, user_id: Optional[int] = None, extra_filters: Optional[dict] = None + ) -> Optional[T]: + """Pop an item from the queue (fetch and return, but do not remove). + - If use_done: only non-done items + - If use_priority: by priority + - Otherwise: random + - extra_filters: additional query filters to apply + """ + query = {} + if user_id: + query["user_id"] = user_id + if self.use_done: + query["done"] = False + if extra_filters: + query.update(extra_filters) + cursor = self.collection.find(query) + if self.use_priority: + cursor = cursor.sort("priority", -1) # Highest priority first + elif not self.use_priority: + # Random order: sample one + docs = await cursor.to_list(length=None) + import random + + if not docs: + return None + doc = random.choice(docs) + return self.item_model.model_validate(doc) + doc = await cursor.to_list(length=1) + if not doc: + return None + return self.item_model.model_validate(doc[0]) + class QueueManager: def __init__(self, settings: QueueManagerSettings, single_user_mode: Optional[bool] = None): @@ -132,8 +196,6 @@ def __init__(self, settings: QueueManagerSettings, single_user_mode: Optional[bo self.db = get_database() # Fetch MongoDB database instance self.queues: Dict[str, Queue] = {} if single_user_mode is None: - from botspot.core.dependency_manager import get_dependency_manager - single_user_mode = not self.settings.enabled self.single_user_mode = single_user_mode @@ -210,6 +272,8 @@ def initialize(settings: QueueManagerSettings) -> QueueManager: raise ConfigurationError("MongoDB is required for queue_manager component") single_user_mode = deps.botspot_settings.single_user_mode.enabled + logger.info(f"Queue Manager initialized with {single_user_mode=}") + return QueueManager(settings, single_user_mode=single_user_mode) diff --git a/botspot/utils/send_safe.py b/botspot/utils/send_safe.py index 530773e..5e1764f 100644 --- a/botspot/utils/send_safe.py +++ b/botspot/utils/send_safe.py @@ -6,6 +6,7 @@ from aiogram import Bot from aiogram.enums import ParseMode from aiogram.types import BufferedInputFile, Chat, Message, User +from aiogram.exceptions import TelegramBadRequest from loguru import logger from botspot.utils.text_utils import MAX_TELEGRAM_MESSAGE_LENGTH, escape_md, split_long_message @@ -406,6 +407,40 @@ async def answer_safe( ) +async def delete_safe( + message: Message, + deps: Optional["DependencyManager"] = None, + **kwargs: Any, +) -> None: + """Delete a message with safe deletion""" + try: + await message.delete() + except TelegramBadRequest as e: + # Error processing message: + # user: None + # timestamp: 2025-06-11_14-26-56 + # error: Telegram server says - Bad Request: message can't be deleted for everyone + # traceback: + # File "/app/outstanding_items_bot/routers/router_triage.py", line 159, in handle_triage_callback + # await callback_query.message.delete() + # File "/usr/local/lib/python3.12/site-packages/aiogram/methods/base.py", line 84, in emit + # return await bot(self) + # ^^^^^^^^^^^^^^^ + # File "/usr/local/lib/python3.12/site-packages/aiogram/client/bot.py", line 478, in __call__ + # return await self.session(self, method, timeout=request_timeout) + # ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + # File "/usr/local/lib/python3.12/site-packages/aiogram/client/session/base.py", line 254, in __call__ + # return cast(TelegramType, await middleware(bot, method)) + # ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + # File "/usr/local/lib/python3.12/site-packages/aiogram/client/session/aiohttp.py", line 185, in make_request + # response = self.check_response( + # ^^^^^^^^^^^^^^^^^^^^ + # File "/usr/local/lib/python3.12/site-packages/aiogram/client/session/base.py", line 120, in check_response + # raise TelegramBadRequest(method=method, message=description) + # aiogram.exceptions.TelegramBadRequest: Telegram server says - Bad Request: message can't be deleted for everyone + logger.warning(f"Failed to delete message: {e}") + + if __name__ == "__main__": async def message_handler(message: Message): diff --git a/pyproject.toml b/pyproject.toml index 6c2d957..6470881 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "botspot" -version = "0.10.22" +version = "0.10.25" description = "" authors = ["Petr Lavrov "] readme = "README.md"