加入GPT Index

ChuanhuChatbot.py

@@ -6,9 +6,11 @@ import sys
 import argparse
 from utils import *
 from presets import *
+from overwrites import *
+from chat_func import *
 
 logging.basicConfig(
-    level=logging.INFO,
+    level=logging.DEBUG,
     format="%(asctime)s [%(levelname)s] [%(filename)s:%(lineno)d] %(message)s",
 )
 
@@ -49,6 +51,7 @@ else:
                 authflag = True
 
 gr.Chatbot.postprocess = postprocess
+PromptHelper.compact_text_chunks = compact_text_chunks
 
 with open("custom.css", "r", encoding="utf-8") as f:
     customCSS = f.read()
@@ -165,7 +168,7 @@ with gr.Blocks(
                         label="实时传输回答", value=True, visible=enable_streaming_option
                     )
                     use_websearch_checkbox = gr.Checkbox(label="使用在线搜索", value=False)
-                    index_files = gr.File(label="上传索引文件", type="file", multiple=True)
+                    index_files = gr.Files(label="上传索引文件", type="file", multiple=True)
 
                 with gr.Tab(label="Prompt"):
                     systemPromptTxt = gr.Textbox(
@@ -286,6 +289,7 @@ with gr.Blocks(
             use_streaming_checkbox,
             model_select_dropdown,
             use_websearch_checkbox,
+            index_files
         ],
         [chatbot, history, status_display, token_count],
         show_progress=True,
@@ -306,6 +310,7 @@ with gr.Blocks(
             use_streaming_checkbox,
             model_select_dropdown,
             use_websearch_checkbox,
+            index_files
         ],
         [chatbot, history, status_display, token_count],
         show_progress=True,

chat_func.py

@@ -0,0 +1,447 @@
+# -*- coding:utf-8 -*-
+from __future__ import annotations
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Tuple, Type
+import logging
+import json
+import gradio as gr
+
+# import openai
+import os
+import traceback
+import requests
+
+# import markdown
+import csv
+import mdtex2html
+from pypinyin import lazy_pinyin
+from presets import *
+from llama_func import *
+from utils import *
+import tiktoken
+from tqdm import tqdm
+import colorama
+import os
+from llama_index import (
+    GPTSimpleVectorIndex,
+    GPTTreeIndex,
+    GPTKeywordTableIndex,
+    GPTListIndex,
+)
+from llama_index import SimpleDirectoryReader, download_loader
+from llama_index import (
+    Document,
+    LLMPredictor,
+    PromptHelper,
+    QuestionAnswerPrompt,
+    RefinePrompt,
+)
+from langchain.llms import OpenAIChat, OpenAI
+from duckduckgo_search import ddg
+import datetime
+
+# 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."
+API_URL = "https://api.openai.com/v1/chat/completions"
+HISTORY_DIR = "history"
+TEMPLATES_DIR = "templates"
+
+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
+
+    # 获取环境变量中的代理设置
+    http_proxy = os.environ.get("HTTP_PROXY") or os.environ.get("http_proxy")
+    https_proxy = os.environ.get("HTTPS_PROXY") or os.environ.get("https_proxy")
+
+    # 如果存在代理设置，使用它们
+    proxies = {}
+    if http_proxy:
+        logging.info(f"Using HTTP proxy: {http_proxy}")
+        proxies["http"] = http_proxy
+    if https_proxy:
+        logging.info(f"Using HTTPS proxy: {https_proxy}")
+        proxies["https"] = https_proxy
+
+    # 如果有代理，使用代理发送请求，否则使用默认设置发送请求
+    if proxies:
+        response = requests.post(
+            API_URL,
+            headers=headers,
+            json=payload,
+            stream=True,
+            timeout=timeout,
+            proxies=proxies,
+        )
+    else:
+        response = requests.post(
+            API_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,
+):
+    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(""))
+    chatbot.append((parse_text(inputs), ""))
+    user_token_count = 0
+    if len(all_token_counts) == 0:
+        system_prompt_token_count = count_token(construct_system(system_prompt))
+        user_token_count = (
+            count_token(construct_user(inputs)) + system_prompt_token_count
+        )
+    else:
+        user_token_count = count_token(construct_user(inputs))
+    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 = ""
+
+    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(
+                    sum(all_token_counts), stream=True
+                )
+                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] = (parse_text(inputs), parse_text(partial_words))
+                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,
+):
+    logging.info("一次性回答模式")
+    history.append(construct_user(inputs))
+    history.append(construct_assistant(""))
+    chatbot.append((parse_text(inputs), ""))
+    all_token_counts.append(count_token(construct_user(inputs)))
+    try:
+        response = get_response(
+            openai_api_key,
+            system_prompt,
+            history,
+            temperature,
+            top_p,
+            False,
+            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)
+    content = response["choices"][0]["message"]["content"]
+    history[-1] = construct_assistant(content)
+    chatbot[-1] = (parse_text(inputs), parse_text(content))
+    total_token_count = response["usage"]["total_tokens"]
+    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
+
+
+def predict(
+    openai_api_key,
+    system_prompt,
+    history,
+    inputs,
+    chatbot,
+    all_token_counts,
+    top_p,
+    temperature,
+    stream=False,
+    selected_model=MODELS[0],
+    use_websearch_checkbox=False,
+    files = None,
+    should_check_token_count=True,
+):  # repetition_penalty, top_k
+    logging.info("输入为：" + colorama.Fore.BLUE + f"{inputs}" + colorama.Style.RESET_ALL)
+    if files:
+        msg = "构建索引中……（这可能需要比较久的时间）"
+        logging.info(msg)
+        yield chatbot, history, msg, all_token_counts
+        index = construct_index(openai_api_key, file_src=files)
+        msg = "索引构建完成，获取回答中……"
+        yield chatbot, history, msg, all_token_counts
+        history, chatbot, status_text = chat_ai(openai_api_key, index, inputs, history, chatbot)
+        yield chatbot, history, status_text, all_token_counts
+        return
+    if use_websearch_checkbox:
+        results = ddg(inputs, max_results=3)
+        web_results = []
+        for idx, result in enumerate(results):
+            logging.info(f"搜索结果{idx + 1}：{result}")
+            web_results.append(f'[{idx+1}]"{result["body"]}"\nURL: {result["href"]}')
+        web_results = "\n\n".join(web_results)
+        inputs = (
+            replace_today(WEBSEARCH_PTOMPT_TEMPLATE)
+            .replace("{query}", inputs)
+            .replace("{web_results}", web_results)
+        )
+    if len(openai_api_key) != 51:
+        status_text = standard_error_msg + no_apikey_msg
+        logging.info(status_text)
+        chatbot.append((parse_text(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, history, status_text, all_token_counts
+        return
+    if stream:
+        yield chatbot, history, "开始生成回答……", all_token_counts
+    if stream:
+        logging.info("使用流式传输")
+        iter = stream_predict(
+            openai_api_key,
+            system_prompt,
+            history,
+            inputs,
+            chatbot,
+            all_token_counts,
+            top_p,
+            temperature,
+            selected_model,
+        )
+        for chatbot, history, status_text, all_token_counts in iter:
+            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,
+        )
+        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 stream:
+        max_token = max_token_streaming
+    else:
+        max_token = max_token_all
+    if sum(all_token_counts) > max_token and should_check_token_count:
+        status_text = f"精简token中{all_token_counts}/{max_token}"
+        logging.info(status_text)
+        yield chatbot, history, status_text, all_token_counts
+        iter = reduce_token_size(
+            openai_api_key,
+            system_prompt,
+            history,
+            chatbot,
+            all_token_counts,
+            top_p,
+            temperature,
+            stream=False,
+            selected_model=selected_model,
+            hidden=True,
+        )
+        for chatbot, history, status_text, all_token_counts in iter:
+            status_text = f"Token 达到上限，已自动降低Token计数至 {status_text}"
+            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],
+):
+    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,
+    )
+    logging.info("重试完毕")
+    for x in iter:
+        yield x
+
+
+def reduce_token_size(
+    openai_api_key,
+    system_prompt,
+    history,
+    chatbot,
+    token_count,
+    top_p,
+    temperature,
+    stream=False,
+    selected_model=MODELS[0],
+    hidden=False,
+):
+    logging.info("开始减少token数量……")
+    iter = predict(
+        openai_api_key,
+        system_prompt,
+        history,
+        summarize_prompt,
+        chatbot,
+        token_count,
+        top_p,
+        temperature,
+        stream=stream,
+        selected_model=selected_model,
+        should_check_token_count=False,
+    )
+    logging.info(f"chatbot: {chatbot}")
+    for chatbot, history, status_text, previous_token_count in iter:
+        history = history[-2:]
+        token_count = previous_token_count[-1:]
+        if hidden:
+            chatbot.pop()
+        yield chatbot, history, construct_token_message(
+            sum(token_count), stream=stream
+        ), token_count
+    logging.info("减少token数量完毕")

