diff --git a/gateway/model_pool.py b/gateway/model_pool.py index c5508b2..a1bc67d 100644 --- a/gateway/model_pool.py +++ b/gateway/model_pool.py @@ -13,7 +13,7 @@ import json import re import threading from pathlib import Path -from typing import Any, Dict, Optional +from typing import Any, Dict, List, Optional _POOL_PATH = Path(__file__).resolve().parent.parent / "config" / "model_pool.json" @@ -28,6 +28,7 @@ ENTRY_FIELDS = { "id", "name", "tier", "backend", "base_url", "model", "api_key", "price_in", "price_out", "temperature", "max_tokens", "enabled", "provider", "in_hit_price", # 代理层扩展（D-P3）：usage 归一化 / 按命中价选上游 + "capabilities", # 能力位（T-X6）：vision/tools/context_window } # 单价默认值（$/1M tokens）；local 档为 0 @@ -41,6 +42,45 @@ def _empty_pool() -> Dict[str, Any]: } +def _normalize_capabilities(raw: Any) -> Dict[str, Any]: + """能力位归一（T-X6）：缺省全兼容。 + + - vision/tools：缺省 True（未声明 = 不设限，老条目行为不变）； + - context_window：缺省 0 = 不按上下文长度过滤。 + """ + if not isinstance(raw, dict): + return {"vision": True, "tools": True, "context_window": 0} + try: + ctx = int(raw.get("context_window", 0) or 0) + except (TypeError, ValueError): + ctx = 0 + return {"vision": bool(raw.get("vision", True)), + "tools": bool(raw.get("tools", True)), + "context_window": max(0, ctx)} + + +def filter_by_capabilities(entries: List[Dict[str, Any]], *, + need_vision: bool = False, need_tools: bool = False, + min_context_tokens: int = 0) -> List[Dict[str, Any]]: + """能力位硬过滤（T-X6，采纳 cortiq capabilities 过滤）。 + + 依次校验 vision/tools 声明位与 context_window >= 提示+输出预估 + （context_window=0 视为不限）。返回保持原顺序的合格条目子集。 + """ + out: List[Dict[str, Any]] = [] + for e in entries: + cap = e.get("capabilities") or _normalize_capabilities(None) + if need_vision and not cap.get("vision", True): + continue + if need_tools and not cap.get("tools", True): + continue + ctx = int(cap.get("context_window") or 0) + if min_context_tokens > 0 and 0 < ctx < min_context_tokens: + continue + out.append(e) + return out + + class PoolError(ValueError): """池条目/角色配置非法。""" @@ -203,6 +243,7 @@ class PoolStore: max_tokens = int(entry.get("max_tokens", 4096)) except (TypeError, ValueError): max_tokens = 4096 + capabilities = _normalize_capabilities(entry.get("capabilities")) return { "id": eid, "name": str(entry.get("name") or entry.get("model") or eid), @@ -220,6 +261,8 @@ class PoolStore: "provider": (str(entry["provider"]) if entry.get("provider") else "openai"), "in_hit_price": (float(entry["in_hit_price"]) if entry.get("in_hit_price") is not None else round(price_in / 30.0, 6)), + # 能力位（T-X6）：缺省全兼容（不破坏既有条目/选型行为） + "capabilities": capabilities, } @staticmethod diff --git a/gateway/proxy/routes.py b/gateway/proxy/routes.py index 6b95d61..bcfbb22 100644 --- a/gateway/proxy/routes.py +++ b/gateway/proxy/routes.py @@ -274,6 +274,41 @@ def _budget_headers(budget_mode: str) -> Dict[str, str]: return {"X-Budget-Mode": budget_mode} if budget_mode and budget_mode != "normal" else {} +def _request_needs(body: dict) -> Dict[str, Any]: + """从请求体推断能力需求（T-X6）：多模态图片 -> vision；带 tools -> tools； + 上下文需求 = 提示字符/3 + min(max_tokens, 4096)（与预扣估算同口径）。""" + vision = False + for m in body.get("messages") or []: + content = m.get("content") + if isinstance(content, list): + for part in content: + if isinstance(part, dict) and ("image_url" in part + or part.get("type") == "image_url"): + vision = True + prompt_chars = sum(len(str(m.get("content") or "")) + if not isinstance(m.get("content"), list) + else sum(len(str(p.get("text") or "")) for p in m["content"] + if isinstance(p, dict)) + for m in (body.get("messages") or [])) + est_out = min(int(body.get("max_tokens") or 1024), 4096) + return {"vision": vision, + "tools": bool(body.get("tools")), + "min_context_tokens": prompt_chars // 3 + est_out} + + +def _entry_meets(entry: Dict[str, Any], needs: Dict[str, Any]) -> bool: + """条目能力位是否满足请求需求（capabilities 缺省全兼容）。""" + cap = entry.get("capabilities") or {} + if needs["vision"] and not cap.get("vision", True): + return False + if needs["tools"] and not cap.get("tools", True): + return False + ctx = int(cap.get("context_window") or 0) + if needs["min_context_tokens"] > 0 and 0 < ctx < needs["min_context_tokens"]: + return False + return True + + def _fallback_chain(pool, entry: Dict[str, Any], max_total: int = 3) -> List[Dict[str, Any]]: """上游有序降级链（T-X2，采纳 cortiq tier 链思路）。 @@ -367,6 +402,21 @@ async def _run_chat(body: dict, headers: Dict[str, str], ctx: Dict[str, Any], if isinstance(body.get("stream_options"), dict) else False is_stream = bool(body.get("stream")) + # ---- 能力位硬过滤（T-X6）：请求需要 vision/tools/上下文而条目不满足时重定向 ---- + needs = _request_needs(body) + if not _entry_meets(entry, needs): + from gateway.model_pool import filter_by_capabilities + cap_entries = filter_by_capabilities( + [e for e in pool.list().get("entries", []) + if e.get("enabled") and e.get("backend") not in ("mock",) + and e.get("base_url")], + need_vision=needs["vision"], need_tools=needs["tools"], + min_context_tokens=needs["min_context_tokens"]) + if cap_entries: + entry = cap_entries[0] + model = str(entry.get("model") or model) + body = {**body, "model": model} + # ---- 预算四档（T-X1）：接近日上限渐进降档；>100% 仍由 try_hold 硬拒 ---- budget_mode = "normal" try: diff --git a/tests/test_model_pool_capabilities.py b/tests/test_model_pool_capabilities.py new file mode 100644 index 0000000..d6c365c --- /dev/null +++ b/tests/test_model_pool_capabilities.py @@ -0,0 +1,94 @@ +"""模型池能力位过滤测试（T-X6）：归一化缺省 / 硬过滤 / 请求需求推断。""" +import pytest + +from gateway.model_pool import PoolStore, filter_by_capabilities +from gateway.proxy.routes import _entry_meets, _request_needs + +pytest.importorskip("fastapi") + + +def _make_pool(tmp_path) -> PoolStore: + store = PoolStore(path=tmp_path / "pool.json") + store.upsert({"id": "text-only", "name": "纯文本", "tier": "budget", + "backend": "openai", "base_url": "http://a", "model": "text-m", + "capabilities": {"vision": False, "tools": True, + "context_window": 8192}, + "enabled": True}) + store.upsert({"id": "vision", "name": "多模态", "tier": "premium", + "backend": "openai", "base_url": "http://b", "model": "vision-m", + "capabilities": {"vision": True, "tools": True, + "context_window": 32768}, + "enabled": True}) + store.upsert({"id": "legacy", "name": "老条目", "tier": "local", + "backend": "llama_server", "base_url": "http://c", + "model": "legacy-m", "enabled": True}) # 无 capabilities 字段 + return store + + +def test_capabilities_default_full_compat(tmp_path): + """未声明能力位的老条目：缺省全兼容（vision/tools True、context 0 不限）。""" + pool = _make_pool(tmp_path) + e = pool.find_by_model("legacy-m") + assert e["capabilities"] == {"vision": True, "tools": True, "context_window": 0} + + +def test_capabilities_invalid_shapes_fall_back(tmp_path): + pool = _make_pool(tmp_path) + pool.upsert({"id": "weird", "name": "怪", "tier": "budget", "backend": "openai", + "base_url": "http://d", "model": "weird-m", + "capabilities": {"vision": "是", "context_window": "abc"}, + "enabled": True}) + cap = pool.find_by_model("weird-m")["capabilities"] + assert cap["vision"] is True # 非布尔按缺省 + assert cap["context_window"] == 0 # 非整数回落不限 + + +def test_filter_by_capabilities_hard_rules(): + entries = [{"capabilities": {"vision": False, "tools": True, "context_window": 8192}}, + {"capabilities": {"vision": True, "tools": True, "context_window": 32768}}, + {"capabilities": {"vision": True, "tools": False, "context_window": 0}}] + out = filter_by_capabilities(entries, need_vision=True) + assert len(out) == 2 + out = filter_by_capabilities(entries, need_vision=True, need_tools=True) + assert len(out) == 1 + out = filter_by_capabilities(entries, need_vision=True, need_tools=True, + min_context_tokens=16384) + assert len(out) == 1 and out[0]["capabilities"]["context_window"] == 32768 + # context_window=0（不限）不受 min_context 约束 + out = filter_by_capabilities(entries, min_context_tokens=999999) + assert len(out) == 1 and out[0]["capabilities"]["context_window"] == 0 + + +def test_request_needs_detection(): + body = {"messages": [ + {"role": "user", "content": "普通文本问题"}, + {"role": "user", "content": [{"type": "text", "text": "看这张图"}, + {"type": "image_url", + "image_url": {"url": "data:image/png;base64,x"}}]}, + ], "tools": [{"type": "function", "function": {"name": "f"}}], + "max_tokens": 3000} + needs = _request_needs(body) + assert needs["vision"] is True and needs["tools"] is True + assert needs["min_context_tokens"] > 3000 + + plain = _request_needs({"messages": [{"role": "user", "content": "hi"}]}) + assert plain["vision"] is False and plain["tools"] is False + assert plain["min_context_tokens"] > 0 + + +def test_entry_meets_defaults_and_limits(): + legacy = {"capabilities": {"vision": True, "tools": True, "context_window": 0}} + assert _entry_meets(legacy, {"vision": True, "tools": True, + "min_context_tokens": 999999}) + limited = {"capabilities": {"vision": True, "tools": True, "context_window": 4096}} + assert _entry_meets(limited, {"vision": False, "tools": False, + "min_context_tokens": 4096}) + assert not _entry_meets(limited, {"vision": False, "tools": False, + "min_context_tokens": 4097}) + + +def test_pool_list_returns_capabilities(tmp_path): + """list() 打码输出保留 capabilities 字段（前端可见）。""" + pool = _make_pool(tmp_path) + e = next(x for x in pool.list()["entries"] if x["id"] == "vision") + assert e["capabilities"]["vision"] is True diff --git a/任务拆解与执行计划.md b/任务拆解与执行计划.md index 8418616..04c7ad1 100644 --- a/任务拆解与执行计划.md +++ b/任务拆解与执行计划.md @@ -168,3 +168,4 @@ P0 完成后的能力：干净的后端抽象 + 可量化的评测 + 可追溯 | T-X1 | 预算四档渐进干预（外部采纳 ai-model-router）：budget_mode 整数基点判定（80/95/100%）+ optimize/cheap 自动降档 + X-Budget-Mode 上报；黄金用例锁边界 | ✅ 完成 | T-X1 | | T-X2 | 上游有序降级链（外部采纳 cortiq tier 链）：池内候选按档位/单价排序、首 token 前 failover、X-Upstream-Fallback 三元组响应头 + admin/stats failover 聚合 | ✅ 完成 | T-X2 | | T-X3 | 路由决策缓存（外部采纳 cortiq 决策哈希缓存）：DecisionCache（sha256+LRU+TTL60s，4096 条）前置 grader live 模式，embed+线性头去重；观察/审计不跳过、invalidate 同步清空、collect/shadow 不缓存 | ✅ 完成 | T-X3 | +| T-X6 | 模型池能力位过滤（外部采纳 cortiq capabilities）：条目 capabilities{vision,tools,context_window}（缺省全兼容）+ filter_by_capabilities 硬过滤 + 代理请求需求推断（图片→vision/tools/上下文）与重定向 | ✅ 完成 | T-X6 |