chibi-bot 1.6.0b0__py3-none-any.whl

chibi/services/providers/utils.py

@@ -0,0 +1,182 @@
+import inspect
+import json
+from typing import Any, Callable, Coroutine, ParamSpec, Type, TypeAlias, TypeVar
+
+from anthropic.types import (
+    Message as AnthropicMessage,
+)
+from google.genai.types import GenerateContentResponse
+from mistralai import ChatCompletionResponse
+from openai.types import CompletionUsage
+from openai.types.chat import ChatCompletion
+from telegram import Update, constants
+from telegram.ext import ContextTypes
+
+from chibi.config import gpt_settings
+from chibi.models import User
+from chibi.schemas.app import UsageSchema
+from chibi.schemas.suno import SunoGetGenerationDetailsSchema
+from chibi.utils.app import get_builtin_skill_names
+
+T = TypeVar("T")
+P = ParamSpec("P")
+M = TypeVar("M", bound=Callable[..., Coroutine[Any, Any, Any]])
+AsyncFunc: TypeAlias = Callable[P, Coroutine[Any, Any, T]]
+
+
+def decorate_async_methods(decorator: Callable[[M], M]) -> Callable[[Type[T]], Type[T]]:
+    def decorate(cls: Type[T]) -> Type[T]:
+        for attr in cls.__dict__:
+            if inspect.iscoroutinefunction(getattr(cls, attr)):
+                original_func = getattr(cls, attr)
+                decorated_func = decorator(original_func)
+                setattr(cls, attr, decorated_func)
+        return cls
+
+    return decorate
+
+
+def escape_and_truncate(message: str | dict[str, Any] | list[dict[str, Any]] | None, limit: int = 50) -> str:
+    if not message:
+        return "no data"
+
+    if isinstance(message, dict):
+        return json.dumps({k: escape_and_truncate(message=v, limit=limit) for k, v in message.items()})
+
+    if isinstance(message, list):
+        return json.dumps([escape_and_truncate(message=m, limit=limit) for m in message])
+
+    escaped_message = str(message).replace("<", r"\<").replace(">", r"\>")
+    if len(escaped_message) < limit + 20:
+        return escaped_message
+    return f"{escaped_message[:limit]}... (truncated)"
+
+
+async def prepare_system_prompt(base_system_prompt: str, user: User | None = None) -> str:
+    if user is None:
+        return base_system_prompt
+
+    prompt = {
+        "current_working_dir": user.working_dir,
+        "user_id": user.id,
+        "user_info": user.info,
+        "system_prompt": base_system_prompt,
+        "available_builtin_skills": get_builtin_skill_names(),
+        "activated_skills": user.llm_skills,
+    }
+    return json.dumps(prompt)
+
+
+async def send_llm_thoughts(thoughts: str, update: Update | None, context: ContextTypes.DEFAULT_TYPE | None) -> None:
+    if not gpt_settings.show_llm_thoughts:
+        return None
+
+    from chibi.utils.telegram import send_long_message
+
+    if update is None or context is None:
+        return None
+    message = f"💡💭 {thoughts}"
+    await send_long_message(
+        message=message,
+        update=update,
+        context=context,
+        parse_mode=constants.ParseMode.MARKDOWN_V2,
+        reply=False,
+    )
+
+
+def get_usage_from_anthropic_response(response_message: AnthropicMessage) -> UsageSchema:
+    return UsageSchema(
+        completion_tokens=response_message.usage.output_tokens,
+        prompt_tokens=response_message.usage.input_tokens,
+        cache_creation_input_tokens=response_message.usage.cache_creation_input_tokens or 0,
+        cache_read_input_tokens=response_message.usage.cache_read_input_tokens or 0,
+        total_tokens=response_message.usage.output_tokens + response_message.usage.input_tokens,
+    )
+
+
+def get_usage_from_openai_response(response_message: ChatCompletion) -> UsageSchema:
+    if response_message.usage is None:
+        return UsageSchema()
+    response_usage = response_message.usage
+    usage = UsageSchema(
+        completion_tokens=response_usage.completion_tokens,
+        prompt_tokens=response_usage.prompt_tokens,
+        total_tokens=response_usage.total_tokens,
+    )
+    if prompt_cache := response_usage.prompt_tokens_details:
+        usage.cache_read_input_tokens = prompt_cache.cached_tokens or 0
+    return usage
+
+
+def get_usage_from_google_response(response_message: GenerateContentResponse) -> UsageSchema:
+    if not response_message.usage_metadata:
+        return UsageSchema()
+
+    return UsageSchema(
+        total_tokens=response_message.usage_metadata.total_token_count or 0,
+        completion_tokens=response_message.usage_metadata.candidates_token_count or 0,
+        prompt_tokens=response_message.usage_metadata.prompt_token_count or 0,
+        cache_read_input_tokens=response_message.usage_metadata.cached_content_token_count or 0,
+    )
+
+
+def get_usage_from_mistral_response(response_message: ChatCompletionResponse) -> UsageSchema:
+    return UsageSchema(
+        completion_tokens=response_message.usage.completion_tokens or 0,
+        prompt_tokens=response_message.usage.prompt_tokens or 0,
+        cache_creation_input_tokens=0,
+        cache_read_input_tokens=0,
+        total_tokens=response_message.usage.total_tokens or 0,
+    )
+
+
+def get_usage_msg(usage: UsageSchema | CompletionUsage | None) -> str:
+    if usage is None:
+        return ""
+    cache_read = getattr(usage, "cache_read_input_tokens", None)
+    cache_create = getattr(usage, "cache_creation_input_tokens", None)
+    return (
+        f"Tokens used: {getattr(usage, 'total_tokens', None) or 'n/a'} "
+        f"({getattr(usage, 'prompt_tokens', None)} prompt, "
+        f"{getattr(usage, 'completion_tokens', None)} completion, "
+        f"{cache_read or 0} cached read/prompt, "
+        f"{cache_create or 0} cached creation)"
+    )
+
+
+def suno_task_still_processing(task_data_response: SunoGetGenerationDetailsSchema) -> bool:
+    return task_data_response.is_in_progress
+
+
+# def limit_recursion(
+#     max_depth: int = application_settings.max_consecutive_tool_calls,
+# ) -> Callable[[AsyncFunc[P, T]], AsyncFunc[P, T]]:
+#     def decorator(func: AsyncFunc[P, T]) -> AsyncFunc[P, T]:
+#         depth_var: ContextVar[int] = ContextVar(f"{func.__name__}_depth", default=0)
+#
+#         @wraps(func)
+#         async def async_wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
+#             current_depth = depth_var.get()
+#             depth_var.set(current_depth + 1)
+#             if depth_var.get() > max_depth + 1:
+#                 depth_var.set(current_depth)
+#                 class_name = ""
+#                 if args and hasattr(args[0], "__class__"):
+#                     class_name = f"{args[0].__class__.__name__}."
+#                 raise RecursionLimitExceeded(
+#                     provider=class_name,
+#                     model=cast(str, kwargs.get("model", "unknown")),
+#                     detail=f"Recursion depth exceeded: {max_depth} (function: {class_name}{func.__name__})",
+#                     exceeded_limit=max_depth,
+#                 )
+#
+#             try:
+#                 result = await func(*args, **kwargs)
+#                 return result
+#             finally:
+#                 depth_var.set(current_depth)
+#
+#         return async_wrapper
+#
+#     return decorator

