Refactor model definitions and message handling: update Person model to use optional fields, enhance asynchronous context handling in database operations, and streamline prompt generation in MessageHandler for improved clarity and efficiency.

Author: Flagro

bot/models/handlers_input.py

@@ -7,8 +7,8 @@
 
 class Person(BaseModel):
     user_handle: str
-    first_name: str
-    last_name: str
+    first_name: Optional[str] = None
+    last_name: Optional[str] = None
 
 
 class Context(BaseModel):

bot/rp_bot/ai_agent/agent_tools/agent.py

@@ -38,7 +38,7 @@ async def get_streaming_reply(
     ) -> AsyncIterator[AIAgentStreamingResponse]:
         total_text = ""
         async for response in self.models_toolkit.text_model.astream_default(
-            user_input, system_prompt
+            user_input=user_input, system_prompt=system_prompt
         ):
             total_text += response.text_chunk
             if response is not None:

bot/rp_bot/db.py

@@ -48,8 +48,8 @@ def __init__(
 
     async def create_if_not_exists(self, person: Person, context: Context) -> None:
         for model in self.models:
-            await model.create_if_not_exists(person, context)
+            await model.create_if_not_exists(context=context, person=person)
 
     async def update_if_needed(self, person: Person, context: Context) -> None:
         for model in self.models:
-            await model.update_if_needed(person, context)
+            await model.update_if_needed(person=person, context=context)

bot/rp_bot/db_models/chat_modes.py

@@ -25,9 +25,7 @@ def __init__(
         self.chat_modes = db.chat_modes
         self.default_chat_modes = default_chat_modes
 
-    async def create_chat_modes_if_not_exist(
-        self, person: Person, context: Context
-    ) -> None:
+    async def create_if_not_exists(self, context: Context, person: Person) -> None:
         for _, mode in self.default_chat_modes.default_chat_modes.items():
             self.chat_modes.update_one(
                 {"chat_id": context.chat_id, "mode_name": mode.name},
@@ -53,7 +51,9 @@ async def get_chat_mode(self, context: Context) -> ChatModeResponse:
             # find the first mode by default
             chat_data = await self.chat_modes.find_one({"chat_id": chat_id})
         return ChatModeResponse(
-            chat_data["_id"], chat_data["mode_name"], chat_data["mode_description"]
+            id=chat_data["_id"],
+            mode_name=chat_data["mode_name"],
+            mode_description=chat_data["mode_description"],
         )
 
     async def get_mode_name_by_id(self, context: Context, mode_id: str) -> str:

bot/rp_bot/db_models/user_facts.py

@@ -26,7 +26,7 @@ async def get_user_facts(self, context: Context, person: Person) -> List[str]:
         facts = await self.user_facts.find_one(
             {"chat_id": context.chat_id, "user_handle": person.user_handle}
         )
-        return facts.get("facts", [])
+        return facts.get("facts", []) if facts else []
 
     async def get_facts_for_user_handle(
         self, context: Context, user_handle: str
@@ -43,7 +43,7 @@ async def get_facts_for_user_handle(
         facts = await self.user_facts.find_one(
             {"chat_id": context.chat_id, "user_handle": user_handle}
         )
-        return facts.get("facts", [])
+        return facts.get("facts", []) if facts else []
 
     async def add_fact(self, context: Context, person: Person, fact: str) -> None:
         await self.user_facts.update_one(

bot/rp_bot/db_models/user_introductions.py

@@ -22,4 +22,4 @@ async def get_user_introduction(self, context: Context, person: Person) -> str:
         introduction = await self.user_introductions.find_one(
             {"chat_id": context.chat_id, "user_handle": person.user_handle}
         )
-        return introduction.get("introduction", "")
+        return introduction.get("introduction", "") if introduction else ""

bot/rp_bot/localizer.py

@@ -22,7 +22,7 @@ async def get_command_response(
         if context is None:
             language = self.default_language
         else:
-            language = self.db.chats.get_language(context)
+            language = await self.db.chats.get_language(context)
         if text not in self.translations.translations:
             return None
         localizer_translation = self.translations.translations[text]

bot/rp_bot/messages/message_handler.py

@@ -77,7 +77,11 @@ async def get_usage_over_limit_response(self, person: Person) -> CommandResponse
     async def get_prompt_from_transcribed_message(
         self, person: Person, context: Context, transcribed_message: TranscribedMessage
     ) -> str:
-        prompt = await self.prompt_manager.compose_prompt(transcribed_message, context)
+        prompt = await self.prompt_manager.compose_prompt(
+            initiator=person,
+            context=context,
+            user_transcribed_message=transcribed_message,
+        )
         self.logger.info(
             "Using AI to generate a response to the message from "
             f"{person.user_handle} in chat {context.chat_id}"
@@ -155,15 +159,20 @@ async def stream_get_response(
                 "as the agent detected a question or a request for information"
             )
 
-        if not self.is_usage_under_limit(person, context, message):
+        if not await self.is_usage_under_limit(person, context, message):
             yield await self.get_usage_over_limit_response(person)
             return
 
         prompt = await self.get_prompt_from_transcribed_message(
-            person, context, message
+            person,
+            context,
+            TranscribedMessage(
+                message_text=message.message_text,
+                timestamp=message.timestamp,
+            ),
         )
         response_message = ""
-        system_prompt = self.prompt_manager.get_reply_system_prompt(context)
+        system_prompt = await self.prompt_manager.get_reply_system_prompt(context)
         ai_agent = AIAgent(
             person,
             context,

bot/rp_bot/prompt_manager.py

@@ -114,7 +114,7 @@ async def compose_prompt(
     async def get_reply_system_prompt(self, context: Context) -> str:
         chat_name = context.chat_name
         chat_mode_prompt = await self._compose_chat_mode_prompt(context)
-        chat_language = self.db.chats.get_language(context)
+        chat_language = await self.db.chats.get_language(context)
         return (
             "You are a helpful assistant. "
             f"You are currently in the chat: {chat_name}. "