Files
qihuaiflow/backend/python/app/capability_registry.py
T
QihuAI Design d3d8d62b94
CI / gate (push) Has been cancelled
fix(media): repair JSON_CONTAINS auto-routing + reconcile stale backend tests
Real bug:
- capability_registry.resolve_auto_model and chat_media._build_media_route used
  `Column.op("JSON_CONTAINS")(col, val)`, but SQLAlchemy's op() accepts only ONE
  argument -> TypeError on every image/video/audio auto-route, silently failing
  media generation (caught as "后台生成异常"). Replaced with func.json_contains().

Test reconciliation (no production behavior change):
- tests/test_context_engine.py: tiktoken upgrade made the hardcoded len//2 token
  estimates wrong; relax to tokenizer-agnostic assertions and add
  structured_memory=None to the _Summary mock (code now branches on it).
- tests/test_regression_r1r2r3.py: R3 FakeModel lacked `name` (code reads m.name);
  also the assertion assumed a literal '超时' string, but the failover architecture
  catches model-level exceptions and writes a user-understandable error marker
  instead. Align the assertion to the real contract: a media failure/timeout must
  be written back as a user-understandable error message, never left as '生成中…'.

CI gate: 28 backend pytest + 15 frontend vitest all green.
2026-07-23 09:16:31 +08:00

109 lines
4.9 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. 媒体类能力(image/video/audio)若未配置,自动发现第一个启用且具备该能力的模型
4. 其余能力(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
# 媒体类能力(image/video/audio):无显式配置时自动发现任一启用的模型
# 避免用户未在后台配置"能力默认模型"时 Auto 回退到文本模型导致"无法生成图片/视频"
if capability in ("image", "video", "audio"):
from sqlalchemy import text as _text, func
m = db.query(AIModel).filter(
AIModel.is_enabled == True,
func.json_contains(AIModel.capabilities, f'"{capability}"'),
).first()
if m:
return m
return None