chibi/services/task_manager.py

@@ -0,0 +1,93 @@
+import asyncio
+from typing import Any, Coroutine
+
+from loguru import logger
+
+from chibi.utils.app import SingletonMeta
+
+
+class BackgroundTaskManager(metaclass=SingletonMeta):
+    def __init__(self) -> None:
+        """Initialize the task manager."""
+        if not hasattr(self, "_tasks"):
+            self._tasks: set[asyncio.Task] = set()
+            self._shutting_down: bool = False
+
+    async def _wrap_with_timeout(self, coro: Coroutine[Any, Any, Any], timeout: float) -> Any:
+        """Wrap a coroutine with a timeout.
+
+        Args:
+            coro: The coroutine to wrap.
+
+        Returns:
+            The wrapped coroutine result.
+        """
+        try:
+            return await asyncio.wait_for(fut=coro, timeout=timeout)
+        except asyncio.TimeoutError:
+            logger.warning(f"Background task timed out after {timeout}s")
+            raise
+
+    def run_task(self, coro: Coroutine[Any, Any, Any], timeout: float | None = None) -> asyncio.Task | None:
+        """Schedule a coroutine to run in the background.
+
+        Args:
+            coro: the coroutine to run
+            timeout: optional timeout in seconds. If the task doesn't complete
+                     within this time, it will be cancelled with TimeoutError
+        """
+        if self._shutting_down:
+            logger.warning("Task manager is shutting down, refusing new task")
+            coro.close()
+            return None
+
+        # Wrap with timeout if specified
+        if timeout is not None:
+            coro = self._wrap_with_timeout(coro, timeout)
+
+        task = asyncio.create_task(coro)
+        self._tasks.add(task)
+        task.add_done_callback(self._discard_task)
+        return task
+
+    def _discard_task(self, task: asyncio.Task) -> None:
+        """Callback to remove a task from the set when it finishes."""
+        try:
+            exc = task.exception()
+            if exc:
+                logger.error(
+                    f"Background task '{task.get_name()}' failed: {exc.__class__.__name__} ({str(exc) or 'no details'})"
+                )
+        except asyncio.CancelledError:
+            pass
+        except Exception as e:
+            logger.error(f"Error checking background task result: {e}")
+        finally:
+            self._tasks.discard(task)
+
+    async def shutdown(self, *args: Any) -> None:
+        """
+        Wait for all background tasks to complete with a timeout.
+        If tasks do not finish within 15 seconds, they are cancelled.
+        """
+        logger.info("Shutting down background tasks...")
+        self._shutting_down = True
+        tasks_to_wait = list(self._tasks)
+        if not tasks_to_wait:
+            logger.info("No background tasks to wait, we're good.")
+            return None
+
+        logger.info(f"Waiting for {len(tasks_to_wait)} background tasks to complete...")
+        try:
+            await asyncio.wait_for(asyncio.gather(*tasks_to_wait, return_exceptions=True), timeout=5.0)
+            logger.info("All background tasks completed.")
+
+        except asyncio.TimeoutError:
+            logger.warning("Timeout reached. Cancelling remaining background tasks...")
+            remaining = [t for t in tasks_to_wait if not t.done()]
+            for task in remaining:
+                task.cancel()
+            logger.info(f"Cancelled {len(remaining)} remaining tasks.")
+
+
+task_manager = BackgroundTaskManager()

