feat: 加入 Azure OpenAI 支持

config_example.json

@@ -8,6 +8,12 @@
     "minimax_api_key": "", // 你的 MiniMax API Key，用于 MiniMax 对话模型
     "minimax_group_id": "", // 你的 MiniMax Group ID，用于 MiniMax 对话模型
 
+    //== Azure ==
+    "azure_openai_api_key": "", // 你的 Azure OpenAI API Key，用于 Azure OpenAI 对话模型
+    "azure_api_base_url": "", // 你的 Azure Base URL
+    "azure_openai_api_version": "2023-05-15", // 你的 Azure OpenAI API 版本
+    "azure_deployment_name": "", // 你的 Azure DEPLOYMENT NAME
+
     //== 基础配置 ==
     "language": "auto", // 界面语言，可选"auto", "zh-CN", "en-US", "ja-JP", "ko-KR"
     "users": [], // 用户列表，[[用户名1, 密码1], [用户名2, 密码2], ...]

modules/config.py

@@ -39,6 +39,12 @@ if os.path.exists("config.json"):
 else:
     config = {}
 
+def load_config_to_environ(key_list):
+    global config
+    for key in key_list:
+        if key in config:
+            os.environ[key.upper()] = os.environ.get(key.upper(), config[key])
+
 sensitive_id = config.get("sensitive_id", "")
 sensitive_id = os.environ.get("SENSITIVE_ID", sensitive_id)
 
@@ -97,6 +103,8 @@ os.environ["MINIMAX_API_KEY"] = minimax_api_key
 minimax_group_id = config.get("minimax_group_id", "")
 os.environ["MINIMAX_GROUP_ID"] = minimax_group_id
 
+load_config_to_environ(["azure_openai_api_key", "azure_api_base_url", "azure_openai_api_version", "azure_deployment_name"])
+
 
 usage_limit = os.environ.get("USAGE_LIMIT", config.get("usage_limit", 120))
 

modules/models/azure.py

@@ -0,0 +1,17 @@
+from langchain.chat_models import AzureChatOpenAI
+import os
+
+from .base_model import Base_Chat_Langchain_Client
+
+# load_config_to_environ(["azure_openai_api_key", "azure_api_base_url", "azure_openai_api_version", "azure_deployment_name"])
+
+class Azure_OpenAI_Client(Base_Chat_Langchain_Client):
+    def setup_model(self):
+        # inplement this to setup the model then return it
+        return AzureChatOpenAI(
+            openai_api_base=os.environ["AZURE_API_BASE_URL"],
+            openai_api_version=os.environ["AZURE_OPENAI_API_VERSION"],
+            deployment_name=os.environ["AZURE_DEPLOYMENT_NAME"],
+            openai_api_key=os.environ["AZURE_OPENAI_API_KEY"],
+            openai_api_type="azure",
+        )

modules/models/base_model.py

@@ -29,6 +29,8 @@ from langchain.input import print_text
 from langchain.schema import AgentAction, AgentFinish, LLMResult
 from threading import Thread, Condition
 from collections import deque
+from langchain.chat_models.base import BaseChatModel
+from langchain.schema import HumanMessage, AIMessage, SystemMessage, BaseMessage
 
 from ..presets import *
 from ..index_func import *
@@ -36,6 +38,7 @@ from ..utils import *
 from .. import shared
 from ..config import retrieve_proxy
 
+
 class CallbackToIterator:
     def __init__(self):
         self.queue = deque()
@@ -52,7 +55,8 @@ class CallbackToIterator:
 
     def __next__(self):
         with self.cond:
-            while not self.queue and not self.finished:  # Wait for a value to be added to the queue.
+            # Wait for a value to be added to the queue.
+            while not self.queue and not self.finished:
                 self.cond.wait()
             if not self.queue:
                 raise StopIteration()
@@ -63,6 +67,7 @@ class CallbackToIterator:
             self.finished = True
             self.cond.notify()  # Wake up the generator if it's waiting.
 
