chibi-bot 1.6.0b0__py3-none-any.whl

chibi/storage/dynamodb.py

@@ -0,0 +1,257 @@
+import asyncio
+import time
+from typing import Any
+
+import boto3
+from botocore.exceptions import ClientError
+
+from chibi.models import Message, User
+from chibi.storage.abstract import Database
+
+
+class DynamoDBStorage(Database):
+    """DynamoDB storage backend implementing Database interface.
+
+    Uses two DynamoDB tables:
+      - users table (PK=user_id)
+      - messages table (PK=user_id, SK=message_id)
+    """
+
+    def __init__(
+        self,
+        users_table_name: str,
+        messages_table_name: str,
+        aws_region: str,
+        aws_access_key_id: str | None = None,
+        aws_secret_access_key: str | None = None,
+    ) -> None:
+        """Initialize the DynamoDBStorage.
+
+        Args:
+            users_table_name: Name of the DynamoDB table for users.
+            messages_table_name: Name of the DynamoDB table for messages.
+            aws_region: AWS region for the DynamoDB tables.
+            aws_access_key_id: AWS access key ID. Defaults to None.
+            aws_secret_access_key: AWS secret access key. Defaults to None.
+        """
+        session = boto3.Session(
+            aws_access_key_id=aws_access_key_id,
+            aws_secret_access_key=aws_secret_access_key,
+            region_name=aws_region,
+        )
+        self.dynamodb = session.resource("dynamodb")
+        self.users_table = self.dynamodb.Table(users_table_name)
+        self.messages_table = self.dynamodb.Table(messages_table_name)
+
+    @classmethod
+    async def create(
+        cls,
+        region: str,
+        access_key: str | None,
+        secret_access_key: str | None,
+        users_table: str,
+        messages_table: str,
+    ) -> "DynamoDBStorage":
+        """Create and initializes an instance of DynamoDBStorage.
+
+        Args:
+            region: AWS region.
+            access_key: AWS access key ID.
+            secret_access_key: AWS secret access key.
+            users_table: Name of the users table.
+            messages_table: Name of the messages table.
+
+        Returns:
+            DynamoDBStorage: An instance of the DynamoDBStorage class.
+        """
+        instance = cls(
+            users_table_name=users_table,
+            messages_table_name=messages_table,
+            aws_region=region,
+            aws_access_key_id=access_key,
+            aws_secret_access_key=secret_access_key,
+        )
+        await instance.connect()
+        return instance
+
+    async def connect(self) -> None:
+        """Ensure users and messages tables exist (creates if missing).
+
+        Raises:
+            ClientError: If there's an issue with DynamoDB operations other than
+                ResourceNotFoundException when checking/creating tables.
+        """
+        client = self.dynamodb.meta.client
+        # users table
+        try:
+            client.describe_table(TableName=self.users_table.table_name)
+        except ClientError as e:
+            code = e.response.get("Error", {}).get("Code")
+            if code == "ResourceNotFoundException":
+                client.create_table(
+                    TableName=self.users_table.table_name,
+                    KeySchema=[{"AttributeName": "user_id", "KeyType": "HASH"}],
+                    AttributeDefinitions=[{"AttributeName": "user_id", "AttributeType": "S"}],
+                    BillingMode="PAY_PER_REQUEST",
+                )
+                client.get_waiter("table_exists").wait(TableName=self.users_table.table_name)
+            else:
+                raise
+        # messages table
+        try:
+            client.describe_table(TableName=self.messages_table.table_name)
+        except ClientError as e:
+            code = e.response.get("Error", {}).get("Code")
+            if code == "ResourceNotFoundException":
+                client.create_table(
+                    TableName=self.messages_table.table_name,
+                    KeySchema=[
+                        {"AttributeName": "user_id", "KeyType": "HASH"},
+                        {"AttributeName": "message_id", "KeyType": "RANGE"},
+                    ],
+                    AttributeDefinitions=[
+                        {"AttributeName": "user_id", "AttributeType": "S"},
+                        {"AttributeName": "message_id", "AttributeType": "S"},
+                    ],
+                    BillingMode="PAY_PER_REQUEST",
+                )
+                client.get_waiter("table_exists").wait(TableName=self.messages_table.table_name)
+            else:
+                raise
+
+    async def get_user(self, user_id: int) -> User | None:
+        """Retrieve a User by ID.
+
+        Args:
+            user_id: The ID of the user to retrieve.
+
+        Returns:
+            The User object if found, otherwise None.
+        """
+
+        def _sync() -> User | None:
+            try:
+                resp = self.users_table.get_item(Key={"user_id": str(user_id)})
+                item = resp.get("Item")
+                if not item or "data" not in item:
+                    return None
+                return User.model_validate_json(item["data"])
+            except ClientError:
+                return None
+
+        return await asyncio.to_thread(_sync)
+
+    async def create_user(self, user_id: int) -> User:
+        """Create a new User record.
+
+        Args:
+            user_id: The ID for the new user.
+
+        Returns:
+            The newly created User object.
+        """
+        user = User(id=user_id)
+        await self.save_user(user)
+        return user
+
+    async def save_user(self, user: User) -> None:
+        """Persist the User object to DynamoDB.
+
+        Args:
+            user: The User object to persist.
+        """
+        await asyncio.to_thread(
+            self.users_table.put_item,
+            Item={"user_id": str(user.id), "data": user.model_dump_json()},
+        )
+
+    async def get_or_create_user(self, user_id: int) -> User:
+        """Retrieve a User by ID or creates if missing; loads messages.
+
+        Args:
+            user_id: The ID of the user to retrieve or create.
+
+        Returns:
+            The retrieved or newly created User object, with messages loaded.
+        """
+        user = await self.get_user(user_id)
+        if not user:
+            user = await self.create_user(user_id)
+
+        # load messages
+        def _load() -> list[dict[str, Any]]:
+            resp = self.messages_table.query(
+                KeyConditionExpression="user_id = :u",
+                ExpressionAttributeValues={":u": str(user_id)},
+            )
+            return resp.get("Items", [])
+
+        items = await asyncio.to_thread(_load)
+        now_ts = int(time.time())
+        msgs: list[Message] = []
+        for it in items:
+            exp = it.get("expire_at")
+            if exp is None or exp >= now_ts:
+                msgs.append(
+                    Message(
+                        id=int(it["message_id"]),
+                        role=it.get("role", ""),
+                        content=it.get("content", ""),
+                        expire_at=exp,
+                    )
+                )
+        msgs.sort(key=lambda m: m.id)
+        user.messages = msgs
+        return user
+
+    async def add_message(self, user: User, message: Message, ttl: int | None = None) -> None:
+        """Add a Message record with optional TTL in seconds.
+
+        Args:
+            user: The user to whom the message belongs.
+            message: The message object to add.
+            ttl: Optional Time To Live for the message in seconds.
+        """
+        expire_at: int | None = None
+        if ttl is not None:
+            expire_at = int(time.time()) + ttl
+
+        item: dict[str, Any] = {
+            "user_id": str(user.id),
+            "message_id": str(message.id),
+            "role": message.role,
+            "content": message.content,
+        }
+        if expire_at is not None:
+            item["expire_at"] = expire_at
+
+        await asyncio.to_thread(self.messages_table.put_item, Item=item)
+
+    async def get_messages(self, user: User) -> list[dict[str, str]]:
+        """Retrieve non-expired messages as simple dicts.
+
+        Args:
+            user: The user whose messages are to be retrieved.
+
+        Returns:
+            A list of non-expired messages, where each message is a dictionary (excluding 'expire_at' and 'id').
+        """
+        user_ref = await self.get_or_create_user(user.id)
+        return [msg.model_dump(exclude={"expire_at", "id"}) for msg in user_ref.messages]
+
+    async def drop_messages(self, user: User) -> None:
+        """Delete all messages for a user.
+
+        Args:
+            user: The user whose messages are to be deleted.
+        """
+
+        def _sync() -> None:
+            resp = self.messages_table.query(
+                KeyConditionExpression="user_id = :u",
+                ExpressionAttributeValues={":u": str(user.id)},
+            )
+            for it in resp.get("Items", []):
+                self.messages_table.delete_item(Key={"user_id": it["user_id"], "message_id": it["message_id"]})
+
+        await asyncio.to_thread(_sync)

