62 lines
1.8 KiB
Python
62 lines
1.8 KiB
Python
import os
|
||
import traceback
|
||
|
||
from langchain.chat_models import BaseChatModel
|
||
from pydantic import SecretStr
|
||
|
||
from src import config
|
||
from src.utils import get_docker_safe_url
|
||
|
||
|
||
def load_chat_model(fully_specified_name: str, **kwargs) -> BaseChatModel:
|
||
"""
|
||
Load a chat model from a fully specified name.
|
||
"""
|
||
provider, model = fully_specified_name.split("/", maxsplit=1)
|
||
|
||
assert provider != "custom", "[弃用] 自定义模型已移除，请在 src/config/static/models.py 中配置"
|
||
|
||
model_info = config.model_names.get(provider)
|
||
if not model_info:
|
||
raise ValueError(f"Unknown model provider: {provider}")
|
||
|
||
env_var = model_info.env
|
||
|
||
api_key = os.getenv(env_var, env_var)
|
||
|
||
base_url = get_docker_safe_url(model_info.base_url)
|
||
|
||
if provider in ["deepseek", "dashscope"]:
|
||
from langchain_deepseek import ChatDeepSeek
|
||
|
||
return ChatDeepSeek(
|
||
model=model,
|
||
api_key=SecretStr(api_key),
|
||
base_url=base_url,
|
||
api_base=base_url,
|
||
stream_usage=True,
|
||
)
|
||
|
||
# elif provider == "together":
|
||
# from langchain_together import ChatTogether
|
||
|
||
# return ChatTogether(
|
||
# model=model,
|
||
# api_key=SecretStr(api_key),
|
||
# base_url=base_url,
|
||
# stream_usage=True,
|
||
# )
|
||
|
||
else:
|
||
try: # 其他模型，默认使用OpenAIBase, like openai, zhipuai
|
||
from langchain_openai import ChatOpenAI
|
||
|
||
return ChatOpenAI(
|
||
model=model,
|
||
api_key=SecretStr(api_key),
|
||
base_url=base_url,
|
||
stream_usage=True,
|
||
)
|
||
except Exception as e:
|
||
raise ValueError(f"Model provider {provider} load failed, {e} \n {traceback.format_exc()}")
|