Add optional user attribution for group chats

src/simulacrum.py

@@ -12,9 +12,10 @@ def __init__(self, context_file):
         self.last_completion_tokens = None
         self.warned_about_cost = False
 
-    async def chat(self, user_input, image_url):
+    async def chat(self, user_input, user_name, image_url):
         self.context.load()
         if user_input:
+            user_input = self._apply_attribution(user_input, user_name)
             self.context.add_message("user", user_input, image_url)
         completion = await ChatExecutor(self.context).execute()
         content = completion.content.strip()
@@ -71,3 +72,11 @@ def _set_stats(self, completion):
         self.last_cost = completion.cost
         self.last_prompt_tokens = completion.prompt_tokens
         self.last_completion_tokens = completion.completion_tokens
+
+    def _apply_attribution(self, input, name):
+        attribute_messages = self.context.vars.get("attribute_messages")
+        if attribute_messages and name:
+            if isinstance(attribute_messages, dict) and name in attribute_messages:
+                name = attribute_messages[name]
+            input = f"{name}: {input}"
+        return input

src/telegram/telegram_bot.py

@@ -147,7 +147,7 @@ async def do_nothing(self, *_):
         pass
 
     async def _chat(self, ctx, user_message, image_url=None):
-        response = await self.sim.chat(user_message, image_url)
+        response = await self.sim.chat(user_message, ctx.user_name, image_url)
         response = response.translate(str.maketrans("*_", "_*"))
         await ctx.send_message(response)
         await self._warn_high_cost(ctx)

src/telegram/telegram_context.py

@@ -20,6 +20,10 @@ def chat_id(self):
     def message(self):
         return self.update.message
 
+    @property
+    def user_name(self):
+        return self.update.message.from_user.first_name
+
     async def get_image_url(self):
         if not self.message.photo:
             return None