AgentCoreService/user_manager.py

#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
import logging
from typing import Dict, Optional, Tuple, Any, List
import json
import time
import os
import abc

import pymysql.cursors

import configs
from database import MySQLManager

class UserManager(abc.ABC):
    @abc.abstractmethod
    def get_user_profile(self, user_id) -> Dict:
        pass

    @abc.abstractmethod
    def save_user_profile(self, user_id, profile: Dict) -> None:
        pass

    @abc.abstractmethod
    def list_all_users(self):
        pass

    @staticmethod
    def get_default_profile(**kwargs) -> Dict:
        default_profile = {
            "name": "",
            "nickname": "",
            "preferred_nickname": "",
            "age": 0,
            "region": '',
            "interests": [],
            "family_members": {},
            "health_conditions": [],
            "medications": [],
            "reminder_preferences": {
                "medication": True,
                "health": True,
                "weather": True,
                "news": False
            },
            "interaction_style": "standard",  # standard, verbose, concise
            "interaction_frequency": "medium",  # low, medium, high
            "last_topics": [],
            "created_at": int(time.time() * 1000),
            "human_intervention_history": []
        }
        for key, value in kwargs.items():
            if key in default_profile:
                default_profile[key] = value
        return default_profile

class UserRelationManager(abc.ABC):
    @abc.abstractmethod
    def list_staffs(self):
        pass

    @abc.abstractmethod
    def list_users(self, staff_id: str, page: int = 1, page_size: int = 100):
        pass

    @abc.abstractmethod
    def list_staff_users(self) -> List[Dict]:
        pass

class LocalUserManager(UserManager):
    def get_user_profile(self, user_id) -> Dict:
        """加载用户个人资料，如不存在则创建默认资料。主要用于本地调试"""
        try:
            with open(f"user_profiles/{user_id}.json", "r", encoding="utf-8") as f:
                return json.load(f)
        except FileNotFoundError:
            # 创建默认用户资料
            default_profile = self.get_default_profile()
            self.save_user_profile(user_id, default_profile)
            return default_profile

    def save_user_profile(self, user_id, profile: Dict) -> None:
        if not user_id:
            raise Exception("Invalid user_id: {}".format(user_id))
        with open(f"user_profiles/{user_id}.json", "w", encoding="utf-8") as f:
            json.dump(profile, f, ensure_ascii=False, indent=2)

    def list_all_users(self):
        user_ids = []
        for root, dirs, files in os.walk('user_profiles/'):
            for file in files:
                if file.endswith('.json'):
                    user_ids.append(os.path.splitext(file)[0])
        return user_ids


class MySQLUserManager(UserManager):
    def __init__(self, db_config, table_name):
        self.db = MySQLManager(db_config)
        self.table_name = table_name

    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}"
        data = self.db.select(sql, pymysql.cursors.DictCursor)
        if not data:
            logging.error(f"user[{user_id}] not found")
            return {}
        data = data[0]
        if not data['profile_data_v1']:
            logging.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 json.loads(data['profile_data_v1'])

    def save_user_profile(self, user_id, profile: Dict) -> None:
        if not user_id:
            raise Exception("Invalid user_id: {}".format(user_id))
        sql = f"UPDATE {self.table_name} SET profile_data_v1 = %s WHERE third_party_user_id = {user_id}"
        self.db.execute(sql, (json.dumps(profile),))

    def list_all_users(self):
        sql = f"SELECT third_party_user_id FROM {self.table_name}"
        data = self.db.select(sql, pymysql.cursors.DictCursor)
        return [user['third_party_user_id'] for user in data]


class MySQLUserRelationManager(UserRelationManager):
    def __init__(self, agent_db_config, wecom_db_config,
                 agent_staff_table, agent_user_table,
                 staff_table, relation_table, user_table):
        # FIXME(zhoutian): 因为现在数据库表不统一，需要从两个库读取
        self.agent_db = MySQLManager(agent_db_config)
        self.wecom_db = MySQLManager(wecom_db_config)
        self.agent_staff_table = agent_staff_table
        self.staff_table = staff_table
        self.relation_table = relation_table
        self.agent_user_table = agent_user_table
        self.user_table = user_table

    def list_staffs(self):
        return []

    def list_users(self, staff_id: str, page: int = 1, page_size: int = 100):
        return []

    def list_staff_users(self):
        sql = f"SELECT third_party_user_id, wxid FROM {self.agent_staff_table} WHERE status = 1"
        agent_staff_data = self.agent_db.select(sql, pymysql.cursors.DictCursor)
        if not agent_staff_data:
            return []
        ret = []
        for agent_staff in agent_staff_data:
            wxid = agent_staff['wxid']
            sql = f"SELECT id FROM {self.staff_table} WHERE carrier_id = '{wxid}'"
            staff_data = self.wecom_db.select(sql, pymysql.cursors.DictCursor)
            if not staff_data:
                logging.error(f"staff[{wxid}] not found in wecom database")
                continue
            staff_id = staff_data[0]['id']
            sql = f"SELECT user_id FROM {self.relation_table} WHERE staff_id = '{staff_id}' AND is_delete = 0"
            user_data = self.wecom_db.select(sql, pymysql.cursors.DictCursor)
            if not user_data:
                logging.warning(f"staff[{wxid}] has no user")
                continue
            user_ids = tuple(user['user_id'] for user in user_data)
            sql = f"SELECT union_id FROM {self.user_table} WHERE id IN {str(user_ids)}"
            user_data = self.wecom_db.select(sql, pymysql.cursors.DictCursor)
            if not user_data:
                logging.error(f"staff[{wxid}] users not found in wecom database")
                continue
            user_union_ids = tuple(user['union_id'] for user in user_data)
            batch_size = 100
            n_batches = (len(user_union_ids) + batch_size - 1) // batch_size
            agent_user_data = []
            for i in range(n_batches):
                idx_begin = i * batch_size
                idx_end = min((i + 1) * batch_size, len(user_union_ids))
                batch_union_ids = user_union_ids[idx_begin:idx_end]
                sql = f"SELECT third_party_user_id, wxid FROM {self.agent_user_table} WHERE wxid IN {str(batch_union_ids)}"
                batch_agent_user_data = self.agent_db.select(sql, pymysql.cursors.DictCursor)
                if len(agent_user_data) != len(batch_union_ids):
                    logging.error(f"staff[{wxid}] some users not found in agent database")
                agent_user_data.extend(batch_agent_user_data)
            staff_user_pairs = [
                {
                    'staff_id': agent_staff['third_party_user_id'],
                    'user_id': agent_user['third_party_user_id']
                }
                for agent_user in agent_user_data
            ]
            ret.extend(staff_user_pairs)
        return ret


if __name__ == '__main__':
    config = configs.get()
    user_db_config = config['storage']['user']
    user_manager = MySQLUserManager(user_db_config['mysql'], user_db_config['table'])
    user_profile = user_manager.get_user_profile('7881301263964433')
    print(user_profile)

    wecom_db_config = config['storage']['user_relation']
    user_relation_manager = MySQLUserRelationManager(
        user_db_config['mysql'], wecom_db_config['mysql'],
        config['storage']['staff']['table'],
        user_db_config['table'],
        wecom_db_config['table']['staff'],
        wecom_db_config['table']['relation'],
        wecom_db_config['table']['user']
    )
    all_staff_users = user_relation_manager.list_staff_users()
    print(all_staff_users)