cathub/LocalChatGPT

1# -*- coding:utf-8 -*-2from __future__ import annotations3from typing import TYPE_CHECKING, Any, Callable, Dict, List, Tuple, Type4import logging5import json6import os7import datetime8import hashlib9import csv10import requests11import re12import html13import sys14import subprocess15 16import gradio as gr17from pypinyin import lazy_pinyin18import tiktoken19import mdtex2html20from markdown import markdown21from pygments import highlight22from pygments.lexers import get_lexer_by_name23from pygments.formatters import HtmlFormatter24 25from modules.presets import *26import modules.shared as shared27 28logging.basicConfig(29    level=logging.INFO,30    format="%(asctime)s [%(levelname)s] [%(filename)s:%(lineno)d] %(message)s",31)32 33if TYPE_CHECKING:34    from typing import TypedDict35 36    class DataframeData(TypedDict):37        headers: List[str]38        data: List[List[str | int | bool]]39 40 41def count_token(message):42    encoding = tiktoken.get_encoding("cl100k_base")43    input_str = f"role: {message['role']}, content: {message['content']}"44    length = len(encoding.encode(input_str))45    return length46 47 48def markdown_to_html_with_syntax_highlight(md_str):49    def replacer(match):50        lang = match.group(1) or "text"51        code = match.group(2)52 53        try:54            lexer = get_lexer_by_name(lang, stripall=True)55        except ValueError:56            lexer = get_lexer_by_name("text", stripall=True)57 58        formatter = HtmlFormatter()59        highlighted_code = highlight(code, lexer, formatter)60 61        return f'<pre><code class="{lang}">{highlighted_code}</code></pre>'62 63    code_block_pattern = r"```(\w+)?\n([\s\S]+?)\n```"64    md_str = re.sub(code_block_pattern, replacer, md_str, flags=re.MULTILINE)65 66    html_str = markdown(md_str)67    return html_str68 69 70def normalize_markdown(md_text: str) -> str:71    lines = md_text.split("\n")72    normalized_lines = []73    inside_list = False74 75    for i, line in enumerate(lines):76        if re.match(r"^(\d+\.|-|\*|\+)\s", line.strip()):77            if not inside_list and i > 0 and lines[i - 1].strip() != "":78                normalized_lines.append("")79            inside_list = True80            normalized_lines.append(line)81        elif inside_list and line.strip() == "":82            if i < len(lines) - 1 and not re.match(83                r"^(\d+\.|-|\*|\+)\s", lines[i + 1].strip()84            ):85                normalized_lines.append(line)86            continue87        else:88            inside_list = False89            normalized_lines.append(line)90 91    return "\n".join(normalized_lines)92 93 94def convert_mdtext(md_text):95    code_block_pattern = re.compile(r"```(.*?)(?:```|$)", re.DOTALL)96    inline_code_pattern = re.compile(r"`(.*?)`", re.DOTALL)97    code_blocks = code_block_pattern.findall(md_text)98    non_code_parts = code_block_pattern.split(md_text)[::2]99 100    result = []101    for non_code, code in zip(non_code_parts, code_blocks + [""]):102        if non_code.strip():103            non_code = normalize_markdown(non_code)104            if inline_code_pattern.search(non_code):105                result.append(markdown(non_code, extensions=["tables"]))106            else:107                result.append(mdtex2html.convert(non_code, extensions=["tables"]))108        if code.strip():109            # _, code = detect_language(code)  # 暂时去除代码高亮功能，因为在大段代码的情况下会出现问题110            # code = code.replace("\n\n", "\n") # 暂时去除代码中的空行，因为在大段代码的情况下会出现问题111            code = f"\n```{code}\n\n```"112            code = markdown_to_html_with_syntax_highlight(code)113            result.append(code)114    result = "".join(result)115    result += ALREADY_CONVERTED_MARK116    return result117 118 119def convert_asis(userinput):120    return (121        f'<p style="white-space:pre-wrap;">{html.escape(userinput)}</p>'122        + ALREADY_CONVERTED_MARK123    )124 125 126def detect_converted_mark(userinput):127    if userinput.endswith(ALREADY_CONVERTED_MARK):128        return True129    else:130        return False131 132 133def detect_language(code):134    if code.startswith("\n"):135        first_line = ""136    else:137        first_line = code.strip().split("\n", 1)[0]138    language = first_line.lower() if first_line else ""139    code_without_language = code[len(first_line) :].lstrip() if first_line else code140    return language, code_without_language141 142 143def construct_text(role, text):144    return {"role": role, "content": text}145 146 147def construct_user(text):148    return construct_text("user", text)149 150 151def construct_system(text):152    return construct_text("system", text)153 154 155def construct_assistant(text):156    return construct_text("assistant", text)157 158 159def construct_token_message(token, stream=False):160    return f"Token 计数: {token}"161 162 163def delete_first_conversation(history, previous_token_count):164    if history:165        del history[:2]166        del previous_token_count[0]167    return (168        history,169        previous_token_count,170        construct_token_message(sum(previous_token_count)),171    )172 173 174def delete_last_conversation(chatbot, history, previous_token_count):175    if len(chatbot) > 0 and standard_error_msg in chatbot[-1][1]:176        logging.info("由于包含报错信息，只删除chatbot记录")177        chatbot.pop()178        return chatbot, history179    if len(history) > 0:180        logging.info("删除了一组对话历史")181        history.pop()182        history.pop()183    if len(chatbot) > 0:184        logging.info("删除了一组chatbot对话")185        chatbot.pop()186    if len(previous_token_count) > 0:187        logging.info("删除了一组对话的token计数记录")188        previous_token_count.pop()189    return (190        chatbot,191        history,192        previous_token_count,193        construct_token_message(sum(previous_token_count)),194    )195 196 197def save_file(filename, system, history, chatbot):198    logging.info("保存对话历史中……")199    os.makedirs(HISTORY_DIR, exist_ok=True)200    if filename.endswith(".json"):201        json_s = {"system": system, "history": history, "chatbot": chatbot}202        print(json_s)203        with open(os.path.join(HISTORY_DIR, filename), "w") as f:204            json.dump(json_s, f)205    elif filename.endswith(".md"):206        md_s = f"system: \n- {system} \n"207        for data in history:208            md_s += f"\n{data['role']}: \n- {data['content']} \n"209        with open(os.path.join(HISTORY_DIR, filename), "w", encoding="utf8") as f:210            f.write(md_s)211    logging.info("保存对话历史完毕")212    return os.path.join(HISTORY_DIR, filename)213 214 215def save_chat_history(filename, system, history, chatbot):216    if filename == "":217        return218    if not filename.endswith(".json"):219        filename += ".json"220    return save_file(filename, system, history, chatbot)221 222 223def export_markdown(filename, system, history, chatbot):224    if filename == "":225        return226    if not filename.endswith(".md"):227        filename += ".md"228    return save_file(filename, system, history, chatbot)229 230 231def load_chat_history(filename, system, history, chatbot):232    logging.info("加载对话历史中……")233    if type(filename) != str:234        filename = filename.name235    try:236        with open(os.path.join(HISTORY_DIR, filename), "r") as f:237            json_s = json.load(f)238        try:239            if type(json_s["history"][0]) == str:240                logging.info("历史记录格式为旧版，正在转换……")241                new_history = []242                for index, item in enumerate(json_s["history"]):243                    if index % 2 == 0:244                        new_history.append(construct_user(item))245                    else:246                        new_history.append(construct_assistant(item))247                json_s["history"] = new_history248                logging.info(new_history)249        except:250            # 没有对话历史251            pass252        logging.info("加载对话历史完毕")253        return filename, json_s["system"], json_s["history"], json_s["chatbot"]254    except FileNotFoundError:255        logging.info("没有找到对话历史文件，不执行任何操作")256        return filename, system, history, chatbot257 258 259def sorted_by_pinyin(list):260    return sorted(list, key=lambda char: lazy_pinyin(char)[0][0])261 262 263def get_file_names(dir, plain=False, filetypes=[".json"]):264    logging.info(f"获取文件名列表，目录为{dir}，文件类型为{filetypes}，是否为纯文本列表{plain}")265    files = []266    try:267        for type in filetypes:268            files += [f for f in os.listdir(dir) if f.endswith(type)]269    except FileNotFoundError:270        files = []271    files = sorted_by_pinyin(files)272    if files == []:273        files = [""]274    if plain:275        return files276    else:277        return gr.Dropdown.update(choices=files)278 279 280def get_history_names(plain=False):281    logging.info("获取历史记录文件名列表")282    return get_file_names(HISTORY_DIR, plain)283 284 285def load_template(filename, mode=0):286    logging.info(f"加载模板文件{filename}，模式为{mode}（0为返回字典和下拉菜单，1为返回下拉菜单，2为返回字典）")287    lines = []288    logging.info("Loading template...")289    if filename.endswith(".json"):290        with open(os.path.join(TEMPLATES_DIR, filename), "r", encoding="utf8") as f:291            lines = json.load(f)292        lines = [[i["act"], i["prompt"]] for i in lines]293    else:294        with open(295            os.path.join(TEMPLATES_DIR, filename), "r", encoding="utf8"296        ) as csvfile:297            reader = csv.reader(csvfile)298            lines = list(reader)299        lines = lines[1:]300    if mode == 1:301        return sorted_by_pinyin([row[0] for row in lines])302    elif mode == 2:303        return {row[0]: row[1] for row in lines}304    else:305        choices = sorted_by_pinyin([row[0] for row in lines])306        return {row[0]: row[1] for row in lines}, gr.Dropdown.update(307            choices=choices, value=choices[0]308        )309 310 311def get_template_names(plain=False):312    logging.info("获取模板文件名列表")313    return get_file_names(TEMPLATES_DIR, plain, filetypes=[".csv", "json"])314 315 316def get_template_content(templates, selection, original_system_prompt):317    logging.info(f"应用模板中，选择为{selection}，原始系统提示为{original_system_prompt}")318    try:319        return templates[selection]320    except:321        return original_system_prompt322 323 324def reset_state():325    logging.info("重置状态")326    return [], [], [], construct_token_message(0)327 328 329def reset_textbox():330    logging.debug("重置文本框")331    return gr.update(value="")332 333 334def reset_default():335    newurl = shared.state.reset_api_url()336    os.environ.pop("HTTPS_PROXY", None)337    os.environ.pop("https_proxy", None)338    return gr.update(value=newurl), gr.update(value=""), "API URL 和代理已重置"339 340 341def change_api_url(url):342    shared.state.set_api_url(url)343    msg = f"API地址更改为了{url}"344    logging.info(msg)345    return msg346 347 348def change_proxy(proxy):349    os.environ["HTTPS_PROXY"] = proxy350    msg = f"代理更改为了{proxy}"351    logging.info(msg)352    return msg353 354 355def hide_middle_chars(s):356    if s is None:357        return ""358    if len(s) <= 8:359        return s360    else:361        head = s[:4]362        tail = s[-4:]363        hidden = "*" * (len(s) - 8)364        return head + hidden + tail365 366 367def submit_key(key):368    key = key.strip()369    msg = f"API密钥更改为了{hide_middle_chars(key)}"370    logging.info(msg)371    return key, msg372 373 374def replace_today(prompt):375    today = datetime.datetime.today().strftime("%Y-%m-%d")376    return prompt.replace("{current_date}", today)377 378 379def get_geoip():380    try:381        response = requests.get("https://ipapi.co/json/", timeout=5)382        data = response.json()383    except:384        data = {"error": True, "reason": "连接ipapi失败"}385    if "error" in data.keys():386        logging.warning(f"无法获取IP地址信息。\n{data}")387        if data["reason"] == "RateLimited":388            return (389                f"获取IP地理位置失败，因为达到了检测IP的速率限制。聊天功能可能仍然可用。"390            )391        else:392            return f"获取IP地理位置失败。原因：{data['reason']}。你仍然可以使用聊天功能。"393    else:394        country = data["country_name"]395        if country == "China":396            text = "**您的IP区域：中国。请立即检查代理设置，在不受支持的地区使用API可能导致账号被封禁。**"397        else:398            text = f"您的IP区域：{country}。"399        logging.info(text)400        return text401 402 403def find_n(lst, max_num):404    n = len(lst)405    total = sum(lst)406 407    if total < max_num:408        return n409 410    for i in range(len(lst)):411        if total - lst[i] < max_num:412            return n - i - 1413        total = total - lst[i]414    return 1415 416 417def start_outputing():418    logging.debug("显示取消按钮，隐藏发送按钮")419    return gr.Button.update(visible=False), gr.Button.update(visible=True)420 421 422def end_outputing():423    return (424        gr.Button.update(visible=True),425        gr.Button.update(visible=False),426    )427 428 429def cancel_outputing():430    logging.info("中止输出……")431    shared.state.interrupt()432 433 434def transfer_input(inputs):435    # 一次性返回，降低延迟436    textbox = reset_textbox()437    outputing = start_outputing()438    return (439        inputs,440        gr.update(value=""),441        gr.Button.update(visible=True),442        gr.Button.update(visible=False),443    )444 445 446def get_proxies():447    # 获取环境变量中的代理设置448    http_proxy = os.environ.get("HTTP_PROXY") or os.environ.get("http_proxy")449    https_proxy = os.environ.get("HTTPS_PROXY") or os.environ.get("https_proxy")450 451    # 如果存在代理设置，使用它们452    proxies = {}453    if http_proxy:454        logging.info(f"使用 HTTP 代理: {http_proxy}")455        proxies["http"] = http_proxy456    if https_proxy:457        logging.info(f"使用 HTTPS 代理: {https_proxy}")458        proxies["https"] = https_proxy459 460    if proxies == {}:461        proxies = None462 463    return proxies464 465def run(command, desc=None, errdesc=None, custom_env=None, live=False):466    if desc is not None:467        print(desc)468    if live:469        result = subprocess.run(command, shell=True, env=os.environ if custom_env is None else custom_env)470        if result.returncode != 0:471            raise RuntimeError(f"""{errdesc or 'Error running command'}.472Command: {command}473Error code: {result.returncode}""")474 475        return ""476    result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True, env=os.environ if custom_env is None else custom_env)477    if result.returncode != 0:478        message = f"""{errdesc or 'Error running command'}.479Command: {command}480Error code: {result.returncode}481stdout: {result.stdout.decode(encoding="utf8", errors="ignore") if len(result.stdout)>0 else '<empty>'}482stderr: {result.stderr.decode(encoding="utf8", errors="ignore") if len(result.stderr)>0 else '<empty>'}483"""484        raise RuntimeError(message)485    return result.stdout.decode(encoding="utf8", errors="ignore")486 487def versions_html():488    git = os.environ.get('GIT', "git")489    python_version = ".".join([str(x) for x in sys.version_info[0:3]])490    try:491        commit_hash = run(f"{git} rev-parse HEAD").strip()492    except Exception:493        commit_hash = "<none>"494    if commit_hash != "<none>":495        short_commit = commit_hash[0:7]496        commit_info = f"<a style=\"text-decoration:none\" href=\"https://github.com/GaiZhenbiao/ChuanhuChatGPT/commit/{short_commit}\">{short_commit}</a>"497    else:498        commit_info = "unknown \U0001F615"499    return f"""500Python: <span title="{sys.version}">{python_version}</span>501 • 502Gradio: {gr.__version__}503 • 504Commit: {commit_info}505"""506 507def add_source_numbers(lst, source_name = "Source", use_source = True):508    if use_source:509        return [f'[{idx+1}]\t "{item[0]}"\n{source_name}: {item[1]}' for idx, item in enumerate(lst)]510    else:511        return [f'[{idx+1}]\t "{item}"' for idx, item in enumerate(lst)]512 513def add_details(lst):514    nodes = []515    for index, txt in enumerate(lst):516        brief = txt[:25].replace("\n", "")517        nodes.append(518            f"<details><summary>{brief}...</summary><p>{txt}</p></details>"519        )520    return nodes521