lollms-webui/endpoints/lollms_discussion.py
2024-01-09 23:26:41 +01:00

193 lines
7.3 KiB
Python

"""
project: lollms_webui
file: lollms_discussion.py
author: ParisNeo
description:
This module contains a set of FastAPI routes that provide information about the Lord of Large Language and Multimodal Systems (LoLLMs) Web UI
application. These routes allow users to manipulate the discussion elements.
"""
from fastapi import APIRouter, Request
from lollms_webui import LOLLMSWebUI
from pydantic import BaseModel
from starlette.responses import StreamingResponse
from lollms.types import MSG_TYPE
from lollms.utilities import detect_antiprompt, remove_text_from_string, trace_exception
from ascii_colors import ASCIIColors
from api.db import DiscussionsDB, Discussion
from safe_store.text_vectorizer import TextVectorizer, VectorizationMethod, VisualizationMethod
import tqdm
class GenerateRequest(BaseModel):
text: str
class DatabaseSelectionParameters(BaseModel):
name: str
class EditTitleParameters(BaseModel):
client_id: str
title: str
id: int
class MakeTitleParameters(BaseModel):
id: int
class DeleteDiscussionParameters(BaseModel):
client_id: str
id: int
# ----------------------- Defining router and main class ------------------------------
router = APIRouter()
lollmsElfServer:LOLLMSWebUI = LOLLMSWebUI.get_instance()
@router.get("/list_discussions")
def list_discussions():
discussions = lollmsElfServer.db.get_discussions()
return discussions
@router.get("/list_databases")
async def list_databases():
"""List all the personal databases in the LoLLMs server."""
# Retrieve the list of database names
databases = [f.name for f in lollmsElfServer.lollms_paths.personal_databases_path.iterdir() if f.suffix == ".db"]
# Return the list of database names
return databases
@router.post("/select_database")
def select_database(data:DatabaseSelectionParameters):
if not data.name.endswith(".db"):
data.name += ".db"
print(f'Selecting database {data.name}')
# Create database object
lollmsElfServer.db = DiscussionsDB(lollmsElfServer.lollms_paths.personal_databases_path/data.name)
ASCIIColors.info("Checking discussions database... ",end="")
lollmsElfServer.db.create_tables()
lollmsElfServer.db.add_missing_columns()
lollmsElfServer.config.db_path = data.name
ASCIIColors.success("ok")
if lollmsElfServer.config.auto_save:
lollmsElfServer.config.save_config()
if lollmsElfServer.config.data_vectorization_activate and lollmsElfServer.config.use_discussions_history:
try:
ASCIIColors.yellow("0- Detected discussion vectorization request")
folder = lollmsElfServer.lollms_paths.personal_databases_path/"vectorized_dbs"
folder.mkdir(parents=True, exist_ok=True)
lollmsElfServer.long_term_memory = TextVectorizer(
vectorization_method=VectorizationMethod.TFIDF_VECTORIZER,#=VectorizationMethod.BM25_VECTORIZER,
database_path=folder/lollmsElfServer.config.db_path,
data_visualization_method=VisualizationMethod.PCA,#VisualizationMethod.PCA,
save_db=True
)
ASCIIColors.yellow("1- Exporting discussions")
lollmsElfServer.info("Exporting discussions")
discussions = lollmsElfServer.db.export_all_as_markdown_list_for_vectorization()
ASCIIColors.yellow("2- Adding discussions to vectorizer")
lollmsElfServer.info("Adding discussions to vectorizer")
index = 0
nb_discussions = len(discussions)
for (title,discussion) in tqdm(discussions):
lollmsElfServer.socketio.emit('update_progress',{'value':int(100*(index/nb_discussions))})
index += 1
if discussion!='':
skill = lollmsElfServer.learn_from_discussion(title, discussion)
lollmsElfServer.long_term_memory.add_document(title, skill, chunk_size=lollmsElfServer.config.data_vectorization_chunk_size, overlap_size=lollmsElfServer.config.data_vectorization_overlap_size, force_vectorize=False, add_as_a_bloc=False)
ASCIIColors.yellow("3- Indexing database")
lollmsElfServer.info("Indexing database",True, None)
lollmsElfServer.long_term_memory.index()
ASCIIColors.yellow("Ready")
except Exception as ex:
lollmsElfServer.error(f"Couldn't vectorize the database:{ex}")
return {"status":False}
return {"status":True}
@router.post("/export_discussion")
def export_discussion():
return {"discussion_text":lollmsElfServer.get_discussion_to()}
@router.post("/edit_title")
def edit_title(data:EditTitleParameters):
client_id = data.client_id
title = data.title
discussion_id = data.id
lollmsElfServer.connections[client_id]["current_discussion"] = Discussion(discussion_id, lollmsElfServer.db)
lollmsElfServer.connections[client_id]["current_discussion"].rename(title)
return {'status':True}
@router.post("/make_title")
def make_title(data:MakeTitleParameters):
ASCIIColors.info("Making title")
discussion_id = data.id
discussion = Discussion(discussion_id, lollmsElfServer.db)
title = lollmsElfServer.make_discussion_title(discussion)
discussion.rename(title)
return {'status':True, 'title':title}
@router.get("/export")
def export():
return lollmsElfServer.db.export_to_json()
@router.post("/make_title")
def delete_discussion(data: DeleteDiscussionParameters):
client_id = data.client_id
discussion_id = data.id
lollmsElfServer.connections[client_id]["current_discussion"] = Discussion(discussion_id, lollmsElfServer.db)
lollmsElfServer.connections[client_id]["current_discussion"].delete_discussion()
lollmsElfServer.connections[client_id]["current_discussion"] = None
return {'status':True}
# ----------------------------- import/export --------------------
@router.post("/export_multiple_discussions")
async def export_multiple_discussions(request: Request):
"""
Opens code in vs code.
:param request: The HTTP request object.
:return: A JSON response with the status of the operation.
"""
try:
data = (await request.json())
discussion_ids = data["discussion_ids"]
export_format = data["export_format"]
if export_format=="json":
discussions = lollmsElfServer.db.export_discussions_to_json(discussion_ids)
elif export_format=="markdown":
discussions = lollmsElfServer.db.export_discussions_to_markdown(discussion_ids)
else:
discussions = lollmsElfServer.db.export_discussions_to_markdown(discussion_ids)
return discussions
except Exception as ex:
trace_exception(ex)
lollmsElfServer.error(ex)
return {"status":False,"error":str(ex)}
@router.post("/import_multiple_discussions")
async def import_multiple_discussions(request: Request):
"""
Opens code in vs code.
:param request: The HTTP request object.
:return: A JSON response with the status of the operation.
"""
try:
discussions = (await request.json())["jArray"]
lollmsElfServer.db.import_from_json(discussions)
return discussions
except Exception as ex:
trace_exception(ex)
lollmsElfServer.error(ex)
return {"status":False,"error":str(ex)}