feat: Added support for custom models

modules/config.py

@@ -5,9 +5,11 @@ import logging
 import sys
 import commentjson as json
 import colorama
+from collections import defaultdict
 
 from . import shared
 from . import presets
+from .presets import i18n
 
 
 __all__ = [
@@ -100,14 +102,25 @@ else:
     sensitive_id = config.get("sensitive_id", "")
     sensitive_id = os.environ.get("SENSITIVE_ID", sensitive_id)
 
+if "extra_model_metadata" in config:
+    presets.MODEL_METADATA.update(config["extra_model_metadata"])
+    logging.info(i18n("已添加 {extra_model_quantity} 个额外的模型元数据").format(extra_model_quantity=len(config["extra_model_metadata"])))
+
+_model_metadata = {}
+for k, v in presets.MODEL_METADATA.items():
+    temp_dict = presets.DEFAULT_METADATA.copy()
+    temp_dict.update(v)
+    _model_metadata[k] = temp_dict
+presets.MODEL_METADATA = _model_metadata
+
 if "available_models" in config:
     presets.MODELS = config["available_models"]
-    logging.info(f"已设置可用模型：{config['available_models']}")
+    logging.info(i18n("已设置可用模型：{available_models}").format(available_models=config["available_models"]))
 
 # 模型配置
 if "extra_models" in  config:
     presets.MODELS.extend(config["extra_models"])
-    logging.info(f"已添加额外的模型：{config['extra_models']}")
+    logging.info(i18n("已添加额外的模型：{extra_models}").format(extra_models=config["extra_models"]))
 
 HIDE_MY_KEY = config.get("hide_my_key", False)
 

modules/models/Claude.py

@@ -11,7 +11,7 @@ class Claude_Client(BaseLLMModel):
         self.api_secret = api_secret
         if None in [self.api_secret]:
             raise Exception("请在配置文件或者环境变量中设置Claude的API Secret")
-        self.claude_client = Anthropic(api_key=self.api_secret)
+        self.claude_client = Anthropic(api_key=self.api_secret, base_url=self.api_host)
 
     def _get_claude_style_history(self):
         history = []

modules/models/DALLE3.py

@@ -7,8 +7,11 @@ from ..config import retrieve_proxy, sensitive_id
 
 class OpenAI_DALLE3_Client(BaseLLMModel):
     def __init__(self, model_name, api_key, user_name="") -> None:
-        super().__init__(model_name=model_name, user=user_name)
-        self.api_key = api_key
+        super().__init__(model_name=model_name, user=user_name, config={"api_key": api_key})
+        if self.api_host is not None:
+            self.chat_completion_url, self.images_completion_url, self.openai_api_base, self.balance_api_url, self.usage_api_url = shared.format_openai_host(self.api_host)
+        else:
+            self.api_host, self.chat_completion_url, self.images_completion_url, self.openai_api_base, self.balance_api_url, self.usage_api_url = shared.state.api_host, shared.state.chat_completion_url, shared.state.images_completion_url, shared.state.openai_api_base, shared.state.balance_api_url, shared.state.usage_api_url
         self._refresh_header()
 
     def _get_dalle3_prompt(self):
@@ -24,7 +27,7 @@ class OpenAI_DALLE3_Client(BaseLLMModel):
             "Authorization": f"Bearer {self.api_key}"
         }
         payload = {
-            "model": "dall-e-3",
+            "model": self.model_name,
             "prompt": prompt,
             "n": 1,
             "size": "1024x1024",
@@ -35,13 +38,13 @@ class OpenAI_DALLE3_Client(BaseLLMModel):
         else:
             timeout = TIMEOUT_ALL
 
-        if shared.state.images_completion_url != IMAGES_COMPLETION_URL:
-            logging.debug(f"使用自定义API URL: {shared.state.images_completion_url}")
+        if self.images_completion_url != IMAGES_COMPLETION_URL:
+            logging.debug(f"使用自定义API URL: {self.images_completion_url}")
 
         with retrieve_proxy():
             try:
                 response = requests.post(
-                    shared.state.images_completion_url,
+                    self.images_completion_url,
                     headers=headers,
                     json=payload,
                     stream=stream,

modules/models/GoogleGemini.py

@@ -17,8 +17,7 @@ from .base_model import BaseLLMModel
 
 class GoogleGeminiClient(BaseLLMModel):
     def __init__(self, model_name, api_key, user_name="") -> None:
-        super().__init__(model_name=model_name, user=user_name)
-        self.api_key = api_key
+        super().__init__(model_name=model_name, user=user_name, config={"api_key": api_key})
         if "vision" in model_name.lower():
             self.multimodal = True
         else:

modules/models/GooglePaLM.py

@@ -4,8 +4,7 @@ import google.generativeai as palm
 
 class Google_PaLM_Client(BaseLLMModel):
     def __init__(self, model_name, api_key, user_name="") -> None:
-        super().__init__(model_name=model_name, user=user_name)
-        self.api_key = api_key
+        super().__init__(model_name=model_name, user=user_name, config={"api_key": api_key})
 
     def _get_palm_style_input(self):
         new_history = []
@@ -20,7 +19,7 @@ class Google_PaLM_Client(BaseLLMModel):
         palm.configure(api_key=self.api_key)
         messages = self._get_palm_style_input()
         response = palm.chat(context=self.system_prompt, messages=messages,
-                             temperature=self.temperature, top_p=self.top_p)
+                             temperature=self.temperature, top_p=self.top_p, model=self.model_name)
         if response.last is not None:
             return response.last, len(response.last)
         else:

modules/models/Groq.py

@@ -18,10 +18,10 @@ from .base_model import BaseLLMModel
 
 class Groq_Client(BaseLLMModel):
     def __init__(self, model_name, api_key, user_name="") -> None:
-        super().__init__(model_name=model_name, user=user_name)
-        self.api_key = api_key
+        super().__init__(model_name=model_name, user=user_name, api_key=api_key)
         self.client = Groq(
             api_key=os.environ.get("GROQ_API_KEY"),
+            base_url=self.api_host,
         )
 
     def _get_groq_style_input(self):

modules/models/OpenAI.py

@@ -1,280 +0,0 @@
-from __future__ import annotations
-
-import json
-import logging
-import traceback
-
-import colorama
-import requests
-
-from .. import shared
-from ..config import retrieve_proxy, sensitive_id, usage_limit
-from ..index_func import *
-from ..presets import *
-from ..utils import *
-from .base_model import BaseLLMModel
-
-
-class OpenAIClient(BaseLLMModel):
-    def __init__(
-        self,
-        model_name,
-        api_key,
-        system_prompt=INITIAL_SYSTEM_PROMPT,
-        temperature=1.0,
-        top_p=1.0,
-        user_name=""
-    ) -> None:
-        super().__init__(
-            model_name=model_name,
-            temperature=temperature,
-            top_p=top_p,
-            system_prompt=system_prompt,
-            user=user_name
-        )
-        self.api_key = api_key
-        self.need_api_key = True
-        self._refresh_header()
-
-    def get_answer_stream_iter(self):
-        if not self.api_key:
-            raise Exception(NO_APIKEY_MSG)
-        response = self._get_response(stream=True)
-        if response is not None:
-            iter = self._decode_chat_response(response)
-            partial_text = ""
-            for i in iter:
-                partial_text += i
-                yield partial_text
-        else:
-            yield STANDARD_ERROR_MSG + GENERAL_ERROR_MSG
-
-    def get_answer_at_once(self):
-        if not self.api_key:
-            raise Exception(NO_APIKEY_MSG)
-        response = self._get_response()
-        response = json.loads(response.text)
-        content = response["choices"][0]["message"]["content"]
-        total_token_count = response["usage"]["total_tokens"]
-        return content, total_token_count
-
-    def count_token(self, user_input):
-        input_token_count = count_token(construct_user(user_input))
-        if self.system_prompt is not None and len(self.all_token_counts) == 0:
-            system_prompt_token_count = count_token(
-                construct_system(self.system_prompt)
-            )
-            return input_token_count + system_prompt_token_count
-        return input_token_count
-
-    def billing_info(self):
-        try:
-            curr_time = datetime.datetime.now()
-            last_day_of_month = get_last_day_of_month(
-                curr_time).strftime("%Y-%m-%d")
-            first_day_of_month = curr_time.replace(day=1).strftime("%Y-%m-%d")
-            usage_url = f"{shared.state.usage_api_url}?start_date={first_day_of_month}&end_date={last_day_of_month}"
-            try:
-                usage_data = self._get_billing_data(usage_url)
-            except Exception as e:
-                # logging.error(f"获取API使用情况失败: " + str(e))
-                if "Invalid authorization header" in str(e):
-                    return i18n("**获取API使用情况失败**，需在填写`config.json`中正确填写sensitive_id")
-                elif "Incorrect API key provided: sess" in str(e):
-                    return i18n("**获取API使用情况失败**，sensitive_id错误或已过期")
-                return i18n("**获取API使用情况失败**")
-            # rounded_usage = "{:.5f}".format(usage_data["total_usage"] / 100)
-            rounded_usage = round(usage_data["total_usage"] / 100, 5)
-            usage_percent = round(usage_data["total_usage"] / usage_limit, 2)
-            from ..webui import get_html
-
-            # return i18n("**本月使用金额** ") + f"\u3000 ${rounded_usage}"
-            return get_html("billing_info.html").format(
-                    label = i18n("本月使用金额"),
-                    usage_percent = usage_percent,
-                    rounded_usage = rounded_usage,
-                    usage_limit = usage_limit
-                )
-        except requests.exceptions.ConnectTimeout:
-            status_text = (
-                STANDARD_ERROR_MSG + CONNECTION_TIMEOUT_MSG + ERROR_RETRIEVE_MSG
-            )
-            return status_text
-        except requests.exceptions.ReadTimeout:
-            status_text = STANDARD_ERROR_MSG + READ_TIMEOUT_MSG + ERROR_RETRIEVE_MSG
-            return status_text
-        except Exception as e:
-            import traceback
-            traceback.print_exc()
-            logging.error(i18n("获取API使用情况失败:") + str(e))
-            return STANDARD_ERROR_MSG + ERROR_RETRIEVE_MSG
-
-    @shared.state.switching_api_key  # 在不开启多账号模式的时候，这个装饰器不会起作用
-    def _get_response(self, stream=False):
-        openai_api_key = self.api_key
-        system_prompt = self.system_prompt
-        history = self.history
-        logging.debug(colorama.Fore.YELLOW +
-                      f"{history}" + colorama.Fore.RESET)
-        headers = {
-            "Content-Type": "application/json",
-            "Authorization": f"Bearer {openai_api_key}",
-        }
-
-        if system_prompt is not None:
-            history = [construct_system(system_prompt), *history]
-
-        payload = {
-            "model": self.model_name,
-            "messages": history,
-            "temperature": self.temperature,
-            "top_p": self.top_p,
-            "n": self.n_choices,
-            "stream": stream,
-            "presence_penalty": self.presence_penalty,
-            "frequency_penalty": self.frequency_penalty,
-        }
-
-        if self.max_generation_token is not None:
-            payload["max_tokens"] = self.max_generation_token
-        if self.stop_sequence is not None:
-            payload["stop"] = self.stop_sequence
-        if self.logit_bias is not None:
-            payload["logit_bias"] = self.encoded_logit_bias()
-        if self.user_identifier:
-            payload["user"] = self.user_identifier
-
-        if stream:
-            timeout = TIMEOUT_STREAMING
-        else:
-            timeout = TIMEOUT_ALL
-
-        # 如果有自定义的api-host，使用自定义host发送请求，否则使用默认设置发送请求
-        if shared.state.chat_completion_url != CHAT_COMPLETION_URL:
-            logging.debug(f"使用自定义API URL: {shared.state.chat_completion_url}")
-
-        with retrieve_proxy():
-            try:
-                response = requests.post(
-                    shared.state.chat_completion_url,
-                    headers=headers,
-                    json=payload,
-                    stream=stream,
-                    timeout=timeout,
-                )
-            except:
-                traceback.print_exc()
-                return None
-        return response
-
-    def _refresh_header(self):
-        self.headers = {
-            "Content-Type": "application/json",
-            "Authorization": f"Bearer {sensitive_id}",
-        }
-
-
-    def _get_billing_data(self, billing_url):
-        with retrieve_proxy():
-            response = requests.get(
-                billing_url,
-                headers=self.headers,
-                timeout=TIMEOUT_ALL,
-            )
-
-        if response.status_code == 200:
-            data = response.json()
-            return data
-        else:
-            raise Exception(
-                f"API request failed with status code {response.status_code}: {response.text}"
-            )
-
-    def _decode_chat_response(self, response):
-        error_msg = ""
-        for chunk in response.iter_lines():
-            if chunk:
-                chunk = chunk.decode()
-                chunk_length = len(chunk)
-                try:
-                    chunk = json.loads(chunk[6:])
-                except:
-                    print(i18n("JSON解析错误,收到的内容: ") + f"{chunk}")
-                    error_msg += chunk
-                    continue
-                try:
-                    if chunk_length > 6 and "delta" in chunk["choices"][0]:
-                        if "finish_reason" in chunk["choices"][0]:
-                            finish_reason = chunk["choices"][0]["finish_reason"]
-                        else:
-                            finish_reason = chunk["finish_reason"]
-                        if finish_reason == "stop":
-                            break
-                        try:
-                            yield chunk["choices"][0]["delta"]["content"]
-                        except Exception as e:
-                            # logging.error(f"Error: {e}")
-                            continue
-                except:
-                    print(f"ERROR: {chunk}")
-                    continue
-        if error_msg and not error_msg=="data: [DONE]":
-            raise Exception(error_msg)
-
-    def set_key(self, new_access_key):
-        ret = super().set_key(new_access_key)
-        self._refresh_header()
-        return ret
-
-    def _single_query_at_once(self, history, temperature=1.0):
-        timeout = TIMEOUT_ALL
-        headers = {
-            "Content-Type": "application/json",
-            "Authorization": f"Bearer {self.api_key}",
-            "temperature": f"{temperature}",
-        }
-        payload = {
-            "model": self.model_name,
-            "messages": history,
-        }
-        # 如果有自定义的api-host，使用自定义host发送请求，否则使用默认设置发送请求
-        if shared.state.chat_completion_url != CHAT_COMPLETION_URL:
-            logging.debug(f"使用自定义API URL: {shared.state.chat_completion_url}")
-
-        with retrieve_proxy():
-            response = requests.post(
-                shared.state.chat_completion_url,
-                headers=headers,
-                json=payload,
-                stream=False,
-                timeout=timeout,
-            )
-
-        return response
-
-
-    def auto_name_chat_history(self, name_chat_method, user_question, chatbot, single_turn_checkbox):
-        if len(self.history) == 2 and not single_turn_checkbox and not hide_history_when_not_logged_in:
-            user_question = self.history[0]["content"]
-            if name_chat_method == i18n("模型自动总结（消耗tokens）"):
-                ai_answer = self.history[1]["content"]
-                try:
-                    history = [
-                        { "role": "system", "content": SUMMARY_CHAT_SYSTEM_PROMPT},
-                        { "role": "user", "content": f"Please write a title based on the following conversation:\n---\nUser: {user_question}\nAssistant: {ai_answer}"}
-                    ]
-                    response = self._single_query_at_once(history, temperature=0.0)
-                    response = json.loads(response.text)
-                    content = response["choices"][0]["message"]["content"]
-                    filename = replace_special_symbols(content) + ".json"
-                except Exception as e:
-                    logging.info(f"自动命名失败。{e}")
-                    filename = replace_special_symbols(user_question)[:16] + ".json"
-                return self.rename_chat_history(filename, chatbot)
-            elif name_chat_method == i18n("第一条提问"):
-                filename = replace_special_symbols(user_question)[:16] + ".json"
-                return self.rename_chat_history(filename, chatbot)
-            else:
-                return gr.update()
-        else:
-            return gr.update()

modules/models/OpenAIInstruct.py

@@ -8,8 +8,7 @@ from ..config import retrieve_proxy
 
 class OpenAI_Instruct_Client(BaseLLMModel):
     def __init__(self, model_name, api_key, user_name="") -> None:
-        super().__init__(model_name=model_name, user=user_name)
-        self.api_key = api_key
+        super().__init__(model_name=model_name, user=user_name, config={"api_key": api_key})
 
     def _get_instruct_style_input(self):
         return "\n\n".join([item["content"] for item in self.history])

modules/models/OpenAIVision.py

@@ -27,22 +27,19 @@ class OpenAIVisionClient(BaseLLMModel):
         self,
         model_name,
         api_key,
-        system_prompt=INITIAL_SYSTEM_PROMPT,
-        temperature=1.0,
-        top_p=1.0,
         user_name=""
     ) -> None:
         super().__init__(
             model_name=model_name,
-            temperature=temperature,
-            top_p=top_p,
-            system_prompt=system_prompt,
-            user=user_name
+            user=user_name,
+            config={
+                "api_key": api_key
+            }
         )
