mirror of
https://github.com/ParisNeo/lollms-webui.git
synced 2024-12-21 05:13:10 +00:00
198 lines
9.1 KiB
Python
198 lines
9.1 KiB
Python
"""
|
|
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 |