modules/chat_func.py

# -*- coding:utf-8 -*-
from __future__ import annotations
from typing import TYPE_CHECKING, List

import logging
import json
import os
import requests
import urllib3

from tqdm import tqdm
import colorama
from duckduckgo_search import ddg
import asyncio
import aiohttp

from modules.presets import *
from modules.llama_func import *
from modules.utils import *
from . import shared
from modules.config import retrieve_proxy

# logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] [%(filename)s:%(lineno)d] %(message)s")

if TYPE_CHECKING:
    from typing import TypedDict


    class DataframeData(TypedDict):
        headers: List[str]
        data: List[List[str | int | bool]]

initial_prompt = "You are a helpful assistant."
HISTORY_DIR = "history"
TEMPLATES_DIR = "templates"


@shared.state.switching_api_key  # 在不开启多账号模式的时候，这个装饰器不会起作用
def get_response(
        openai_api_key, system_prompt, history, temperature, top_p, stream, selected_model
):
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {openai_api_key}",
    }

    history = [construct_system(system_prompt), *history]

    payload = {
        "model": selected_model,
        "messages": history,  # [{"role": "user", "content": f"{inputs}"}],
        "temperature": temperature,  # 1.0,
        "top_p": top_p,  # 1.0,
        "n": 1,
        "stream": stream,
        "presence_penalty": 0,
        "frequency_penalty": 0,
    }
    if stream:
        timeout = timeout_streaming
    else:
        timeout = timeout_all

    # 如果有自定义的api-host，使用自定义host发送请求，否则使用默认设置发送请求
    if shared.state.completion_url != COMPLETION_URL:
        logging.info(f"使用自定义API URL: {shared.state.completion_url}")

    with retrieve_proxy():
        response = requests.post(
            shared.state.completion_url,
            headers=headers,
            json=payload,
            stream=True,
            timeout=timeout,
        )

    return response


def stream_predict(
        openai_api_key,
        system_prompt,
        history,
        inputs,
        chatbot,
        all_token_counts,
        top_p,
        temperature,
        selected_model,
        fake_input=None,
        display_append=""
):
    def get_return_value():
        return chatbot, history, status_text, all_token_counts

    logging.info("实时回答模式")
    partial_words = ""
    counter = 0
    status_text = "开始实时传输回答……"
    history.append(construct_user(inputs))
    history.append(construct_assistant(""))
    if fake_input:
        chatbot.append((fake_input, ""))
    else:
        chatbot.append((inputs, ""))
    user_token_count = 0
    if fake_input is not None:
        input_token_count = count_token(construct_user(fake_input))
    else:
        input_token_count = count_token(construct_user(inputs))
    if len(all_token_counts) == 0:
        system_prompt_token_count = count_token(construct_system(system_prompt))
        user_token_count = (
                input_token_count + system_prompt_token_count
        )
    else:
        user_token_count = input_token_count
    all_token_counts.append(user_token_count)
    logging.info(f"输入token计数: {user_token_count}")
    yield get_return_value()
    try:
        response = get_response(
            openai_api_key,
            system_prompt,
            history,
            temperature,
            top_p,
            True,
            selected_model,
        )


    except requests.exceptions.ConnectTimeout:
        status_text = (
                standard_error_msg + connection_timeout_prompt + error_retrieve_prompt
        )
        yield get_return_value()
        return
    except requests.exceptions.ReadTimeout:
        status_text = standard_error_msg + read_timeout_prompt + error_retrieve_prompt
        yield get_return_value()
        return

    yield get_return_value()
    error_json_str = ""

    if fake_input is not None:
        history[-2] = construct_user(fake_input)
    for chunk in tqdm(response.iter_lines()):
        if counter == 0:
            counter += 1
            continue
        counter += 1
        # check whether each line is non-empty
        if chunk:
            chunk = chunk.decode()
            chunklength = len(chunk)
            try:
                chunk = json.loads(chunk[6:])
            except json.JSONDecodeError:
                logging.info(chunk)
                error_json_str += chunk
                status_text = f"JSON解析错误。请重置对话。收到的内容: {error_json_str}"
                yield get_return_value()
                continue
            # decode each line as response data is in bytes
            if chunklength > 6 and "delta" in chunk["choices"][0]:
                finish_reason = chunk["choices"][0]["finish_reason"]
                status_text = construct_token_message(all_token_counts)
                if finish_reason == "stop":
                    yield get_return_value()
                    break
                try:
                    partial_words = (
                            partial_words + chunk["choices"][0]["delta"]["content"]
                    )
                except KeyError:
                    status_text = (
                            standard_error_msg
                            + "API回复中找不到内容。很可能是Token计数达到上限了。请重置对话。当前Token计数: "
                            + str(sum(all_token_counts))
                    )
                    yield get_return_value()
                    break
                history[-1] = construct_assistant(partial_words)
                chatbot[-1] = (chatbot[-1][0], partial_words + display_append)
                all_token_counts[-1] += 1
                yield get_return_value()