llama_func.py

@@ -0,0 +1,201 @@
+import os
+from llama_index import (
+    GPTSimpleVectorIndex,
+    GPTTreeIndex,
+    GPTKeywordTableIndex,
+    GPTListIndex,
+)
+from llama_index import SimpleDirectoryReader, download_loader
+from llama_index import (
+    Document,
+    LLMPredictor,
+    PromptHelper,
+    QuestionAnswerPrompt,
+    RefinePrompt,
+)
+from langchain.llms import OpenAIChat, OpenAI
+from googlesearch import search as google_search
+from baidusearch.baidusearch import search as baidu_search
+from duckduckgo_search import ddg
+import colorama
+
+import logging
+import sys
+
+from presets import *
+from utils import *
+
+
+def get_documents(file_src):
+    documents = []
+    index_name = ""
+    logging.debug("Loading documents...")
+    logging.debug(f"file_src: {file_src}")
+    for file in file_src:
+        logging.debug(f"file: {file.name}")
+        index_name += file.name
+        if os.path.splitext(file.name)[1] == ".pdf":
+            logging.debug("Loading PDF...")
+            CJKPDFReader = download_loader("CJKPDFReader")
+            loader = CJKPDFReader()
+            documents += loader.load_data(file=file.name)
+        elif os.path.splitext(file.name)[1] == ".docx":
+            logging.debug("Loading DOCX...")
+            DocxReader = download_loader("DocxReader")
+            loader = DocxReader()
+            documents += loader.load_data(file=file.name)
+        elif os.path.splitext(file.name)[1] == ".epub":
+            logging.debug("Loading EPUB...")
+            EpubReader = download_loader("EpubReader")
+            loader = EpubReader()
+            documents += loader.load_data(file=file.name)
+        else:
+            logging.debug("Loading text file...")
+            with open(file.name, "r", encoding="utf-8") as f:
+                text = add_space(f.read())
+                documents += [Document(text)]
+    index_name = sha1sum(index_name)
+    return documents, index_name
+
+
+def construct_index(
+    api_key,
+    file_src,
+    max_input_size=4096,
+    num_outputs=1,
+    max_chunk_overlap=20,
+    chunk_size_limit=600,
+    embedding_limit=None,
+    separator=" ",
+    num_children=10,
+    max_keywords_per_chunk=10,
+):
+    os.environ["OPENAI_API_KEY"] = api_key
+    chunk_size_limit = None if chunk_size_limit == 0 else chunk_size_limit
+    embedding_limit = None if embedding_limit == 0 else embedding_limit
+    separator = " " if separator == "" else separator
+
+    llm_predictor = LLMPredictor(
+        llm=OpenAI(model_name="gpt-3.5-turbo-0301", openai_api_key=api_key)
+    )
+    prompt_helper = PromptHelper(
+        max_input_size,
+        num_outputs,
+        max_chunk_overlap,
+        embedding_limit,
+        chunk_size_limit,
+        separator=separator,
+    )
+    documents, index_name = get_documents(file_src)
+    if os.path.exists(f"./index/{index_name}.json"):
+        logging.info("找到了缓存的索引文件，加载中……")
+        return GPTSimpleVectorIndex.load_from_disk(f"./index/{index_name}.json")
+    else:
+        try:
+            logging.debug("构建索引中……")
+            index = GPTSimpleVectorIndex(
+                documents, llm_predictor=llm_predictor, prompt_helper=prompt_helper
+            )
+            os.makedirs("./index", exist_ok=True)
+            index.save_to_disk(f"./index/{index_name}.json")
+            return index
+        except Exception as e:
+            print(e)
+            return None
+
+
+def chat_ai(
+    api_key,
+    index,
+    question,
+    context,
+    chatbot,
+):
+    os.environ["OPENAI_API_KEY"] = api_key
+
+    logging.info(f"Question: {question}")
+
+    response, status_text = ask_ai(
+        api_key,
+        index,
+        question,
+        replace_today(PROMPT_TEMPLATE),
+        REFINE_TEMPLATE,
+        SIM_K,
+        INDEX_QUERY_TEMPRATURE,
+        context,
+    )
+    if response is None:
+        status_text = "查询失败，请换个问法试试"
+        return context, chatbot
+    response = response
+
+    context.append({"role": "user", "content": question})
+    context.append({"role": "assistant", "content": response})
+    chatbot.append((question, response))
+
+    os.environ["OPENAI_API_KEY"] = ""
+    return context, chatbot, status_text
+
+
+def ask_ai(
+    api_key,
+    index,
+    question,
+    prompt_tmpl,
+    refine_tmpl,
+    sim_k=1,
+    temprature=0,
+    prefix_messages=[],
+):
+    os.environ["OPENAI_API_KEY"] = api_key
+
+    logging.debug("Index file found")
+    logging.debug("Querying index...")
+    llm_predictor = LLMPredictor(
+        llm=OpenAI(
+            temperature=temprature,
+            model_name="gpt-3.5-turbo-0301",
+            prefix_messages=prefix_messages,
+        )
+    )
+
+    response = None  # Initialize response variable to avoid UnboundLocalError
+    qa_prompt = QuestionAnswerPrompt(prompt_tmpl)
+    rf_prompt = RefinePrompt(refine_tmpl)
+    response = index.query(
+        question,
+        llm_predictor=llm_predictor,
+        similarity_top_k=sim_k,
+        text_qa_template=qa_prompt,
+        refine_template=rf_prompt,
+        response_mode="compact",
+    )
+
+    if response is not None:
+        logging.info(f"Response: {response}")
+        ret_text = response.response
+        ret_text += "\n----------\n"
+        nodes = []
+        for index, node in enumerate(response.source_nodes):
+            brief = node.source_text[:25].replace("\n", "")
+            nodes.append(
+                f"<details><summary>[{index+1}]\t{brief}...</summary><p>{node.source_text}</p></details>"
+            )
+        ret_text += "\n\n".join(nodes)
+        logging.info(
+            f"Response: {colorama.Fore.BLUE}{ret_text}{colorama.Style.RESET_ALL}"
+        )
+        os.environ["OPENAI_API_KEY"] = ""
+        return ret_text, f"查询消耗了{llm_predictor.last_token_usage} tokens"
+    else:
+        logging.warning("No response found, returning None")
+        os.environ["OPENAI_API_KEY"] = ""
+        return None
+
+
+def add_space(text):
+    punctuations = {"，": "， ", "。": "。 ", "？": "？ ", "！": "！ ", "：": "： ", "；": "； "}
+    for cn_punc, en_punc in punctuations.items():
+        text = text.replace(cn_punc, en_punc)
+    return text

overwrites.py

@@ -0,0 +1,97 @@
+from __future__ import annotations
+import os
+
+import llama_index
+
+from llama_index import (
+    LLMPredictor,
+    GPTTreeIndex,
+    Document,
+    GPTSimpleVectorIndex,
+    SimpleDirectoryReader,
+    RefinePrompt,
+    QuestionAnswerPrompt,
+    GPTListIndex,
+    PromptHelper,
+)
+from pathlib import Path
+from docx import Document as DocxDocument
+from tqdm import tqdm
+import re
+from langchain.llms import OpenAIChat, OpenAI
+from llama_index.composability import ComposableGraph
+from IPython.display import Markdown, display
+import json
+from llama_index import Prompt
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Tuple, Type
+
+import logging
+import sys
+
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Tuple, Type
+import logging
+import json
+import gradio as gr
+
+# import openai
+import os
+import traceback
+import requests
+
+# import markdown
+import csv
+import mdtex2html
+from pypinyin import lazy_pinyin
+from presets import *
+from llama_func import *
+import tiktoken
+from tqdm import tqdm
+import colorama
+import os
+from llama_index import (
+    GPTSimpleVectorIndex,
+    GPTTreeIndex,
+    GPTKeywordTableIndex,
+    GPTListIndex,
+)
+from llama_index import SimpleDirectoryReader, download_loader
+from llama_index import (
+    Document,
+    LLMPredictor,
+    PromptHelper,
+    QuestionAnswerPrompt,
+    RefinePrompt,
+)
+from langchain.llms import OpenAIChat, OpenAI
+from duckduckgo_search import ddg
+import datetime
+
+def compact_text_chunks(self, prompt: Prompt, text_chunks: List[str]) -> List[str]:
+    logging.debug("Compacting text chunks...🚀🚀🚀")
+    combined_str = [c.strip() for c in text_chunks if c.strip()]
+    combined_str = [f"[{index+1}] {c}" for index, c in enumerate(combined_str)]
+    combined_str = "\n\n".join(combined_str)
+    # resplit based on self.max_chunk_overlap
+    text_splitter = self.get_text_splitter_given_prompt(prompt, 1, padding=1)
+    return text_splitter.split_text(combined_str)
+
+
+def postprocess(
+    self, y: List[Tuple[str | None, str | None]]
+) -> List[Tuple[str | None, str | None]]:
+    """
+    Parameters:
+        y: List of tuples representing the message and response pairs. Each message and response should be a string, which may be in Markdown format.
+    Returns:
+        List of tuples representing the message and response. Each message and response will be a string of HTML.
+    """
+    if y is None:
+        return []
+    for i, (message, response) in enumerate(y):
+        y[i] = (
+            # None if message is None else markdown.markdown(message),
+            # None if response is None else markdown.markdown(response),
+            None if message is None else message,
+            None if response is None else mdtex2html.convert(response, extensions=['fenced_code','codehilite','tables']),
+        )
+    return y

