Upload 19 files

app.py

@@ -25,9 +25,11 @@ else:
     dockerflag = False
 
 authflag = False
+auth_list = []
 
-if dockerflag:
+if not my_api_key:
     my_api_key = os.environ.get("my_api_key")
+if dockerflag:
     if my_api_key == "empty":
         logging.error("Please give a api key!")
         sys.exit(1)
@@ -35,6 +37,7 @@ if dockerflag:
     username = os.environ.get("USERNAME")
     password = os.environ.get("PASSWORD")
     if not (isinstance(username, type(None)) or isinstance(password, type(None))):
+        auth_list.append((os.environ.get("USERNAME"), os.environ.get("PASSWORD")))
         authflag = True
 else:
     if (
@@ -45,12 +48,15 @@ else:
         with open("api_key.txt", "r") as f:
             my_api_key = f.read().strip()
     if os.path.exists("auth.json"):
+        authflag = True
         with open("auth.json", "r", encoding='utf-8') as f:
             auth = json.load(f)
-            username = auth["username"]
-            password = auth["password"]
-            if username != "" and password != "":
-                authflag = True
+            for _ in auth:
+                if auth[_]["username"] and auth[_]["password"]:
+                    auth_list.append((auth[_]["username"], auth[_]["password"]))
+                else:
+                    logging.error("请检查auth.json文件中的用户名和密码！")
+                    sys.exit(1)
 
 gr.Chatbot.postprocess = postprocess
 PromptHelper.compact_text_chunks = compact_text_chunks
@@ -75,19 +81,19 @@ with gr.Blocks(css=customCSS, theme=small_and_beautiful_theme) as demo:
         with gr.Column(scale=4):
             status_display = gr.Markdown(get_geoip(), elem_id="status_display")
 
-    with gr.Row(scale=1).style(equal_height=True):
+    with gr.Row().style(equal_height=True):
         with gr.Column(scale=5):
-            with gr.Row(scale=1):
+            with gr.Row():
                 chatbot = gr.Chatbot(elem_id="chuanhu_chatbot").style(height="100%")
-            with gr.Row(scale=1):
+            with gr.Row():
                 with gr.Column(scale=12):
                     user_input = gr.Textbox(
-                        show_label=False, placeholder="在这里输入", interactive=True
+                        show_label=False, placeholder="在这里输入"
                     ).style(container=False)
                 with gr.Column(min_width=70, scale=1):
                     submitBtn = gr.Button("发送", variant="primary")
                     cancelBtn = gr.Button("取消", variant="secondary", visible=False)
-            with gr.Row(scale=1):
+            with gr.Row():
                 emptyBtn = gr.Button(
                     "🧹 新的对话",
                 )
@@ -107,7 +113,7 @@ with gr.Blocks(css=customCSS, theme=small_and_beautiful_theme) as demo:
                         visible=not HIDE_MY_KEY,
                         label="API-Key",
                     )
-                    usageTxt = gr.Markdown(get_usage(my_api_key), elem_id="usage_display")
+                    usageTxt = gr.Markdown("**发送消息** 或 **提交key** 以显示额度", elem_id="usage_display")
                     model_select_dropdown = gr.Dropdown(
                         label="选择模型", choices=MODELS, multiselect=False, value=MODELS[0]
                     )
@@ -207,7 +213,7 @@ with gr.Blocks(css=customCSS, theme=small_and_beautiful_theme) as demo:
                             label="Temperature",
                         )
 
