Merge pull request #1 from pha123661/feat/tele

Add telementry

app/app.py

@@ -8,10 +8,12 @@
 import string
 import sys
 import time
+from datetime import datetime
 from argparse import ArgumentParser
 from asyncio import Semaphore
 
 import aiohttp
+import motor.motor_asyncio
 from aiohttp import web
 from aiohttp.web_runner import TCPSite
 from async_lru import alru_cache
@@ -30,6 +32,8 @@
 HF_API_HEADER: dict = None
 API_URL = "https://api-inference.huggingface.co/models/liswei/EmojiLMSeq2SeqLoRA"
 
+MONGO_CLIENT_STRING = os.getenv('MONGO_CLIENT', None)
+
 if CHANNEL_SECRET is None or CHANNEL_ACCESS_TOKEN is None or HF_API_TOKEN_LIST is None:
     print(
         "Please set LINE_CHANNEL_SECRET, LINE_CHANNEL_ACCESS_TOKEN and HF_API_TOKEN environment variables.")
@@ -53,6 +57,15 @@ def __init__(self, line_bot_api: AsyncMessagingApi, parser: WebhookParser, worke
         self.parser = parser
         self.semaphore = Semaphore(workers)
 
+        self.datacol = motor.motor_asyncio.AsyncIOMotorClient(
+            MONGO_CLIENT_STRING)["data"]["data"]
+        self.usercol = motor.motor_asyncio.AsyncIOMotorClient(
+            MONGO_CLIENT_STRING)["analysis"]["user_status"]
+        self.groupcol = motor.motor_asyncio.AsyncIOMotorClient(
+            MONGO_CLIENT_STRING)["analysis"]["group_status"]
+        self.msgcol = motor.motor_asyncio.AsyncIOMotorClient(
+            MONGO_CLIENT_STRING)["analysis"]["emoji_status"]
+
         self.last_query_time = time.time()
         self.last_query_time_lock = asyncio.Lock()
         self.keep_alive_interval = keep_alive_interval
@@ -92,17 +105,26 @@ async def __call__(self, request):
             if isinstance(event, JoinEvent):
                 logger.info(f'加入群組 {event.source.group_id}')
                 await self.send_help_message(event)
+                await self.groupcol.find_one_and_update({"_id": event.source.group_id}, {"$set": {"leave": False}, "$setOnInsert": {"first_use": datetime.fromtimestamp(event.timestamp/1000)}}, upsert=True)
             elif isinstance(event, FollowEvent):
                 logger.info(f'加入好友 {event.source.user_id}')
             elif isinstance(event, LeaveEvent):
                 logger.warning(f'幹被踢了啦 {event.source.group_id}')
+                await self.groupcol.find_one_and_update({"_id": event.source.group_id}, {"$set": {"leave": True}}, upsert=True)
             elif isinstance(event, UnfollowEvent):
                 logger.warning(f'幹被封鎖了啦 {event.source.user_id}')
+                await self.usercol.find_one_and_update({"_id": event.source.user_id}, {"$set": {"block": True, "last_block": datetime.fromtimestamp(event.timestamp/1000)}}, upsert=True)
             elif isinstance(event, MessageEvent) and isinstance(event.message, TextMessageContent):
                 await self.handle_text_message(event)
 
         return web.Response(text="OK\n")
 
+    async def update_emoji_count(self, emoji_list):
+        for emojis in emoji_list:
+            emojis = set(emojis)
+            for emoji in emojis:
+                await self.msgcol.find_one_and_update({"_id": emoji}, {"$inc": {"usage_count": 1}})
+
     async def send_help_message(self, event: MessageEvent):
         await self.line_bot_api.reply_message(
             ReplyMessageRequest(
@@ -119,8 +141,15 @@ async def handle_text_message(self, event: MessageEvent):
         if input_text == f"{self.BOT_NAME}幫幫我":
             logger.info(f"幫幫我 by {event.source.user_id}")
             await self.send_help_message(event)
+            await self.usercol.find_one_and_update(
+                {"_id": event.source.user_id},
+                {
+                    "$inc": {"help_count": 1},
+                    "$setOnInsert": {"first_use": datetime.fromtimestamp(event.timestamp/1000)}
+                },
+                upsert=True
+            )
             return
-
         if input_text.startswith(f"@{self.BOT_NAME}") or input_text.endswith(f"＠{self.BOT_NAME}"):
             input_text = input_text[len(f"@{self.BOT_NAME}"):]
         elif input_text.endswith(f"@{self.BOT_NAME}") or input_text.endswith(f"＠{self.BOT_NAME}"):
@@ -130,7 +159,7 @@ async def handle_text_message(self, event: MessageEvent):
 
         async with self.semaphore:
             await asyncio.sleep(0.1)
-            output = await generate_output(self.INPUT_TASK_PREFIX, input_text)
+            output, out_emoji_list = await generate_output(self.INPUT_TASK_PREFIX, input_text)
 
         last_query_time = time.time()
         async with self.last_query_time_lock:
@@ -144,6 +173,36 @@ async def handle_text_message(self, event: MessageEvent):
             )
         )
 
+        document = {
+            "Input": input_text,
+            "Output": output,
+            "User_ID": event.source.user_id,
+            "Creat_Time": event.timestamp
+        }
+        await self.datacol.insert_one(document)
+
+        if event.source.type == "group":
+            await self.groupcol.find_one_and_update(
+                {"_id": event.source.group_id},
+                {
+                    "$inc": {"msg_count": 1},
+                    "$set": {"last_use": datetime.fromtimestamp(event.timestamp/1000)},
+                    "$setOnInsert": {"first_use": datetime.fromtimestamp(event.timestamp/1000)}
+                },
+                upsert=True
+            )
+
+        await self.usercol.find_one_and_update(
+            {"_id": event.source.user_id},
+            {
+                "$inc": {"msg_count": 1},
+                "$set": {"last_use": datetime.fromtimestamp(event.timestamp/1000)},
+                "$setOnInsert": {"first_use": datetime.fromtimestamp(event.timestamp/1000)}
+            },
+            upsert=True
+        )
+        await self.update_emoji_count(out_emoji_list)
+
 
 async def generate_output(prefix, input_text):
     sentences_limit = 35
@@ -181,7 +240,7 @@ async def generate_output(prefix, input_text):
 
     output = "".join(output_list)
 
-    return output
+    return output, out_emoji_list
 
 
 @alru_cache(maxsize=1024)