presets.py

@@ -1,4 +1,23 @@
 # -*- coding:utf-8 -*-
+# 错误信息
+standard_error_msg = "☹️发生了错误："  # 错误信息的标准前缀
+error_retrieve_prompt = "请检查网络连接，或者API-Key是否有效。"  # 获取对话时发生错误
+connection_timeout_prompt = "连接超时，无法获取对话。"  # 连接超时
+read_timeout_prompt = "读取超时，无法获取对话。"  # 读取超时
+proxy_error_prompt = "代理错误，无法获取对话。"  # 代理错误
+ssl_error_prompt = "SSL错误，无法获取对话。"  # SSL 错误
+no_apikey_msg = "API key长度不是51位，请检查是否输入正确。"  # API key 长度不足 51 位
+
+max_token_streaming = 3500  # 流式对话时的最大 token 数
+timeout_streaming = 30  # 流式对话时的超时时间
+max_token_all = 3500  # 非流式对话时的最大 token 数
+timeout_all = 200  # 非流式对话时的超时时间
+enable_streaming_option = True  # 是否启用选择选择是否实时显示回答的勾选框
+HIDE_MY_KEY = False  # 如果你想在UI中隐藏你的 API 密钥，将此值设置为 True
+
+SIM_K = 5
+INDEX_QUERY_TEMPRATURE = 1.0
+
 title = """<h1 align="left" style="min-width:200px; margin-top:0;">川虎ChatGPT 🚀</h1>"""
 description = """\
 <div align="center" style="margin:16px 0">
@@ -12,6 +31,7 @@ description = """\
 """
 
 summarize_prompt = "你是谁？我们刚才聊了什么？"  # 总结对话时的 prompt
+
 MODELS = [
     "gpt-3.5-turbo",
     "gpt-3.5-turbo-0301",
@@ -21,7 +41,8 @@ MODELS = [
     "gpt-4-32k-0314",
 ]  # 可选的模型
 
-websearch_prompt = """\
+
+WEBSEARCH_PTOMPT_TEMPLATE = """\
 Web search results:
 
 {web_results}
@@ -31,18 +52,29 @@ Instructions: Using the provided web search results, write a comprehensive reply
 Query: {query}
 Reply in 中文"""
 