chibi/services/user.py

@@ -0,0 +1,269 @@
+import datetime
+import json
+from copy import deepcopy
+from datetime import timezone
+from io import BytesIO
+from typing import TYPE_CHECKING, Any, Optional
+
+from aiocache import cached
+from telegram import Update
+from telegram.ext import ContextTypes
+
+from chibi.config import gpt_settings
+from chibi.models import Message, User
+from chibi.schemas.app import ChatResponseSchema, ModelChangeSchema
+from chibi.services.lock_manager import LockManager
+from chibi.storage.abstract import Database
+from chibi.storage.database import inject_database
+
+if TYPE_CHECKING:
+    from chibi.services.providers.provider import Provider
+    from chibi.services.providers.tools import ToolResponse
+
+
+@inject_database
+async def get_chibi_user(db: Database, user_id: int) -> User:
+    return await db.get_or_create_user(user_id=user_id)
+
+
+@inject_database
+async def set_active_model(db: Database, user_id: int, model: ModelChangeSchema) -> None:
+    user = await db.get_or_create_user(user_id=user_id)
+    if model.image_generation:
+        user.selected_image_model_name = model.name
+        user.selected_image_provider_name = model.provider
+    else:
+        user.selected_gpt_model_name = model.name
+        user.selected_gpt_provider_name = model.provider
+    await db.save_user(user)
+
+
+@inject_database
+async def reset_chat_history(db: Database, user_id: int) -> None:
+    user = await db.get_or_create_user(user_id=user_id)
+    await db.drop_messages(user=user)
+
+
+@inject_database
+async def emergency_summarization(db: Database, user_id: int) -> None:
+    user = await db.get_or_create_user(user_id=user_id)
+
+    chat_history = await db.get_messages(user=user)
+    chat_history_string = str(msg for msg in chat_history if not any((msg.get("tool_calls"), msg.get("tool_call_id"))))
+    user_messages: list[Message] = [Message(role="user", content=chat_history_string)]
+
+    response, _ = await user.active_gpt_provider.get_chat_response(
+        messages=user_messages,
+        user=user,
+        system_prompt="Summarize this conversation, keeping the most important and useful information using English.",
+    )
+    initial_message = Message(role="user", content="What we were talking about?")
+    answer_message = Message(role="assistant", content=response.answer)
+    await reset_chat_history(user_id=user_id)
+    await db.add_message(user=user, message=initial_message, ttl=gpt_settings.messages_ttl)
+    await db.add_message(user=user, message=answer_message, ttl=gpt_settings.messages_ttl)
+
+
+@inject_database
+async def get_llm_chat_completion_answer(
+    db: Database,
+    user_id: int,
+    update: Update,
+    context: ContextTypes.DEFAULT_TYPE,
+    user_text_message: str | None = None,
+    user_voice_message: BytesIO | None = None,
+    tool_message: Optional["ToolResponse"] = None,
+) -> ChatResponseSchema:
+    user = await db.get_or_create_user(user_id=user_id)
+    lock = await LockManager().get_lock(key=user_id)
+
+    if not user_text_message and not user_voice_message and not tool_message:
+        raise ValueError("No prompt data provided")
+
+    if user_voice_message and not user.stt_provider:
+        raise ValueError("Can't compute voice message: no STT provide available.")
+
+    prompt: dict[str, Any]
+
+    if tool_message:
+        prompt = {
+            "type": "tool response",
+            "desc": "background task is done",
+            "tool_name": tool_message.tool_name,
+            "tool_response": tool_message.model_dump(),
+            "datetime_now": datetime.datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S %Z%z"),
+        }
+    else:
+        user_message = (
+            await user.stt_provider.transcribe(audio=user_voice_message) if user_voice_message else user_text_message
+        )
+        assert user_message
+        prompt = {
+            "prompt": user_message,
+            "datetime_now": datetime.datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S %Z%z"),
+            "type": "user message",
+            "transcribed_from_voice_message": bool(user_voice_message),
+        }
+
+    async with lock:
+        conversation_messages: list[Message] = await db.get_conversation_messages(user=user)
+        new_message_to_llm = Message(role="user", content=json.dumps(prompt))
+        conversation_messages.append(new_message_to_llm)
+
+        chat_response, new_messages = await user.active_gpt_provider.get_chat_response(
+            messages=conversation_messages,
+            user=user,
+            model=user.selected_gpt_model_name,
+            update=update,
+            context=context,
+        )
+        await db.add_message(user=user, message=new_message_to_llm, ttl=gpt_settings.messages_ttl)
+        for message in new_messages:
+            await db.add_message(user=user, message=message, ttl=gpt_settings.messages_ttl)
+        return chat_response
+
+
+@inject_database
+async def check_history_and_summarize(db: Database, user_id: int) -> bool:
+    user = await db.get_or_create_user(user_id=user_id)
+    messages = await db.get_messages(user=user)
+    # Roughly estimating how many tokens the current conversation history will comprise. It is possible to calculate
+    # this accurately, but the modules that can be used for this need to be separately built for armv7, which is
+    # difficult to do right now (but will be done further, I hope).
+    if len(str(messages)) / 4 >= gpt_settings.max_history_tokens:
+        await emergency_summarization(user_id=user_id)
+        return True
+    return False
+
+
+@inject_database
+async def generate_image(
+    db: Database, user_id: int, prompt: str, model: str | None = None, provider_name: str | None = None
+) -> list[str] | list[BytesIO]:
+    user = await db.get_or_create_user(user_id=user_id)
+
+    if provider_name:
+        provider = user.providers.get(provider_name)
+        selected_model = model
+    elif user.selected_image_provider_name:
+        provider = user.active_image_provider
+        selected_model = model or user.selected_image_model_name
+    else:
+        provider = user.active_image_provider
+        selected_model = None
+    if not provider:
+        raise ValueError(f"User {user_id}: no image provider available.")
+    images = await provider.get_images(prompt=prompt, model=selected_model)
+    if user_id not in gpt_settings.image_generations_whitelist:
+        await db.count_image(user_id)
+    return images
+
+
+@cached(ttl=300)
+@inject_database
+async def get_user_cached_models(db: Database, user_id: int, image_generation: bool = False) -> list[ModelChangeSchema]:
+    user = await db.get_or_create_user(user_id=user_id)
+    return await user.get_available_models(image_generation=image_generation)
+
+
+@inject_database
+async def get_models_available(db: Database, user_id: int, image_generation: bool = False) -> list[ModelChangeSchema]:
+    user = await db.get_or_create_user(user_id=user_id)
+    user_models = await get_user_cached_models(user_id=user_id, image_generation=image_generation)
+
+    if not user_models:
+        return []
+
+    available_models = deepcopy(user_models)
+    if image_generation:
+        active_model = user.selected_image_model_name or user.active_image_provider.default_image_model
+    else:
+        active_model = user.selected_gpt_model_name or user.active_gpt_provider.default_image_model
+
+    for model in available_models:
+        if model.name == active_model:
+            model.display_name = f"🟢 {model.display_name}️"
+    return available_models
+
+
+@inject_database
+async def user_has_reached_images_generation_limit(db: Database, user_id: int) -> bool:
+    user = await db.get_or_create_user(user_id=user_id)
+    return user.has_reached_image_limits
+
+
+@inject_database
+async def set_api_key(db: Database, user_id: int, api_key: str, provider_name: str) -> None:
+    user = await db.get_or_create_user(user_id=user_id)
+    user.tokens[provider_name] = api_key
+    await db.save_user(user)
+    return None
+
+
+@inject_database
+async def get_info(db: Database, user_id: int) -> str:
+    user = await db.get_or_create_user(user_id=user_id)
+    return user.info
+
+
+@inject_database
+async def set_info(db: Database, user_id: int, new_info: str) -> None:
+    user = await db.get_or_create_user(user_id=user_id)
+    user.info = new_info
+    await db.save_user(user)
+
+
+@inject_database
+async def activate_llm_skill(db: Database, user_id: int, skill_name: str, skill_payload: str) -> None:
+    user = await db.get_or_create_user(user_id=user_id)
+    user.llm_skills[skill_name] = skill_payload
+    await db.save_user(user)
+
+
+@inject_database
+async def deactivate_llm_skill(db: Database, user_id: int, skill_name: str) -> None:
+    user = await db.get_or_create_user(user_id=user_id)
+    if skill_name not in user.llm_skills.keys():
+        raise ValueError(f"The skill {skill_name} seems never been activated")
+    user.llm_skills.pop(skill_name)
+    await db.save_user(user)
+
+
+@inject_database
+async def set_working_dir(db: Database, user_id: int, new_wd: str) -> None:
+    user = await db.get_or_create_user(user_id=user_id)
+    user.working_dir = new_wd
+    await db.save_user(user)
+
+
+@inject_database
+async def get_cwd(db: Database, user_id: int) -> str:
+    user = await db.get_or_create_user(user_id=user_id)
+    return user.working_dir
+
+
+@inject_database
+async def get_moderation_provider(db: Database, user_id: int) -> "Provider":
+    user = await db.get_or_create_user(user_id=user_id)
+    return user.moderation_provider
+
+
+@inject_database
+async def drop_tool_call_history(db: Database, user_id: int) -> None:
+    user = await db.get_or_create_user(user_id=user_id)
+    chat_history: list[Message] = await db.get_conversation_messages(user=user)
+    await reset_chat_history(user_id=user_id)
+    for message in chat_history:
+        if message.role == "tool":
+            continue
+        message.tool_calls = None
+        message.tool_call_id = None
+        await db.add_message(user=user, message=message, ttl=gpt_settings.messages_ttl)
+
+
+@inject_database
+async def summarize_history(db: Database, user_id: int) -> None:
+    user = await db.get_or_create_user(user_id=user_id)
+    chat_history: list[Message] = await db.get_conversation_messages(user=user)
+    await reset_chat_history(user_id=user_id)
+    await db.add_message(user=user, message=chat_history[0], ttl=gpt_settings.messages_ttl)

