live-digital-avatar-chat/helper.py

import os
import sqlite3
import sys
from enum import Enum
from multiprocessing import Queue
from settings import sqlite_file
from queue import Empty


class MessageType(Enum):
    ENTER_LIVE_ROOM = 1
    FOLLOW = 2
    LIKE = 3
    GIFT = 4
    CHAT = 5


def resource_path(relative_path, base_dir='.'):
    # PyInstaller打包后会把文件放到临时文件夹 _MEIPASS
    try:
        base_path = sys._MEIPASS
    except AttributeError:
        base_path = os.path.abspath(base_dir)
    return os.path.join(base_path, relative_path)


class LiveChatConfig:
    def __init__(self):
        self.conn = sqlite3.connect(resource_path(sqlite_file, '..'))
        self.conn.execute("PRAGMA journal_mode=WAL;")

    def _query_config(self, key):
        cursor = self.conn.cursor()
        cursor.execute('select value from config where key=?', (key,))
        result = cursor.fetchone()
        result = result[0] if result is not None else None
        cursor.close()
        return result

    @property
    def enter_live_room_prob(self):
        result = self._query_config('reply_prob_enter_live_room')
        return int(result) if result is not None else None

    @property
    def follow_prob(self):
        result = self._query_config('reply_prob_follow')
        return int(result) if result is not None else None

    @property
    def like_prob(self):
        result = self._query_config('reply_prob_like')
        return int(result) if result is not None else None

    @property
    def gift_prob(self):
        result = self._query_config('reply_prob_gift')
        return int(result) if result is not None else None

    @property
    def chat_prob(self):
        result = self._query_config('reply_prob_chat')
        return int(result) if result is not None else None

    @property
    def enter_live_room_prompt(self):
        return self._query_config('enter_live_room_prompt')

    @property
    def follow_prompt(self):
        return self._query_config('follow_prompt')

    @property
    def like_prompt(self):
        return self._query_config('like_prompt')

    @property
    def gift_prompt(self):
        return self._query_config('gift_prompt')

    @property
    def chat_prompt(self):
        return self._query_config('chat_prompt')

    @property
    def product_related_prompt(self):
        return self._query_config('product_related_prompt')

    @property
    def backend_token(self):
        return self._query_config('backend_token')

    @property
    def refine_system_message_prompt(self):
        return self._query_config('refine_system_message_prompt')

    @property
    def live_id(self):
        cursor = self.conn.cursor()
        cursor.execute("select value from live_config where key='live_id'")
        live_id = cursor.fetchone()[0]
        cursor.close()
        return live_id

    @property
    def livetalking_address(self):
        cursor = self.conn.cursor()
        cursor.execute("select value from live_config where key='livetalking_address'")
        livetalking_address = cursor.fetchone()[0]
        cursor.close()
        return livetalking_address

    @property
    def ollama_address(self):
        cursor = self.conn.cursor()
        cursor.execute("select value from live_config where key='ollama_address'")
        ollama_address = cursor.fetchone()[0]
        cursor.close()
        return ollama_address

    @property
    def system_messages(self) -> list:
        results = []
        cursor = self.conn.cursor()
        cursor.execute('select message from system_message')
        rows = cursor.fetchall()
        for message in rows:
            results.append(message[0])
        cursor.close()
        return results

    @property
    def prohibited_words(self) -> dict:
        results = {}
        cursor = self.conn.cursor()
        cursor.execute('select word, substitutes from prohibited_word')
        rows = cursor.fetchall()
        for word, substitutes in rows:
            results[word] = substitutes
        cursor.close()
        return results

    def update_chat_enable_status(self):
        cursor = self.conn.cursor()
        cursor.execute("update live_config set value = 1 where key = 'chat_enable_status'")
        self.conn.commit()
        cursor.close()


class PromptQueue:
    def __init__(self, maxsize=0):
        self.queue = Queue(maxsize)
        self.maxsize = maxsize

    def put(self, item):
        if self.queue.full():
            try:
                self.queue.get_nowait()  # 丢掉最旧的数据
            except:
                pass
        self.queue.put(item)

    def get(self, block=True, timeout=None):
        try:
            return self.queue.get(block, timeout)
        except Empty:
            return None

    def qsize(self):
        return self.queue.qsize()

    def empty(self):
        return self.queue.empty()

    def full(self):
        return self.queue.full()