chibi/storage/local.py

@@ -0,0 +1,70 @@
+import os
+import pickle
+import time
+from typing import Optional
+
+from loguru import logger
+
+from chibi.models import Message, User
+from chibi.storage.abstract import Database
+
+
+class LocalStorage(Database):
+    def __init__(self, storage_path: str):
+        self.storage_path = storage_path
+        logger.info("Local storage initialized.")
+
+    def _get_storage_filename(self, user_id: int) -> str:
+        return os.path.join(self.storage_path, f"{user_id}.pkl")
+
+    async def save_user(self, user: User) -> None:
+        filename = self._get_storage_filename(user.id)
+        with open(filename, "wb") as f:
+            pickle.dump(user.model_dump(), f)
+
+    async def create_user(self, user_id: int) -> User:
+        user = User(id=user_id)
+        await self.save_user(user=user)
+        return user
+
+    async def get_user(self, user_id: int) -> User | None:
+        filename = self._get_storage_filename(user_id)
+        user = None
+        if not os.path.exists(filename):
+            return None
+        with open(filename, "rb") as f:
+            data = pickle.load(f)
+        if isinstance(data, dict):
+            user = User(**data)
+        if isinstance(data, User):
+            user = User(**data.dict())
+        if user:
+            current_time = time.time()
+            if hasattr(user, "images"):
+                user.images = [img for img in user.images if img.expire_at > current_time]
+            return user
+
+        return None
+
+    async def add_message(self, user: User, message: Message, ttl: Optional[int] = None) -> None:
+        user_refreshed = await self.get_or_create_user(user_id=user.id)
+        expire_at = time.time() + ttl if ttl else None
+
+        message.expire_at = expire_at
+        user_refreshed.messages.append(message)
+        await self.save_user(user_refreshed)
+
+    async def get_messages(self, user: User) -> list[dict[str, str]]:
+        user_refreshed = await self.get_or_create_user(user_id=user.id)
+        current_time = time.time()
+
+        msgs = [
+            msg.model_dump(exclude={"expire_at", "id"})
+            for msg in user_refreshed.messages
+            if msg.expire_at is None or msg.expire_at > current_time
+        ]
+        return msgs
+
+    async def drop_messages(self, user: User) -> None:
+        user.messages = []
+        await self.save_user(user=user)

