231 lines
8.6 KiB
Python
231 lines
8.6 KiB
Python
import nltk
|
||
import sys
|
||
import os
|
||
|
||
sys.path.append(os.path.dirname(os.path.dirname(__file__)))
|
||
|
||
from configs.model_config import LLM_MODEL, NLTK_DATA_PATH
|
||
from configs.server_config import OPEN_CROSS_DOMAIN, HTTPX_DEFAULT_TIMEOUT
|
||
from configs import VERSION
|
||
import argparse
|
||
import uvicorn
|
||
from fastapi import Body
|
||
from fastapi.middleware.cors import CORSMiddleware
|
||
from starlette.responses import RedirectResponse
|
||
from server.chat import (chat, knowledge_base_chat, openai_chat,
|
||
search_engine_chat)
|
||
from server.knowledge_base.kb_api import list_kbs, create_kb, delete_kb
|
||
from server.knowledge_base.kb_doc_api import (list_files, upload_doc, delete_doc,
|
||
update_doc, download_doc, recreate_vector_store,
|
||
search_docs, DocumentWithScore)
|
||
from server.utils import BaseResponse, ListResponse, FastAPI, MakeFastAPIOffline, fschat_controller_address
|
||
import httpx
|
||
from typing import List
|
||
|
||
nltk.data.path = [NLTK_DATA_PATH] + nltk.data.path
|
||
|
||
|
||
async def document():
|
||
return RedirectResponse(url="/docs")
|
||
|
||
|
||
def create_app():
|
||
app = FastAPI(
|
||
title="Langchain-Chatchat API Server",
|
||
version=VERSION
|
||
)
|
||
MakeFastAPIOffline(app)
|
||
# Add CORS middleware to allow all origins
|
||
# 在config.py中设置OPEN_DOMAIN=True，允许跨域
|
||
# set OPEN_DOMAIN=True in config.py to allow cross-domain
|
||
if OPEN_CROSS_DOMAIN:
|
||
app.add_middleware(
|
||
CORSMiddleware,
|
||
allow_origins=["*"],
|
||
allow_credentials=True,
|
||
allow_methods=["*"],
|
||
allow_headers=["*"],
|
||
)
|
||
|
||
app.get("/",
|
||
response_model=BaseResponse,
|
||
summary="swagger 文档")(document)
|
||
|
||
# Tag: Chat
|
||
app.post("/chat/fastchat",
|
||
tags=["Chat"],
|
||
summary="与llm模型对话(直接与fastchat api对话)")(openai_chat)
|
||
|
||
app.post("/chat/chat",
|
||
tags=["Chat"],
|
||
summary="与llm模型对话(通过LLMChain)")(chat)
|
||
|
||
app.post("/chat/knowledge_base_chat",
|
||
tags=["Chat"],
|
||
summary="与知识库对话")(knowledge_base_chat)
|
||
|
||
app.post("/chat/search_engine_chat",
|
||
tags=["Chat"],
|
||
summary="与搜索引擎对话")(search_engine_chat)
|
||
|
||
# Tag: Knowledge Base Management
|
||
app.get("/knowledge_base/list_knowledge_bases",
|
||
tags=["Knowledge Base Management"],
|
||
response_model=ListResponse,
|
||
summary="获取知识库列表")(list_kbs)
|
||
|
||
app.post("/knowledge_base/create_knowledge_base",
|
||
tags=["Knowledge Base Management"],
|
||
response_model=BaseResponse,
|
||
summary="创建知识库"
|
||
)(create_kb)
|
||
|
||
app.post("/knowledge_base/delete_knowledge_base",
|
||
tags=["Knowledge Base Management"],
|
||
response_model=BaseResponse,
|
||
summary="删除知识库"
|
||
)(delete_kb)
|
||
|
||
app.get("/knowledge_base/list_files",
|
||
tags=["Knowledge Base Management"],
|
||
response_model=ListResponse,
|
||
summary="获取知识库内的文件列表"
|
||
)(list_files)
|
||
|
||
app.post("/knowledge_base/search_docs",
|
||
tags=["Knowledge Base Management"],
|
||
response_model=List[DocumentWithScore],
|
||
summary="搜索知识库"
|
||
)(search_docs)
|
||
|
||
app.post("/knowledge_base/upload_doc",
|
||
tags=["Knowledge Base Management"],
|
||
response_model=BaseResponse,
|
||
summary="上传文件到知识库"
|
||
)(upload_doc)
|
||
|
||
app.post("/knowledge_base/delete_doc",
|
||
tags=["Knowledge Base Management"],
|
||
response_model=BaseResponse,
|
||
summary="删除知识库内指定文件"
|
||
)(delete_doc)
|
||
|
||
app.post("/knowledge_base/update_doc",
|
||
tags=["Knowledge Base Management"],
|
||
response_model=BaseResponse,
|
||
summary="更新现有文件到知识库"
|
||
)(update_doc)
|
||
|
||
app.get("/knowledge_base/download_doc",
|
||
tags=["Knowledge Base Management"],
|
||
summary="下载对应的知识文件")(download_doc)
|
||
|
||
app.post("/knowledge_base/recreate_vector_store",
|
||
tags=["Knowledge Base Management"],
|
||
summary="根据content中文档重建向量库，流式输出处理进度。"
|
||
)(recreate_vector_store)
|
||
|
||
# LLM模型相关接口
|
||
@app.post("/llm_model/list_models",
|
||
tags=["LLM Model Management"],
|
||
summary="列出当前已加载的模型")
|
||
def list_models(
|
||
controller_address: str = Body(None, description="Fastchat controller服务器地址", examples=[fschat_controller_address()])
|
||
) -> BaseResponse:
|
||
'''
|
||
从fastchat controller获取已加载模型列表
|
||
'''
|
||
try:
|
||
controller_address = controller_address or fschat_controller_address()
|
||
r = httpx.post(controller_address + "/list_models")
|
||
return BaseResponse(data=r.json()["models"])
|
||
except Exception as e:
|
||
return BaseResponse(
|
||
code=500,
|
||
data=[],
|
||
msg=f"failed to get available models from controller: {controller_address}。错误信息是： {e}")
|
||
|
||
@app.post("/llm_model/stop",
|
||
tags=["LLM Model Management"],
|
||
summary="停止指定的LLM模型（Model Worker)",
|
||
)
|
||
def stop_llm_model(
|
||
model_name: str = Body(..., description="要停止的LLM模型名称", examples=[LLM_MODEL]),
|
||
controller_address: str = Body(None, description="Fastchat controller服务器地址", examples=[fschat_controller_address()])
|
||
) -> BaseResponse:
|
||
'''
|
||
向fastchat controller请求停止某个LLM模型。
|
||
注意：由于Fastchat的实现方式，实际上是把LLM模型所在的model_worker停掉。
|
||
'''
|
||
try:
|
||
controller_address = controller_address or fschat_controller_address()
|
||
r = httpx.post(
|
||
controller_address + "/release_worker",
|
||
json={"model_name": model_name},
|
||
)
|
||
return r.json()
|
||
except Exception as e:
|
||
return BaseResponse(
|
||
code=500,
|
||
msg=f"failed to stop LLM model {model_name} from controller: {controller_address}。错误信息是： {e}")
|
||
|
||
@app.post("/llm_model/change",
|
||
tags=["LLM Model Management"],
|
||
summary="切换指定的LLM模型（Model Worker)",
|
||
)
|
||
def change_llm_model(
|
||
model_name: str = Body(..., description="当前运行模型", examples=[LLM_MODEL]),
|
||
new_model_name: str = Body(..., description="要切换的新模型", examples=[LLM_MODEL]),
|
||
controller_address: str = Body(None, description="Fastchat controller服务器地址", examples=[fschat_controller_address()])
|
||
):
|
||
'''
|
||
向fastchat controller请求切换LLM模型。
|
||
'''
|
||
try:
|
||
controller_address = controller_address or fschat_controller_address()
|
||
r = httpx.post(
|
||
controller_address + "/release_worker",
|
||
json={"model_name": model_name, "new_model_name": new_model_name},
|
||
timeout=HTTPX_DEFAULT_TIMEOUT, # wait for new worker_model
|
||
)
|
||
return r.json()
|
||
except Exception as e:
|
||
return BaseResponse(
|
||
code=500,
|
||
msg=f"failed to switch LLM model from controller: {controller_address}。错误信息是： {e}")
|
||
|
||
return app
|
||
|
||
|
||
app = create_app()
|
||
|
||
|
||
def run_api(host, port, **kwargs):
|
||
if kwargs.get("ssl_keyfile") and kwargs.get("ssl_certfile"):
|
||
uvicorn.run(app,
|
||
host=host,
|
||
port=port,
|
||
ssl_keyfile=kwargs.get("ssl_keyfile"),
|
||
ssl_certfile=kwargs.get("ssl_certfile"),
|
||
)
|
||
else:
|
||
uvicorn.run(app, host=host, port=port)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
parser = argparse.ArgumentParser(prog='langchain-ChatGLM',
|
||
description='About langchain-ChatGLM, local knowledge based ChatGLM with langchain'
|
||
' ｜ 基于本地知识库的 ChatGLM 问答')
|
||
parser.add_argument("--host", type=str, default="0.0.0.0")
|
||
parser.add_argument("--port", type=int, default=7861)
|
||
parser.add_argument("--ssl_keyfile", type=str)
|
||
parser.add_argument("--ssl_certfile", type=str)
|
||
# 初始化消息
|
||
args = parser.parse_args()
|
||
args_dict = vars(args)
|
||
run_api(host=args.host,
|
||
port=args.port,
|
||
ssl_keyfile=args.ssl_keyfile,
|
||
ssl_certfile=args.ssl_certfile,
|
||
)
|