-        self.image_token = 0
-        self.api_key = api_key
-        self.need_api_key = True
-        self.max_generation_token = 4096
+        if self.api_host is not None:
+            self.chat_completion_url, self.images_completion_url, self.openai_api_base, self.balance_api_url, self.usage_api_url = shared.format_openai_host(self.api_host)
+        else:
+            self.api_host, self.chat_completion_url, self.images_completion_url, self.openai_api_base, self.balance_api_url, self.usage_api_url = shared.state.api_host, shared.state.chat_completion_url, shared.state.images_completion_url, shared.state.openai_api_base, shared.state.balance_api_url, shared.state.usage_api_url
         self._refresh_header()
 
     def get_answer_stream_iter(self):
@@ -176,7 +173,7 @@ class OpenAIVisionClient(BaseLLMModel):
             "stream": stream,
             "presence_penalty": self.presence_penalty,
             "frequency_penalty": self.frequency_penalty,
-            "max_tokens": 4096
+            "max_tokens": self.max_generation_token
         }
 
         if self.stop_sequence:
@@ -296,3 +293,29 @@ class OpenAIVisionClient(BaseLLMModel):
             )
 
         return response
+
+    def auto_name_chat_history(self, name_chat_method, user_question, chatbot, single_turn_checkbox):
+        if len(self.history) == 2 and not single_turn_checkbox and not hide_history_when_not_logged_in:
+            user_question = self.history[0]["content"]
+            if name_chat_method == i18n("模型自动总结（消耗tokens）"):
+                ai_answer = self.history[1]["content"]
+                try:
+                    history = [
+                        { "role": "system", "content": SUMMARY_CHAT_SYSTEM_PROMPT},
+                        { "role": "user", "content": f"Please write a title based on the following conversation:\n---\nUser: {user_question}\nAssistant: {ai_answer}"}
+                    ]
+                    response = self._single_query_at_once(history, temperature=0.0)
+                    response = json.loads(response.text)
+                    content = response["choices"][0]["message"]["content"]
+                    filename = replace_special_symbols(content) + ".json"
+                except Exception as e:
+                    logging.info(f"自动命名失败。{e}")
+                    filename = replace_special_symbols(user_question)[:16] + ".json"
+                return self.rename_chat_history(filename, chatbot)
+            elif name_chat_method == i18n("第一条提问"):
+                filename = replace_special_symbols(user_question)[:16] + ".json"
+                return self.rename_chat_history(filename, chatbot)
+            else:
+                return gr.update()
+        else:
+            return gr.update()

