lollms-webui/events/lollms_generation_events.py

"""
project: lollms
file: lollms_generation_events.py 
author: ParisNeo
description: 
    This module contains a set of Socketio routes that provide information about the Lord of Large Language and Multimodal Systems (LoLLMs) Web UI
    application. These routes are specific to text generation operation

"""
from fastapi import APIRouter, Request
from fastapi import HTTPException
from pydantic import BaseModel
import pkg_resources
from lollms.server.elf_server import LOLLMSElfServer
from fastapi.responses import FileResponse
from lollms.binding import BindingBuilder, InstallOption
from ascii_colors import ASCIIColors
from lollms.personality import MSG_TYPE, AIPersonality
from lollms.types import MSG_TYPE, SENDER_TYPES
from lollms.utilities import load_config, trace_exception, gc
from lollms.utilities import find_first_available_file_index, convert_language_name
from lollms.security import forbid_remote_access
from lollms_webui import LOLLMSWebUI
from pathlib import Path
from typing import List
import socketio
import threading
from datetime import datetime
import os

router = APIRouter()
lollmsElfServer = LOLLMSWebUI.get_instance()


# ----------------------------------- events -----------------------------------------
def add_events(sio:socketio):
    forbid_remote_access(lollmsElfServer)
    @sio.on('generate_msg')
    def handle_generate_msg(sid, data):        
        client_id = sid
        lollmsElfServer.cancel_gen = False
        client = lollmsElfServer.session.get_client(client_id)

        client.generated_text=""
        client.cancel_generation=False
        client.continuing=False
        client.first_chunk=True
        
       
        if not lollmsElfServer.model:
            ASCIIColors.error("Model not selected. Please select a model")
            lollmsElfServer.error("Model not selected. Please select a model", client_id=client_id)
            return

        if not lollmsElfServer.busy:
            if lollmsElfServer.session.get_client(client_id).discussion is None:
                if lollmsElfServer.db.does_last_discussion_have_messages():
                    lollmsElfServer.session.get_client(client_id).discussion = lollmsElfServer.db.create_discussion()
                else:
                    lollmsElfServer.session.get_client(client_id).discussion = lollmsElfServer.db.load_last_discussion()

            prompt = data["prompt"]
            try:
                nb_tokens = len(lollmsElfServer.model.tokenize(prompt))
            except:
                nb_tokens = None
            created_at = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            ump = lollmsElfServer.config.discussion_prompt_separator +lollmsElfServer.config.user_name.strip() if lollmsElfServer.config.use_user_name_in_discussions else lollmsElfServer.personality.user_message_prefix
            message = lollmsElfServer.session.get_client(client_id).discussion.add_message(
                message_type    = MSG_TYPE.MSG_TYPE_FULL.value,
                sender_type     = SENDER_TYPES.SENDER_TYPES_USER.value,
                sender          = ump.replace(lollmsElfServer.config.discussion_prompt_separator,"").replace(":",""),
                content=prompt,
                metadata=None,
                parent_message_id=lollmsElfServer.message_id,
                created_at=created_at,
                nb_tokens=nb_tokens
            )

            ASCIIColors.green("Starting message generation by "+lollmsElfServer.personality.name)
            client.generation_thread = threading.Thread(target=lollmsElfServer.start_message_generation, args=(message, message.id, client_id))
            client.generation_thread.start()
            
            # lollmsElfServer.sio.sleep(0.01)
            ASCIIColors.info("Started generation task")
            lollmsElfServer.busy=True
            #tpe = threading.Thread(target=lollmsElfServer.start_message_generation, args=(message, message_id, client_id))
            #tpe.start()
        else:
            lollmsElfServer.error("I am busy. Come back later.", client_id=client_id)
    @sio.on('generate_msg_with_internet')
    def generate_msg_with_internet(sid, data):        
        client_id = sid
        lollmsElfServer.cancel_gen = False
        client = lollmsElfServer.session.get_client(client_id)

        client.generated_text=""
        client.cancel_generation=False
        client.continuing=False
        client.first_chunk=True
        

        
        if not lollmsElfServer.model:
            ASCIIColors.error("Model not selected. Please select a model")
            lollmsElfServer.error("Model not selected. Please select a model", client_id=client_id)
            return

        if not lollmsElfServer.busy:
            if lollmsElfServer.session.get_client(client_id).discussion is None:
                if lollmsElfServer.db.does_last_discussion_have_messages():
                    lollmsElfServer.session.get_client(client_id).discussion = lollmsElfServer.db.create_discussion()
                else:
                    lollmsElfServer.session.get_client(client_id).discussion = lollmsElfServer.db.load_last_discussion()

            prompt = data["prompt"]
            try:
                nb_tokens = len(lollmsElfServer.model.tokenize(prompt))
            except:
                nb_tokens = None
            created_at = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            ump = lollmsElfServer.config.discussion_prompt_separator +lollmsElfServer.config.user_name.strip() if lollmsElfServer.config.use_user_name_in_discussions else lollmsElfServer.personality.user_message_prefix
            message = lollmsElfServer.session.get_client(client_id).discussion.add_message(
                message_type    = MSG_TYPE.MSG_TYPE_FULL.value,
                sender_type     = SENDER_TYPES.SENDER_TYPES_USER.value,
                sender          = ump.replace(lollmsElfServer.config.discussion_prompt_separator,"").replace(":",""),
                content=prompt,
                metadata=None,
                parent_message_id=lollmsElfServer.message_id,
                created_at=created_at,
                nb_tokens=nb_tokens
            )

            ASCIIColors.green("Starting message generation by "+lollmsElfServer.personality.name)
            
            client.generation_thread = threading.Thread(target=lollmsElfServer.start_message_generation, args=(message, message.id, client_id, False, None, True))
            client.generation_thread.start()
            
            # lollmsElfServer.sio.sleep(0.01)
            ASCIIColors.info("Started generation task")
            lollmsElfServer.busy=True
            #tpe = threading.Thread(target=lollmsElfServer.start_message_generation, args=(message, message_id, client_id))
            #tpe.start()
        else:
            lollmsElfServer.error("I am busy. Come back later.", client_id=client_id)




    @sio.on('generate_msg_from')
    def handle_generate_msg_from(sid, data):
        client_id = sid
        client = lollmsElfServer.session.get_client(client_id)
        lollmsElfServer.cancel_gen = False
        client.continuing=False
        client.first_chunk=True
        
        if lollmsElfServer.session.get_client(client_id).discussion is None:
            ASCIIColors.warning("Please select a discussion")
            lollmsElfServer.error("Please select a discussion first", client_id=client_id)
            return
        id_ = data['id']
        generation_type = data.get('msg_type',None)
        if id_==-1:
            message = lollmsElfServer.session.get_client(client_id).discussion.current_message
        else:
            message = lollmsElfServer.session.get_client(client_id).discussion.load_message(id_)
        if message is None:
            return            
        client.generation_thread = threading.Thread(target=lollmsElfServer.start_message_generation, args=(message, message.id, client_id, False, generation_type))
        client.generation_thread.start()

    @sio.on('continue_generate_msg_from')
    def handle_continue_generate_msg_from(sid, data):
        client_id = sid
        client = lollmsElfServer.session.get_client(client_id)
        lollmsElfServer.cancel_gen = False
        client.continuing=True
        client.first_chunk=True
        
        if lollmsElfServer.session.get_client(client_id).discussion is None:
            ASCIIColors.yellow("Please select a discussion")
            lollmsElfServer.error("Please select a discussion", client_id=client_id)
            return
        id_ = data['id']
        if id_==-1:
            message = lollmsElfServer.session.get_client(client_id).discussion.current_message
        else:
            message = lollmsElfServer.session.get_client(client_id).discussion.load_message(id_)

        client.generated_text=message.content
        client.generation_thread = threading.Thread(target=lollmsElfServer.start_message_generation, args=(message, message.id, client_id, True))
        client.generation_thread.start()

    #add functions to lollm
    lollmsElfServer.handle_generate_msg                 = handle_generate_msg
    lollmsElfServer.generate_msg_with_internet          = generate_msg_with_internet
    lollmsElfServer.handle_continue_generate_msg_from   = handle_continue_generate_msg_from