Langchain-Chatchat/server/chat/chat.py

from fastapi import Body
from fastapi.responses import StreamingResponse
from langchain.agents import initialize_agent, AgentType
from langchain_core.callbacks import AsyncCallbackManager, BaseCallbackManager
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnableBranch
from server.agent.agent_factory import initialize_glm3_agent
from server.agent.tools_factory.tools_registry import all_tools
from server.utils import wrap_done, get_ChatOpenAI
from langchain.chains import LLMChain
from typing import AsyncIterable, Dict
import asyncio
import json
from langchain.prompts.chat import ChatPromptTemplate
from typing import List, Union
from server.chat.utils import History
from langchain.prompts import PromptTemplate
from server.utils import get_prompt_template
from server.memory.conversation_db_buffer_memory import ConversationBufferDBMemory
from server.db.repository import add_message_to_db
from server.callback_handler.agent_callback_handler import Status, CustomAsyncIteratorCallbackHandler


def create_models_from_config(configs, callbacks):
    if configs is None:
        configs = {}
    models = {}
    prompts = {}
    for model_type, model_configs in configs.items():
        for model_name, params in model_configs.items():
            callback = callbacks if params.get('callbacks', False) else None
            model_instance = get_ChatOpenAI(
                model_name=model_name,
                temperature=params.get('temperature', 0.5),
                max_tokens=params.get('max_tokens', 1000),
                callbacks=callback
            )
            models[model_type] = model_instance
            prompt_name = params.get('prompt_name', 'default')
            prompt_template = get_prompt_template(type=model_type, name=prompt_name)
            prompts[model_type] = prompt_template
    return models, prompts


# 在这里写构建逻辑
def create_models_chains(history, history_len, prompts, models, tools, callbacks, conversation_id, metadata):
    memory = None
    chat_prompt = None
    if history:
        history = [History.from_data(h) for h in history]
        input_msg = History(role="user", content=prompts["llm_model"]).to_msg_template(False)
        chat_prompt = ChatPromptTemplate.from_messages(
            [i.to_msg_template() for i in history] + [input_msg])
    elif conversation_id and history_len > 0:
        memory = ConversationBufferDBMemory(conversation_id=conversation_id, llm=models["llm_model"],
                                            message_limit=history_len)
    else:
        input_msg = History(role="user", content=prompts["llm_model"]).to_msg_template(False)
        chat_prompt = ChatPromptTemplate.from_messages([input_msg])

    chain = LLMChain(prompt=chat_prompt, llm=models["llm_model"], memory=memory)
    classifier_chain = (
            PromptTemplate.from_template(prompts["preprocess_model"])
            | models["preprocess_model"]
            | StrOutputParser()
    )
    if "chatglm3" in models["action_model"].model_name.lower():
        agent_executor = initialize_glm3_agent(
            llm=models["action_model"],
            tools=tools,
            prompt=prompts["action_model"],
            input_variables=["input", "intermediate_steps", "history"],
            memory=memory,
            callback_manager=BaseCallbackManager(handlers=callbacks),
            verbose=True,
        )
    else:
        agent_executor = initialize_agent(
            llm=models["action_model"],
            tools=tools,
            callbacks=callbacks,
            agent=AgentType.STRUCTURED_CHAT_ZERO_SHOT_REACT_DESCRIPTION,
            memory=memory,
            verbose=True,
        )
    agent_use = False
    if agent_use:
        branch = RunnableBranch(
            (lambda x: "1" in x["topic"].lower(), agent_executor),
            chain
        )
        full_chain = ({"topic": classifier_chain, "input": lambda x: x["input"]} | branch)
    else:
        full_chain = ({"input": lambda x: x["input"]} | chain)
    return full_chain


async def chat(query: str = Body(..., description="用户输入", examples=["恼羞成怒"]),
               metadata: dict = Body({}, description="附件，可能是图像或者其他功能", examples=[]),
               conversation_id: str = Body("", description="对话框ID"),
               history_len: int = Body(-1, description="从数据库中取历史消息的数量"),
               history: Union[int, List[History]] = Body([],
                                                         description="历史对话，设为一个整数可以从数据库中读取历史消息",
                                                         examples=[[
                                                             {"role": "user",
                                                              "content": "我们来玩成语接龙，我先来，生龙活虎"},
                                                             {"role": "assistant", "content": "虎头虎脑"}]]
                                                         ),
               stream: bool = Body(False, description="流式输出"),
               model_config: Dict = Body({}, description="LLM 模型配置。"),
               tool_config: Dict = Body({}, description="工具配置"),
               ):
    async def chat_iterator() -> AsyncIterable[str]:
        message_id = add_message_to_db(chat_type="llm_chat", query=query, conversation_id=conversation_id) if conversation_id else  None
        callback = CustomAsyncIteratorCallbackHandler()
        callbacks = [callback]
        models, prompts = create_models_from_config(callbacks=callbacks, configs=model_config)
        tools = [tool for tool in all_tools if tool.name in tool_config]
        # 构建完整的Chain
        full_chain = create_models_chains(prompts=prompts,
                                          models=models,
                                          conversation_id=conversation_id,
                                          tools=tools,
                                          callbacks=callbacks,
                                          history=history,
                                          history_len=history_len,
                                          metadata=metadata)

        # 执行完整的Chain
        task = asyncio.create_task(wrap_done(full_chain.ainvoke({"input": query}, callbacks=callbacks), callback.done))
        if stream:
            async for chunk in callback.aiter():
                data = json.loads(chunk)
                if data["status"] == Status.start:
                    continue
                elif data["status"] == Status.agent_action:
                    tool_info = {
                        "tool_name": data["tool_name"],
                        "tool_input": data["tool_input"]
                    }
                    yield json.dumps({"agent_action": tool_info, "message_id": message_id}, ensure_ascii=False)
                elif data["status"] == Status.agent_finish:
                    yield json.dumps({"agent_finish": data["agent_finish"], "message_id": message_id},
                                     ensure_ascii=False)
                else:
                    yield json.dumps({"text": data["llm_token"], "message_id": message_id}, ensure_ascii=False)
        else:
            text = ""
            agent_finish = ""
            tool_info = None
            async for chunk in callback.aiter():
                data = json.loads(chunk)
                if data["status"] == Status.agent_action:
                    tool_info = {
                        "tool_name": data["tool_name"],
                        "tool_input": data["tool_input"]
                    }
                if data["status"] == Status.agent_finish:
                    agent_finish = data["agent_finish"]
                else:
                    text += data["llm_token"]
            if tool_info:
                yield json.dumps(
                    {"text": text, "agent_action": tool_info, "agent_finish": agent_finish, "message_id": message_id},
                    ensure_ascii=False)
            else:
                yield json.dumps(
                    {"text": text, "message_id": message_id},
                    ensure_ascii=False)
        await task

    return EventSourceResponse(chat_iterator())