ai/openai_api.py

# 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)