Refactor AI class to remove reply methods and integrate AIAgent for handling replies in MessageHandler

Author: Flagro

bot/rp_bot/ai_agent/agent_tools/agent.py

@@ -1,3 +1,4 @@
+from typing import AsyncIterator
 from omnimodkit import ModelsToolkit
 from motor.motor_asyncio import AsyncIOMotorDatabase
 from .agent_toolkit import AIAgentToolkit
@@ -18,3 +19,19 @@ def __init__(
         self.toolkit = AIAgentToolkit(
             person, context, message, db, models_toolkit, prompt_manager
         )
+        self.models_toolkit = models_toolkit
+
+    async def get_reply(self, user_input: str, system_prompt: str) -> str:
+        response = await self.models_toolkit.text_model.arun_default(
+            user_input, system_prompt
+        )
+        return response.text
+
+    async def get_streaming_reply(
+        self, user_input: str, system_prompt: str
+    ) -> AsyncIterator[str]:
+        async for response in self.models_toolkit.text_model.astream_default(
+            user_input, system_prompt
+        ):
+            if response is not None:
+                yield response.text_chunk

bot/rp_bot/ai_agent/ai.py

@@ -74,21 +74,6 @@ async def generate_image(
         response = await self.models_toolkit.image_generation_model.arun_default(prompt)
         return response.image_url
 
-    async def get_reply(self, user_input: str, system_prompt: str) -> str:
-        response = await self.models_toolkit.text_model.arun_default(
-            user_input, system_prompt
-        )
-        return response.text
-
-    async def get_streaming_reply(
-        self, user_input: str, system_prompt: str
-    ) -> AsyncIterator[str]:
-        async for response in self.models_toolkit.text_model.astream_default(
-            user_input, system_prompt
-        ):
-            if response is not None:
-                yield response.text_chunk
-
     async def get_price(
         self,
         message: Message,

bot/rp_bot/messages/message_handler.py

@@ -5,6 +5,7 @@
 from ...models.handlers_input import Person, Context, Message, TranscribedMessage
 from ..auth import AllowedUser, BotAdmin, NotBanned
 from ..rp_bot_handlers import RPBotMessageHandler
+from ..ai_agent.agent_tools.agent import AIAgent
 
 
 class MessageHandler(RPBotMessageHandler):
@@ -144,7 +145,15 @@ async def get_response(
             person, context, transcribed_message
         )
         system_prompt = self.prompt_manager.get_reply_system_prompt(context)
-        response_message = await self.ai.get_reply(prompt, system_prompt)
+        ai_agent = AIAgent(
+            person,
+            context,
+            message,
+            self.db,
+            self.ai.models_toolkit,
+            self.prompt_manager,
+        )
+        response_message = await ai_agent.get_reply(prompt, system_prompt)
         result = CommandResponse(
             text="message_response", kwargs={"response_text": response_message}
         )
@@ -167,7 +176,15 @@ async def stream_get_response(
         )
         response_message = ""
         system_prompt = self.prompt_manager.get_reply_system_prompt(context)
-        async for response_message_chunk in self.ai.get_streaming_reply(
+        ai_agent = AIAgent(
+            person,
+            context,
+            message,
+            self.db,
+            self.ai.models_toolkit,
+            self.prompt_manager,
+        )
+        async for response_message_chunk in ai_agent.get_streaming_reply(
             prompt, system_prompt
         ):
             if not response_message_chunk: