Refactor message handling in MessageHandler: streamline the logic for adding messages to the dialog by ensuring both user and agent messages are recorded, improving the accuracy of conversation tracking and response generation.

Author: Flagro

bot/rp_bot/messages/message_handler.py

@@ -55,7 +55,6 @@ async def get_response(
     async def stream_get_response(
         self, person: Person, context: Context, message: Message, args: List[str]
     ) -> AsyncIterator[CommandResponse]:
-        # Prepare and analyze the message (formerly prepare_get_reply logic)
         conversation_tracker_enabled = (
             await self.db.chats.get_conversation_tracker_state(context=context)
         )
@@ -149,15 +148,12 @@ async def stream_get_response(
             )
             yield response
 
-        await self.db.dialogs.add_message_to_dialog(
-            context=context,
-            person=person,
-            transcribed_message=TranscribedMessage(
-                message_text=message.message_text,
-                timestamp=message.timestamp,
-            ),
-        )
         if agent_response:
+            await self.db.dialogs.add_message_to_dialog(
+                context=context,
+                person=person,
+                transcribed_message=agent_response.transcribed_user_message,
+            )
             await self.db.user_usage.add_usage_points(
                 person=person, points=agent_response.total_price
             )
@@ -182,6 +178,14 @@ async def stream_get_response(
                 f"with usage of {agent_response.total_price}"
             )
         else:
+            await self.db.dialogs.add_message_to_dialog(
+                context=context,
+                person=person,
+                transcribed_message=TranscribedMessage(
+                    message_text=message.message_text,
+                    timestamp=message.timestamp,
+                ),
+            )
             self.logger.info(
                 f"No response generated for the message from {person.user_handle} in chat {context.chat_id}"
             )