diff --git a/helper.py b/helper.py index 672d61f..adf428e 100644 --- a/helper.py +++ b/helper.py @@ -94,6 +94,30 @@ class LiveChatConfig: 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 = [] diff --git a/liveMan.py b/liveMan.py index 875d722..eb26d06 100644 --- a/liveMan.py +++ b/liveMan.py @@ -21,7 +21,6 @@ from py_mini_racer import MiniRacer from protobuf.douyin import * from message_processor import * from helper import resource_path, PromptQueue, MessageType -from settings import live_talking_host, Backend @contextmanager @@ -87,7 +86,7 @@ def generateMsToken(length=107): class DouyinLiveWebFetcher: - def __init__(self, live_id, ws_open_event, queue: PromptQueue): + def __init__(self, ws_open_event, queue: PromptQueue): """ 直播间弹幕抓取对象 :param live_id: 直播间的直播id,打开直播间web首页的链接如:https://live.douyin.com/261378947940, @@ -95,7 +94,7 @@ class DouyinLiveWebFetcher: """ self.__ttwid = None self.__room_id = None - self.live_id = live_id + self.live_id = LiveChatConfig().live_id self.live_url = "https://live.douyin.com/" self.user_agent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) " \ "Chrome/120.0.0.0 Safari/537.36" @@ -329,7 +328,7 @@ class DouyinLiveWebReply: } def _gen(): - response = requests.post(f'{Backend.backend_host}{Backend.ollama_uri}', json=payload, + response = requests.post(f'{self.live_chat_config.ollama_address}/live-digital-avatar-manage/ollama/generate', json=payload, headers={'Authorization': f'Bearer {self.live_chat_config.backend_token}'}, stream=True) buffer = '' @@ -352,7 +351,7 @@ class DouyinLiveWebReply: if stream: return _gen() else: - response = requests.post(f'{Backend.backend_host}{Backend.ollama_uri}', json=payload, + response = requests.post(f'{self.live_chat_config.ollama_address}/live-digital-avatar-manage/ollama/generate', json=payload, headers={'Authorization': f'Bearer {self.live_chat_config.backend_token}'}, timeout=10).content.decode()[5:] response = json.loads(response) @@ -368,7 +367,7 @@ class DouyinLiveWebReply: async def post_to_human(self, text: str): async with httpx.AsyncClient() as client: await client.post( - f'{live_talking_host}/human', + f'{self.live_chat_config.livetalking_address}/human', json={ "type": "echo", "sessionid": self.session_id, @@ -384,7 +383,7 @@ class DouyinLiveWebReply: live_chat_config.update_chat_enable_status() while True: try: - is_speaking = requests.post(f'{live_talking_host}/is_speaking', json={'sessionid': self.session_id}, + is_speaking = requests.post(f'{self.live_chat_config.livetalking_address}/is_speaking', json={'sessionid': self.session_id}, timeout=5).json()['data'] if is_speaking: time.sleep(0.1) diff --git a/main.py b/main.py index 233febe..36929f0 100644 --- a/main.py +++ b/main.py @@ -1,18 +1,10 @@ from liveMan import DouyinLiveWebFetcher, DouyinLiveWebReply -from argparse import ArgumentParser from helper import PromptQueue from multiprocessing import Process, Event, freeze_support -def parse_args(): - parser = ArgumentParser() - parser.add_argument('--live_id', type=str, required=True, help='直播间id') - args = parser.parse_args() - return args - - -def fetch_user_chat_content(live_id, ws_open_event, queue): - fetcher = DouyinLiveWebFetcher(live_id, ws_open_event, queue) +def fetch_user_chat_content(ws_open_event, queue): + fetcher = DouyinLiveWebFetcher(ws_open_event, queue) fetcher.start() @@ -23,12 +15,11 @@ def reply_user_chat_content(queue): if __name__ == '__main__': freeze_support() - args = parse_args() queue = PromptQueue(10) ws_open_event = Event() - fetch_process = Process(target=fetch_user_chat_content, args=(args.live_id, ws_open_event, queue)) + fetch_process = Process(target=fetch_user_chat_content, args=(ws_open_event, queue)) reply_process = Process(target=reply_user_chat_content, args=(queue,)) fetch_process.start() ws_open_event.wait() diff --git a/settings.py b/settings.py index 2dbdd88..9ba3857 100644 --- a/settings.py +++ b/settings.py @@ -1,8 +1 @@ sqlite_file = 'live_chat.db' - -live_talking_host = 'http://192.168.10.96:8010' - -class Backend: - backend_host = 'http://192.168.10.25:9909' - ollama_uri = '/live-digital-avatar-manage/ollama/generate' - login_uri = '/live-digital-avatar-manage/auth/login'