chibi/storage/redis.py

@@ -0,0 +1,91 @@
+from typing import Optional
+from urllib.parse import urlparse
+
+from loguru import logger
+from redis.asyncio import Redis, from_url
+
+from chibi.models import Message, User
+from chibi.storage.abstract import Database
+
+
+class RedisStorage(Database):
+    def __init__(self, url: str, password: str | None = None, db: int = 0) -> None:
+        self.redis: Redis
+        self.url = url
+        self.password = password
+        self.db = db
+        logger.info("Redis storage initialized.")
+
+    @classmethod
+    async def create(cls, url: str, password: str | None = None, db: int = 0) -> "RedisStorage":
+        instance = cls(url, password, db)
+        await instance.connect()
+        return instance
+
+    async def connect(self) -> None:
+        redis_dsn = self._combine_redis_dsn(base_dsn=self.url, password=self.password)
+        self.redis = await from_url(redis_dsn)
+
+    def _combine_redis_dsn(self, base_dsn: str, password: str | None) -> str:
+        if not password:
+            return base_dsn
+
+        parsed_dsn = urlparse(base_dsn)
+        password_in_dsn = parsed_dsn.password or None
+
+        if password_in_dsn:
+            logger.warning(
+                "Redis password specified twice: in the REDIS_PASSWORD and REDIS environment variables. "
+                "Trying to use the password from the Redis DSN..."
+            )
+            return base_dsn
+
+        if host := parsed_dsn.hostname:
+            return base_dsn.replace(host, f":{password}@{host}")
+
+        raise ValueError("Incorrect Redis DSN string provided.")
+
+    async def save_user(self, user: User) -> None:
+        user_key = f"user:{user.id}"
+        user_data = user.model_dump_json()
+        await self.redis.set(user_key, user_data)
+
+    async def create_user(self, user_id: int) -> User:
+        user = User(id=user_id)
+        user_key = f"user:{user_id}"
+        await self.redis.set(user_key, user.model_dump_json())
+        return user
+
+    async def get_user(self, user_id: int) -> Optional[User]:
+        user_key = f"user:{user_id}"
+        user_data = await self.redis.get(user_key)
+        if not user_data:
+            return None
+
+        user = User.model_validate_json(user_data)
+        message_keys_pattern = f"user:{user.id}:message:*"
+        message_keys = set(await self.redis.keys(message_keys_pattern))
+        user_messages = [Message.model_validate_json(await self.redis.get(message_key)) for message_key in message_keys]
+        user.messages = sorted(user_messages, key=lambda msg: msg.id)
+
+        return user
+
+    async def add_message(self, user: User, message: Message, ttl: Optional[int] = None) -> None:
+        message_key = f"user:{user.id}:message:{message.id}"
+        await self.redis.set(name=message_key, value=message.model_dump_json(exclude={"expire_at"}))
+        if ttl:
+            await self.redis.expire(name=message_key, time=ttl)
+
+    async def get_messages(self, user: User) -> list[dict[str, str]]:
+        user_refreshed = await self.get_or_create_user(user_id=user.id)
+        return [msg.model_dump(exclude={"expire_at", "id"}) for msg in user_refreshed.messages]
+
+    async def drop_messages(self, user: User) -> None:
+        message_keys_pattern = f"user:{user.id}:message:*"
+        message_keys = await self.redis.keys(message_keys_pattern)
+
+        for message_key in message_keys:
+            await self.redis.delete(message_key)
+
+    async def close(self) -> None:
+        await self.redis.aclose()