-                    with gr.Accordion("网络设置", open=False):
+                    with gr.Accordion("网络设置", open=False, visible=False):
                         apiurlTxt = gr.Textbox(
                             show_label=True,
                             placeholder=f"在这里输入API地址...",
@@ -226,7 +232,7 @@ with gr.Blocks(css=customCSS, theme=small_and_beautiful_theme) as demo:
                         changeProxyBtn = gr.Button("🔄 设置代理地址")
 
     gr.Markdown(description)
-
+    gr.HTML(footer.format(versions=versions_html()), elem_id="footer")
     chatgpt_predict_args = dict(
         fn=predict,
         inputs=[
@@ -264,13 +270,14 @@ with gr.Blocks(css=customCSS, theme=small_and_beautiful_theme) as demo:
     )
 
     transfer_input_args = dict(
-        fn=transfer_input, inputs=[user_input], outputs=[user_question, user_input], show_progress=True
+        fn=transfer_input, inputs=[user_input], outputs=[user_question, user_input, submitBtn, cancelBtn], show_progress=True
     )
 
     get_usage_args = dict(
         fn=get_usage, inputs=[user_api_key], outputs=[usageTxt], show_progress=False
     )
 
+
     # Chatbot
     cancelBtn.click(cancel_outputing, [], [])
 
@@ -287,8 +294,7 @@ with gr.Blocks(css=customCSS, theme=small_and_beautiful_theme) as demo:
     )
     emptyBtn.click(**reset_textbox_args)
 
-    retryBtn.click(**reset_textbox_args)
-    retryBtn.click(
+    retryBtn.click(**start_outputing_args).then(
         retry,
         [
             user_api_key,
@@ -304,7 +310,7 @@ with gr.Blocks(css=customCSS, theme=small_and_beautiful_theme) as demo:
         ],
         [chatbot, history, status_display, token_count],
         show_progress=True,
-    )
+    ).then(**end_outputing_args)
     retryBtn.click(**get_usage_args)
 
     delFirstBtn.click(
@@ -330,7 +336,7 @@ with gr.Blocks(css=customCSS, theme=small_and_beautiful_theme) as demo:
             token_count,
             top_p,
             temperature,
-            gr.State(0),
+            gr.State(sum(token_count.value[-4:])),
             model_select_dropdown,
             language_select_dropdown,
         ],
@@ -341,6 +347,7 @@ with gr.Blocks(css=customCSS, theme=small_and_beautiful_theme) as demo:
 
     # ChatGPT
     keyTxt.change(submit_key, keyTxt, [user_api_key, status_display]).then(**get_usage_args)
+    keyTxt.submit(**get_usage_args)
 
     # Template
     templateRefreshBtn.click(get_template_names, None, [templateFileSelectDropdown])
@@ -417,7 +424,7 @@ if __name__ == "__main__":
             demo.queue(concurrency_count=CONCURRENT_COUNT).launch(
                 server_name="0.0.0.0",
                 server_port=7860,
-                auth=(username, password),
+                auth=auth_list,
                 favicon_path="./assets/favicon.ico",
             )
         else:
@@ -432,7 +439,7 @@ if __name__ == "__main__":
         if authflag:
             demo.queue(concurrency_count=CONCURRENT_COUNT).launch(
                 share=False,
-                auth=(username, password),
+                auth=auth_list,
                 favicon_path="./assets/favicon.ico",
                 inbrowser=True,
             )

assets/custom.css

@@ -3,6 +3,21 @@
     --chatbot-color-dark: #121111;
 }
 
+/* 覆盖gradio的页脚信息QAQ */
+footer {
+    display: none !important;
+}
+#footer{
+    text-align: center;
+}
+#footer div{
+    display: inline-block;
+}
+#footer .versions{
+    font-size: 85%;
+    opacity: 0.85;
+}
+
 /* status_display */
 #status_display {
     display: flex;
@@ -22,14 +37,45 @@
 
 /* usage_display */
 #usage_display {
-    height: 1em;
-}
-#usage_display p{
-    padding: 0 1em;
+    position: relative;
+    margin: 0;
+    box-shadow: var(--block-shadow);
+    border-width: var(--block-border-width);
+    border-color: var(--block-border-color);
+    border-radius: var(--block-radius);
+    background: var(--block-background-fill);
+    width: 100%;
+    line-height: var(--line-sm);
+    min-height: 2em;
+}
+#usage_display p, #usage_display span {
+    margin: 0;
+    padding: .5em 1em;
     font-size: .85em;
-    font-family: monospace;
     color: var(--body-text-color-subdued);
 }
