diff --git a/src/models/chat_model.py b/src/models/chat_model.py index 49c1ca85..548369be 100644 --- a/src/models/chat_model.py +++ b/src/models/chat_model.py @@ -4,8 +4,11 @@ from src.utils import logger, get_docker_safe_url class OpenAIBase(): def __init__(self, api_key, base_url, model_name): + self.api_key = api_key + self.base_url = base_url self.client = OpenAI(api_key=api_key, base_url=base_url) self.model_name = model_name + # logger.debug(f"{self.get_models()=}") def predict(self, message, stream=False): if isinstance(message, str): @@ -34,6 +37,13 @@ class OpenAIBase(): stream=False, ) return response.choices[0].message + + def get_models(self): + try: + return self.client.models.list() + except Exception as e: + logger.error(f"Error getting models: {e}") + return [] class OpenModel(OpenAIBase): @@ -60,7 +70,8 @@ class GeneralResponse: self.is_full = False -class Qianfan: +class Qianfan(OpenAIBase): + """弃用""" def __init__(self, model_name="ernie_speed") -> None: import qianfan @@ -99,7 +110,7 @@ class Qianfan: -class DashScope: +class DashScope(OpenAIBase): def __init__(self, model_name="qwen-max-latest") -> None: self.model_name = model_name @@ -143,6 +154,4 @@ class DashScope: if __name__ == "__main__": - model = SiliconFlow() - for a in model.predict("你好", stream=True): - print(a.content, end="") \ No newline at end of file + pass \ No newline at end of file diff --git a/src/routers/chat_router.py b/src/routers/chat_router.py index 16f97ba3..35aa5157 100644 --- a/src/routers/chat_router.py +++ b/src/routers/chat_router.py @@ -21,6 +21,7 @@ def chat_post( history: list = Body(...), cur_res_id: str = Body(...)): + meta["server_model_name"] = startup.model.model_name history_manager = HistoryManager(history) logger.debug(f"Received query: {query} with meta: {meta}") diff --git a/web/src/components/RefsComponent.vue b/web/src/components/RefsComponent.vue index 9a081fa8..e79504cc 100644 --- a/web/src/components/RefsComponent.vue +++ b/web/src/components/RefsComponent.vue @@ -3,7 +3,7 @@
- {{ msg.model_name }} + {{ msg.meta.server_model_name }}