423 lines
15 KiB
Python
423 lines
15 KiB
Python
|
# coding=utf-8
|
|||
|
|
import json
|
|||
|
|
import time
|
|||
|
|
import uuid
|
|||
|
|
from contextlib import asynccontextmanager
|
|||
|
|
from copy import deepcopy
|
|||
|
|
from typing import List, Union
|
|||
|
|
from typing import Literal, Optional
|
|||
|
|
from colorama import init, Fore
|
|||
|
|
import requests
|
|||
|
|
import torch
|
|||
|
|
import uvicorn
|
|||
|
|
from fastapi import FastAPI
|
|||
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|||
|
|
from loguru import logger
|
|||
|
|
from pydantic import BaseModel, Field
|
|||
|
|
from sse_starlette.sse import EventSourceResponse
|
|||
|
|
from transformers import AutoTokenizer, AutoModel
|
|||
|
|
|
|||
|
|
from api import get_complete_docs
|
|||
|
|
from tool_using.tool_register import get_tools, dispatch_tool
|
|||
|
|
from utils import process_response, generate_chatglm3, generate_stream_chatglm3
|
|||
|
|
|
|||
|
|
SYSTEM_PROMPT = """你是一个农业领域的AI助手，你需要通过已知信息和历史对话来回答人类的问题。
|
|||
|
|
特别注意：已知信息不一定是正确答案，还要根据对话的历史问题和答案进行思考回答，如果找不到答案，请用自己的知识进行回答，不可以乱回答！！！
|
|||
|
|
直接给出答案即可，答案使用中文。\n"""
|
|||
|
|
|
|||
|
|
|
|||
|
|
@asynccontextmanager
|
|||
|
|
async def lifespan(app: FastAPI): # collects GPU memory
|
|||
|
|
yield
|
|||
|
|
if torch.cuda.is_available():
|
|||
|
|
torch.cuda.empty_cache()
|
|||
|
|
torch.cuda.ipc_collect()
|
|||
|
|
|
|||
|
|
|
|||
|
|
app = FastAPI(lifespan=lifespan)
|
|||
|
|
|
|||
|
|
app.add_middleware(
|
|||
|
|
CORSMiddleware,
|
|||
|
|
allow_origins=["*"],
|
|||
|
|
allow_credentials=True,
|
|||
|
|
allow_methods=["*"],
|
|||
|
|
allow_headers=["*"],
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
class ClearHistory(BaseModel):
|
|||
|
|
code: int
|
|||
|
|
message: str
|
|||
|
|
|
|||
|
|
|
|||
|
|
class ModelCard(BaseModel):
|
|||
|
|
id: str
|
|||
|
|
object: str = "model"
|
|||
|
|
created: int = Field(default_factory=lambda: int(time.time()))
|
|||
|
|
owned_by: str = "owner"
|
|||
|
|
root: Optional[str] = None
|
|||
|
|
parent: Optional[str] = None
|
|||
|
|
permission: Optional[list] = None
|
|||
|
|
|
|||
|
|
|
|||
|
|
class ModelList(BaseModel):
|
|||
|
|
object: str = "list"
|
|||
|
|
data: List[ModelCard] = []
|
|||
|
|
|
|||
|
|
|
|||
|
|
class FunctionCallResponse(BaseModel):
|
|||
|
|
name: Optional[str] = None
|
|||
|
|
arguments: Optional[str] = None
|
|||
|
|
|
|||
|
|
|
|||
|
|
class ChatMessage(BaseModel):
|
|||
|
|
role: Literal["user", "assistant", "system", "function"]
|
|||
|
|
content: str = None
|
|||
|
|
name: Optional[str] = None
|
|||
|
|
function_call: Optional[FunctionCallResponse] = None
|
|||
|
|
|
|||
|
|
|
|||
|
|
class DeltaMessage(BaseModel):
|
|||
|
|
role: Optional[Literal["user", "assistant", "system"]] = None
|
|||
|
|
content: Optional[str] = None
|
|||
|
|
function_call: Optional[FunctionCallResponse] = None
|
|||
|
|
|
|||
|
|
|
|||
|
|
class ChatCompletionRequest(BaseModel):
|
|||
|
|
model: str
|
|||
|
|
messages: List[ChatMessage]
|
|||
|
|
kb_name: str
|
|||
|
|
stream: Optional[bool] = False
|
|||
|
|
functions: Optional[Union[dict, List[dict]]] = None
|
|||
|
|
# Additional parameters
|
|||
|
|
temperature: Optional[float] = 0.8
|
|||
|
|
top_p: Optional[float] = 0.8
|
|||
|
|
max_tokens: Optional[int] = None
|
|||
|
|
max_length: Optional[int] = None
|
|||
|
|
repetition_penalty: Optional[float] = 1.1
|
|||
|
|
top_k: int = 3 # 召回的知识库文档数量
|
|||
|
|
score_threshold: float = 1.0 # 知识库严谨度
|
|||
|
|
system_prompt: str = SYSTEM_PROMPT
|
|||
|
|
|
|||
|
|
|
|||
|
|
class ChatCompletionResponseChoice(BaseModel):
|
|||
|
|
index: int
|
|||
|
|
message: ChatMessage
|
|||
|
|
finish_reason: Literal["stop", "length", "function_call", "finish"]
|
|||
|
|
|
|||
|
|
|
|||
|
|
class ChatCompletionResponseStreamChoice(BaseModel):
|
|||
|
|
index: int
|
|||
|
|
delta: DeltaMessage
|
|||
|
|
finish_reason: Optional[Literal["stop", "length", "function_call", "finish"]]
|
|||
|
|
|
|||
|
|
|
|||
|
|
class UsageInfo(BaseModel):
|
|||
|
|
prompt_tokens: int = 0
|
|||
|
|
total_tokens: int = 0
|
|||
|
|
completion_tokens: Optional[int] = 0
|
|||
|
|
|
|||
|
|
|
|||
|
|
class History(BaseModel):
|
|||
|
|
id: str
|
|||
|
|
question: str
|
|||
|
|
messages: List[ChatMessage]
|
|||
|
|
query: str = None
|
|||
|
|
wiki_content: List[str] = None
|
|||
|
|
|
|||
|
|
|
|||
|
|
class ChatCompletionResponse(BaseModel):
|
|||
|
|
model: str
|
|||
|
|
object: Literal["chat.completion", "chat.completion.chunk"]
|
|||
|
|
choices: List[Union[ChatCompletionResponseChoice, ChatCompletionResponseStreamChoice]]
|
|||
|
|
created: Optional[int] = Field(default_factory=lambda: int(time.time()))
|
|||
|
|
usage: Optional[UsageInfo] = None
|
|||
|
|
history: History = None
|
|||
|
|
|
|||
|
|
|
|||
|
|
def get_contexts(query, top_k, score_threshold, kb_name):
|
|||
|
|
return get_complete_docs(query, top_k, score_threshold, kb_name)
|
|||
|
|
|
|||
|
|
|
|||
|
|
from fastapi import Depends, HTTPException, Request
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 验证
|
|||
|
|
async def verify_token(request: Request):
|
|||
|
|
auth_header = request.headers.get('Authorization')
|
|||
|
|
if auth_header:
|
|||
|
|
token_type, _, token = auth_header.partition(' ')
|
|||
|
|
if (
|
|||
|
|
token_type.lower() == "bearer"
|
|||
|
|
and token == "sk-chatglm3-6b"
|
|||
|
|
): # 这里配置你的token
|
|||
|
|
return True
|
|||
|
|
raise HTTPException(
|
|||
|
|
status_code=401,
|
|||
|
|
detail="Invalid authorization credentials",
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
|
|||
|
|
@app.get("/v1/clear/history")
|
|||
|
|
async def clear_history():
|
|||
|
|
try:
|
|||
|
|
if torch.cuda.is_available():
|
|||
|
|
torch.cuda.empty_cache()
|
|||
|
|
torch.cuda.ipc_collect()
|
|||
|
|
except:
|
|||
|
|
raise HTTPException(status_code=500, detail="Internal 500 Error")
|
|||
|
|
return ClearHistory(code=200, message="success")
|
|||
|
|
|
|||
|
|
|
|||
|
|
@app.get("/v1/models", response_model=ModelList)
|
|||
|
|
async def list_models():
|
|||
|
|
model_card = ModelCard(id="chatglm3-6b-32k")
|
|||
|
|
return ModelList(data=[model_card])
|
|||
|
|
|
|||
|
|
|
|||
|
|
@app.get("/v1/tool", response_model=ChatCompletionResponse)
|
|||
|
|
def run_conversation(query: str, stream=False, functions=None):
|
|||
|
|
params = dict(model="chatglm3-6b", messages=[{"role": "user", "content": query}], kb_name="ceshi", stream=stream)
|
|||
|
|
if functions:
|
|||
|
|
params["functions"] = functions
|
|||
|
|
else:
|
|||
|
|
params["functions"] = get_tools()
|
|||
|
|
response = requests.post(f"{base_url}/v1/chat/completions", json=params, stream=False)
|
|||
|
|
max_retry = 1
|
|||
|
|
response = ChatCompletionResponse(**response.json())
|
|||
|
|
print(response)
|
|||
|
|
for _ in range(max_retry):
|
|||
|
|
if not stream:
|
|||
|
|
if response.choices[0].message.function_call:
|
|||
|
|
function_call = response.choices[0].message.function_call
|
|||
|
|
logger.info(f"Function Call Response: {function_call.json(ensure_ascii=False)}")
|
|||
|
|
function_args = json.loads(function_call.arguments)
|
|||
|
|
tool_response = dispatch_tool(function_call.name, function_args)
|
|||
|
|
logger.info(f"Tool Call Response: {tool_response}")
|
|||
|
|
msg = response.choices[0].message
|
|||
|
|
params["messages"].append({
|
|||
|
|
'role': msg.role,
|
|||
|
|
'content': msg.content
|
|||
|
|
})
|
|||
|
|
params["messages"].append({
|
|||
|
|
"role": "function",
|
|||
|
|
"name": function_call.name,
|
|||
|
|
"content": tool_response, # 调用函数返回结果
|
|||
|
|
})
|
|||
|
|
else:
|
|||
|
|
reply = response.choices[0].message.content
|
|||
|
|
logger.info(f"Final Reply: \n{reply}")
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
else:
|
|||
|
|
output = ""
|
|||
|
|
for chunk in response:
|
|||
|
|
content = chunk.choices[0].delta.content or ""
|
|||
|
|
print(Fore.BLUE + content, end="", flush=True)
|
|||
|
|
output += content
|
|||
|
|
|
|||
|
|
if chunk.choices[0].finish_reason == "stop":
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
elif chunk.choices[0].finish_reason == "function_call":
|
|||
|
|
print("\n")
|
|||
|
|
|
|||
|
|
function_call = chunk.choices[0].delta.function_call
|
|||
|
|
logger.info(f"Function Call Response: {function_call.model_dump()}")
|
|||
|
|
|
|||
|
|
function_args = json.loads(function_call.arguments)
|
|||
|
|
tool_response = dispatch_tool(function_call.name, function_args)
|
|||
|
|
logger.info(f"Tool Call Response: {tool_response}")
|
|||
|
|
|
|||
|
|
params["messages"].append(
|
|||
|
|
{
|
|||
|
|
"role": "assistant",
|
|||
|
|
"content": output
|
|||
|
|
}
|
|||
|
|
)
|
|||
|
|
params["messages"].append(
|
|||
|
|
{
|
|||
|
|
"role": "function",
|
|||
|
|
"name": function_call.name,
|
|||
|
|
"content": tool_response, # 调用函数返回结果
|
|||
|
|
}
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
break
|
|||
|
|
print(params)
|
|||
|
|
res = requests.post(f"{base_url}/v1/chat/completions", json=params)
|
|||
|
|
return ChatCompletionResponse(**res.json())
|
|||
|
|
|
|||
|
|
|
|||
|
|
def init(request):
|
|||
|
|
# TODO 如果用户上传的有SYSTEM，则重新处理
|
|||
|
|
r = deepcopy(request)
|
|||
|
|
# 向量库查询语句
|
|||
|
|
query = "\n".join([m.content for m in r.messages])
|
|||
|
|
# 用户问题
|
|||
|
|
question = r.messages[-1].content
|
|||
|
|
# 召回的文档数量
|
|||
|
|
top_k = r.top_k
|
|||
|
|
# 分数匹配
|
|||
|
|
score_threshold = r.score_threshold
|
|||
|
|
# 知识库名称
|
|||
|
|
kb_name = r.kb_name
|
|||
|
|
# 向量库查询内容
|
|||
|
|
contexts = get_contexts(query, top_k, score_threshold, kb_name)
|
|||
|
|
# 添加系统提示词
|
|||
|
|
system_prompt = r.system_prompt or SYSTEM_PROMPT
|
|||
|
|
r.messages.insert(0, ChatMessage(role='system', content=system_prompt))
|
|||
|
|
# 重构消息
|
|||
|
|
new_messages = r.messages
|
|||
|
|
# 知识库模板
|
|||
|
|
wiki_content = "\n".join(contexts)
|
|||
|
|
prompt = f"已知信息：\n{wiki_content}\n问题：\n{question}\n"
|
|||
|
|
new_messages[-1].content = prompt
|
|||
|
|
return query, question, contexts, prompt, new_messages
|
|||
|
|
|
|||
|
|
|
|||
|
|
@app.post("/v1/chat/completions", response_model=ChatCompletionResponse)
|
|||
|
|
async def create_chat_completion(request: ChatCompletionRequest):
|
|||
|
|
# TODO agent调用 需要判断function是否存在
|
|||
|
|
global model, tokenizer
|
|||
|
|
if len(request.messages) < 1 or request.messages[-1].role == "assistant":
|
|||
|
|
raise HTTPException(status_code=400, detail="Invalid request")
|
|||
|
|
# TODO 添加wiki模板提示词
|
|||
|
|
query, question, contexts, prompt, new_messages = init(request)
|
|||
|
|
###########################################################
|
|||
|
|
gen_params = dict(
|
|||
|
|
messages=request.messages if request.functions is not None else new_messages,
|
|||
|
|
temperature=request.temperature,
|
|||
|
|
top_p=request.top_p,
|
|||
|
|
max_tokens=request.max_tokens or 1024,
|
|||
|
|
max_length=request.max_length or 8192,
|
|||
|
|
echo=False,
|
|||
|
|
stream=request.stream,
|
|||
|
|
repetition_penalty=request.repetition_penalty,
|
|||
|
|
functions=request.functions,
|
|||
|
|
system_prompt=request.system_prompt or SYSTEM_PROMPT,
|
|||
|
|
top_k=request.top_k or 3,
|
|||
|
|
score_threshold=request.score_threshold or 1.0,
|
|||
|
|
kb_name=request.kb_name,
|
|||
|
|
)
|
|||
|
|
logger.debug(f"==== request ====\n{gen_params}")
|
|||
|
|
# 流式输出
|
|||
|
|
if request.stream:
|
|||
|
|
generate = predict(request.model, gen_params, request.messages, question, query, contexts)
|
|||
|
|
return EventSourceResponse(generate, media_type="text/event-stream")
|
|||
|
|
#############################################################################################
|
|||
|
|
# 非流式输出
|
|||
|
|
response = generate_chatglm3(model, tokenizer, gen_params)
|
|||
|
|
usage = UsageInfo()
|
|||
|
|
|
|||
|
|
function_call, finish_reason = None, "stop"
|
|||
|
|
if request.functions:
|
|||
|
|
try:
|
|||
|
|
function_call = process_response(response["text"], use_tool=True)
|
|||
|
|
except:
|
|||
|
|
logger.warning("Failed to parse tool call")
|
|||
|
|
|
|||
|
|
if isinstance(function_call, dict):
|
|||
|
|
finish_reason = "function_call"
|
|||
|
|
function_call = FunctionCallResponse(**function_call)
|
|||
|
|
|
|||
|
|
message = ChatMessage(
|
|||
|
|
role="assistant",
|
|||
|
|
content=response["text"],
|
|||
|
|
function_call=function_call if isinstance(function_call, FunctionCallResponse) else None,
|
|||
|
|
)
|
|||
|
|
# 添加历史记录
|
|||
|
|
request.messages.append(message)
|
|||
|
|
h = History(id=str(uuid.uuid4()), messages=request.messages, question=question, query=query, wiki_content=contexts)
|
|||
|
|
choice_data = ChatCompletionResponseChoice(
|
|||
|
|
index=0,
|
|||
|
|
message=message,
|
|||
|
|
finish_reason=finish_reason,
|
|||
|
|
)
|
|||
|
|
task_usage = UsageInfo.parse_obj(response["usage"])
|
|||
|
|
for usage_key, usage_value in task_usage.dict().items():
|
|||
|
|
setattr(usage, usage_key, getattr(usage, usage_key) + usage_value)
|
|||
|
|
|
|||
|
|
return ChatCompletionResponse(model=request.model, choices=[choice_data], object="chat.completion", usage=usage,
|
|||
|
|
history=h)
|
|||
|
|
|
|||
|
|
|
|||
|
|
async def predict(model_id: str, params: dict, messages: List[ChatMessage], question, query, contexts):
|
|||
|
|
global model, tokenizer
|
|||
|
|
all_text = ""
|
|||
|
|
# 第一条消息没有content
|
|||
|
|
choice_data = ChatCompletionResponseStreamChoice(
|
|||
|
|
index=0,
|
|||
|
|
delta=DeltaMessage(role="assistant"),
|
|||
|
|
finish_reason=None
|
|||
|
|
)
|
|||
|
|
chunk = ChatCompletionResponse(model=model_id, choices=[choice_data], object="chat.completion.chunk")
|
|||
|
|
yield "{}".format(chunk.json(exclude_unset=True, ensure_ascii=False))
|
|||
|
|
|
|||
|
|
previous_text = ""
|
|||
|
|
f = None
|
|||
|
|
for new_response in generate_stream_chatglm3(model, tokenizer, params):
|
|||
|
|
decoded_unicode = new_response["text"]
|
|||
|
|
delta_text = decoded_unicode[len(previous_text):]
|
|||
|
|
previous_text = decoded_unicode
|
|||
|
|
|
|||
|
|
finish_reason = new_response["finish_reason"]
|
|||
|
|
if len(delta_text) == 0 and finish_reason != "function_call":
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
function_call = None
|
|||
|
|
if finish_reason == "function_call":
|
|||
|
|
try:
|
|||
|
|
function_call = process_response(decoded_unicode, use_tool=True)
|
|||
|
|
except:
|
|||
|
|
print("Failed to parse tool call")
|
|||
|
|
|
|||
|
|
if isinstance(function_call, dict):
|
|||
|
|
function_call = FunctionCallResponse(**function_call)
|
|||
|
|
f = function_call
|
|||
|
|
|
|||
|
|
delta = DeltaMessage(
|
|||
|
|
content=delta_text,
|
|||
|
|
role="assistant",
|
|||
|
|
function_call=function_call if isinstance(function_call, FunctionCallResponse) else None,
|
|||
|
|
)
|
|||
|
|
all_text += delta_text
|
|||
|
|
choice_data = ChatCompletionResponseStreamChoice(
|
|||
|
|
index=0,
|
|||
|
|
delta=delta,
|
|||
|
|
finish_reason=finish_reason
|
|||
|
|
)
|
|||
|
|
# 这里添加历史记录
|
|||
|
|
chunk = ChatCompletionResponse(model=model_id, choices=[choice_data], object="chat.completion.chunk")
|
|||
|
|
yield "{}".format(chunk.json(exclude_unset=True, ensure_ascii=False))
|
|||
|
|
messages.append(ChatMessage(role='assistant', content=all_text, name='', function_call=f))
|
|||
|
|
choice_data = ChatCompletionResponseStreamChoice(
|
|||
|
|
index=0,
|
|||
|
|
delta=DeltaMessage(),
|
|||
|
|
finish_reason="stop"
|
|||
|
|
)
|
|||
|
|
chunk = ChatCompletionResponse(model=model_id, choices=[choice_data], object="chat.completion.chunk")
|
|||
|
|
yield "{}".format(chunk.json(exclude_unset=True, ensure_ascii=False))
|
|||
|
|
h = History(id=str(uuid.uuid4()), question=question, messages=messages, query=query, wiki_content=contexts)
|
|||
|
|
yield '[DONE]'
|
|||
|
|
choice_data = ChatCompletionResponseChoice(
|
|||
|
|
index=0,
|
|||
|
|
message=ChatMessage(role='assistant', content=all_text, function_call=f),
|
|||
|
|
finish_reason="finish"
|
|||
|
|
)
|
|||
|
|
chunk = ChatCompletionResponse(model=model_id, choices=[choice_data], object="chat.completion.chunk", history=h)
|
|||
|
|
yield "{}".format(chunk.json(exclude_unset=True, ensure_ascii=False))
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm3-6b", trust_remote_code=True)
|
|||
|
|
model = AutoModel.from_pretrained("THUDM/chatglm3-6b", trust_remote_code=True).cuda()
|
|||
|
|
# 多显卡支持，使用下面两行代替上面一行，将num_gpus改为你实际的显卡数量
|
|||
|
|
# from utils import load_model_on_gpus
|
|||
|
|
# model = load_model_on_gpus("THUDM/chatglm3-6b", num_gpus=2)
|
|||
|
|
model = model.eval()
|
|||
|
|
host = "localhost"
|
|||
|
|
port = 8000
|
|||
|
|
base_url = f"http://{host}:{port}"
|
|||
|
|
uvicorn.run(app, host=host, port=port, workers=1)
|