-# 错误信息
-standard_error_msg = "☹️发生了错误："  # 错误信息的标准前缀
-error_retrieve_prompt = "请检查网络连接，或者API-Key是否有效。"  # 获取对话时发生错误
-connection_timeout_prompt = "连接超时，无法获取对话。"  # 连接超时
-read_timeout_prompt = "读取超时，无法获取对话。"  # 读取超时
-proxy_error_prompt = "代理错误，无法获取对话。"  # 代理错误
-ssl_error_prompt = "SSL错误，无法获取对话。"  # SSL 错误
-no_apikey_msg = "API key长度不是51位，请检查是否输入正确。"  # API key 长度不足 51 位
+PROMPT_TEMPLATE = """\
+Context information is below.
+---------------------
+{context_str}
+---------------------
+Using the provided context information, write a comprehensive reply to the given query.
+Make sure to cite results using [number] notation after the reference.
+If the provided context information refer to multiple subjects with the same name, write separate answers for each subject.
+Use prior knowledge only if the given context didn't provide enough information.
+Today is {current_date}.
+Answer the question: {query_str}
+Reply in 中文
+"""
 
-max_token_streaming = 3500  # 流式对话时的最大 token 数
-timeout_streaming = 30  # 流式对话时的超时时间
-max_token_all = 3500  # 非流式对话时的最大 token 数
-timeout_all = 200  # 非流式对话时的超时时间
-enable_streaming_option = True  # 是否启用选择选择是否实时显示回答的勾选框
-HIDE_MY_KEY = False  # 如果你想在UI中隐藏你的 API 密钥，将此值设置为 True
+REFINE_TEMPLATE = """\
+The original question is as follows: {query_str}
+We have provided an existing answer: {existing_answer}
+We have the opportunity to refine the existing answer
+(only if needed) with some more context below.
+------------
+{context_msg}
+------------
+Given the new context, refine the original answer to better
+Answer in the same language as the question, such as English, 中文, 日本語, Español, Français, or Deutsch.
+If the context isn't useful, return the original answer.
+"""