modules/models/XMChat.py

@@ -26,6 +26,8 @@ class XMChat(BaseLLMModel):
         self.image_path = None
         self.xm_history = []
         self.url = "https://xmbot.net/web"
+        if self.api_host is not None:
+            self.url = self.api_host
         self.last_conv_id = None
 
     def reset(self, remain_system_prompt=False):

modules/models/base_model.py

@@ -159,6 +159,14 @@ class ModelType(Enum):
 
     @classmethod
     def get_type(cls, model_name: str):
+        # 1. get model type from model metadata (if exists)
+        model_type = MODEL_METADATA[model_name]["model_type"]
+        if model_type is not None:
+            for member in cls:
+                if member.name == model_type:
+                    return member
+
+        # 2. infer model type from model name
         model_type = None
         model_name_lower = model_name.lower()
         if "gpt" in model_name_lower:
@@ -249,66 +257,56 @@ class BaseLLMModel:
     def __init__(
         self,
         model_name,
-        system_prompt=INITIAL_SYSTEM_PROMPT,
-        temperature=1.0,
-        top_p=1.0,
-        n_choices=1,
-        stop=[],
-        max_generation_token=None,
-        presence_penalty=0,
-        frequency_penalty=0,
-        logit_bias=None,
         user="",
-        single_turn=False,
+        config=None,
     ) -> None:
+
+        if config is not None:
+            temp = MODEL_METADATA[model_name].copy()
+            keys_with_diff_values = {key: temp[key] for key in temp if key in DEFAULT_METADATA and temp[key] != DEFAULT_METADATA[key]}
+            config.update(keys_with_diff_values)
+            temp.update(config)
+            config = temp
+        else:
+            config = MODEL_METADATA[model_name]
+
+        self.model_name = config["model_name"]
+        self.multimodal = config["multimodal"]
+        self.description = config["description"]
+        self.token_upper_limit = config["token_limit"]
+        self.system_prompt = config["system"]
+        self.api_key = config["api_key"]
+        self.api_host = config["api_host"]
+
+        self.interrupted = False
+        self.need_api_key = self.api_key is not None
         self.history = []
         self.all_token_counts = []
         self.model_type = ModelType.get_type(model_name)
-        try:
-            self.model_name = MODEL_METADATA[model_name]["model_name"]
-        except:
-            self.model_name = model_name
-        try:
-            self.multimodal = MODEL_METADATA[model_name]["multimodal"]
-        except:
-            self.multimodal = False
-        if max_generation_token is None:
-            try:
-                max_generation_token = MODEL_METADATA[model_name]["max_generation"]
-            except:
-                pass
-        try:
-            self.token_upper_limit = MODEL_METADATA[model_name]["token_limit"]
-        except KeyError:
-            self.token_upper_limit = DEFAULT_TOKEN_LIMIT
-        self.interrupted = False
-        self.system_prompt = system_prompt
-        self.api_key = None
-        self.need_api_key = False
         self.history_file_path = get_first_history_name(user)
         self.user_name = user
         self.chatbot = []
 
