57 lines
1.8 KiB
Python
57 lines
1.8 KiB
Python
from langchain.chat_models import BaseChatModel
|
|
from pydantic import SecretStr
|
|
|
|
from yuxi.services.model_cache import model_cache
|
|
from yuxi.utils import get_docker_safe_url
|
|
from yuxi.utils.logging_config import logger
|
|
|
|
|
|
def load_chat_model(fully_specified_name: str, **kwargs) -> BaseChatModel:
|
|
if not fully_specified_name:
|
|
raise ValueError("model spec 不能为空")
|
|
|
|
info = model_cache.get_model_info(fully_specified_name)
|
|
if not info:
|
|
available_specs = model_cache.get_all_specs("chat")
|
|
available_ids = [item.spec for item in available_specs[:10]]
|
|
raise ValueError(
|
|
f"Unknown model spec: '{fully_specified_name}'. "
|
|
f"Available chat models ({len(available_specs)}): {available_ids}"
|
|
)
|
|
|
|
if info.model_type != "chat":
|
|
raise ValueError(f"Model {fully_specified_name} is not a chat model (type={info.model_type})")
|
|
|
|
api_key = info.api_key
|
|
base_url = get_docker_safe_url(info.base_url)
|
|
|
|
logger.debug(f"Loading model {fully_specified_name} with provider_type={info.provider_type}")
|
|
|
|
if info.provider_type == "anthropic":
|
|
from langchain_anthropic import ChatAnthropic
|
|
|
|
return ChatAnthropic(
|
|
model=info.model_id,
|
|
api_key=SecretStr(api_key),
|
|
base_url=base_url,
|
|
**kwargs,
|
|
)
|
|
if info.provider_type == "gemini":
|
|
from langchain_google_genai import ChatGoogleGenerativeAI
|
|
|
|
return ChatGoogleGenerativeAI(
|
|
model=info.model_id,
|
|
google_api_key=SecretStr(api_key),
|
|
**kwargs,
|
|
)
|
|
|
|
from langchain_openai import ChatOpenAI
|
|
|
|
return ChatOpenAI(
|
|
model=info.model_id,
|
|
api_key=SecretStr(api_key),
|
|
base_url=base_url,
|
|
stream_usage=True,
|
|
**kwargs,
|
|
)
|