requirements.txt

@@ -6,4 +6,6 @@ socksio
 tqdm
 colorama
 duckduckgo_search
-Pygments
+Pygments
+llama_index
+langchain

utils.py

@@ -18,8 +18,25 @@ from presets import *
 import tiktoken
 from tqdm import tqdm
 import colorama
+import os
+from llama_index import (
+    GPTSimpleVectorIndex,
+    GPTTreeIndex,
+    GPTKeywordTableIndex,
+    GPTListIndex,
+)
+from llama_index import SimpleDirectoryReader, download_loader
+from llama_index import (
+    Document,
+    LLMPredictor,
+    PromptHelper,
+    QuestionAnswerPrompt,
+    RefinePrompt,
+)
+from langchain.llms import OpenAIChat, OpenAI
 from duckduckgo_search import ddg
 import datetime
+import hashlib
 
 # logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] [%(filename)s:%(lineno)d] %(message)s")
 
@@ -37,27 +54,6 @@ HISTORY_DIR = "history"
 TEMPLATES_DIR = "templates"
 
 
-def postprocess(
-    self, y: List[Tuple[str | None, str | None]]
-) -> List[Tuple[str | None, str | None]]:
-    """
-    Parameters:
-        y: List of tuples representing the message and response pairs. Each message and response should be a string, which may be in Markdown format.
-    Returns:
-        List of tuples representing the message and response. Each message and response will be a string of HTML.
-    """
-    if y is None:
-        return []
-    for i, (message, response) in enumerate(y):
-        y[i] = (
-            # None if message is None else markdown.markdown(message),
-            # None if response is None else markdown.markdown(response),
-            None if message is None else message,
-            None if response is None else mdtex2html.convert(response, extensions=['fenced_code','codehilite','tables']),
-        )
-    return y
-
-
 def count_token(message):
     encoding = tiktoken.get_encoding("cl100k_base")
     input_str = f"role: {message['role']}, content: {message['content']}"
