62 lines
2.2 KiB
Python
62 lines
2.2 KiB
Python
from typing import Sequence
|
|
|
|
import openai
|
|
from openai import AsyncOpenAI
|
|
from sqlalchemy.orm import joinedload
|
|
from sqlmodel import Session, select
|
|
|
|
from nyahome.database import AiiModel
|
|
|
|
|
|
def apply_get_models(session: Session) -> list[dict]:
|
|
"""
|
|
从数据库中获取可用的 AI 模型列表。
|
|
|
|
Args:
|
|
session: 数据库连接对象。
|
|
|
|
Returns:
|
|
|
|
"""
|
|
aii_models: Sequence[AiiModel] = session.exec(select(AiiModel).options(joinedload(AiiModel.aii_provider))).all() # type: ignore[arg-type]
|
|
|
|
final_model_list = []
|
|
for aii_model in aii_models:
|
|
final_model_list.append({
|
|
"id": aii_model.id,
|
|
"model_name": aii_model.model_name,
|
|
"max_context_length": aii_model.max_context_length,
|
|
"provider_id": aii_model.aii_provider_id,
|
|
"provider_name": aii_model.aii_provider.name,
|
|
"base_url": aii_model.aii_provider.base_url,
|
|
"reasonable": bool(aii_model.reasonable), # 数据库中的 reasonable 字段可能为 None，在这里归一为 False
|
|
})
|
|
|
|
return final_model_list
|
|
|
|
|
|
async def s_list_remote_provider_models(base_url: str, api_key: str) -> list[dict]:
|
|
client = AsyncOpenAI(base_url=base_url, api_key=api_key)
|
|
try:
|
|
models = await client.models.list()
|
|
final_model_list = []
|
|
async for model in models:
|
|
# model 实际上是 pydantic 模型，因此拥有 BaseModel 的所有方法。
|
|
# model.model_dump() 的示例结果：
|
|
# {'id': 'xxx', 'created': None, 'object': 'model', 'owned_by': 'xxx'}
|
|
final_model_list.append(model.model_dump())
|
|
return final_model_list
|
|
except Exception as e:
|
|
raise TypeError(f"获取模型提供商 {base_url} 的可用模型列表失败。") from e
|
|
|
|
|
|
async def s_check_remote_model(model_name: str, base_url: str, api_key: str) -> bool:
|
|
client = AsyncOpenAI(base_url=base_url, api_key=api_key)
|
|
try:
|
|
await client.models.retrieve(model_name)
|
|
return True
|
|
except openai.NotFoundError:
|
|
return False
|
|
except Exception as e:
|
|
raise TypeError(f"从模型提供商 {base_url} 检测模型 {model_name} 可用性时遇到未知错误") from e
|