mirror of
https://github.com/csunny/DB-GPT.git
synced 2026-01-21 14:05:13 +00:00
655 lines
22 KiB
Python
655 lines
22 KiB
Python
#!/usr/bin/env python3
|
||
# -*- coding: utf-8 -*-
|
||
|
||
import argparse
|
||
import datetime
|
||
import json
|
||
import os
|
||
import shutil
|
||
import sys
|
||
import time
|
||
import uuid
|
||
from urllib.parse import urljoin
|
||
|
||
import gradio as gr
|
||
import requests
|
||
from langchain import PromptTemplate
|
||
|
||
ROOT_PATH = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||
sys.path.append(ROOT_PATH)
|
||
|
||
from pilot.commands.command_mange import CommandRegistry
|
||
|
||
from pilot.scene.base_chat import BaseChat
|
||
|
||
from pilot.configs.config import Config
|
||
from pilot.configs.model_config import (
|
||
DATASETS_DIR,
|
||
KNOWLEDGE_UPLOAD_ROOT_PATH,
|
||
LLM_MODEL_CONFIG,
|
||
LOGDIR,
|
||
VECTOR_SEARCH_TOP_K,
|
||
)
|
||
from pilot.connections.mysql import MySQLOperator
|
||
from pilot.conversation import (
|
||
SeparatorStyle,
|
||
conv_qa_prompt_template,
|
||
conv_templates,
|
||
conversation_sql_mode,
|
||
conversation_types,
|
||
default_conversation,
|
||
)
|
||
from pilot.plugins import scan_plugins
|
||
from pilot.prompts.auto_mode_prompt import AutoModePrompt
|
||
from pilot.prompts.generator import PromptGenerator
|
||
from pilot.server.gradio_css import code_highlight_css
|
||
from pilot.server.gradio_patch import Chatbot as grChatbot
|
||
from pilot.server.vectordb_qa import KnownLedgeBaseQA
|
||
from pilot.source_embedding.knowledge_embedding import KnowledgeEmbedding
|
||
from pilot.utils import build_logger, server_error_msg
|
||
from pilot.vector_store.extract_tovec import (
|
||
get_vector_storelist,
|
||
knownledge_tovec_st,
|
||
load_knownledge_from_doc,
|
||
)
|
||
|
||
from pilot.commands.command import execute_ai_response_json
|
||
from pilot.scene.base import ChatScene
|
||
from pilot.scene.chat_factory import ChatFactory
|
||
|
||
logger = build_logger("webserver", LOGDIR + "webserver.log")
|
||
headers = {"User-Agent": "dbgpt Client"}
|
||
|
||
no_change_btn = gr.Button.update()
|
||
enable_btn = gr.Button.update(interactive=True)
|
||
disable_btn = gr.Button.update(interactive=True)
|
||
|
||
enable_moderation = False
|
||
models = []
|
||
dbs = []
|
||
vs_list = ["新建知识库"] + get_vector_storelist()
|
||
autogpt = False
|
||
vector_store_client = None
|
||
vector_store_name = {"vs_name": ""}
|
||
|
||
priority = {"vicuna-13b": "aaa"}
|
||
|
||
# 加载插件
|
||
CFG = Config()
|
||
CHAT_FACTORY = ChatFactory()
|
||
|
||
DB_SETTINGS = {
|
||
"user": CFG.LOCAL_DB_USER,
|
||
"password": CFG.LOCAL_DB_PASSWORD,
|
||
"host": CFG.LOCAL_DB_HOST,
|
||
"port": CFG.LOCAL_DB_PORT,
|
||
}
|
||
|
||
|
||
def get_simlar(q):
|
||
docsearch = knownledge_tovec_st(os.path.join(DATASETS_DIR, "plan.md"))
|
||
docs = docsearch.similarity_search_with_score(q, k=1)
|
||
|
||
contents = [dc.page_content for dc, _ in docs]
|
||
return "\n".join(contents)
|
||
|
||
|
||
def gen_sqlgen_conversation(dbname):
|
||
mo = MySQLOperator(**DB_SETTINGS)
|
||
|
||
message = ""
|
||
|
||
schemas = mo.get_schema(dbname)
|
||
for s in schemas:
|
||
message += s["schema_info"] + ";"
|
||
return f"数据库{dbname}的Schema信息如下: {message}\n"
|
||
|
||
|
||
def get_database_list():
|
||
mo = MySQLOperator(**DB_SETTINGS)
|
||
return mo.get_db_list()
|
||
|
||
|
||
get_window_url_params = """
|
||
function() {
|
||
const params = new URLSearchParams(window.location.search);
|
||
url_params = Object.fromEntries(params);
|
||
console.log(url_params);
|
||
gradioURL = window.location.href
|
||
if (!gradioURL.endsWith('?__theme=dark')) {
|
||
window.location.replace(gradioURL + '?__theme=dark');
|
||
}
|
||
return url_params;
|
||
}
|
||
"""
|
||
|
||
|
||
def load_demo(url_params, request: gr.Request):
|
||
logger.info(f"load_demo. ip: {request.client.host}. params: {url_params}")
|
||
|
||
# dbs = get_database_list()
|
||
dropdown_update = gr.Dropdown.update(visible=True)
|
||
if dbs:
|
||
gr.Dropdown.update(choices=dbs)
|
||
|
||
state = default_conversation.copy()
|
||
|
||
unique_id = uuid.uuid1()
|
||
state.conv_id = str(unique_id)
|
||
|
||
return (
|
||
state,
|
||
dropdown_update,
|
||
gr.Chatbot.update(visible=True),
|
||
gr.Textbox.update(visible=True),
|
||
gr.Button.update(visible=True),
|
||
gr.Row.update(visible=True),
|
||
gr.Accordion.update(visible=True),
|
||
)
|
||
|
||
|
||
def get_conv_log_filename():
|
||
t = datetime.datetime.now()
|
||
name = os.path.join(LOGDIR, f"{t.year}-{t.month:02d}-{t.day:02d}-conv.json")
|
||
return name
|
||
|
||
|
||
def regenerate(state, request: gr.Request):
|
||
logger.info(f"regenerate. ip: {request.client.host}")
|
||
state.messages[-1][-1] = None
|
||
state.skip_next = False
|
||
return (state, state.to_gradio_chatbot(), "") + (disable_btn,) * 5
|
||
|
||
|
||
def clear_history(request: gr.Request):
|
||
logger.info(f"clear_history. ip: {request.client.host}")
|
||
state = None
|
||
return (state, [], "") + (disable_btn,) * 5
|
||
|
||
|
||
def add_text(state, text, request: gr.Request):
|
||
logger.info(f"add_text. ip: {request.client.host}. len: {len(text)}")
|
||
if len(text) <= 0:
|
||
state.skip_next = True
|
||
return (state, state.to_gradio_chatbot(), "") + (no_change_btn,) * 5
|
||
|
||
""" Default support 4000 tokens, if tokens too lang, we will cut off """
|
||
text = text[:4000]
|
||
state.append_message(state.roles[0], text)
|
||
state.append_message(state.roles[1], None)
|
||
state.skip_next = False
|
||
### TODO
|
||
state.last_user_input = text
|
||
return (state, state.to_gradio_chatbot(), "") + (disable_btn,) * 5
|
||
|
||
|
||
def post_process_code(code):
|
||
sep = "\n```"
|
||
if sep in code:
|
||
blocks = code.split(sep)
|
||
if len(blocks) % 2 == 1:
|
||
for i in range(1, len(blocks), 2):
|
||
blocks[i] = blocks[i].replace("\\_", "_")
|
||
code = sep.join(blocks)
|
||
return code
|
||
|
||
|
||
def get_chat_mode(mode, sql_mode, db_selector) -> ChatScene:
|
||
if mode == conversation_types["default_knownledge"] and not db_selector:
|
||
return ChatScene.ChatKnowledge
|
||
elif mode == conversation_types["custome"] and not db_selector:
|
||
return ChatScene.ChatNewKnowledge
|
||
elif sql_mode == conversation_sql_mode["auto_execute_ai_response"] and db_selector:
|
||
return ChatScene.ChatWithDb
|
||
|
||
elif mode == conversation_types["auto_execute_plugin"] and not db_selector:
|
||
return ChatScene.ChatExecution
|
||
else:
|
||
return ChatScene.ChatNormal
|
||
|
||
|
||
def http_bot(
|
||
state, mode, sql_mode, db_selector, temperature, max_new_tokens, request: gr.Request
|
||
):
|
||
logger.info(f"User message send!{state.conv_id},{sql_mode},{db_selector}")
|
||
start_tstamp = time.time()
|
||
scene: ChatScene = get_chat_mode(mode, sql_mode, db_selector)
|
||
print(f"当前对话模式:{scene.value}")
|
||
model_name = CFG.LLM_MODEL
|
||
|
||
if ChatScene.ChatWithDb == scene:
|
||
logger.info("基于DB对话走新的模式！")
|
||
chat_param = {
|
||
"chat_session_id": state.conv_id,
|
||
"db_name": db_selector,
|
||
"user_input": state.last_user_input,
|
||
}
|
||
chat: BaseChat = CHAT_FACTORY.get_implementation(scene.value, **chat_param)
|
||
chat.call()
|
||
state.messages[-1][-1] = f"{chat.current_ai_response()}"
|
||
yield (state, state.to_gradio_chatbot()) + (enable_btn,) * 5
|
||
|
||
else:
|
||
dbname = db_selector
|
||
# TODO 这里的请求需要拼接现有知识库, 使得其根据现有知识库作答, 所以prompt需要继续优化
|
||
if state.skip_next:
|
||
# This generate call is skipped due to invalid inputs
|
||
yield (state, state.to_gradio_chatbot()) + (no_change_btn,) * 5
|
||
return
|
||
|
||
if len(state.messages) == state.offset + 2:
|
||
query = state.messages[-2][1]
|
||
|
||
template_name = "conv_one_shot"
|
||
new_state = conv_templates[template_name].copy()
|
||
# prompt 中添加上下文提示, 根据已有知识对话, 上下文提示是否也应该放在第一轮, 还是每一轮都添加上下文?
|
||
# 如果用户侧的问题跨度很大, 应该每一轮都加提示。
|
||
if db_selector:
|
||
new_state.append_message(
|
||
new_state.roles[0], gen_sqlgen_conversation(dbname) + query
|
||
)
|
||
new_state.append_message(new_state.roles[1], None)
|
||
else:
|
||
new_state.append_message(new_state.roles[0], query)
|
||
new_state.append_message(new_state.roles[1], None)
|
||
|
||
new_state.conv_id = uuid.uuid4().hex
|
||
state = new_state
|
||
|
||
prompt = state.get_prompt()
|
||
skip_echo_len = len(prompt.replace("</s>", " ")) + 1
|
||
if mode == conversation_types["default_knownledge"] and not db_selector:
|
||
query = state.messages[-2][1]
|
||
knqa = KnownLedgeBaseQA()
|
||
state.messages[-2][1] = knqa.get_similar_answer(query)
|
||
prompt = state.get_prompt()
|
||
state.messages[-2][1] = query
|
||
skip_echo_len = len(prompt.replace("</s>", " ")) + 1
|
||
|
||
if mode == conversation_types["custome"] and not db_selector:
|
||
persist_dir = os.path.join(
|
||
KNOWLEDGE_UPLOAD_ROOT_PATH, vector_store_name["vs_name"] + ".vectordb"
|
||
)
|
||
print("向量数据库持久化地址: ", persist_dir)
|
||
knowledge_embedding_client = KnowledgeEmbedding(
|
||
file_path="",
|
||
model_name=LLM_MODEL_CONFIG["sentence-transforms"],
|
||
local_persist=False,
|
||
vector_store_config={
|
||
"vector_store_name": vector_store_name["vs_name"],
|
||
"vector_store_path": KNOWLEDGE_UPLOAD_ROOT_PATH,
|
||
},
|
||
)
|
||
query = state.messages[-2][1]
|
||
docs = knowledge_embedding_client.similar_search(query, 1)
|
||
context = [d.page_content for d in docs]
|
||
prompt_template = PromptTemplate(
|
||
template=conv_qa_prompt_template,
|
||
input_variables=["context", "question"],
|
||
)
|
||
result = prompt_template.format(context="\n".join(context), question=query)
|
||
state.messages[-2][1] = result
|
||
prompt = state.get_prompt()
|
||
state.messages[-2][1] = query
|
||
skip_echo_len = len(prompt.replace("</s>", " ")) + 1
|
||
# Make requests
|
||
payload = {
|
||
"model": model_name,
|
||
"prompt": prompt,
|
||
"temperature": float(temperature),
|
||
"max_new_tokens": int(max_new_tokens),
|
||
"stop": state.sep
|
||
if state.sep_style == SeparatorStyle.SINGLE
|
||
else state.sep2,
|
||
}
|
||
logger.info(f"Requert: \n{payload}")
|
||
|
||
# 流式输出
|
||
state.messages[-1][-1] = "▌"
|
||
yield (state, state.to_gradio_chatbot()) + (disable_btn,) * 5
|
||
|
||
try:
|
||
# Stream output
|
||
response = requests.post(
|
||
urljoin(CFG.MODEL_SERVER, "generate_stream"),
|
||
headers=headers,
|
||
json=payload,
|
||
stream=True,
|
||
timeout=20,
|
||
)
|
||
for chunk in response.iter_lines(decode_unicode=False, delimiter=b"\0"):
|
||
if chunk:
|
||
data = json.loads(chunk.decode())
|
||
if data["error_code"] == 0:
|
||
output = data["text"][skip_echo_len:].strip()
|
||
output = post_process_code(output)
|
||
state.messages[-1][-1] = output + "▌"
|
||
yield (state, state.to_gradio_chatbot()) + (disable_btn,) * 5
|
||
else:
|
||
output = data["text"] + f" (error_code: {data['error_code']})"
|
||
state.messages[-1][-1] = output
|
||
yield (state, state.to_gradio_chatbot()) + (
|
||
disable_btn,
|
||
disable_btn,
|
||
disable_btn,
|
||
enable_btn,
|
||
enable_btn,
|
||
)
|
||
return
|
||
|
||
except requests.exceptions.RequestException as e:
|
||
state.messages[-1][-1] = server_error_msg + f" (error_code: 4)"
|
||
yield (state, state.to_gradio_chatbot()) + (
|
||
disable_btn,
|
||
disable_btn,
|
||
disable_btn,
|
||
enable_btn,
|
||
enable_btn,
|
||
)
|
||
return
|
||
|
||
state.messages[-1][-1] = state.messages[-1][-1][:-1]
|
||
yield (state, state.to_gradio_chatbot()) + (enable_btn,) * 5
|
||
|
||
# 记录运行日志
|
||
finish_tstamp = time.time()
|
||
logger.info(f"{output}")
|
||
|
||
with open(get_conv_log_filename(), "a") as fout:
|
||
data = {
|
||
"tstamp": round(finish_tstamp, 4),
|
||
"type": "chat",
|
||
"model": model_name,
|
||
"start": round(start_tstamp, 4),
|
||
"finish": round(start_tstamp, 4),
|
||
"state": state.dict(),
|
||
"ip": request.client.host,
|
||
}
|
||
fout.write(json.dumps(data) + "\n")
|
||
|
||
|
||
block_css = (
|
||
code_highlight_css
|
||
+ """
|
||
pre {
|
||
white-space: pre-wrap; /* Since CSS 2.1 */
|
||
white-space: -moz-pre-wrap; /* Mozilla, since 1999 */
|
||
white-space: -pre-wrap; /* Opera 4-6 */
|
||
white-space: -o-pre-wrap; /* Opera 7 */
|
||
word-wrap: break-word; /* Internet Explorer 5.5+ */
|
||
}
|
||
#notice_markdown th {
|
||
display: none;
|
||
}
|
||
"""
|
||
)
|
||
|
||
|
||
def change_sql_mode(sql_mode):
|
||
if sql_mode in ["直接执行结果"]:
|
||
return gr.update(visible=True)
|
||
else:
|
||
return gr.update(visible=False)
|
||
|
||
|
||
def change_mode(mode):
|
||
if mode in ["默认知识库对话", "LLM原生对话"]:
|
||
return gr.update(visible=False)
|
||
else:
|
||
return gr.update(visible=True)
|
||
|
||
|
||
def change_tab():
|
||
autogpt = True
|
||
|
||
|
||
def build_single_model_ui():
|
||
notice_markdown = """
|
||
# DB-GPT
|
||
|
||
[DB-GPT](https://github.com/csunny/DB-GPT) 是一个开源的以数据库为基础的GPT实验项目，使用本地化的GPT大模型与您的数据和环境进行交互，无数据泄露风险，100% 私密，100% 安全。
|
||
"""
|
||
learn_more_markdown = """
|
||
### Licence
|
||
The service is a research preview intended for non-commercial use only. subject to the model [License](https://github.com/facebookresearch/llama/blob/main/MODEL_CARD.md) of Vicuna-13B
|
||
"""
|
||
|
||
state = gr.State()
|
||
|
||
gr.Markdown(notice_markdown, elem_id="notice_markdown")
|
||
|
||
with gr.Accordion("参数", open=False, visible=False) as parameter_row:
|
||
temperature = gr.Slider(
|
||
minimum=0.0,
|
||
maximum=1.0,
|
||
value=0.7,
|
||
step=0.1,
|
||
interactive=True,
|
||
label="Temperature",
|
||
)
|
||
|
||
max_output_tokens = gr.Slider(
|
||
minimum=0,
|
||
maximum=1024,
|
||
value=512,
|
||
step=64,
|
||
interactive=True,
|
||
label="最大输出Token数",
|
||
)
|
||
tabs = gr.Tabs()
|
||
with tabs:
|
||
tab_sql = gr.TabItem("SQL生成与诊断", elem_id="SQL")
|
||
with tab_sql:
|
||
# TODO A selector to choose database
|
||
with gr.Row(elem_id="db_selector"):
|
||
db_selector = gr.Dropdown(
|
||
label="请选择数据库",
|
||
choices=dbs,
|
||
value=dbs[0] if len(models) > 0 else "",
|
||
interactive=True,
|
||
show_label=True,
|
||
).style(container=False)
|
||
|
||
sql_mode = gr.Radio(["直接执行结果", "不执行结果"], show_label=False, value="不执行结果")
|
||
sql_vs_setting = gr.Markdown("自动执行模式下, DB-GPT可以具备执行SQL、从网络读取知识自动化存储学习的能力")
|
||
sql_mode.change(fn=change_sql_mode, inputs=sql_mode, outputs=sql_vs_setting)
|
||
|
||
tab_qa = gr.TabItem("知识问答", elem_id="QA")
|
||
with tab_qa:
|
||
mode = gr.Radio(
|
||
["LLM原生对话", "默认知识库对话", "新增知识库对话"], show_label=False, value="LLM原生对话"
|
||
)
|
||
vs_setting = gr.Accordion("配置知识库", open=False)
|
||
mode.change(fn=change_mode, inputs=mode, outputs=vs_setting)
|
||
with vs_setting:
|
||
vs_name = gr.Textbox(label="新知识库名称", lines=1, interactive=True)
|
||
vs_add = gr.Button("添加为新知识库")
|
||
with gr.Column() as doc2vec:
|
||
gr.Markdown("向知识库中添加文件")
|
||
with gr.Tab("上传文件"):
|
||
files = gr.File(
|
||
label="添加文件",
|
||
file_types=[".txt", ".md", ".docx", ".pdf"],
|
||
file_count="multiple",
|
||
allow_flagged_uploads=True,
|
||
show_label=False,
|
||
)
|
||
|
||
load_file_button = gr.Button("上传并加载到知识库")
|
||
with gr.Tab("上传文件夹"):
|
||
folder_files = gr.File(
|
||
label="添加文件夹",
|
||
accept_multiple_files=True,
|
||
file_count="directory",
|
||
show_label=False,
|
||
)
|
||
load_folder_button = gr.Button("上传并加载到知识库")
|
||
|
||
with gr.Blocks():
|
||
chatbot = grChatbot(elem_id="chatbot", visible=False).style(height=550)
|
||
with gr.Row():
|
||
with gr.Column(scale=20):
|
||
textbox = gr.Textbox(
|
||
show_label=False,
|
||
placeholder="Enter text and press ENTER",
|
||
visible=False,
|
||
).style(container=False)
|
||
with gr.Column(scale=2, min_width=50):
|
||
send_btn = gr.Button(value="发送", visible=False)
|
||
|
||
with gr.Row(visible=False) as button_row:
|
||
regenerate_btn = gr.Button(value="重新生成", interactive=False)
|
||
clear_btn = gr.Button(value="清理", interactive=False)
|
||
|
||
gr.Markdown(learn_more_markdown)
|
||
btn_list = [regenerate_btn, clear_btn]
|
||
regenerate_btn.click(regenerate, state, [state, chatbot, textbox] + btn_list).then(
|
||
http_bot,
|
||
[state, mode, sql_mode, db_selector, temperature, max_output_tokens],
|
||
[state, chatbot] + btn_list,
|
||
)
|
||
clear_btn.click(clear_history, None, [state, chatbot, textbox] + btn_list)
|
||
|
||
textbox.submit(
|
||
add_text, [state, textbox], [state, chatbot, textbox] + btn_list
|
||
).then(
|
||
http_bot,
|
||
[state, mode, sql_mode, db_selector, temperature, max_output_tokens],
|
||
[state, chatbot] + btn_list,
|
||
)
|
||
|
||
send_btn.click(
|
||
add_text, [state, textbox], [state, chatbot, textbox] + btn_list
|
||
).then(
|
||
http_bot,
|
||
[state, mode, sql_mode, db_selector, temperature, max_output_tokens],
|
||
[state, chatbot] + btn_list,
|
||
)
|
||
vs_add.click(
|
||
fn=save_vs_name, show_progress=True, inputs=[vs_name], outputs=[vs_name]
|
||
)
|
||
load_file_button.click(
|
||
fn=knowledge_embedding_store,
|
||
show_progress=True,
|
||
inputs=[vs_name, files],
|
||
outputs=[vs_name],
|
||
)
|
||
load_folder_button.click(
|
||
fn=knowledge_embedding_store,
|
||
show_progress=True,
|
||
inputs=[vs_name, folder_files],
|
||
outputs=[vs_name],
|
||
)
|
||
return state, chatbot, textbox, send_btn, button_row, parameter_row
|
||
|
||
|
||
def build_webdemo():
|
||
with gr.Blocks(
|
||
title="数据库智能助手",
|
||
# theme=gr.themes.Base(),
|
||
theme=gr.themes.Default(),
|
||
css=block_css,
|
||
) as demo:
|
||
url_params = gr.JSON(visible=False)
|
||
(
|
||
state,
|
||
chatbot,
|
||
textbox,
|
||
send_btn,
|
||
button_row,
|
||
parameter_row,
|
||
) = build_single_model_ui()
|
||
|
||
if args.model_list_mode == "once":
|
||
demo.load(
|
||
load_demo,
|
||
[url_params],
|
||
[
|
||
state,
|
||
chatbot,
|
||
textbox,
|
||
send_btn,
|
||
button_row,
|
||
parameter_row,
|
||
],
|
||
_js=get_window_url_params,
|
||
)
|
||
else:
|
||
raise ValueError(f"Unknown model list mode: {args.model_list_mode}")
|
||
return demo
|
||
|
||
|
||
def save_vs_name(vs_name):
|
||
vector_store_name["vs_name"] = vs_name
|
||
return vs_name
|
||
|
||
|
||
def knowledge_embedding_store(vs_id, files):
|
||
# vs_path = os.path.join(VS_ROOT_PATH, vs_id)
|
||
if not os.path.exists(os.path.join(KNOWLEDGE_UPLOAD_ROOT_PATH, vs_id)):
|
||
os.makedirs(os.path.join(KNOWLEDGE_UPLOAD_ROOT_PATH, vs_id))
|
||
for file in files:
|
||
filename = os.path.split(file.name)[-1]
|
||
shutil.move(
|
||
file.name, os.path.join(KNOWLEDGE_UPLOAD_ROOT_PATH, vs_id, filename)
|
||
)
|
||
knowledge_embedding_client = KnowledgeEmbedding(
|
||
file_path=os.path.join(KNOWLEDGE_UPLOAD_ROOT_PATH, vs_id, filename),
|
||
model_name=LLM_MODEL_CONFIG["text2vec"],
|
||
local_persist=False,
|
||
vector_store_config={
|
||
"vector_store_name": vector_store_name["vs_name"],
|
||
"vector_store_path": KNOWLEDGE_UPLOAD_ROOT_PATH,
|
||
},
|
||
)
|
||
knowledge_embedding_client.knowledge_embedding()
|
||
|
||
logger.info("knowledge embedding success")
|
||
return os.path.join(KNOWLEDGE_UPLOAD_ROOT_PATH, vs_id, vs_id + ".vectordb")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument("--host", type=str, default="0.0.0.0")
|
||
parser.add_argument("--port", type=int)
|
||
parser.add_argument("--concurrency-count", type=int, default=10)
|
||
parser.add_argument(
|
||
"--model-list-mode", type=str, default="once", choices=["once", "reload"]
|
||
)
|
||
parser.add_argument("--share", default=False, action="store_true")
|
||
|
||
args = parser.parse_args()
|
||
logger.info(f"args: {args}")
|
||
# 配置初始化
|
||
cfg = Config()
|
||
|
||
dbs = cfg.local_db.get_database_list()
|
||
|
||
cfg.set_plugins(scan_plugins(cfg, cfg.debug_mode))
|
||
|
||
# 加载插件可执行命令
|
||
command_categories = [
|
||
"pilot.commands.audio_text",
|
||
"pilot.commands.image_gen",
|
||
]
|
||
# 排除禁用命令
|
||
command_categories = [
|
||
x for x in command_categories if x not in cfg.disabled_command_categories
|
||
]
|
||
command_registry = CommandRegistry()
|
||
for command_category in command_categories:
|
||
command_registry.import_commands(command_category)
|
||
|
||
cfg.command_registry = command_registry
|
||
|
||
logger.info(args)
|
||
demo = build_webdemo()
|
||
demo.queue(
|
||
concurrency_count=args.concurrency_count, status_update_rate=10, api_open=False
|
||
).launch(
|
||
server_name=args.host,
|
||
server_port=args.port,
|
||
share=args.share,
|
||
max_threads=200,
|
||
)
|