qa.py

# -*- coding: utf-8 -*-
import sys
import time
from datetime import datetime
from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate
from typing import Awaitable
import asyncio
from langchain.callbacks import AsyncIteratorCallbackHandler
from src.llm.ernie_with_sdk import ChatERNIESerLLM
from qianfan import ChatCompletion
from src.pgdb.chat.c_db import UPostgresDB
from src.pgdb.chat.crud import CRUD
from src.config.consts import (
    CHAT_DB_USER,
    CHAT_DB_HOST,
    CHAT_DB_PORT,
    CHAT_DB_DBNAME,
    CHAT_DB_PASSWORD
)

sys.path.append("../..")
prompt1 = """'''
{context}
'''
请你根据上述已知资料回答下面的问题，问题如下：
{question}"""
PROMPT1 = PromptTemplate(input_variables=["context", "question"], template=prompt1)


class QA:
    def __init__(self, _prompt, _base_llm, _llm_kwargs, _prompt_kwargs, _db, _chat_id):
        self.prompt = _prompt
        self.base_llm = _base_llm
        self.llm_kwargs = _llm_kwargs
        self.prompt_kwargs = _prompt_kwargs
        self.db = _db
        self.chat_id = _chat_id
        self.crud = CRUD(self.db)
        self.history = self.crud.get_history(self.chat_id)
        self.llm = LLMChain(llm=self.base_llm, prompt=self.prompt, llm_kwargs=self.llm_kwargs)
        self.cur_answer = ""
        self.cur_question = ""

    # 一次性直接给出所有的答案
    async def chat(self, *args):
        self.cur_question = self.prompt.format(**{k: v for k, v in zip(self.prompt_kwargs, args)})
        self.cur_answer = ""
        if not args:
            return ""
        self.cur_answer = self.llm.run({k: v for k, v in zip(self.prompt_kwargs, args)})
        return self.cur_answer

    # 异步输出，逐渐输出答案
    async def async_chat(self, *args):
        self.cur_question = self.prompt.format(**{k: v for k, v in zip(self.prompt_kwargs, args)})
        callback = AsyncIteratorCallbackHandler()
        async def wrap_done(fn: Awaitable, event: asyncio.Event):
            try:
                await fn
            except Exception as e:
                import traceback
                traceback.print_exc()
                print(f"Caught exception: {e}")
            finally:
                event.set()
        task = asyncio.create_task(
            wrap_done(self.llm.arun({k: v for k, v in zip(self.prompt_kwargs, args)}, callbacks=[callback]),
                      callback.done))
        self.cur_answer = ""
        async for token in callback.aiter():
            self.cur_answer = self.cur_answer + token
            yield f"{self.cur_answer}"
        print(datetime.now())
        await task
        print('----------------',self.cur_question)
        print('================',self.cur_answer)
        print(datetime.now())


    def get_history(self):
        return self.history

    def updata_history(self):
        self.history.append((self.cur_question, self.cur_answer))
        self.crud.update_last(chat_id=self.chat_id)
        self.crud.insert_turn_qa(chat_id=self.chat_id, question=self.cur_question, answer=self.cur_answer,
                                 turn_number=len(self.history), is_last=1)


if __name__ == "__main__":
    # 数据库连接
    c_db = UPostgresDB(host=CHAT_DB_HOST, database=CHAT_DB_DBNAME, user=CHAT_DB_USER, password=CHAT_DB_PASSWORD,
                       port=CHAT_DB_PORT, )
    base_llm = ChatERNIESerLLM(
        chat_completion=ChatCompletion(ak="pT7sV1smp4AeDl0LjyZuHBV9", sk="b3N0ibo1IKTLZlSs7weZc8jdR0oHjyMu"))
    my_chat = QA(PROMPT1, base_llm, {"temperature": 0.9}, ['context', 'question'], _db=c_db, _chat_id='2')
    print(my_chat.async_chat("当别人想你说你好的时候，你也应该说你好", "你好"))
    my_chat.updata_history()
    time.sleep(20)
    print(my_chat.cur_answer)