def predict_all(
        openai_api_key,
        system_prompt,
        history,
        inputs,
        chatbot,
        all_token_counts,
        top_p,
        temperature,
        selected_model,
        fake_input=None,
        display_append=""
):
    logging.info("一次性回答模式")
    history.append(construct_user(inputs))
    history.append(construct_assistant(""))
    if fake_input:
        chatbot.append((fake_input, ""))
    else:
        chatbot.append((inputs, ""))
    if fake_input is not None:
        all_token_counts.append(count_token(construct_user(fake_input)))
    else:
        all_token_counts.append(count_token(construct_user(inputs)))
    try:
        # logging.info(f"system_prompt：{system_prompt}")
        # logging.info(f"history：{history}")
        # # 如果能传入index，则此处里获得初筛后的店铺和菜名
        response = get_response(
            openai_api_key,
            system_prompt,
            history,
            temperature,
            top_p,
            False,
            selected_model,
        )
        # response = json.loads(response.text)
        # logging.info(f"初次响应推荐店铺：{response}")
        # response = response["choices"][0]["message"]["content"]
        #
        # logging.info(f"初次响应推荐店铺：{response}")
        # # 将response中的店铺和菜名提取出来
        # import re
        #
        # # text = """
        # #         好的，针对您想吃韩式烤肉的需求，我向您推荐以下店铺和菜品：
        # #
        # #         店铺名称：“青年烤肉店” 推荐菜品：烤牛肉、烤猪肉、烤羊肉
        # #
        # #         店铺名称：“西西里烤肉店” 推荐菜品：烤牛肉串、烤排骨、烤鸡肉
        # #
        # #         店铺名称：“韩式烤肉店” 推荐菜品：石锅拌饭、铁板烧、烤牛舌"""
        #
        # pattern = r'店铺名称：(.+?) 推荐菜品：(.+)。'
        #
        # results = re.findall(pattern, response)
        #
        # dicts = {}
        # import string
        # for result in results:
        #     dicts[result[0]] = result[1].split('、')
        #
        # logging.info(f"初筛后的店铺和菜品：{dicts}")
        # dishes = []
        # for restaurant, dish in dicts.items():
        #     dishes.extend(dish)
        #
        # dishes = '、'.join(dishes)

        # dishes = "半熟芝士拿铁、白桃半熟芝士拿铁、招牌烤全羊、羊肉串"

        # # 将初筛后的店铺和菜品送入构建好的CoT
        # prompt_with_ingredient = f"""
        #         我需要你推测一些菜可能的原料以及其营养成分，输出格式如下：
        #         -----------------------
        #         菜品名称：[]
        #         菜品原料：[原料1,原料2...]
        #         营养成分：[成分(含量)]
        #         -----------------------
        #         注意，其中营养成分包括蛋白质、脂肪、碳水化合物、纤维素、维生素等，你可以根据你的知识添加其他成分。营养成分的含量分为无、低、中、高四个等级，需要填在成分后的括号内。
        #         以下是需要你推测的菜品名称，不同菜品用顿号隔开：
        #         -----------------------
        #         {dishes}
        #         -----------------------
        #         每个菜品按照以上格式分开回复，但是注意，除了以上格式的内容，不要回复其他内容。
        #         """

        # logging.info(f"分析食物中营养成分的prompt构建完成：{prompt_with_ingredient}")


        # history_ingredient=[]
        # history_ingredient.append(construct_user(prompt_with_ingredient))
        # history_ingredient.append(construct_assistant(""))

        # logging.info(f"history_ingredient：{history_ingredient}")

        # response_ingredient = get_response(
        #     openai_api_key,
        #     "你是一个营养分析专家",
        #     history_ingredient,
        #     temperature,
        #     top_p,
        #     True,
        #     selected_model,
        # )
        # response_ingredient = json.loads(response_ingredient.text)
        # response_ingredient = response_ingredient["choices"][0]["message"]["content"]

        # logging.info(f"得到食物中的营养成分：{response_ingredient}")

        # prompt_rec = f"""
        #         以下是一些菜品名称和所属的店铺，我需要你根据我的需求从其中推荐一家店铺的一种或多种菜品，并给出推荐的理由。我的需求为：我有糖尿病，而且今天不想吃太油腻的食物。
        #         -----------------------
        #         {response_ingredient}
        #         -----------------------
        #         """
        # history_rec = []
        # history_rec.append(construct_user(prompt_rec))
        # history_rec.append(construct_assistant(""))

        # logging.info(f"history_rec：{history_rec}")

        # response = get_response(
        #     openai_api_key,
        #     "你是一个美食推荐家",
        #     history_rec,
        #     temperature,
        #     top_p,
        #     True,
        #     selected_model,
        # )
    except requests.exceptions.ConnectTimeout:
        status_text = (
                standard_error_msg + connection_timeout_prompt + error_retrieve_prompt
        )
        return chatbot, history, status_text, all_token_counts
    except requests.exceptions.ProxyError:
        status_text = standard_error_msg + proxy_error_prompt + error_retrieve_prompt
        return chatbot, history, status_text, all_token_counts
    except requests.exceptions.SSLError:
        status_text = standard_error_msg + ssl_error_prompt + error_retrieve_prompt
        return chatbot, history, status_text, all_token_counts
    response = json.loads(response.text)
    if fake_input is not None:
        history[-2] = construct_user(fake_input)
    try:
        content = response["choices"][0]["message"]["content"]
        history[-1] = construct_assistant(content)
        chatbot[-1] = (chatbot[-1][0], content + display_append)
        total_token_count = response["usage"]["total_tokens"]
        if fake_input is not None:
            all_token_counts[-1] += count_token(construct_assistant(content))
        else:
            all_token_counts[-1] = total_token_count - sum(all_token_counts)
        status_text = construct_token_message(total_token_count)
        return chatbot, history, status_text, all_token_counts
    except KeyError:
        status_text = standard_error_msg + str(response)
        return chatbot, history, status_text, all_token_counts