chibi/utils/app.py

@@ -0,0 +1,249 @@
+from functools import wraps
+from pathlib import Path
+from typing import Any, Callable
+
+import httpx
+from loguru import logger
+from telegram import Update
+from telegram.ext import ContextTypes
+
+from chibi.config import application_settings, gpt_settings, telegram_settings
+from chibi.constants import SETTING_DISABLED, SETTING_ENABLED, SETTING_SET, SETTING_UNSET
+from chibi.exceptions import (
+    NoApiKeyProvidedError,
+    NoModelSelectedError,
+    NoProviderSelectedError,
+    NoResponseError,
+    NotAuthorizedError,
+    RecursionLimitExceeded,
+    ServiceRateLimitError,
+    ServiceResponseError,
+)
+from chibi.utils.telegram import chat_data, send_message, user_data
+
+
+class SingletonMeta(type):
+    _instances: dict[type, Any] = {}
+
+    def __call__(cls, *args: Any, **kwargs: Any) -> Any:
+        if cls not in cls._instances:
+            instance = super().__call__(*args, **kwargs)
+            cls._instances[cls] = instance
+        return cls._instances[cls]
+
+
+async def run_heartbeat(context: ContextTypes.DEFAULT_TYPE) -> None:
+    """Send a heartbeat GET request to a configured monitoring URL.
+
+    This function is designed to be called periodically by a scheduler
+    (python-telegram-bot's JobQueue) to signal the bot's operational status
+    to an external monitoring service (e.g., Healthchecks.io, Uptime Kuma, etc.).
+
+    Args:
+        context: The callback context provided by the JobQueue.
+    """
+    if not application_settings.heartbeat_url:
+        return None
+
+    transport = httpx.AsyncHTTPTransport(
+        retries=application_settings.heartbeat_retry_calls,
+        proxy=application_settings.heartbeat_proxy,
+    )
+
+    async with httpx.AsyncClient(transport=transport, proxy=application_settings.heartbeat_proxy) as client:
+        try:
+            await context.bot.get_me()
+
+            result = await client.get(application_settings.heartbeat_url)
+        except Exception as error:
+            logger.error(f"Uptime Checker failed with an Exception: {error}")
+            return
+        if result.is_error:
+            logger.error(f"Uptime Checker failed, status_code: {result.status_code}, msg: {result.text}")
+
+
+def _provider_statuses() -> list[str]:
+    """Prepare a provider clients statuses data for logging.
+
+    Returns:
+        list of string containing the provider clients statuses data.
+    """
+    from chibi.services.providers import RegisteredProviders
+
+    statuses = [
+        "<magenta>Provider clients:</magenta>",
+    ]
+    for provider_name in RegisteredProviders.all.keys():
+        status = SETTING_SET if provider_name in RegisteredProviders.available else SETTING_UNSET
+        statuses.append(f"{provider_name} client: {status}")
+    return statuses
+
+
+def log_application_settings() -> None:
+    mode = "<yellow>PUBLIC</yellow>" if gpt_settings.public_mode else "<cyan>PRIVATE</cyan>"
+    storage = "<red>REDIS</red>" if application_settings.redis else "<yellow>LOCAL</yellow>"
+    proxy = f"<cyan>{telegram_settings.proxy}</cyan>" if telegram_settings.proxy else SETTING_UNSET
+    users_whitelist = (
+        f"<cyan>{','.join(telegram_settings.users_whitelist)}</cyan>"
+        if telegram_settings.users_whitelist
+        else SETTING_UNSET
+    )
+    groups_whitelist = (
+        f"<cyan>{telegram_settings.groups_whitelist}</cyan>" if telegram_settings.groups_whitelist else SETTING_UNSET
+    )
+    models_whitelist = (
+        f"<cyan>{', '.join(gpt_settings.models_whitelist)}</cyan>" if gpt_settings.models_whitelist else SETTING_UNSET
+    )
+    images_whitelist = (
+        f"<cyan>{','.join(gpt_settings.image_generations_whitelist)}</cyan>"
+        if gpt_settings.image_generations_whitelist
+        else SETTING_UNSET
+    )
+
+    messages = [
+        "<magenta>General Settings:</magenta>",
+        f"Application is initialized in the {mode} mode using {storage} storage.",
+        f"Proxy is {proxy}",
+        "<magenta>LLM Settings:</magenta>",
+        f"Bot name is <cyan>{telegram_settings.bot_name}</cyan>",
+        # f"Initial assistant prompt: <cyan>{gpt_settings.assistant_prompt}</cyan>",
+        f"Messages TTL: <cyan>{gpt_settings.max_conversation_age_minutes} minutes</cyan>",
+        f"Maximum conversation history size: <cyan>{gpt_settings.max_history_tokens}</cyan> tokens",
+        f"Maximum answer size: <cyan>{gpt_settings.max_tokens}</cyan> tokens",
+        f"Images generation limit: <cyan>{gpt_settings.image_generations_monthly_limit}</cyan>",
+        f"Filesystem access: {SETTING_ENABLED if gpt_settings.filesystem_access else SETTING_DISABLED}",
+        "<magenta>Whitelists:</magenta>",
+        f"Images limit whitelist: {images_whitelist}",
+        f"Users whitelist: {users_whitelist}",
+        f"Groups whitelist: {groups_whitelist}",
+        f"Models whitelist: {models_whitelist}",
+        "<magenta>Heartbeat:</magenta>",
+        f"Heartbeat mechanism: {SETTING_SET if application_settings.heartbeat_url else SETTING_UNSET}",
+    ]
+    messages += _provider_statuses()
+
+    for message in messages:
+        logger.opt(colors=True).info(message)
+
+    if application_settings.redis_password:
+        logger.opt(colors=True).warning(
+            "`REDIS_PASSWORD` environment variable is <red>deprecated</red>. Use `REDIS` instead, i.e. "
+            "`redis://:password@localhost:6379/0`"
+        )
+
+
+def handle_gpt_exceptions(func: Callable[..., Any]) -> Callable[..., Any]:
+    """Decorator handling openai module's exceptions.
+
+    If the specific exception occurred, handles it and sends the corresponding message.
+
+    Args:
+        func: async function that may rise openai exception.
+
+    Returns:
+        Wrapper function object.
+    """
+
+    @wraps(func)
+    async def wrapper(*args: Any, **kwargs: Any) -> Any:
+        update: Update = kwargs.get("update") or args[1]
+        context: ContextTypes.DEFAULT_TYPE = kwargs.get("context") or args[2]
+        error_msg_prefix = f"{user_data(update)} didn't get a GPT answer in the {chat_data(update)}"
+        try:
+            return await func(*args, **kwargs)
+        except NoApiKeyProvidedError as e:
+            logger.error(f"{error_msg_prefix}: {e}")
+            await send_message(
+                update=update,
+                context=context,
+                text="Oops! It looks like you didn't set the API key for this provider.",
+            )
+            return None
+
+        except NotAuthorizedError as e:
+            logger.error(f"{error_msg_prefix}: {e}")
+            await send_message(
+                update=update,
+                context=context,
+                text=(
+                    "We encountered an authorization problem when interacting with a remote service.\n"
+                    f"Please check your {e.provider} API key."
+                ),
+            )
+            return None
+
+        except ServiceResponseError as e:
+            logger.error(f"{error_msg_prefix}: {e}")
+
+            await send_message(
+                update=update,
+                context=context,
+                text=(
+                    f"😲Lol... we got an unexpected response from the {e.provider} service! \n"
+                    f"Please, try again a bit later."
+                ),
+            )
+            return None
+
+        except NoResponseError as e:
+            logger.error(f"{error_msg_prefix}: {e}")
+            return None
+
+        except ServiceRateLimitError as e:
+            logger.error(f"{error_msg_prefix}: {e}")
+            await send_message(
+                update=update,
+                context=context,
+                text=f"Rate Limit exceeded for {e.provider}. We should back off a bit.",
+            )
+            return None
+
+        except NoModelSelectedError as e:
+            logger.error(f"{error_msg_prefix}: {e}")
+
+            await send_message(
+                update=update,
+                context=context,
+                text="Please, select your model first.",
+            )
+            return None
+
+        except NoProviderSelectedError as e:
+            logger.error(f"{error_msg_prefix}: {e}")
+
+            await send_message(
+                update=update,
+                context=context,
+                text="Please, select your provider first.",
+            )
+            return None
+
+        except RecursionLimitExceeded as e:
+            logger.error(f"{error_msg_prefix}: {e}")
+
+            await send_message(
+                update=update,
+                context=context,
+                text=(
+                    f"{e.provider} ({e.model}) exceeded the limit on the maximum number of consecutive tool calls "
+                    f"({e.exceeded_limit}) and was stopped. The model has likely entered an infinite loop of tool "
+                    f"calls. Please check the logs. If the model was functioning as intended, you should either "
+                    f"rephrase the task or increase the value of the `MAX_CONSECUTIVE_TOOL_CALLS` setting."
+                ),
+            )
+
+        except Exception as e:
+            logger.exception(f"{error_msg_prefix}: {e!r}")
+            msg = (
+                "I'm sorry, but there seems to be a little hiccup with your request at the moment 😥 Would you mind "
+                "trying again later? Don't worry, I'll be here to assist you whenever you're ready! 😼"
+            )
+            await send_message(update=update, context=context, text=msg)
+            # raise
+
+    return wrapper
+
+
+def get_builtin_skill_names() -> list[str]:
+    path = Path(application_settings.skills_dir)
+    return [f.name for f in path.iterdir() if f.is_file()]