更优雅的更换自定义api方式

ChuanhuChatbot.py

@@ -164,12 +164,12 @@ with gr.Blocks(css=customCSS, theme=small_and_beautiful_theme) as demo:
                         )
 
                     with gr.Accordion("网络设置", open=False):
-                        apiurlTxt = gr.Textbox(
+                        apihostTxt = gr.Textbox(
                             show_label=True,
-                            placeholder=f"在这里输入API地址...",
-                            label="API地址",
-                            value="https://api.openai.com/v1/chat/completions",
-                            lines=2,
+                            placeholder=f"在这里输入API-Host...",
+                            label="API-Host",
+                            value="api.openai.com",
+                            lines=1,
                         )
                         changeAPIURLBtn = gr.Button("🔄 切换API地址")
                         proxyTxt = gr.Textbox(
@@ -343,11 +343,11 @@ with gr.Blocks(css=customCSS, theme=small_and_beautiful_theme) as demo:
 
     # Advanced
     default_btn.click(
-        reset_default, [], [apiurlTxt, proxyTxt, status_display], show_progress=True
+        reset_default, [], [apihostTxt, proxyTxt, status_display], show_progress=True
     )
     changeAPIURLBtn.click(
-        change_api_url,
-        [apiurlTxt],
+        change_api_host,
+        [apihostTxt],
         [status_display],
         show_progress=True,
     )

modules/chat_func.py

@@ -13,9 +13,7 @@ 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 *
@@ -63,13 +61,13 @@ def get_response(
         timeout = timeout_all
 
 
-    # 如果有自定义的api-url，使用自定义url发送请求，否则使用默认设置发送请求
-    if shared.state.api_url != API_URL:
-        logging.info(f"使用自定义API URL: {shared.state.api_url}")
+    # 如果有自定义的api-host，使用自定义host发送请求，否则使用默认设置发送请求
+    if shared.state.completion_url != COMPLETION_URL:
+        logging.info(f"使用自定义API URL: {shared.state.completion_url}")
 
     with retrieve_proxy():
         response = requests.post(
-            shared.state.api_url,
+            shared.state.completion_url,
             headers=headers,
             json=payload,
             stream=True,
@@ -270,6 +268,11 @@ def predict(
     reply_language="中文",
     should_check_token_count=True,
 ):  # repetition_penalty, top_k
+    from llama_index.indices.query.vector_store import GPTVectorStoreIndexQuery
+    from llama_index.indices.query.schema import QueryBundle
+    from langchain.llms import OpenAIChat
+
+    
     logging.info("输入为：" + colorama.Fore.BLUE + f"{inputs}" + colorama.Style.RESET_ALL)
     if should_check_token_count:
         yield chatbot+[(inputs, "")], history, "开始生成回答……", all_token_counts

modules/config.py

@@ -110,4 +110,4 @@ def retrieve_proxy(proxy=None):
 
 
 ## 处理advance pdf
-advance_pdf = config.get("advance_pdf", {})
+advance_pdf = config.get("advance_pdf", {})

modules/llama_func.py

@@ -1,7 +1,6 @@
 import os
 import logging
 
-from llama_index import GPTSimpleVectorIndex, ServiceContext
 from llama_index import download_loader
 from llama_index import (
     Document,
@@ -10,8 +9,6 @@ from llama_index import (
     QuestionAnswerPrompt,
     RefinePrompt,
 )
-from langchain.llms import OpenAI
-from langchain.chat_models import ChatOpenAI
 import colorama
 import PyPDF2
 from tqdm import tqdm
@@ -89,6 +86,9 @@ def construct_index(
         embedding_limit=None,
         separator=" "
 ):
+    from langchain.chat_models import ChatOpenAI
+    from llama_index import GPTSimpleVectorIndex, ServiceContext
+    
     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
@@ -120,98 +120,7 @@ def construct_index(
             logging.error("索引构建失败！", e)
             print(e)
             return None
-
-
-def chat_ai(
-        api_key,
-        index,
-        question,
-        context,
-        chatbot,
-        reply_language,
-):
-    os.environ["OPENAI_API_KEY"] = api_key
-
-    logging.info(f"Question: {question}")
-
-    response, chatbot_display, status_text = ask_ai(
-        api_key,
-        index,
-        question,
-        replace_today(PROMPT_TEMPLATE),
-        REFINE_TEMPLATE,
-        SIM_K,
-        INDEX_QUERY_TEMPRATURE,
-        context,
-        reply_language,
-    )
-    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, chatbot_display))
-
-    os.environ["OPENAI_API_KEY"] = ""
-    return context, chatbot, status_text
-
-
-def ask_ai(
-        api_key,
-        index,
-        question,
-        prompt_tmpl,
-        refine_tmpl,
-        sim_k=5,
-        temprature=0,
-        prefix_messages=[],
-        reply_language="中文",
-):
-    os.environ["OPENAI_API_KEY"] = api_key
-
-    logging.debug("Index file found")
-    logging.debug("Querying index...")
-    llm_predictor = LLMPredictor(
-        llm=ChatOpenAI(
-            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.replace("{reply_language}", reply_language))
-    rf_prompt = RefinePrompt(refine_tmpl.replace("{reply_language}", reply_language))
-    response = index.query(
-        question,
-        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
-        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>"
-            )
-        new_response = ret_text + "\n----------\n" + "\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, new_response, 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 = {"，": "， ", "。": "。 ", "？": "？ ", "！": "！ ", "：": "： ", "；": "； "}

modules/openai_func.py

@@ -37,7 +37,7 @@ def get_billing_data(openai_api_key, billing_url):
 
 def get_usage(openai_api_key):
     try:
-        balance_data=get_billing_data(openai_api_key, BALANCE_API_URL)
+        balance_data=get_billing_data(openai_api_key, shared.state.balance_api_url)
         logging.debug(balance_data)
         try:
             balance = balance_data["total_available"] if balance_data["total_available"] else 0
@@ -51,7 +51,7 @@ def get_usage(openai_api_key):
         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}"
+            usage_url = f"{shared.state.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:

modules/presets.py

@@ -3,7 +3,8 @@ import gradio as gr
 
 # ChatGPT 设置
 initial_prompt = "You are a helpful assistant."
-API_URL = "https://api.openai.com/v1/chat/completions"
+API_HOST = "api.openai.com"
+COMPLETION_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"

modules/shared.py

@@ -1,8 +1,10 @@
-from modules.presets import API_URL
-
+from modules.presets import COMPLETION_URL, BALANCE_API_URL, USAGE_API_URL, API_HOST
+import os
 class State:
     interrupted = False
-    api_url = API_URL
+    completion_url = COMPLETION_URL
+    balance_api_url = BALANCE_API_URL
+    usage_api_url = USAGE_API_URL
 
     def interrupt(self):
         self.interrupted = True
@@ -10,15 +12,21 @@ class State:
     def recover(self):
         self.interrupted = False
 
-    def set_api_url(self, api_url):
-        self.api_url = api_url
+    def set_api_host(self, api_host):
+        self.completion_url = f"https://{api_host}/v1/chat/completions"
+        self.balance_api_url = f"https://{api_host}/dashboard/billing/credit_grants"
+        self.usage_api_url = f"https://{api_host}/dashboard/billing/usage"
+        os.environ["OPENAI_API_BASE"] = f"https://{api_host}/v1"
 
-    def reset_api_url(self):
-        self.api_url = API_URL
-        return self.api_url
+    def reset_api_host(self):
+        self.completion_url = COMPLETION_URL
+        self.balance_api_url = BALANCE_API_URL
+        self.usage_api_url = USAGE_API_URL
+        os.environ["OPENAI_API_BASE"] = f"https://{API_HOST}/v1"
+        return API_HOST
 
     def reset_all(self):
         self.interrupted = False
-        self.api_url = API_URL
+        self.completion_url = COMPLETION_URL
 
 state = State()

modules/utils.py

@@ -328,14 +328,14 @@ def reset_textbox():
 
 
 def reset_default():
-    newurl = shared.state.reset_api_url()
+    default_host = shared.state.reset_api_host()
     retrieve_proxy("")
-    return gr.update(value=newurl), gr.update(value=""), "API URL 和代理已重置"
+    return gr.update(value=default_host), gr.update(value=""), "API-Host 和代理已重置"
 
 
-def change_api_url(url):
-    shared.state.set_api_url(url)
-    msg = f"API地址更改为了{url}"
+def change_api_host(host):
+    shared.state.set_api_host(host)
+    msg = f"API-Host更改为了{host}"
     logging.info(msg)
     return msg