-        self.default_single_turn = single_turn
-        self.default_temperature = temperature
-        self.default_top_p = top_p
-        self.default_n_choices = n_choices
-        self.default_stop_sequence = stop
-        self.default_max_generation_token = max_generation_token
-        self.default_presence_penalty = presence_penalty
-        self.default_frequency_penalty = frequency_penalty
-        self.default_logit_bias = logit_bias
+        self.default_single_turn = config["single_turn"]
+        self.default_temperature = config["temperature"]
+        self.default_top_p = config["top_p"]
+        self.default_n_choices = config["n_choices"]
+        self.default_stop_sequence = config["stop"]
+        self.default_max_generation_token = config["max_generation"]
+        self.default_presence_penalty = config["presence_penalty"]
+        self.default_frequency_penalty = config["frequency_penalty"]
+        self.default_logit_bias = config["logit_bias"]
         self.default_user_identifier = user
 
-        self.single_turn = single_turn
-        self.temperature = temperature
-        self.top_p = top_p
-        self.n_choices = n_choices
-        self.stop_sequence = stop
-        self.max_generation_token = max_generation_token
-        self.presence_penalty = presence_penalty
-        self.frequency_penalty = frequency_penalty
-        self.logit_bias = logit_bias
+        self.single_turn = self.default_single_turn
+        self.temperature = self.default_temperature
+        self.top_p = self.default_top_p
+        self.n_choices = self.default_n_choices
+        self.stop_sequence = self.default_stop_sequence
+        self.max_generation_token = self.default_max_generation_token
+        self.presence_penalty = self.default_presence_penalty
+        self.frequency_penalty = self.default_frequency_penalty
+        self.logit_bias = self.default_logit_bias
         self.user_identifier = user
 
         self.metadata = {}