+.progress-bar {
+    background-color: var(--input-background-fill);;
+    margin: 0 1em;
+    height: 20px;
+    border-radius: 10px;
+    overflow: hidden;
+}
+.progress {
+    background-color: var(--block-title-background-fill);;
+    height: 100%;
+    border-radius: 10px;
+    text-align: right;
+    transition: width 0.5s ease-in-out;
+}
+.progress-text {
+    /* color: white; */
+    color: var(--color-accent) !important;
+    font-size: 1em !important;
+    font-weight: bold;
+    padding-right: 10px;
+    line-height: 20px;
+}
 /* list */
 ol:not(.options), ul:not(.options) {
     padding-inline-start: 2em !important;
@@ -64,6 +110,7 @@ ol:not(.options), ul:not(.options) {
         background-color: var(--neutral-950) !important;
     }
 }
+
 /* 对话气泡 */
 [class *= "message"] {
     border-radius: var(--radius-xl) !important;

modules/chat_func.py

@@ -13,6 +13,9 @@ import colorama
 from duckduckgo_search import ddg
 import asyncio
 import aiohttp
+from llama_index.indices.query.vector_store import GPTVectorStoreIndexQuery
+from llama_index.indices.query.schema import QueryBundle
+from langchain.llms import OpenAIChat
 
 from modules.presets import *
 from modules.llama_func import *
@@ -58,39 +61,21 @@ def get_response(
     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"使用 HTTP 代理: {http_proxy}")
-        proxies["http"] = http_proxy
-    if https_proxy:
-        logging.info(f"使用 HTTPS 代理: {https_proxy}")
-        proxies["https"] = https_proxy
+    proxies = get_proxies()
 
     # 如果有自定义的api-url，使用自定义url发送请求，否则使用默认设置发送请求
     if shared.state.api_url != API_URL:
         logging.info(f"使用自定义API URL: {shared.state.api_url}")
-    if proxies:
-        response = requests.post(
-            shared.state.api_url,
-            headers=headers,
-            json=payload,
-            stream=True,
-            timeout=timeout,
-            proxies=proxies,
-        )
-    else:
-        response = requests.post(
-            shared.state.api_url,
-            headers=headers,
-            json=payload,
-            stream=True,
-            timeout=timeout,
-        )
+
+    response = requests.post(
+        shared.state.api_url,
+        headers=headers,
+        json=payload,
+        stream=True,
+        timeout=timeout,
+        proxies=proxies,
+    )
+
     return response
 
 
@@ -121,13 +106,17 @@ def stream_predict(
     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 = (
-            count_token(construct_user(inputs)) + system_prompt_token_count
+            input_token_count + system_prompt_token_count
         )
     else:
-        user_token_count = count_token(construct_user(inputs))
+        user_token_count = input_token_count
     all_token_counts.append(user_token_count)
     logging.info(f"输入token计数: {user_token_count}")
     yield get_return_value()
@@ -155,6 +144,8 @@ def stream_predict(
     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
@@ -219,7 +210,10 @@ def predict_all(
         chatbot.append((fake_input, ""))
     else:
         chatbot.append((inputs, ""))
-    all_token_counts.append(count_token(construct_user(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:
         response = get_response(
             openai_api_key,
@@ -242,13 +236,22 @@ def predict_all(
         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] = (chatbot[-1][0], content+display_append)
-    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
+    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(
@@ -268,40 +271,59 @@ def predict(
     should_check_token_count=True,
 ):  # repetition_penalty, top_k
     logging.info("输入为：" + colorama.Fore.BLUE + f"{inputs}" + colorama.Style.RESET_ALL)
-    yield chatbot+[(inputs, "")], history, "开始生成回答……", all_token_counts
+    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
-        history, chatbot, status_text = chat_ai(openai_api_key, index, inputs, history, chatbot, reply_language)
-        yield chatbot, history, status_text, all_token_counts
-        return
-
-    old_inputs = ""
-    link_references = []
-    if use_websearch:
+        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)
+        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
-        web_results = []
+        reference_results = []
         for idx, result in enumerate(search_results):
             logging.info(f"搜索结果{idx + 1}：{result}")
             domain_name = urllib3.util.parse_url(result["href"]).host
-            web_results.append(f'[{idx+1}]"{result["body"]}"\nURL: {result["href"]}')
-            link_references.append(f"{idx+1}. [{domain_name}]({result['href']})\n")
-        link_references = "\n\n" + "".join(link_references)
+            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(web_results))
+            .replace("{web_results}", "\n\n".join(reference_results))
             .replace("{reply_language}", reply_language )
         )
     else:
-        link_references = ""
+        display_reference = ""
 
     if len(openai_api_key) != 51:
         status_text = standard_error_msg + no_apikey_msg
@@ -334,7 +356,7 @@ def predict(
             temperature,
             selected_model,
             fake_input=old_inputs,
-            display_append=link_references
+            display_append=display_reference
         )
         for chatbot, history, status_text, all_token_counts in iter:
             if shared.state.interrupted:
@@ -354,7 +376,7 @@ def predict(
             temperature,
             selected_model,
             fake_input=old_inputs,
-            display_append=link_references
+            display_append=display_reference
         )
         yield chatbot, history, status_text, all_token_counts
 
@@ -367,10 +389,15 @@ def predict(
             + 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 = max_token_streaming
+        max_token = MODEL_SOFT_TOKEN_LIMIT[selected_model]["streaming"]
     else:
-        max_token = max_token_all
+        max_token = MODEL_SOFT_TOKEN_LIMIT[selected_model]["all"]
 
     if sum(all_token_counts) > max_token and should_check_token_count:
         status_text = f"精简token中{all_token_counts}/{max_token}"
@@ -460,6 +487,7 @@ def reduce_token_size(
     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

modules/llama_func.py

@@ -1,7 +1,7 @@
 import os
 import logging
 
-from llama_index import GPTSimpleVectorIndex
+from llama_index import GPTSimpleVectorIndex, ServiceContext
 from llama_index import download_loader
 from llama_index import (
     Document,
@@ -11,19 +11,32 @@ from llama_index import (
     RefinePrompt,
 )
 from langchain.llms import OpenAI
+from langchain.chat_models import ChatOpenAI
 import colorama
+import PyPDF2
+from tqdm import tqdm
 
 from modules.presets import *
 from modules.utils import *
 
 def get_index_name(file_src):
-    index_name = []
-    for file in file_src:
-        index_name.append(os.path.basename(file.name))
-    index_name = sorted(index_name)
-    index_name = "".join(index_name)
-    index_name = sha1sum(index_name)
-    return index_name
+    file_paths = [x.name for x in file_src]
+    file_paths.sort(key=lambda x: os.path.basename(x))
+
+    md5_hash = hashlib.md5()
+    for file_path in file_paths:
+        with open(file_path, "rb") as f:
+            while chunk := f.read(8192):
+                md5_hash.update(chunk)
+
+    return md5_hash.hexdigest()
+
+def block_split(text):
+    blocks = []
+    while len(text) > 0:
+        blocks.append(Document(text[:1000]))
+        text = text[1000:]
+    return blocks
 
 def get_documents(file_src):
     documents = []
@@ -33,9 +46,12 @@ def get_documents(file_src):
         logging.info(f"loading file: {file.name}")
         if os.path.splitext(file.name)[1] == ".pdf":
             logging.debug("Loading PDF...")
-            CJKPDFReader = download_loader("CJKPDFReader")
-            loader = CJKPDFReader()
-            text_raw = loader.load_data(file=file.name)[0].text
+            pdftext = ""
+            with open(file.name, 'rb') as pdfFileObj:
+                pdfReader = PyPDF2.PdfReader(pdfFileObj)
+                for page in tqdm(pdfReader.pages):
+                    pdftext += page.extract_text()
+            text_raw = pdftext
         elif os.path.splitext(file.name)[1] == ".docx":
             logging.debug("Loading DOCX...")
             DocxReader = download_loader("DocxReader")
@@ -51,7 +67,10 @@ def get_documents(file_src):
             with open(file.name, "r", encoding="utf-8") as f:
                 text_raw = f.read()
         text = add_space(text_raw)
+        # text = block_split(text)
+        # documents += text
         documents += [Document(text)]
+    logging.debug("Documents loaded.")
     return documents
 
 
@@ -59,13 +78,11 @@ def construct_index(
         api_key,
         file_src,
         max_input_size=4096,
-        num_outputs=1,
+        num_outputs=5,
         max_chunk_overlap=20,
         chunk_size_limit=600,
         embedding_limit=None,
-        separator=" ",
-        num_children=10,
-        max_keywords_per_chunk=10,
+        separator=" "
 ):
     os.environ["OPENAI_API_KEY"] = api_key
     chunk_size_limit = None if chunk_size_limit == 0 else chunk_size_limit
@@ -73,16 +90,9 @@ def construct_index(
     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,
+        llm=ChatOpenAI(model_name="gpt-3.5-turbo-0301", openai_api_key=api_key)
     )
+    prompt_helper = PromptHelper(max_input_size = max_input_size, num_output = num_outputs, max_chunk_overlap = max_chunk_overlap, embedding_limit=embedding_limit, chunk_size_limit=600, separator=separator)
     index_name = get_index_name(file_src)
     if os.path.exists(f"./index/{index_name}.json"):
         logging.info("找到了缓存的索引文件，加载中……")
@@ -90,14 +100,19 @@ def construct_index(
     else:
         try:
             documents = get_documents(file_src)
-            logging.debug("构建索引中……")
-            index = GPTSimpleVectorIndex(
-                documents, llm_predictor=llm_predictor, prompt_helper=prompt_helper
+            logging.info("构建索引中……")
+            service_context = ServiceContext.from_defaults(llm_predictor=llm_predictor, prompt_helper=prompt_helper, chunk_size_limit=chunk_size_limit)
+            index = GPTSimpleVectorIndex.from_documents(
+                documents,  service_context=service_context
             )
+            logging.debug("索引构建完成！")
             os.makedirs("./index", exist_ok=True)
             index.save_to_disk(f"./index/{index_name}.json")
+            logging.debug("索引已保存至本地!")
             return index
+
         except Exception as e:
+            logging.error("索引构建失败！", e)
             print(e)
             return None
 
@@ -144,7 +159,7 @@ def ask_ai(
         question,
         prompt_tmpl,
         refine_tmpl,
-        sim_k=1,
+        sim_k=5,
         temprature=0,
         prefix_messages=[],
         reply_language="中文",
@@ -154,7 +169,7 @@ def ask_ai(
     logging.debug("Index file found")
     logging.debug("Querying index...")
     llm_predictor = LLMPredictor(
-        llm=OpenAI(
+        llm=ChatOpenAI(
             temperature=temprature,
             model_name="gpt-3.5-turbo-0301",
             prefix_messages=prefix_messages,
@@ -166,7 +181,6 @@ def ask_ai(
     rf_prompt = RefinePrompt(refine_tmpl.replace("{reply_language}", reply_language))
     response = index.query(
         question,
-        llm_predictor=llm_predictor,
         similarity_top_k=sim_k,
         text_qa_template=qa_prompt,
         refine_template=rf_prompt,

modules/openai_func.py

@@ -1,70 +1,82 @@
 import requests
 import logging
-from modules.presets import timeout_all, BALANCE_API_URL,standard_error_msg,connection_timeout_prompt,error_retrieve_prompt,read_timeout_prompt
-from modules import shared
-import os
+from modules.presets import (
+    timeout_all,
+    USAGE_API_URL,
+    BALANCE_API_URL,
+    standard_error_msg,
+    connection_timeout_prompt,
+    error_retrieve_prompt,
+    read_timeout_prompt
+)
 
+from modules import shared
+from modules.utils import get_proxies
+import os, datetime
 
-def get_usage_response(openai_api_key):
+def get_billing_data(openai_api_key, billing_url):
     headers = {
         "Content-Type": "application/json",
-        "Authorization": f"Bearer {openai_api_key}",
+        "Authorization": f"Bearer {openai_api_key}"
     }
-
+    
     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"使用 HTTP 代理: {http_proxy}")
-        proxies["http"] = http_proxy
-    if https_proxy:
-        logging.info(f"使用 HTTPS 代理: {https_proxy}")
-        proxies["https"] = https_proxy
-
-    # 如果有代理，使用代理发送请求，否则使用默认设置发送请求
-    """
-    暂不支持修改
-    if shared.state.balance_api_url != BALANCE_API_URL:
-        logging.info(f"使用自定义BALANCE API URL: {shared.state.balance_api_url}")
-    """
-    if proxies:
-        response = requests.get(
-            BALANCE_API_URL,
-            headers=headers,
-            timeout=timeout,
-            proxies=proxies,
-        )
+    proxies = get_proxies()
+    response = requests.get(
+        billing_url,
+        headers=headers,
+        timeout=timeout,
+        proxies=proxies,
+    )
+    
+    if response.status_code == 200:
+        data = response.json()
+        return data
     else:
-        response = requests.get(
-            BALANCE_API_URL,
-            headers=headers,
-            timeout=timeout,
-        )
-    return response
+        raise Exception(f"API request failed with status code {response.status_code}: {response.text}")
+    
 
 def get_usage(openai_api_key):
     try:
-        response=get_usage_response(openai_api_key=openai_api_key)
-        logging.debug(response.json())
+        balance_data=get_billing_data(openai_api_key, BALANCE_API_URL)
+        logging.debug(balance_data)
         try:
-            balance = response.json().get("total_available") if response.json().get(
-                "total_available") else 0
-            total_used = response.json().get("total_used") if response.json().get(
-                "total_used") else 0
+            balance = balance_data["total_available"] if balance_data["total_available"] else 0
+            total_used = balance_data["total_used"] if balance_data["total_used"] else 0
+            usage_percent = round(total_used / (total_used+balance) * 100, 2)
         except Exception as e:
             logging.error(f"API使用情况解析失败:"+str(e))
             balance = 0
             total_used=0
-        return f"**API使用情况**（已用/余额）\u3000{total_used}$ / {balance}$"
+            return f"**API使用情况解析失败**"
+        if balance == 0:
+            last_day_of_month = datetime.datetime.now().strftime("%Y-%m-%d")
+            first_day_of_month = datetime.datetime.now().replace(day=1).strftime("%Y-%m-%d")
+            usage_url = f"{USAGE_API_URL}?start_date={first_day_of_month}&end_date={last_day_of_month}"
+            try:
+                usage_data = get_billing_data(openai_api_key, usage_url)
+            except Exception as e:
+                logging.error(f"获取API使用情况失败:"+str(e))
+                return f"**获取API使用情况失败**"
+            return f"**本月使用金额** \u3000 ${usage_data['total_usage'] / 100}"
+        
+        # return f"**免费额度**（已用/余额）\u3000${total_used} / ${balance}"
+        return f"""\
+        <b>免费额度使用情况</b>
+        <div class="progress-bar">
+            <div class="progress" style="width: {usage_percent}%;">
+                <span class="progress-text">{usage_percent}%</span>
+            </div>
+        </div>
+        <div style="display: flex; justify-content: space-between;"><span>已用 ${total_used}</span><span>可用 ${balance}</span></div>
+        """
+    
     except requests.exceptions.ConnectTimeout:
         status_text = standard_error_msg + connection_timeout_prompt + error_retrieve_prompt
         return status_text
     except requests.exceptions.ReadTimeout:
         status_text = standard_error_msg + read_timeout_prompt + error_retrieve_prompt
         return status_text
+    except Exception as e:
+        logging.error(f"获取API使用情况失败:"+str(e))
+        return standard_error_msg + error_retrieve_prompt

modules/presets.py

@@ -5,6 +5,7 @@ import gradio as gr
 initial_prompt = "You are a helpful assistant."
 API_URL = "https://api.openai.com/v1/chat/completions"
 BALANCE_API_URL="https://api.openai.com/dashboard/billing/credit_grants"
+USAGE_API_URL="https://api.openai.com/dashboard/billing/usage"
 HISTORY_DIR = "history"
 TEMPLATES_DIR = "templates"
 
@@ -18,9 +19,7 @@ ssl_error_prompt = "SSL错误，无法获取对话。"  # SSL 错误
 no_apikey_msg = "API key长度不是51位，请检查是否输入正确。"  # API key 长度不足 51 位
 no_input_msg = "请输入对话内容。"  # 未输入对话内容
 
-max_token_streaming = 3500  # 流式对话时的最大 token 数
 timeout_streaming = 10  # 流式对话时的超时时间
-max_token_all = 3500  # 非流式对话时的最大 token 数
 timeout_all = 200  # 非流式对话时的超时时间
 enable_streaming_option = True  # 是否启用选择选择是否实时显示回答的勾选框
 HIDE_MY_KEY = False  # 如果你想在UI中隐藏你的 API 密钥，将此值设置为 True
@@ -41,6 +40,10 @@ description = """\
 </div>
 """
 
+footer = """\
+<div class="versions">{versions}</div>
+"""
+
 summarize_prompt = "你是谁？我们刚才聊了什么？"  # 总结对话时的 prompt
 
 MODELS = [
@@ -52,8 +55,36 @@ MODELS = [
     "gpt-4-32k-0314",
 ]  # 可选的模型
 
+MODEL_SOFT_TOKEN_LIMIT = {
+    "gpt-3.5-turbo": {
+        "streaming": 3500,
+        "all": 3500
+    },
+    "gpt-3.5-turbo-0301": {
+        "streaming": 3500,
+        "all": 3500
+    },
+    "gpt-4": {
+        "streaming": 7500,
+        "all": 7500
+    },
+    "gpt-4-0314": {
+        "streaming": 7500,
+        "all": 7500
+    },
+    "gpt-4-32k": {
+        "streaming": 31000,
+        "all": 31000
+    },
+    "gpt-4-32k-0314": {
+        "streaming": 31000,
+        "all": 31000
+    }
+}
+
 REPLY_LANGUAGES = [
-    "中文",
+    "简体中文",
+    "繁體中文",
     "English",
     "日本語",
     "Español",

modules/utils.py

@@ -10,6 +10,8 @@ import csv
 import requests
 import re
 import html
+import sys
+import subprocess
 
 import gradio as gr
 from pypinyin import lazy_pinyin
@@ -115,7 +117,11 @@ def convert_mdtext(md_text):
 
 
 def convert_asis(userinput):
-    return f"<p style=\"white-space:pre-wrap;\">{html.escape(userinput)}</p>"+ALREADY_CONVERTED_MARK
+    return (
+        f'<p style="white-space:pre-wrap;">{html.escape(userinput)}</p>'
+        + ALREADY_CONVERTED_MARK
+    )
+
 
 def detect_converted_mark(userinput):
     if userinput.endswith(ALREADY_CONVERTED_MARK):
@@ -153,6 +159,7 @@ def construct_assistant(text):
 def construct_token_message(token, stream=False):
     return f"Token 计数: {token}"
 
+
 def delete_first_conversation(history, previous_token_count):
     if history:
         del history[:2]
@@ -346,6 +353,8 @@ def change_proxy(proxy):
 
 
 def hide_middle_chars(s):
+    if s is None:
+        return ""
     if len(s) <= 8:
         return s
     else:
@@ -362,20 +371,14 @@ def submit_key(key):
     return key, msg
 
 
-def sha1sum(filename):
-    sha1 = hashlib.sha1()
-    sha1.update(filename.encode("utf-8"))
-    return sha1.hexdigest()
-
-
 def replace_today(prompt):
     today = datetime.datetime.today().strftime("%Y-%m-%d")
     return prompt.replace("{current_date}", today)
 
 
 def get_geoip():
-    response = requests.get("https://ipapi.co/json/", timeout=5)
     try:
+        response = requests.get("https://ipapi.co/json/", timeout=5)
         data = response.json()
     except:
         data = {"error": True, "reason": "连接ipapi失败"}
@@ -383,7 +386,7 @@ def get_geoip():
         logging.warning(f"无法获取IP地址信息。\n{data}")
         if data["reason"] == "RateLimited":
             return (
-                f"获取IP地理位置失败，因为达到了检测IP的速率限制。聊天功能可能仍然可用，但请注意，如果您的IP地址在不受支持的地区，您可能会遇到问题。"
+                f"获取IP地理位置失败，因为达到了检测IP的速率限制。聊天功能可能仍然可用。"
             )
         else:
             return f"获取IP地理位置失败。原因：{data['reason']}。你仍然可以使用聊天功能。"
@@ -427,8 +430,91 @@ def cancel_outputing():
     logging.info("中止输出……")
     shared.state.interrupt()
 
+
 def transfer_input(inputs):
     # 一次性返回，降低延迟
     textbox = reset_textbox()
     outputing = start_outputing()
-    return inputs, gr.update(value="")
+    return (
+        inputs,
+        gr.update(value=""),
+        gr.Button.update(visible=True),
+        gr.Button.update(visible=False),
+    )
+
+
+def get_proxies():
+    # 获取环境变量中的代理设置
+    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"使用 HTTP 代理: {http_proxy}")
+        proxies["http"] = http_proxy
+    if https_proxy:
+        logging.info(f"使用 HTTPS 代理: {https_proxy}")
+        proxies["https"] = https_proxy
+
+    if proxies == {}:
+        proxies = None
+
+    return proxies
+
+def run(command, desc=None, errdesc=None, custom_env=None, live=False):
+    if desc is not None:
+        print(desc)
+    if live:
+        result = subprocess.run(command, shell=True, env=os.environ if custom_env is None else custom_env)
+        if result.returncode != 0:
+            raise RuntimeError(f"""{errdesc or 'Error running command'}.
+Command: {command}
+Error code: {result.returncode}""")
+
+        return ""
+    result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True, env=os.environ if custom_env is None else custom_env)
+    if result.returncode != 0:
+        message = f"""{errdesc or 'Error running command'}.
+Command: {command}
+Error code: {result.returncode}
+stdout: {result.stdout.decode(encoding="utf8", errors="ignore") if len(result.stdout)>0 else '<empty>'}
+stderr: {result.stderr.decode(encoding="utf8", errors="ignore") if len(result.stderr)>0 else '<empty>'}
+"""
+        raise RuntimeError(message)
+    return result.stdout.decode(encoding="utf8", errors="ignore")
+
+def versions_html():
+    git = os.environ.get('GIT', "git")
+    python_version = ".".join([str(x) for x in sys.version_info[0:3]])
+    try:
+        commit_hash = run(f"{git} rev-parse HEAD").strip()
+    except Exception:
+        commit_hash = "<none>"
+    if commit_hash != "<none>":
+        short_commit = commit_hash[0:7]
+        commit_info = f"<a style=\"text-decoration:none\" href=\"https://github.com/GaiZhenbiao/ChuanhuChatGPT/commit/{short_commit}\">{short_commit}</a>"
+    else:
+        commit_info = "unknown \U0001F615"
+    return f"""
+Python: <span title="{sys.version}">{python_version}</span>
+ • 
+Gradio: {gr.__version__}
+ • 
+Commit: {commit_info}
+"""
+
+def add_source_numbers(lst, source_name = "Source", use_source = True):
+    if use_source:
+        return [f'[{idx+1}]\t "{item[0]}"\n{source_name}: {item[1]}' for idx, item in enumerate(lst)]
+    else:
+        return [f'[{idx+1}]\t "{item}"' for idx, item in enumerate(lst)]
+
+def add_details(lst):
+    nodes = []
+    for index, txt in enumerate(lst):
+        brief = txt[:25].replace("\n", "")
+        nodes.append(
+            f"<details><summary>{brief}...</summary><p>{txt}</p></details>"
+        )
+    return nodes