chibi/storage/abstract.py

@@ -0,0 +1,54 @@
+import time
+from abc import ABC, abstractmethod
+from typing import Optional
+
+from openai.types.chat import (
+    ChatCompletionAssistantMessageParam,
+    ChatCompletionFunctionMessageParam,
+    ChatCompletionToolMessageParam,
+    ChatCompletionUserMessageParam,
+)
+
+from chibi.models import ImageMeta, Message, User
+
+CHAT_COMPLETION_CLASSES = {
+    "assistant": ChatCompletionAssistantMessageParam,
+    "function": ChatCompletionFunctionMessageParam,
+    "tool": ChatCompletionToolMessageParam,
+    "user": ChatCompletionUserMessageParam,
+}
+
+
+class Database(ABC):
+    async def get_or_create_user(self, user_id: int) -> User:
+        if user := await self.get_user(user_id=user_id):
+            return user
+        return await self.create_user(user_id=user_id)
+
+    @abstractmethod
+    async def save_user(self, user: User) -> None: ...
+
+    @abstractmethod
+    async def create_user(self, user_id: int) -> User: ...
+
+    @abstractmethod
+    async def get_user(self, user_id: int) -> User | None: ...
+
+    @abstractmethod
+    async def add_message(self, user: User, message: Message, ttl: Optional[int] = None) -> None: ...
+
+    @abstractmethod
+    async def get_messages(self, user: User) -> list[dict[str, str]]: ...
+
+    @abstractmethod
+    async def drop_messages(self, user: User) -> None: ...
+
+    async def get_conversation_messages(self, user: User) -> list[Message]:
+        messages = await self.get_messages(user=user)
+        return [Message(**msg) for msg in messages]
+
+    async def count_image(self, user_id: int) -> None:
+        user = await self.get_or_create_user(user_id=user_id)
+        expire_at = time.time() + 60 * 750  # ~ 1 month
+        user.images.append(ImageMeta(expire_at=expire_at))
+        await self.save_user(user)