+
 def get_action_description(text):
     match = re.search('```(.*?)```', text, re.S)
     json_text = match.group(1)
@@ -76,6 +81,7 @@ def get_action_description(text):
     else:
         return ""
 
+
 class ChuanhuCallbackHandler(BaseCallbackHandler):
 
     def __init__(self, callback) -> None:
@@ -117,6 +123,10 @@ class ChuanhuCallbackHandler(BaseCallbackHandler):
         """Run on new LLM token. Only available when streaming is enabled."""
         self.callback(token)
 
+    def on_chat_model_start(self, serialized: Dict[str, Any], messages: List[List[BaseMessage]],  **kwargs: Any) -> Any:
+        """Run when a chat model starts running."""
+        pass
+
 
 class ModelType(Enum):
     Unknown = -1
@@ -130,6 +140,7 @@ class ModelType(Enum):
     Minimax = 7
     ChuanhuAgent = 8
     GooglePaLM = 9
+    LangchainChat = 10
 
     @classmethod
     def get_type(cls, model_name: str):
@@ -155,6 +166,8 @@ class ModelType(Enum):
             model_type = ModelType.ChuanhuAgent
         elif "palm" in model_name_lower:
             model_type = ModelType.GooglePaLM
+        elif "azure" or "api" in model_name_lower:
+            model_type = ModelType.LangchainChat
         else:
             model_type = ModelType.Unknown
         return model_type
@@ -164,7 +177,7 @@ class BaseLLMModel:
     def __init__(
         self,
         model_name,
-        system_prompt="",
+        system_prompt=INITIAL_SYSTEM_PROMPT,
         temperature=1.0,
         top_p=1.0,
         n_choices=1,
@@ -204,7 +217,8 @@ class BaseLLMModel:
         conversations are stored in self.history, with the most recent question, in OpenAI format
         should return a generator, each time give the next word (str) in the answer
         """
-        logging.warning("stream predict not implemented, using at once predict instead")
+        logging.warning(
+            "stream predict not implemented, using at once predict instead")
         response, _ = self.get_answer_at_once()
         yield response
 
@@ -215,7 +229,8 @@ class BaseLLMModel:
         the answer (str)
         total token count (int)
         """
-        logging.warning("at once predict not implemented, using stream predict instead")
+        logging.warning(
+            "at once predict not implemented, using stream predict instead")
         response_iter = self.get_answer_stream_iter()
         count = 0
         for response in response_iter:
@@ -276,9 +291,11 @@ class BaseLLMModel:
             self.history[-2] = construct_user(fake_input)
         chatbot[-1] = (chatbot[-1][0], ai_reply + display_append)
         if fake_input is not None:
-            self.all_token_counts[-1] += count_token(construct_assistant(ai_reply))
+            self.all_token_counts[-1] += count_token(
+                construct_assistant(ai_reply))
         else:
-            self.all_token_counts[-1] = total_token_count - sum(self.all_token_counts)
+            self.all_token_counts[-1] = total_token_count - \
+                sum(self.all_token_counts)
         status_text = self.token_message()
         return chatbot, status_text
 
@@ -302,10 +319,13 @@ class BaseLLMModel:
             from langchain.chat_models import ChatOpenAI
             from langchain.callbacks import StdOutCallbackHandler
             prompt_template = "Write a concise summary of the following:\n\n{text}\n\nCONCISE SUMMARY IN " + language + ":"
-            PROMPT = PromptTemplate(template=prompt_template, input_variables=["text"])
+            PROMPT = PromptTemplate(
+                template=prompt_template, input_variables=["text"])
             llm = ChatOpenAI()
-            chain = load_summarize_chain(llm, chain_type="map_reduce", return_intermediate_steps=True, map_prompt=PROMPT, combine_prompt=PROMPT)
-            summary = chain({"input_documents": list(index.docstore.__dict__["_dict"].values())}, return_only_outputs=True)["output_text"]
+            chain = load_summarize_chain(
+                llm, chain_type="map_reduce", return_intermediate_steps=True, map_prompt=PROMPT, combine_prompt=PROMPT)
+            summary = chain({"input_documents": list(index.docstore.__dict__[
+                            "_dict"].values())}, return_only_outputs=True)["output_text"]
             print(i18n("总结") + f": {summary}")
             chatbot.append([i18n("上传了")+str(len(files))+"个文件", summary])
         return chatbot, status
@@ -326,9 +346,12 @@ class BaseLLMModel:
             msg = "索引获取成功，生成回答中……"
             logging.info(msg)
             with retrieve_proxy():
-                retriever = VectorStoreRetriever(vectorstore=index, search_type="similarity_score_threshold",search_kwargs={"k":6, "score_threshold": 0.5})
-                relevant_documents = retriever.get_relevant_documents(real_inputs)
-            reference_results = [[d.page_content.strip("�"), os.path.basename(d.metadata["source"])] for d in relevant_documents]
+                retriever = VectorStoreRetriever(vectorstore=index, search_type="similarity_score_threshold", search_kwargs={
+                                                 "k": 6, "score_threshold": 0.5})
+                relevant_documents = retriever.get_relevant_documents(
+                    real_inputs)
+            reference_results = [[d.page_content.strip("�"), os.path.basename(
+                d.metadata["source"])] for d in relevant_documents]
             reference_results = add_source_numbers(reference_results)
             display_append = add_details(reference_results)
             display_append = "\n\n" + "".join(display_append)
@@ -355,7 +378,8 @@ class BaseLLMModel:
                 )
             reference_results = add_source_numbers(reference_results)
             # display_append = "<ol>\n\n" + "".join(display_append) + "</ol>"
