mirror of
https://github.com/ParisNeo/lollms-webui.git
synced 2024-12-26 23:41:03 +00:00
261 lines
10 KiB
Python
261 lines
10 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
|
|
|
|
"""
|
|
|
|
import os
|
|
import threading
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
from typing import List
|
|
|
|
import pkg_resources
|
|
import socketio
|
|
from ascii_colors import ASCIIColors
|
|
from fastapi import APIRouter, HTTPException, Request
|
|
from fastapi.responses import FileResponse
|
|
from lollms.binding import BindingBuilder, InstallOption
|
|
from lollms.personality import AIPersonality
|
|
from lollms.security import forbid_remote_access
|
|
from lollms.server.elf_server import LOLLMSElfServer
|
|
from lollms.types import MSG_OPERATION_TYPE, SENDER_TYPES
|
|
from lollms.utilities import (convert_language_name,
|
|
find_first_available_file_index, gc, load_config,
|
|
trace_exception)
|
|
from pydantic import BaseModel
|
|
|
|
from lollms_webui import LOLLMSWebUI
|
|
|
|
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, use_threading=True):
|
|
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_OPERATION_TYPE.MSG_OPERATION_TYPE_SET_CONTENT.value,
|
|
sender_type=SENDER_TYPES.SENDER_TYPES_USER.value,
|
|
sender=ump.replace(
|
|
lollmsElfServer.config.discussion_prompt_separator, ""
|
|
).replace(":", ""),
|
|
content=prompt,
|
|
steps=[],
|
|
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
|
|
)
|
|
if use_threading:
|
|
client.generation_thread = threading.Thread(
|
|
target=lollmsElfServer.start_message_generation,
|
|
args=(message, message.id, client_id),
|
|
)
|
|
client.generation_thread.start()
|
|
else:
|
|
lollmsElfServer.start_message_generation(message, message.id, client_id)
|
|
|
|
# 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_OPERATION_TYPE.MSG_OPERATION_TYPE_SET_CONTENT.value,
|
|
sender_type=SENDER_TYPES.SENDER_TYPES_USER.value,
|
|
sender=ump.replace(
|
|
lollmsElfServer.config.discussion_prompt_separator, ""
|
|
).replace(":", ""),
|
|
content=prompt,
|
|
steps=[],
|
|
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
|
|
)
|