chibi/storage/database.py

@@ -0,0 +1,86 @@
+import asyncio
+from functools import wraps
+from typing import Awaitable, Callable, Concatenate, Optional, ParamSpec, TypeVar, cast
+
+from chibi.config.app import application_settings
+from chibi.storage.abstract import Database
+from chibi.storage.dynamodb import DynamoDBStorage
+from chibi.storage.local import LocalStorage
+from chibi.storage.redis import RedisStorage
+
+R = TypeVar("R")
+P = ParamSpec("P")
+
+
+class DatabaseCache:
+    """
+    Caches a Database instance according to application settings.
+    Supports 'local', 'redis', and 'dynamodb' backends.
+    """
+
+    def __init__(self) -> None:
+        self._cache: Optional[Database] = None
+        self._lock = asyncio.Lock()
+
+    async def get_database(self) -> Database:
+        """Get or create the Database instance based on storage_backend setting.
+
+        Returns:
+            Initialized Database instance.
+        """
+        async with self._lock:
+            if self._cache is not None:
+                return self._cache
+
+            backend = application_settings.storage_backend.lower()
+            if backend == "redis":
+                # RedisStorage.create expects URL and password
+                self._cache = await RedisStorage.create(
+                    url=cast(str, application_settings.redis),
+                    password=application_settings.redis_password,
+                )
+            elif backend == "dynamodb":
+                # DynamoDBStorage.create expects region, access_key, secret_key, tables
+                self._cache = await DynamoDBStorage.create(
+                    region=application_settings.aws_region or "",
+                    access_key=application_settings.aws_access_key_id,
+                    secret_access_key=application_settings.aws_secret_access_key,
+                    users_table=application_settings.ddb_users_table or "",
+                    messages_table=application_settings.ddb_messages_table or "",
+                )
+            else:
+                # default to local storage
+                self._cache = LocalStorage(application_settings.local_data_path)
+
+            return self._cache
+
+    def clear_cache(self) -> None:
+        """
+        Clear the cached Database instance, forcing reinitialization on next use.
+        """
+        self._cache = None
+
+
+_db_provider = DatabaseCache()
+
+
+def inject_database(
+    func: Callable[Concatenate[Database, P], Awaitable[R]],
+) -> Callable[P, Awaitable[R]]:
+    """Decorator to inject the Database instance into async functions.
+
+    Wraps a function with signature func(db, *args, **kwargs) -> Awaitable.
+
+    Args:
+        func: The function to decorate.
+
+    Returns:
+        Function execution wrapper.
+    """
+
+    @wraps(func)
+    async def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
+        db = await _db_provider.get_database()
+        return await func(db, *args, **kwargs)
+
+    return wrapper