gangyan/langchain-chat/server/chat/chat.py

import uuid
from fastapi import Body
from langchain.memory import (
    CombinedMemory,
    ConversationBufferMemory,
    ConversationSummaryMemory,
    ConversationBufferWindowMemory
)
from sse_starlette.sse import EventSourceResponse
from configs import LLM_MODELS, TEMPERATURE, HISTORY_LEN
from server.utils import wrap_done, get_ChatOpenAI
from langchain.chains import LLMChain, ConversationChain
from langchain.callbacks import AsyncIteratorCallbackHandler
from typing import AsyncIterable
import asyncio
import json
from langchain.prompts.chat import ChatPromptTemplate
from typing import List, Optional, Union
from server.chat.utils import History
from langchain.prompts import PromptTemplate
from server.utils import get_prompt_template, get_format_template
from server.memory.conversation_db_buffer_memory import ConversationBufferDBMemory
from server.db.repository import add_message_to_db
from server.callback_handler.conversation_callback_handler import ConversationCallbackHandler
from datetime import datetime
from langchain_core.messages import SystemMessage
import time as t
from configs.basic_config import *
async def chat(query: str = 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_name: str = Body(LLM_MODELS[0], description="LLM 模型名称。"),
               temperature: float = Body(TEMPERATURE, description="LLM 采样温度", ge=0.0, le=2.0),
               max_tokens: Optional[int] = Body(None, description="限制LLM生成Token数量，默认None代表模型最大值"),
               # top_p: float = Body(TOP_P, description="LLM 核采样。勿与temperature同时设置", gt=0.0, lt=1.0),
               prompt_name: str = Body("default", description="使用的prompt模板名称(在configs/prompt_config.py中配置)"),
               ):
    async def chat_iterator() -> AsyncIterable[str]:
        nonlocal history, max_tokens
        callback = AsyncIteratorCallbackHandler()
        callbacks = [callback]
        memory = None
        time = datetime.now().strftime("%Y年%m月%d日")
        # 负责保存llm response到message db
        message_id = str(uuid.uuid1())+"q"
        conversation_callback = ConversationCallbackHandler(conversation_id=conversation_id, message_id=message_id,
                                                            chat_type="llm_chat",
                                                            query=query)
        callbacks.append(conversation_callback)
        logger.info(f"智能对话的入参信息：query：{query}，conversation_id：{conversation_id}，history：{history}，stream：{stream}，model_name：{model_name}，temperature：{temperature}，max_tokens：{max_tokens}prompt_name：{prompt_name}")

        if isinstance(max_tokens, int) and max_tokens <= 0:
            max_tokens = None

        if prompt_name == "Search Summary":
            model = get_ChatOpenAI(
                model_name=LLM_MODELS[0],
                temperature=temperature,
                max_tokens=max_tokens,
                callbacks=callbacks,
            )
            # print ("model info >>>", LLM_MODELS[0])
        else:
            model = get_ChatOpenAI(
                model_name=model_name,
                temperature=temperature,
                max_tokens=max_tokens,
                callbacks=callbacks,
            )
        logger.info(f"当前使用的模型为：{model_name}")
        if history and prompt_name not in ["Search Summary", "get_policy_time"]:
            history = [History.from_data(h) for h in history]
            if prompt_name == "default":
                prompt_template = get_prompt_template("llm_chat", "default_with_history")
            if prompt_name == "Policy History Assistant":
                prompt_template = get_prompt_template("llm_chat", "Policy History Assistant_with_history")
            if prompt_name == "Topic Recommend Assistant":
                prompt_template = get_prompt_template("llm_chat", "Topic Recommend Assistant_with_history")
            if prompt_name == "Abstract Assistant":
                prompt_template = get_prompt_template("llm_chat", "Abstract Assistant_with_history")
            # input_prompt = History(role="system", content=prompt_template).to_msg_template(False)
            # chat_prompt = ChatPromptTemplate.from_messages([input_prompt])
            chat_prompt =  PromptTemplate.from_template(prompt_template)
            # 把history转成memory
            buff_memory = ConversationBufferMemory(human_prefix='user', ai_prefix='assistant', memory_key="history", input_key="input")
            for message in history:
                # 检查消息的角色
                if message.role == 'user':
                    # 添加用户消息
                    buff_memory.chat_memory.add_user_message(message.content)
                elif message.role == 'assistant':
                    # 添加AI消息
                    buff_memory.chat_memory.add_ai_message(message.content)
            background_memory = ConversationBufferMemory(human_prefix='user', ai_prefix='assistant', memory_key="time", input_key="input")
            message = SystemMessage(content = f'当前的时间是：{time}')
            background_memory.chat_memory.add_message(message)
            memory = CombinedMemory(memories=[background_memory, buff_memory])
            chain = ConversationChain(llm=model, verbose=True, memory=memory, prompt=chat_prompt)
        # elif conversation_id and history_len > 0: # 前端要求从数据库取历史消息
        #     # 使用memory 时必须 prompt 必须含有memory.memory_key 对应的变量
        #     prompt = get_prompt_template("llm_chat", "with_history")
        #     chat_prompt = PromptTemplate.from_template(prompt)
        #     # 根据conversation_id 获取message 列表进而拼凑 memory
        #     memory = ConversationBufferDBMemory(conversation_id=conversation_id,
        #                                         llm=model,
        #                                         message_limit=history_len)
        else:
            prompt_template = get_prompt_template("llm_chat", prompt_name)
            input_prompt = History(role="system", content=prompt_template).to_msg_template(False)
            # input_msg = History(role="user", content=query).to_msg_template(False)
            chat_prompt = ChatPromptTemplate.from_messages([input_prompt])
            chain = LLMChain(prompt=chat_prompt, llm=model)

        # print(f'智能对话的chain>>>\n{chain}\n')
        task = asyncio.create_task(wrap_done(
            chain.acall({"input": query, "time": time}),
            callback.done),
        )

        if stream:
            async for token in callback.aiter():
                # Use server-sent-events to stream the response
                yield json.dumps(
                    {"text": token, "message_id": message_id},
                    ensure_ascii=False)
        else:
            answer = ""
            async for token in callback.aiter():
                answer += token
            yield json.dumps(
                {"text": answer, "message_id": message_id},
                ensure_ascii=False)

        await task

    return EventSourceResponse(chat_iterator())