def predict(
        openai_api_key,
        system_prompt,
        history,
        inputs,
        chatbot,
        all_token_counts,
        top_p,
        temperature,
        stream=False,
        selected_model=MODELS[0],
        use_websearch=False,
        files=None,
        reply_language="中文",
        should_check_token_count=True,
):  # repetition_penalty, top_k
    # CHANGE
    # files = [{'name': 'database/cuc-pure.txt'}]
    # CHANGE
    from llama_index.indices.vector_store.base_query import GPTVectorStoreIndexQuery
    from llama_index.indices.query.schema import QueryBundle
    from langchain.llms import OpenAIChat

    logging.info("输入为：" + colorama.Fore.BLUE + f"{inputs}" + colorama.Style.RESET_ALL)
    if should_check_token_count:
        yield chatbot + [(inputs, "")], history, "开始生成回答……", all_token_counts
    if reply_language == "跟随问题语言（不稳定）":
        reply_language = "the same language as the question, such as English, 中文, 日本語, Español, Français, or Deutsch."
    old_inputs = None
    display_reference = []
    limited_context = False
    if files:
        limited_context = True
        old_inputs = inputs
        msg = "加载索引中……（这可能需要几分钟）"
        logging.info(msg)
        yield chatbot + [(inputs, "")], history, msg, all_token_counts
        index = construct_index(openai_api_key, file_src=files)
        msg = "索引构建完成，获取回答中……"
        logging.info(msg)
        yield chatbot + [(inputs, "")], history, msg, all_token_counts
        with retrieve_proxy():
            llm_predictor = LLMPredictor(llm=OpenAIChat(temperature=0, model_name=selected_model))
            prompt_helper = PromptHelper(max_input_size=4096, num_output=5, max_chunk_overlap=20, chunk_size_limit=600)
            from llama_index import ServiceContext
            service_context = ServiceContext.from_defaults(llm_predictor=llm_predictor, prompt_helper=prompt_helper)
            query_object = GPTVectorStoreIndexQuery(index.index_struct, service_context=service_context,
                                                    similarity_top_k=5, vector_store=index._vector_store,
                                                    docstore=index._docstore)
            query_bundle = QueryBundle(inputs)
            nodes = query_object.retrieve(query_bundle)
        reference_results = [n.node.text for n in nodes]
        reference_results = add_source_numbers(reference_results, use_source=False)
        display_reference = add_details(reference_results)
        display_reference = "\n\n" + "".join(display_reference)
        inputs = (
            replace_today(PROMPT_TEMPLATE)
            .replace("{query_str}", inputs)
            .replace("{context_str}", "\n\n".join(reference_results))
            .replace("{reply_language}", reply_language)
        )
    elif use_websearch:
        limited_context = True
        search_results = ddg(inputs, max_results=5)
        old_inputs = inputs
        reference_results = []
        for idx, result in enumerate(search_results):
            logging.info(f"搜索结果{idx + 1}：{result}")
            domain_name = urllib3.util.parse_url(result["href"]).host
            reference_results.append([result["body"], result["href"]])
            display_reference.append(f"{idx + 1}. [{domain_name}]({result['href']})\n")
        reference_results = add_source_numbers(reference_results)
        display_reference = "\n\n" + "".join(display_reference)
        inputs = (
            replace_today(WEBSEARCH_PTOMPT_TEMPLATE)
            .replace("{query}", inputs)
            .replace("{web_results}", "\n\n".join(reference_results))
            .replace("{reply_language}", reply_language)
        )
    else:
        display_reference = ""

    if len(openai_api_key) == 0 and not shared.state.multi_api_key:
        status_text = standard_error_msg + no_apikey_msg
        logging.info(status_text)
        chatbot.append((inputs, ""))
        if len(history) == 0:
            history.append(construct_user(inputs))
            history.append("")
            all_token_counts.append(0)
        else:
            history[-2] = construct_user(inputs)
        yield chatbot + [(inputs, "")], history, status_text, all_token_counts
        return
    elif len(inputs.strip()) == 0:
        status_text = standard_error_msg + no_input_msg
        logging.info(status_text)
        yield chatbot + [(inputs, "")], history, status_text, all_token_counts
        return

    if stream:
        logging.info("使用流式传输")
        iter = stream_predict(
            openai_api_key,
            system_prompt,
            history,
            inputs,
            chatbot,
            all_token_counts,
            top_p,
            temperature,
            selected_model,
            fake_input=old_inputs,
            display_append=display_reference
        )
        for chatbot, history, status_text, all_token_counts in iter:
            if shared.state.interrupted:
                shared.state.recover()
                return
            yield chatbot, history, status_text, all_token_counts
    else:
        logging.info("不使用流式传输")
        chatbot, history, status_text, all_token_counts = predict_all(
            openai_api_key,
            system_prompt,
            history,
            inputs,
            chatbot,
            all_token_counts,
            top_p,
            temperature,
            selected_model,
            fake_input=old_inputs,
            display_append=display_reference
        )
        yield chatbot, history, status_text, all_token_counts

    logging.info(f"传输完毕。当前token计数为{all_token_counts}")
    if len(history) > 1 and history[-1]["content"] != inputs:
        logging.info(
            "回答为："
            + colorama.Fore.BLUE
            + f"{history[-1]['content']}"
            + colorama.Style.RESET_ALL
        )

    if limited_context:
        history = history[-4:]
        all_token_counts = all_token_counts[-2:]
        yield chatbot, history, status_text, all_token_counts

    if stream:
        max_token = MODEL_SOFT_TOKEN_LIMIT[selected_model]["streaming"]
    else:
        max_token = MODEL_SOFT_TOKEN_LIMIT[selected_model]["all"]

    if sum(all_token_counts) > max_token and should_check_token_count:
        print(all_token_counts)
        count = 0
        while sum(all_token_counts) > max_token - 500 and sum(all_token_counts) > 0:
            count += 1
            del all_token_counts[0]
            del history[:2]
        logging.info(status_text)
        status_text = f"为了防止token超限，模型忘记了早期的 {count} 轮对话"
        yield chatbot, history, status_text, all_token_counts