-            display_append = '<div class = "source-a">' + "".join(display_append) + '</div>'
+            display_append = '<div class = "source-a">' + \
+                "".join(display_append) + '</div>'
             real_inputs = (
                 replace_today(WEBSEARCH_PTOMPT_TEMPLATE)
                 .replace("{query}", real_inputs)
@@ -379,14 +403,16 @@ class BaseLLMModel:
 
         status_text = "开始生成回答……"
         logging.info(
-             "用户" + f"{self.user_identifier}" + "的输入为：" + colorama.Fore.BLUE + f"{inputs}" + colorama.Style.RESET_ALL
+            "用户" + f"{self.user_identifier}" + "的输入为：" +
+            colorama.Fore.BLUE + f"{inputs}" + colorama.Style.RESET_ALL
         )
         if should_check_token_count:
             yield chatbot + [(inputs, "")], status_text
         if reply_language == "跟随问题语言（不稳定）":
             reply_language = "the same language as the question, such as English, 中文, 日本語, Español, Français, or Deutsch."
 
-        limited_context, fake_inputs, display_append, inputs, chatbot = self.prepare_inputs(real_inputs=inputs, use_websearch=use_websearch, files=files, reply_language=reply_language, chatbot=chatbot)
+        limited_context, fake_inputs, display_append, inputs, chatbot = self.prepare_inputs(
+            real_inputs=inputs, use_websearch=use_websearch, files=files, reply_language=reply_language, chatbot=chatbot)
         yield chatbot + [(fake_inputs, "")], status_text
 
         if (
@@ -587,7 +613,8 @@ class BaseLLMModel:
         self.history = []
         self.all_token_counts = []
         self.interrupted = False
-        pathlib.Path(os.path.join(HISTORY_DIR, self.user_identifier, new_auto_history_filename(os.path.join(HISTORY_DIR, self.user_identifier)))).touch()
+        pathlib.Path(os.path.join(HISTORY_DIR, self.user_identifier, new_auto_history_filename(
+            os.path.join(HISTORY_DIR, self.user_identifier)))).touch()
         return [], self.token_message([0])
 
     def delete_first_conversation(self):
@@ -630,7 +657,8 @@ class BaseLLMModel:
 
     def auto_save(self, chatbot):
         history_file_path = get_history_filepath(self.user_identifier)
-        save_file(history_file_path, self.system_prompt, self.history, chatbot, self.user_identifier)
+        save_file(history_file_path, self.system_prompt,
+                  self.history, chatbot, self.user_identifier)
 
     def export_markdown(self, filename, chatbot, user_name):
         if filename == "":
@@ -646,7 +674,8 @@ class BaseLLMModel:
             filename = filename.name
         try:
             if "/" not in filename:
-                history_file_path = os.path.join(HISTORY_DIR, user_name, filename)
+                history_file_path = os.path.join(
+                    HISTORY_DIR, user_name, filename)
             else:
                 history_file_path = filename
             with open(history_file_path, "r", encoding="utf-8") as f:
@@ -695,10 +724,10 @@ class BaseLLMModel:
             self.reset()
             return self.system_prompt, gr.update()
         history_file_path = get_history_filepath(self.user_identifier)
-        filename, system_prompt, chatbot = self.load_chat_history(history_file_path, self.user_identifier)
+        filename, system_prompt, chatbot = self.load_chat_history(
+            history_file_path, self.user_identifier)
         return system_prompt, chatbot
 
-
     def like(self):
         """like the last response, implement if needed
         """
@@ -708,3 +737,47 @@ class BaseLLMModel:
         """dislike the last response, implement if needed
         """
         return gr.update()
+
+
+class Base_Chat_Langchain_Client(BaseLLMModel):
+    def __init__(self, model_name, user_name=""):
+        super().__init__(model_name, user=user_name)
+        self.need_api_key = False
+        self.model = self.setup_model()
+
+    def setup_model(self):
+        # inplement this to setup the model then return it
+        pass
+
+    def _get_langchain_style_history(self):
+        history = [SystemMessage(content=self.system_prompt)]
+        for i in self.history:
+            if i["role"] == "user":
+                history.append(HumanMessage(content=i["content"]))
+            elif i["role"] == "assistant":
+                history.append(AIMessage(content=i["content"]))
+        return history
+
+    def get_answer_at_once(self):
+        assert isinstance(
+            self.model, BaseChatModel), "model is not instance of LangChain BaseChatModel"
+        history = self._get_langchain_style_history()
+        response = self.model.generate(history)
+        return response.content, sum(response.content)
+
+    def get_answer_stream_iter(self):
+        it = CallbackToIterator()
+        assert isinstance(
+            self.model, BaseChatModel), "model is not instance of LangChain BaseChatModel"
+        history = self._get_langchain_style_history()
+
+        def thread_func():
+            self.model(messages=history, callbacks=[
+                ChuanhuCallbackHandler(it.callback)])
+            it.finish()
+        t = Thread(target=thread_func)
+        t.start()
+        partial_text = ""
+        for value in it:
+            partial_text += value
+            yield partial_text

modules/models/models.py

@@ -616,6 +616,9 @@ def get_model(
             from .Google_PaLM import Google_PaLM_Client
             access_key = os.environ.get("GOOGLE_PALM_API_KEY")
             model = Google_PaLM_Client(model_name, access_key, user_name=user_name)
+        elif model_type == ModelType.LangchainChat:
+            from .azure import Azure_OpenAI_Client
+            model = Azure_OpenAI_Client(model_name, user_name=user_name)
         elif model_type == ModelType.Unknown:
             raise ValueError(f"未知模型: {model_name}")
         logging.info(msg)

modules/presets.py

@@ -62,6 +62,7 @@ ONLINE_MODELS = [
     "川虎助理 Pro",
     "GooglePaLM",
     "xmchat",
+    "Azure OpenAI",
     "yuanai-1.0-base_10B",
     "yuanai-1.0-translate",
     "yuanai-1.0-dialog",