@@ -102,389 +98,6 @@ def construct_token_message(token, stream=False):
     return f"Token 计数: {token}"
 
 
-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
-
-    # 获取环境变量中的代理设置
-    http_proxy = os.environ.get("HTTP_PROXY") or os.environ.get("http_proxy")
-    https_proxy = os.environ.get("HTTPS_PROXY") or os.environ.get("https_proxy")
-
-    # 如果存在代理设置，使用它们
-    proxies = {}
-    if http_proxy:
-        logging.info(f"Using HTTP proxy: {http_proxy}")
-        proxies["http"] = http_proxy
-    if https_proxy:
-        logging.info(f"Using HTTPS proxy: {https_proxy}")
-        proxies["https"] = https_proxy
-
-    # 如果有代理，使用代理发送请求，否则使用默认设置发送请求
-    if proxies:
-        response = requests.post(
-            API_URL,
-            headers=headers,
-            json=payload,
-            stream=True,
-            timeout=timeout,
-            proxies=proxies,
-        )
-    else:
-        response = requests.post(
-            API_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,
-):
-    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(""))
-    chatbot.append((parse_text(inputs), ""))
-    user_token_count = 0
-    if len(all_token_counts) == 0:
-        system_prompt_token_count = count_token(construct_system(system_prompt))
-        user_token_count = (
-            count_token(construct_user(inputs)) + system_prompt_token_count
-        )
-    else:
-        user_token_count = count_token(construct_user(inputs))
-    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 = ""
-
-    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(
-                    sum(all_token_counts), stream=True
-                )
-                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] = (parse_text(inputs), parse_text(partial_words))
-                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,
-):
-    logging.info("一次性回答模式")
-    history.append(construct_user(inputs))
-    history.append(construct_assistant(""))
-    chatbot.append((parse_text(inputs), ""))
-    all_token_counts.append(count_token(construct_user(inputs)))
-    try:
-        response = get_response(
-            openai_api_key,
-            system_prompt,
-            history,
-            temperature,
-            top_p,
-            False,
-            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)
-    content = response["choices"][0]["message"]["content"]
-    history[-1] = construct_assistant(content)
-    chatbot[-1] = (parse_text(inputs), parse_text(content))
-    total_token_count = response["usage"]["total_tokens"]
-    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
-
-
-def predict(
-    openai_api_key,
-    system_prompt,
-    history,
-    inputs,
-    chatbot,
-    all_token_counts,
-    top_p,
-    temperature,
-    stream=False,
-    selected_model=MODELS[0],
-    use_websearch_checkbox=False,
-    should_check_token_count=True,
-):  # repetition_penalty, top_k
-    logging.info("输入为：" + colorama.Fore.BLUE + f"{inputs}" + colorama.Style.RESET_ALL)
-    if use_websearch_checkbox:
-        results = ddg(inputs, max_results=3)
-        web_results = []
-        for idx, result in enumerate(results):
-            logging.info(f"搜索结果{idx + 1}：{result}")
-            web_results.append(f'[{idx+1}]"{result["body"]}"\nURL: {result["href"]}')
-        web_results = "\n\n".join(web_results)
-        today = datetime.datetime.today().strftime("%Y-%m-%d")
-        inputs = (
-            websearch_prompt.replace("{current_date}", today)
-            .replace("{query}", inputs)
-            .replace("{web_results}", web_results)
-        )
-    if len(openai_api_key) != 51:
-        status_text = standard_error_msg + no_apikey_msg
-        logging.info(status_text)
-        chatbot.append((parse_text(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, history, status_text, all_token_counts
-        return
-    if stream:
-        yield chatbot, history, "开始生成回答……", all_token_counts
-    if stream:
-        logging.info("使用流式传输")
-        iter = stream_predict(
-            openai_api_key,
-            system_prompt,
-            history,
-            inputs,
-            chatbot,
-            all_token_counts,
-            top_p,
-            temperature,
-            selected_model,
-        )
-        for chatbot, history, status_text, all_token_counts in iter:
-            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,
-        )
-        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 stream:
-        max_token = max_token_streaming
-    else:
-        max_token = max_token_all
-    if sum(all_token_counts) > max_token and should_check_token_count:
-        status_text = f"精简token中{all_token_counts}/{max_token}"
-        logging.info(status_text)
-        yield chatbot, history, status_text, all_token_counts
-        iter = reduce_token_size(
-            openai_api_key,
-            system_prompt,
-            history,
-            chatbot,
-            all_token_counts,
-            top_p,
-            temperature,
-            stream=False,
-            selected_model=selected_model,
-            hidden=True,
-        )
-        for chatbot, history, status_text, all_token_counts in iter:
-            status_text = f"Token 达到上限，已自动降低Token计数至 {status_text}"
-            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],
-):
-    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,
-    )
-    logging.info("重试完毕")
-    for x in iter:
-        yield x
-
-
-def reduce_token_size(
-    openai_api_key,
-    system_prompt,
-    history,
-    chatbot,
-    token_count,
-    top_p,
-    temperature,
-    stream=False,
-    selected_model=MODELS[0],
-    hidden=False,
-):
-    logging.info("开始减少token数量……")
-    iter = predict(
-        openai_api_key,
-        system_prompt,
-        history,
-        summarize_prompt,
-        chatbot,
-        token_count,
-        top_p,
-        temperature,
-        stream=stream,
-        selected_model=selected_model,
-        should_check_token_count=False,
-    )
-    logging.info(f"chatbot: {chatbot}")
-    for chatbot, history, status_text, previous_token_count in iter:
-        history = history[-2:]
-        token_count = previous_token_count[-1:]
-        if hidden:
-            chatbot.pop()
-        yield chatbot, history, construct_token_message(
-            sum(token_count), stream=stream
-        ), token_count
-    logging.info("减少token数量完毕")
-
-
 def delete_last_conversation(chatbot, history, previous_token_count):
     if len(chatbot) > 0 and standard_error_msg in chatbot[-1][1]:
         logging.info("由于包含报错信息，只删除chatbot记录")
