Social_Helper.py

import pickle, os, random
from toolbox import CatchException, update_ui, get_conf, get_log_folder, update_ui_lastest_msg
from crazy_functions.crazy_utils import input_clipping
from crazy_functions.crazy_utils import request_gpt_model_in_new_thread_with_ui_alive
from request_llms.bridge_all import predict_no_ui_long_connection
from crazy_functions.json_fns.select_tool import structure_output, select_tool
from pydantic import BaseModel, Field
from loguru import logger
from typing import List


SOCIAL_NETWOK_WORKER_REGISTER = {}

class SocialNetwork():
    def __init__(self):
        self.people = []

class SaveAndLoad():
    def __init__(self, user_name, llm_kwargs, auto_load_checkpoint=True, checkpoint_dir=None) -> None:
        self.user_name = user_name
        self.checkpoint_dir = checkpoint_dir
        if auto_load_checkpoint:
            self.social_network = self.load_from_checkpoint(checkpoint_dir)
        else:
            self.social_network = SocialNetwork()

    def does_checkpoint_exist(self, checkpoint_dir=None):
        import os, glob
        if checkpoint_dir is None: checkpoint_dir = self.checkpoint_dir
        if not os.path.exists(checkpoint_dir): return False
        if len(glob.glob(os.path.join(checkpoint_dir, "social_network.pkl"))) == 0: return False
        return True

    def save_to_checkpoint(self, checkpoint_dir=None):
        if checkpoint_dir is None: checkpoint_dir = self.checkpoint_dir
        with open(os.path.join(checkpoint_dir, 'social_network.pkl'), "wb+") as f:
            pickle.dump(self.social_network, f)
        return

    def load_from_checkpoint(self, checkpoint_dir=None):
        if checkpoint_dir is None: checkpoint_dir = self.checkpoint_dir
        if self.does_checkpoint_exist(checkpoint_dir=checkpoint_dir):
            with open(os.path.join(checkpoint_dir, 'social_network.pkl'), "rb") as f:
                social_network = pickle.load(f)
                return social_network
        else:
            return SocialNetwork()


class Friend(BaseModel):
    friend_name: str = Field(description="name of a friend")
    friend_description: str = Field(description="description of a friend (everything about this friend)")
    friend_relationship: str = Field(description="The relationship with a friend (e.g. friend, family, colleague)")

class FriendList(BaseModel):
    friends_list: List[Friend] = Field(description="The list of friends")


class SocialNetworkWorker(SaveAndLoad):
    def ai_socail_advice(self, prompt, llm_kwargs, plugin_kwargs, chatbot, history, system_prompt, run_gpt_fn, intention_type):
        pass

    def ai_remove_friend(self, prompt, llm_kwargs, plugin_kwargs, chatbot, history, system_prompt, run_gpt_fn, intention_type):
        pass

    def ai_list_friends(self, prompt, llm_kwargs, plugin_kwargs, chatbot, history, system_prompt, run_gpt_fn, intention_type):
        pass

    def ai_add_multi_friends(self, prompt, llm_kwargs, plugin_kwargs, chatbot, history, system_prompt, run_gpt_fn, intention_type):
        friend, err_msg = structure_output(
            txt=prompt,
            prompt="根据提示, 解析多个联系人的身份信息\n\n",
            err_msg=f"不能理解该联系人",
            run_gpt_fn=run_gpt_fn,
            pydantic_cls=FriendList
        )
        if friend.friends_list:
            for f in friend.friends_list: 
                self.add_friend(f)
            msg = f"成功添加{len(friend.friends_list)}个联系人: {str(friend.friends_list)}"
            yield from update_ui_lastest_msg(lastmsg=msg, chatbot=chatbot, history=history, delay=0)


    def run(self, txt, llm_kwargs, plugin_kwargs, chatbot, history, system_prompt, user_request):
        prompt = txt
        run_gpt_fn = lambda inputs, sys_prompt: predict_no_ui_long_connection(inputs=inputs, llm_kwargs=llm_kwargs, history=[], sys_prompt=sys_prompt, observe_window=[])
        self.tools_to_select = {
            "SocialAdvice":{
                "explain_to_llm": "如果用户希望获取社交指导，调用SocialAdvice生成一些社交建议",
                "callback": self.ai_socail_advice,
            },
            "AddFriends":{
                "explain_to_llm": "如果用户给出了联系人，调用AddMultiFriends把联系人添加到数据库",
                "callback": self.ai_add_multi_friends,
            },
            "RemoveFriend":{
                "explain_to_llm": "如果用户希望移除某个联系人，调用RemoveFriend",
                "callback": self.ai_remove_friend,
            },
            "ListFriends":{
                "explain_to_llm": "如果用户列举联系人，调用ListFriends",
                "callback": self.ai_list_friends,
            }
        }

        try:
            Explaination = '\n'.join([f'{k}: {v["explain_to_llm"]}' for k, v in self.tools_to_select.items()])
            class UserSociaIntention(BaseModel):
                intention_type: str = Field(
                    description=
                        f"The type of user intention. You must choose from {self.tools_to_select.keys()}.\n\n" 
                        f"Explaination:\n{Explaination}", 
                    default="SocialAdvice"
                )
            pydantic_cls_instance, err_msg = select_tool(
                prompt=txt,
                run_gpt_fn=run_gpt_fn,
                pydantic_cls=UserSociaIntention
            )
        except Exception as e:
            yield from update_ui_lastest_msg(
                lastmsg=f"无法理解用户意图 {err_msg}", 
                chatbot=chatbot, 
                history=history, 
                delay=0
            )
            return

        intention_type = pydantic_cls_instance.intention_type
        intention_callback = self.tools_to_select[pydantic_cls_instance.intention_type]['callback']
        yield from intention_callback(txt, llm_kwargs, plugin_kwargs, chatbot, history, system_prompt, run_gpt_fn, intention_type)


    def add_friend(self, friend):
        # check whether the friend is already in the social network
        for f in self.social_network.people:
            if f.friend_name == friend.friend_name:
                f.friend_description = friend.friend_description
                f.friend_relationship = friend.friend_relationship
                logger.info(f"Repeated friend, update info: {friend}")
                return
        logger.info(f"Add a new friend: {friend}")
        self.social_network.people.append(friend)
        return


@CatchException
def I人助手(txt, llm_kwargs, plugin_kwargs, chatbot, history, system_prompt, user_request):

    # 1. we retrieve worker from global context
    user_name = chatbot.get_user()
    checkpoint_dir=get_log_folder(user_name, plugin_name='experimental_rag')
    if user_name in SOCIAL_NETWOK_WORKER_REGISTER:
        social_network_worker = SOCIAL_NETWOK_WORKER_REGISTER[user_name]
    else:
        social_network_worker = SOCIAL_NETWOK_WORKER_REGISTER[user_name] = SocialNetworkWorker(
            user_name, 
            llm_kwargs, 
            checkpoint_dir=checkpoint_dir, 
            auto_load_checkpoint=True
        )

    # 2. save
    yield from social_network_worker.run(txt, llm_kwargs, plugin_kwargs, chatbot, history, system_prompt, user_request)
    social_network_worker.save_to_checkpoint(checkpoint_dir)
    yield from update_ui(chatbot=chatbot, history=history) # 刷新界面