Refactor RPBot to replace AI class with ModelsToolkit for handling model interactions and streamline message processing in MessageHandler

Author: Flagro

bot/rp_bot/ai_agent/ai.py

@@ -1,6 +1,7 @@
 from typing import List, Union, Type
 from logging import Logger
 from omnimodkit.ai_config import AIConfig
+from omnimodkit.models_toolkit import ModelsToolkit
 from ..models.base_bot import BaseBot
 from .rp_bot_handlers import (
     RPBotCommandHandler,
@@ -47,11 +48,9 @@ def get_rp_bot(
         default_language=bot_config.default_language,
     )
     prompt_manager = PromptManager(db=db)
-    ai = AI(
+    models_toolkit = ModelsToolkit(
         openai_api_key=openai_api_key,
         ai_config=ai_config,
-        prompt_manager=prompt_manager,
-        db=db,
     )
     auth = Auth(
         allowed_handles=allowed_handles,
@@ -60,7 +59,7 @@ def get_rp_bot(
     )
     return RPBot(
         db=db,
-        ai=ai,
+        models_toolkit=models_toolkit,
         localizer=localizer,
         prompt_manager=prompt_manager,
         auth=auth,
@@ -73,15 +72,15 @@ class RPBot(BaseBot):
     def __init__(
         self,
         db: DB,
-        ai: AI,
+        models_toolkit: ModelsToolkit,
         localizer: Localizer,
         prompt_manager: PromptManager,
         auth: Auth,
         bot_config: BotConfig,
         logger: Logger,
     ):
         self.db = db
-        self.ai = ai
+        self.models_toolkit = models_toolkit
         self.localizer = localizer
         self.prompt_manager = prompt_manager
         self.auth = auth
@@ -101,7 +100,7 @@ def _init_handlers(
         return [
             handler(
                 db=self.db,
-                ai=self.ai,
+                models_toolkit=self.models_toolkit,
                 localizer=self.localizer,
                 prompt_manager=self.prompt_manager,
                 auth=self.auth,

bot/rp_bot/bot.py

@@ -11,23 +11,44 @@
 class MessageHandler(RPBotMessageHandler):
     permission_classes = (AllowedUser, BotAdmin, NotBanned)
 
+    async def estimate_price(self, message: Message) -> float:
+        """
+        Estimates the price of the message
+        """
+        return self.models_toolkit.estimate_price(
+            input_text=message.message_text,
+            input_image=message.in_file_image,
+            input_audio=message.in_file_audio,
+        )
+
     async def _get_user_usage(
         self, input_message: Message, generated_message: str
     ) -> int:
-        return self.ai.get_price(
-            message=input_message, generated_message=generated_message
+        return self.models_toolkit.get_price(
+            input_text=input_message.message_text,
+            output_text=generated_message,
+            input_image=input_message.in_file_image,
+            input_audio=input_message.in_file_audio,
         )
 
     async def _get_transcribed_message(self, message: Message) -> TranscribedMessage:
         # Note that here the responsibility to pass NULL images and Audio is on the
         # outer level bot processing (TG bot or other bot)
         image_description = (
-            await self.ai.describe_image(message.in_file_image)
+            str(
+                await self.models_toolkit.vision_model.arun_default(
+                    message.in_file_image
+                )
+            )
             if message.in_file_image
             else None
         )
         voice_description = (
-            await self.ai.transcribe_audio(message.in_file_audio)
+            str(
+                await self.models_toolkit.audio_recognition_model.arun_default(
+                    message.in_file_audio
+                )
+            )
             if message.in_file_audio
             else None
         )
@@ -66,7 +87,12 @@ async def prepare_get_reply(
         autoengage_state = await self.db.chats.get_autoengage_state(context)
         engage_is_needed = False
         if autoengage_state:
-            engage_is_needed = await self.ai.engage_is_needed(message)
+            prompt = await self.prompt_manager.compose_engage_needed_prompt(
+                message.message_text
+            )
+            engage_is_needed = (
+                await self.models_toolkit.text_model.async_ask_yes_no_question(prompt)
+            )
         if not context.is_bot_mentioned and not engage_is_needed:
             self.logger.info(
                 f"Saving the message from {person.user_handle} in chat {context.chat_id} "
@@ -83,7 +109,7 @@ async def prepare_get_reply(
     async def is_usage_under_limit(
         self, person: Person, context: Context, transcribed_message: TranscribedMessage
     ) -> bool:
-        estimated_usage = await self.ai.estimate_price(transcribed_message)
+        estimated_usage = await self.estimate_price(transcribed_message)
         user_usage = await self.db.user_usage.get_user_usage(person)
         user_limit = await self.db.user_usage.get_user_usage_limit(person)
         return user_usage + estimated_usage < user_limit
@@ -150,7 +176,7 @@ async def get_response(
             context,
             message,
             self.db,
-            self.ai.models_toolkit,
+            self.models_toolkit,
             self.prompt_manager,
         )
         response_message = await ai_agent.get_reply(prompt, system_prompt)
@@ -181,7 +207,7 @@ async def stream_get_response(
             context,
             message,
             self.db,
-            self.ai.models_toolkit,
+            self.models_toolkit,
             self.prompt_manager,
         )
         async for response_message_chunk in ai_agent.get_streaming_reply(

bot/rp_bot/messages/message_handler.py

@@ -1,8 +1,7 @@
 from abc import ABC, abstractmethod
 from typing import List, Optional
-
+from omnimodkit.models_toolkit import ModelsToolkit
 from .db import DB
-from ai_agent.ai import AI
 from .prompt_manager import PromptManager
 from .auth import Auth
 from .localizer import Localizer
@@ -20,7 +19,7 @@ class RPBotHandlerMixin(ABC):
     def __init__(
         self,
         db: DB,
-        ai: AI,
+        models_toolkit: ModelsToolkit,
         localizer: Localizer,
         prompt_manager: PromptManager,
         auth: Auth,
@@ -29,7 +28,7 @@ def __init__(
         **kwargs
     ):
         self.db = db
-        self.ai = ai
+        self.models_toolkit = models_toolkit
         self.localizer = localizer
         self.prompt_manager = prompt_manager
         self.auth = auth
@@ -48,7 +47,10 @@ def get_initialized_permissions(self) -> List[BasePermission]:
         return result
 
     async def get_localized_text(
-        self, text: str, kwargs: Optional[dict] = None, context: Optional[Context] = None
+        self,
+        text: str,
+        kwargs: Optional[dict] = None,
+        context: Optional[Context] = None,
     ) -> Optional[str]:
         return await self.localizer.get_command_response(
             text=text, kwargs=kwargs, context=context