ws_service.py

from datetime import datetime
import websockets
import asyncio
import os
import uuid
import json
import functools
import traceback
import sys
import logging
from multiprocessing import current_process, Process, Queue, queues

from common import MessageType, format_message, timestamp
import startup
user_dict = {}

KEY_TO_USE_DEFAULT = os.getenv("KEY_TO_USE_DEFAULT")
DEFAULT_LLM_API_KEY = os.getenv("DEFAULT_LLM_API_KEY") if KEY_TO_USE_DEFAULT is not None else None
DEFAULT_SERP_API_KEY = os.getenv("DEFAULT_SERP_API_KEY") if KEY_TO_USE_DEFAULT is not None else None

logging.basicConfig(level=logging.WARNING, format='%(asctime)s | %(levelname)-8s | %(module)s:%(funcName)s:%(lineno)d - %(message)s')
logger = logging.getLogger(__name__)

async def handle_message(task_id=None, message=None, alg_msg_queue=None, proxy=None, llm_api_key=None, serpapi_key=None): 
    if "llm_api_key" in message["data"] and len(message["data"]["llm_api_key"].strip()) >= 32:
        llm_api_key = message["data"]["llm_api_key"].strip()
    if KEY_TO_USE_DEFAULT is not None and \
            DEFAULT_LLM_API_KEY is not None and \
            llm_api_key == KEY_TO_USE_DEFAULT:
        # replace with default key
        logger.warning("Using default llm api key")
        llm_api_key = DEFAULT_LLM_API_KEY

    if "serpapi_key" in message["data"] and len(message["data"]["serpapi_key"].strip()) >= 32:
        serpapi_key = message["data"]["serpapi_key"].strip()
    if KEY_TO_USE_DEFAULT is not None and \
            DEFAULT_SERP_API_KEY is not None and \
            serpapi_key == KEY_TO_USE_DEFAULT:
        # replace with default key
        logger.warning("Using default serp api key")
        serpapi_key = DEFAULT_SERP_API_KEY

    idea = message["data"]["idea"].strip() 

    if not llm_api_key:
        alg_msg_queue.put_nowait(format_message(action=MessageType.RunTask.value, msg="Invalid OpenAI key"))
        return
    if not serpapi_key:
        alg_msg_queue.put_nowait(format_message(action=MessageType.RunTask.value, msg="Invalid SerpAPI key"))
        return
    if not idea or len(idea) < 2:
        alg_msg_queue.put_nowait(format_message(action=MessageType.RunTask.value, msg="Invalid task idea"))
        return
    try:
        await startup.startup(idea=idea, task_id=task_id, llm_api_key=llm_api_key, serpapi_key=serpapi_key, proxy=proxy, alg_msg_queue=alg_msg_queue)
        alg_msg_queue.put_nowait(format_message(action=MessageType.RunTask.value, data={'task_id':task_id}, msg="finished"))
    except Exception as e:
        alg_msg_queue.put_nowait(format_message(action=MessageType.RunTask.value, msg=f"{e}"))

        exc_type, exc_value, exc_traceback = sys.exc_info()
        error_message = traceback.format_exception(exc_type, exc_value, exc_traceback)
        logger.error("".join(error_message))

def handle_message_wrapper(task_id=None, message=None, alg_msg_queue=None, proxy=None, llm_api_key=None, serpapi_key=None):
    logger.warning("New task:"+current_process().name)
    asyncio.run(handle_message(task_id, message, alg_msg_queue, proxy, llm_api_key, serpapi_key))

def clear_queue(alg_msg_queue:Queue=None):
    if not Queue:
        return
    try:
        while True:
            alg_msg_queue.get_nowait()
    except queues.Empty:
        pass

# read websocket messages
async def read_msg_worker(websocket=None, alg_msg_queue=None, proxy=None, llm_api_key=None, serpapi_key=None):
    process = None
    async for raw_message in websocket:
        message = json.loads(raw_message)
        if message["action"] == MessageType.Interrupt.value:
            # force interrupt a specific task
            task_id = message["data"]["task_id"]
            if process and process.is_alive() and process.name == task_id:
                logger.warning("Interrupt task:" + process.name)
                process.terminate()
                process = None
            clear_queue(alg_msg_queue=alg_msg_queue)
            alg_msg_queue.put_nowait(format_message(action=MessageType.Interrupt.value, data={'task_id': task_id}))
            alg_msg_queue.put_nowait(format_message(action=MessageType.RunTask.value, data={'task_id': task_id}, msg="finished"))
                
        elif message["action"] == MessageType.RunTask.value:
            # auto interrupt previous task
            if process and process.is_alive():
                logger.warning("Interrupt task:" + process.name)
                process.terminate()
                process = None
                clear_queue(alg_msg_queue=alg_msg_queue)

            task_id = str(uuid.uuid4())
            process = Process(target=handle_message_wrapper, args=(task_id, message, alg_msg_queue, proxy, llm_api_key, serpapi_key))
            process.daemon = True
            process.name = task_id
            process.start()
        
    # auto terminate process
    if process and process.is_alive():
        logger.warning("Interrupt task:" + process.name)
        process.terminate()
        process = None
        clear_queue(alg_msg_queue=alg_msg_queue)
    
    raise websockets.exceptions.ConnectionClosed(0, "websocket closed")

# send
async def send_msg_worker(websocket=None, alg_msg_queue=None):
    while True:
        if alg_msg_queue.empty():
            await asyncio.sleep(0.5)
        else:
            msg = alg_msg_queue.get_nowait()
            print("=====Sending msg=====\n", msg)
            await websocket.send(msg)

async def echo(websocket, proxy=None, llm_api_key=None, serpapi_key=None):
    # audo register
    uid = datetime.strftime(datetime.now(), '%Y%m%d%H%M%S.%f')+'_'+str(uuid.uuid4())
    logger.warning(f"New user registered, uid: {uid}")
    if uid not in user_dict:
        user_dict[uid] = websocket
    else:
        logger.warning(f"Duplicate user, uid: {uid}")
        
    # message handling
    try:
        alg_msg_queue = Queue()
        await asyncio.gather(
            read_msg_worker(websocket=websocket, alg_msg_queue=alg_msg_queue, proxy=proxy, llm_api_key=llm_api_key, serpapi_key=serpapi_key), 
            send_msg_worker(websocket=websocket, alg_msg_queue=alg_msg_queue)
        )
    except websockets.exceptions.ConnectionClosed:
        logger.warning("Websocket closed: remote endpoint going away")
    finally:
        asyncio.current_task().cancel()
        # auto unregister
        logger.warning(f"Auto unregister, uid: {uid}")
       
        if uid in user_dict:
            user_dict.pop(uid)


async def run_service(host: str = "localhost", port: int=9000, proxy: str=None, llm_api_key:str=None, serpapi_key:str=None):
    message_handler = functools.partial(echo, proxy=proxy,llm_api_key=llm_api_key, serpapi_key=serpapi_key)
    async with websockets.serve(message_handler, host, port):
        logger.warning(f"Websocket server started: {host}:{port} {f'[proxy={proxy}]' if proxy else ''}")
        await asyncio.Future()