Browse Source

Update user profile: add gender

StrayWarrior 2 weeks ago
parent
commit
57cd4a1f1f
2 changed files with 18 additions and 4 deletions
  1. 1 0
      prompt_templates.py
  2. 17 4
      user_manager.py

+ 1 - 0
prompt_templates.py

@@ -178,6 +178,7 @@ USER_PROFILE_EXTRACT_PROMPT = """
 已知信息(可能为空):
 - 姓名:{name}
 - 希望的称呼:{preferred_nickname}
+- 性别:{gender}
 - 年龄:{age}
 - 地区:{region}
 - 健康状况:{health_conditions}

+ 17 - 4
user_manager.py

@@ -113,19 +113,32 @@ class MySQLUserManager(UserManager):
         self.staff_table = staff_table
 
     def get_user_profile(self, user_id) -> Dict:
-        sql = f"SELECT name, wxid, profile_data_v1 FROM {self.table_name} WHERE third_party_user_id = {user_id}"
+        sql = f"SELECT name, wxid, profile_data_v1, gender" \
+              f" FROM {self.table_name} WHERE third_party_user_id = {user_id}"
         data = self.db.select(sql, pymysql.cursors.DictCursor)
         if not data:
             logger.error(f"user[{user_id}] not found")
             return {}
         data = data[0]
+        gender_map = {0: '未知', 1: '男', 2: '女', None: '未知'}
+        gender = gender_map[data['gender']]
+        default_profile = self.get_default_profile(nickname=data['name'], gender=gender)
         if not data['profile_data_v1']:
             logger.warning(f"user[{user_id}] profile not found, create a default one")
-            default_profile = self.get_default_profile(nickname=data['name'])
             self.save_user_profile(user_id, default_profile)
             return default_profile
         else:
-            return json.loads(data['profile_data_v1'])
+            profile = json.loads(data['profile_data_v1'])
+            # 资料条目有增加时,需合并更新
+            entry_added = False
+            for key, value in default_profile.items():
+                if key not in profile:
+                    logger.debug(f"user[{user_id}] add profile key[{key}] value[{value}]")
+                    profile[key] = value
+                    entry_added = True
+            if entry_added:
+                self.save_user_profile(user_id, profile)
+            return profile
 
     def save_user_profile(self, user_id, profile: Dict) -> None:
         if not user_id:
@@ -149,7 +162,7 @@ class MySQLUserManager(UserManager):
             return {}
         profile = data[0]
         # 转换性别格式
-        gender_map = {0: '未知', 1: '男', 2: '女'}
+        gender_map = {0: '未知', 1: '男', 2: '女', None: '未知'}
         profile['agent_gender'] = gender_map[profile['agent_gender']]
         return profile