zk
73 lines
2.3 KiB
Python
73 lines
2.3 KiB
Python
"""LLM 模型枚举与实例获取
|
||
|
||
Usage:
|
||
from app.ai.models import LLM
|
||
|
||
llm = LLM.DOUBAO_PRO_256K.create(temperature=0)
|
||
llm = LLM.DEEPSEEK_V3.create()
|
||
"""
|
||
|
||
from enum import Enum
|
||
|
||
from langchain_openai import ChatOpenAI
|
||
|
||
from app.config import settings
|
||
|
||
# 供应商连接配置
|
||
_VOLCENGINE = (lambda: settings.volcengine_api_key, lambda: settings.volcengine_base_url)
|
||
_JIAYU = (lambda: settings.jiayu_api_key, lambda: settings.jiayu_base_url)
|
||
_JIEKOU = (lambda: settings.jiekou_api_key, lambda: settings.jiekou_base_url)
|
||
|
||
_ZM = (lambda: settings.zm_api_key, lambda: settings.zm_base_url)
|
||
|
||
|
||
class LLM(Enum):
|
||
"""所有可用模型，每个枚举值 = (模型名, api_key函数, base_url函数)"""
|
||
|
||
# 火山引擎
|
||
|
||
DOUBAO_PRO_32K = ("doubao-1-5-pro-32k-250115", *_VOLCENGINE)
|
||
DOUBAO_LITE_32K = ("doubao-1-5-lite-32k-250115", *_VOLCENGINE)
|
||
DEEPSEEK_V3 = ("deepseek-v3-250324", *_VOLCENGINE)
|
||
DEEPSEEK_R1 = ("deepseek-r1-250528", *_VOLCENGINE)
|
||
DOUBAO_SEED_MINI = ("doubao-seed-2-0-mini-260215", *_VOLCENGINE)
|
||
DOUBAO_SEED_LITE = ("doubao-seed-2-0-lite-260215", *_VOLCENGINE)
|
||
DOUBAO_SEED_PRO = ("doubao-seed-2-0-pro-260215", *_VOLCENGINE)
|
||
|
||
# jiekou
|
||
GPT_4O = ("gpt-4o", *_JIAYU)
|
||
GPT_4O_MINI = ("gpt-4o-mini", *_JIEKOU)
|
||
GEMINI_FLASH = ("gemini-2.5-flash", *_JIEKOU)
|
||
|
||
# 加鱼
|
||
JIAYU_CLAUDE_SONNET_4_5 = ("claude-sonnet-4.5", *_JIAYU)
|
||
JIAYU_CLAUDE_HAIKU_4_5 = ("claude-haiku-4.5", *_JIAYU)
|
||
JIAYU_DEEPSEEK_3_2 = ("deepseek-3.2", *_JIAYU)
|
||
JIAYU_GLM_5 = ("glm-5", *_JIAYU)
|
||
JIAYU_QWEN3_CODER_NEXT = ("qwen3-coder-next", *_JIAYU)
|
||
JIAYU_MINIMAX_M2_5 = ("minimax-m2.5", *_JIAYU)
|
||
|
||
|
||
# ZM
|
||
ZM_GPT_5_5 = ("gpt-5.5", *_ZM)
|
||
ZM_GPT_5_4 = ("gpt-5.4", *_ZM)
|
||
ZM_GPT_5_4_MINI = ("gpt-5.4-mini", *_ZM)
|
||
ZM_GPT_5_2 = ("gpt-5.2", *_ZM)
|
||
ZM_GPT_4O = ("gpt-5.5", *_ZM)
|
||
|
||
|
||
|
||
def __init__(self, model_name: str, api_key_fn, base_url_fn):
|
||
self.model_name = model_name
|
||
self._api_key_fn = api_key_fn
|
||
self._base_url_fn = base_url_fn
|
||
|
||
def create(self, **kwargs) -> ChatOpenAI:
|
||
"""创建 LLM 实例，kwargs 透传给 ChatOpenAI（temperature, max_tokens 等）"""
|
||
return ChatOpenAI(
|
||
model=self.model_name,
|
||
api_key=self._api_key_fn(),
|
||
base_url=self._base_url_fn(),
|
||
**kwargs,
|
||
)
|