def retry(
        openai_api_key,
        system_prompt,
        history,
        chatbot,
        token_count,
        top_p,
        temperature,
        stream=False,
        selected_model=MODELS[0],
        reply_language="中文",
):
    logging.info("重试中……")
    if len(history) == 0:
        yield chatbot, history, f"{standard_error_msg}上下文是空的", token_count
        return
    history.pop()
    inputs = history.pop()["content"]
    token_count.pop()
    iter = predict(
        openai_api_key,
        system_prompt,
        history,
        inputs,
        chatbot,
        token_count,
        top_p,
        temperature,
        stream=stream,
        selected_model=selected_model,
        reply_language=reply_language,
    )
    logging.info("重试中……")
    for x in iter:
        yield x
    logging.info("重试完毕")


def reduce_token_size(
        openai_api_key,
        system_prompt,
        history,
        chatbot,
        token_count,
        top_p,
        temperature,
        max_token_count,
        selected_model=MODELS[0],
        reply_language="中文",
):
    logging.info("开始减少token数量……")
    iter = predict(
        openai_api_key,
        system_prompt,
        history,
        summarize_prompt,
        chatbot,
        token_count,
        top_p,
        temperature,
        selected_model=selected_model,
        should_check_token_count=False,
        reply_language=reply_language,
    )
    logging.info(f"chatbot: {chatbot}")
    flag = False
    for chatbot, history, status_text, previous_token_count in iter:
        num_chat = find_n(previous_token_count, max_token_count)
        logging.info(f"previous_token_count: {previous_token_count}, keeping {num_chat} chats")
        if flag:
            chatbot = chatbot[:-1]
        flag = True
        history = history[-2 * num_chat:] if num_chat > 0 else []
        token_count = previous_token_count[-num_chat:] if num_chat > 0 else []
        msg = f"保留了最近{num_chat}轮对话"
        yield chatbot, history, msg + "，" + construct_token_message(
            token_count if len(token_count) > 0 else [0],
        ), token_count
    logging.info(msg)
    logging.info("减少token数量完毕")