@@ -643,6 +256,7 @@ def reset_state():
 def reset_textbox():
     return gr.update(value="")
 
+
 def reset_default():
     global API_URL
     API_URL = "https://api.openai.com/v1/chat/completions"
@@ -650,6 +264,7 @@ def reset_default():
     os.environ.pop("https_proxy", None)
     return gr.update(value=API_URL), gr.update(value=""), "API URL 和代理已重置"
 
+
 def change_api_url(url):
     global API_URL
     API_URL = url
@@ -657,22 +272,41 @@ def change_api_url(url):
     logging.info(msg)
     return msg
 
+
 def change_proxy(proxy):
     os.environ["HTTPS_PROXY"] = proxy
     msg = f"代理更改为了{proxy}"
     logging.info(msg)
     return msg
 
+
 def hide_middle_chars(s):
     if len(s) <= 8:
         return s
     else:
         head = s[:4]
         tail = s[-4:]
-        hidden = '*' * (len(s) - 8)
+        hidden = "*" * (len(s) - 8)
         return head + hidden + tail
 
+
 def submit_key(key):
     msg = f"API密钥更改为了{hide_middle_chars(key)}"
     logging.info(msg)
     return key, msg
+
+
+def sha1sum(filename):
+    sha1 = hashlib.sha1()
+    with open(filename, "rb") as f:
+        while True:
+            data = f.read(65536)
+            if not data:
+                break
+            sha1.update(data)
+    return sha1.hexdigest()
+
+
+def replace_today(prompt):
+    today = datetime.datetime.today().strftime("%Y-%m-%d")
+    return prompt.replace("{current_date}", today)