@@ -1073,7 +1071,7 @@ class BaseLLMModel:
             self.reset()
             return (
                 os.path.basename(self.history_file_path),
-                "",
+                self.system_prompt,
                 [],
                 self.single_turn,
                 self.temperature,

modules/models/models.py

@@ -35,28 +35,18 @@ def get_model(
     model = original_model
     chatbot = gr.Chatbot(label=model_name)
     try:
-        if model_type == ModelType.OpenAI:
-            logging.info(f"正在加载OpenAI模型: {model_name}")
-            from .OpenAI import OpenAIClient
+        if model_type == ModelType.OpenAIVision or model_type == ModelType.OpenAI:
+            logging.info(f"正在加载 OpenAI 模型: {model_name}")
+            from .OpenAIVision import OpenAIVisionClient
             access_key = os.environ.get("OPENAI_API_KEY", access_key)
-            model = OpenAIClient(
-                model_name=model_name,
-                api_key=access_key,
-                system_prompt=system_prompt,
-                user_name=user_name,
-            )
+            model = OpenAIVisionClient(
+                model_name, api_key=access_key, user_name=user_name)
         elif model_type == ModelType.OpenAIInstruct:
             logging.info(f"正在加载OpenAI Instruct模型: {model_name}")
             from .OpenAIInstruct import OpenAI_Instruct_Client
             access_key = os.environ.get("OPENAI_API_KEY", access_key)
             model = OpenAI_Instruct_Client(
                 model_name, api_key=access_key, user_name=user_name)
-        elif model_type == ModelType.OpenAIVision:
-            logging.info(f"正在加载OpenAI Vision模型: {model_name}")
-            from .OpenAIVision import OpenAIVisionClient
-            access_key = os.environ.get("OPENAI_API_KEY", access_key)
-            model = OpenAIVisionClient(
-                model_name, api_key=access_key, user_name=user_name)
         elif model_type == ModelType.ChatGLM:
             logging.info(f"正在加载ChatGLM模型: {model_name}")
             from .ChatGLM import ChatGLM_Client

modules/presets.py

@@ -110,6 +110,29 @@ LOCAL_MODELS = [
     "Qwen 14B"
 ]
 
+DEFAULT_METADATA = {
+    "repo_id": None, # HuggingFace repo id, used if this model is meant to be downloaded from HuggingFace then run locally
+    "model_name": None, # api model name, used if this model is meant to be used online
+    "filelist": None, # file list in the repo to download, now only support .gguf file
+    "description": None, # description of the model, displayed in the chat area when no message is present
+    "model_type": None, # model type, used to determine the model's behavior. If not set, the model type is inferred from the model name
+    "multimodal": False, # whether the model is multimodal
+    "api_host": None, # base url for the model's api
+    "api_key": None, # api key for the model's api
+    "system": INITIAL_SYSTEM_PROMPT, # system prompt for the model
+    "token_limit": 4096, # context window size
+    "single_turn": False, # whether the model is single turn
+    "temperature": 1.0,
+    "top_p": 1.0,
+    "n_choices": 1,
+    "stop": [],
+    "max_generation": None, # maximum token limit for a single generation
+    "presence_penalty": 0.0,
+    "frequency_penalty": 0.0,
+    "logit_bias": None,
+    "metadata": {} # additional metadata for the model
+}
+
 # Additional metadata for online and local models
 MODEL_METADATA = {
     "Llama-2-7B":{
@@ -166,11 +189,8 @@ MODEL_METADATA = {
     "GPT4 Vision": {
         "model_name": "gpt-4-turbo",
         "token_limit": 128000,
-        "multimodal": True
-    },
-    "Claude": {
-        "model_name": "Claude",
-        "token_limit": 4096,
+        "multimodal": True,
+        "max_generation": 4096,
     },
     "Claude 3 Haiku": {
         "model_name": "claude-3-haiku-20240307",
@@ -190,6 +210,9 @@ MODEL_METADATA = {
         "max_generation": 4096,
         "multimodal": True
     },
+    "川虎助理": {"model_name": "川虎助理"},
+    "川虎助理 Pro": {"model_name": "川虎助理 Pro"},
+    "DALL-E 3": {"model_name": "dall-e-3"},
     "ERNIE-Bot-turbo": {
         "model_name": "ERNIE-Bot-turbo",
         "token_limit": 1024,
@@ -243,7 +266,19 @@ MODEL_METADATA = {
     "Groq Gemma 7B": {
         "model_name": "gemma-7b-it",
         "token_limit": 8192,
-    }
+    },
+    "GooglePaLM": {"model_name": "models/chat-bison-001"},
+    "xmchat": {"model_name": "xmchat"},
+    "Azure OpenAI": {"model_name": "azure-openai"},
+    "yuanai-1.0-base_10B": {"model_name": "yuanai-1.0-base_10B"},
+    "yuanai-1.0-translate": {"model_name": "yuanai-1.0-translate"},
+    "yuanai-1.0-dialog": {"model_name": "yuanai-1.0-dialog"},
+    "yuanai-1.0-rhythm_poems": {"model_name": "yuanai-1.0-rhythm_poems"},
+    "minimax-abab5-chat": {"model_name": "minimax-abab5-chat"},
+    "midjourney": {"model_name": "midjourney"},
+    "讯飞星火大模型V3.0": {"model_name": "讯飞星火大模型V3.0"},
+    "讯飞星火大模型V2.0": {"model_name": "讯飞星火大模型V2.0"},
+    "讯飞星火大模型V1.5": {"model_name": "讯飞星火大模型V1.5"},
 }
 
 if os.environ.get('HIDE_LOCAL_MODELS', 'false') == 'true':

modules/shared.py

@@ -3,6 +3,19 @@ import os
 import queue
 import openai
 
+def format_openai_host(api_host: str):
+    api_host = api_host.rstrip("/")
+    if not api_host.startswith("http"):
+        api_host = f"https://{api_host}"
+    if api_host.endswith("/v1"):
+        api_host = api_host[:-3]
+    chat_completion_url = f"{api_host}/v1/chat/completions"
+    images_completion_url = f"{api_host}/v1/images/generations"
+    openai_api_base = f"{api_host}/v1"
+    balance_api_url = f"{api_host}/dashboard/billing/credit_grants"
+    usage_api_url = f"{api_host}/dashboard/billing/usage"
+    return chat_completion_url, images_completion_url, openai_api_base, balance_api_url, usage_api_url
+
 class State:
     interrupted = False
     multi_api_key = False
@@ -11,6 +24,7 @@ class State:
     usage_api_url = USAGE_API_URL
     openai_api_base = OPENAI_API_BASE
     images_completion_url = IMAGES_COMPLETION_URL
+    api_host = API_HOST
 
     def interrupt(self):
         self.interrupted = True
@@ -19,23 +33,16 @@ class State:
         self.interrupted = False
 
     def set_api_host(self, api_host: str):
-        api_host = api_host.rstrip("/")
-        if not api_host.startswith("http"):
-            api_host = f"https://{api_host}"
-        if api_host.endswith("/v1"):
-            api_host = api_host[:-3]
-        self.chat_completion_url = f"{api_host}/v1/chat/completions"
-        self.images_completion_url = f"{api_host}/v1/images/generations"
-        self.openai_api_base = f"{api_host}/v1"
-        self.balance_api_url = f"{api_host}/dashboard/billing/credit_grants"
-        self.usage_api_url = f"{api_host}/dashboard/billing/usage"
-        os.environ["OPENAI_API_BASE"] = api_host + "/v1"
+        self.api_host = api_host
+        self.chat_completion_url, self.images_completion_url, self.openai_api_base, self.balance_api_url, self.usage_api_url = format_openai_host(api_host)
+        os.environ["OPENAI_API_BASE"] = self.openai_api_base
 
     def reset_api_host(self):
         self.chat_completion_url = CHAT_COMPLETION_URL
         self.images_completion_url = IMAGES_COMPLETION_URL
         self.balance_api_url = BALANCE_API_URL
         self.usage_api_url = USAGE_API_URL
+        self.api_host = API_HOST
         os.environ["OPENAI_API_BASE"] = f"https://{API_HOST}"
         return API_HOST
 

modules/train_func.py

@@ -144,6 +144,12 @@ def add_to_models():
                 data['extra_models'].append(i)
     else:
         data['extra_models'] = extra_models
+    if 'extra_model_metadata' in data:
+        for i in extra_models:
+            if i not in data['extra_model_metadata']:
+                data['extra_model_metadata'][i] = {"model_name": i, "model_type": "OpenAIVision"}
+    else:
+        data['extra_model_metadata'] = {i: {"model_name": i, "model_type": "OpenAIVision"} for i in extra_models}
     with open('config.json', 'w') as f:
         commentjson.dump(data, f, indent=4)