98 lines
4.3 KiB
Python
98 lines
4.3 KiB
Python
"""
|
|
能力(capability)注册中心
|
|
|
|
集中定义平台支持的全部模型能力，以及"每个能力对应的默认模型"的读写与解析。
|
|
- CAPABILITY_TYPES: 全部能力（id / 中文标签 / 测试类型 / 描述）
|
|
- get_capability_defaults / set_capability_default: 读写 site_settings 中的 model_capability_defaults
|
|
- resolve_auto_model: 根据能力解析平台默认模型（供聊天 / 技能 / 应用 / IM Bot 自动匹配）
|
|
"""
|
|
from typing import Optional
|
|
from sqlalchemy.orm import Session
|
|
from .models import AIModel, SiteSetting
|
|
import json
|
|
|
|
|
|
# ============================================================
|
|
# 平台支持的全部能力类型
|
|
# test_type 用于模型连接测试时选择对应的 API 路径
|
|
# ============================================================
|
|
CAPABILITY_TYPES = [
|
|
{"id": "text", "label": "文本对话", "test_type": "text", "desc": "文本生成与多轮对话"},
|
|
{"id": "thinking", "label": "推理思考", "test_type": "text", "desc": "深度推理 / 思维链"},
|
|
{"id": "vision", "label": "视觉理解", "test_type": "vision", "desc": "图片 / 多模态理解"},
|
|
{"id": "tool", "label": "工具调用", "test_type": "text", "desc": "Function Calling 函数调用"},
|
|
{"id": "image", "label": "图像生成", "test_type": "image", "desc": "文生图 / 图像生成"},
|
|
{"id": "audio", "label": "语音合成", "test_type": "audio", "desc": "TTS 文字转语音"},
|
|
{"id": "asr", "label": "语音识别", "test_type": "asr", "desc": "ASR 语音转文字"},
|
|
{"id": "ocr", "label": "OCR识别", "test_type": "ocr", "desc": "图片文字识别提取"},
|
|
{"id": "embedding", "label": "嵌入向量", "test_type": "embedding", "desc": "文本向量化 / 嵌入"},
|
|
{"id": "rerank", "label": "重排序", "test_type": "rerank", "desc": "语义重排 / Rerank"},
|
|
{"id": "video", "label": "视频生成", "test_type": "video", "desc": "文生视频 / 视频生成"},
|
|
]
|
|
|
|
# 文本类能力（对话 / 视觉 / 推理 / 工具 / OCR）在缺少专属默认模型时，回退到全局默认模型
|
|
TEXT_FAMILY = {"text", "thinking", "vision", "tool", "ocr"}
|
|
|
|
# 仅用于前端展示的固定配色
|
|
CAPABILITY_COLORS = {
|
|
"text": "#5B5FED", "thinking": "#8B5CF6", "vision": "#219B5E", "tool": "#3B82F6",
|
|
"image": "#D9383A", "audio": "#E8922C", "asr": "#0EA5E9", "ocr": "#F59E0B",
|
|
"embedding": "#10B981", "rerank": "#EC4899", "video": "#EF4444",
|
|
}
|
|
|
|
SETTING_KEY = "model_capability_defaults"
|
|
|
|
|
|
def get_capability_defaults(db: Session) -> dict:
|
|
"""读取 {capability: model_id} 映射"""
|
|
row = db.query(SiteSetting).filter(SiteSetting.key == SETTING_KEY).first()
|
|
if not row or not row.value:
|
|
return {}
|
|
try:
|
|
data = json.loads(row.value)
|
|
return data if isinstance(data, dict) else {}
|
|
except Exception:
|
|
return {}
|
|
|
|
|
|
def set_capability_default(db: Session, capability: str, model_id: Optional[str]) -> dict:
|
|
"""设置某个能力的默认模型（model_id 为空则清除）"""
|
|
data = get_capability_defaults(db)
|
|
if model_id:
|
|
data[capability] = model_id
|
|
else:
|
|
data.pop(capability, None)
|
|
row = db.query(SiteSetting).filter(SiteSetting.key == SETTING_KEY).first()
|
|
if not row:
|
|
from uuid import uuid4
|
|
row = SiteSetting(id=str(uuid4()), key=SETTING_KEY, value="{}")
|
|
db.add(row)
|
|
row.value = json.dumps(data, ensure_ascii=False)
|
|
db.commit()
|
|
return data
|
|
|
|
|
|
def resolve_auto_model(capability: str, db: Session) -> Optional[AIModel]:
|
|
"""
|
|
根据能力解析平台默认模型：
|
|
1. 优先使用后台"按能力默认模型"中配置的值
|
|
2. 文本类能力若未单独配置，回退到全局默认模型(is_default)
|
|
3. 其余能力(embedding/rerank/asr 等)若未配置则返回 None
|
|
"""
|
|
data = get_capability_defaults(db)
|
|
mid = data.get(capability)
|
|
if mid:
|
|
m = db.query(AIModel).filter(
|
|
AIModel.id == mid, AIModel.is_enabled == True
|
|
).first()
|
|
if m:
|
|
return m
|
|
# 文本类能力回退到全局默认模型
|
|
if capability in TEXT_FAMILY:
|
|
m = db.query(AIModel).filter(
|
|
AIModel.is_default == True, AIModel.is_enabled == True
|
|
).first()
|
|
if m:
|
|
return m
|
|
return None
|