请先构建前端：cd webapp && npm run build
") + @app.get("/health", response_model=HealthResponse, tags=["system"]) async def health(): return get_router().health() - @app.get("/", response_class=HTMLResponse, tags=["ui"]) - async def index(): - """端云协同 Web 界面（单文件前端，无需构建）。""" - if _INDEX_PATH.exists(): - return HTMLResponse(_INDEX_PATH.read_text(encoding="utf-8")) - return HTMLResponse("缺少 gateway/static/index.html
") - - # ---------------- v2：/chat ---------------- + # Vue SPA（由 StaticFiles mount 在 / 路径提供） @app.post("/chat", tags=["chat"]) async def chat(req: QueryRequest): - try: - result = await get_pipeline().run(req.query) - except ValueError as e: - raise HTTPException(status_code=400, detail=str(e)) - # 记账 + 入队 - if _v2stats is not None: - _v2stats.record(result) - _maybe_enqueue(result) + """立即返回 request_id，协作管线在后台 asyncio.Task 中运行。 + + 完成后结果写入 runs/{id}/workspace.json（由 pipeline 内部完成）， + 状态通过 GET /runs/{id}/status 查询，SSE 通过 /runs/{id}/stream 订阅。 + """ + import uuid + + request_id = uuid.uuid4().hex[:12] + run_dir = Path("runs") / request_id + run_dir.mkdir(parents=True, exist_ok=True) + + # 注册任务（容量满时拒绝） + ok, msg = get_job_store().register(request_id) + if not ok: + raise HTTPException(status_code=503, detail=msg) + + # 提交后台协程 + async def _run(): + job_store = get_job_store() + try: + pipeline = get_pipeline() + # mode=full 时临时跳过快路径，直接走架构师协作（DeepSeek V4） + if req.mode == "full": + old_fast = pipeline.fast_path + pipeline.fast_path = False + try: + result = await pipeline.run(req.query, request_id=request_id) + finally: + pipeline.fast_path = old_fast + else: + result = await pipeline.run(req.query, request_id=request_id) + if _v2stats is not None: + # 按模型池单价分账（条目未命中时 cost_est 保持管线原值） + entry = get_pool().find_by_model(result.model_used or "") + if entry is not None: + result.cost_est = compute_cost( + entry, result.api_input_tokens, result.api_output_tokens) + _v2stats.record(result) + _maybe_enqueue(result) + job_store._task_done(request_id, result) + except Exception as exc: + import traceback + traceback.print_exc() + job_store._task_failed(request_id, "pipeline_error", detail=str(exc)) + # 写一份失败状态到 workspace.json（pipeline 异常时） + try: + import json + fail_ws = { + "request_id": request_id, + "query": req.query, + "status": "failed", + "error": str(exc), + "brief": None, "plan": [], "progress": [], + "issues": [], "decisions": [], "archive": [], + "meta": {"state": "failed", "round": 0, "rounds_cap": 6, + "api_input_tokens": 0, "api_output_tokens": 0}, + } + (run_dir / "workspace.json").write_text( + json.dumps(fail_ws, ensure_ascii=False, indent=2), encoding="utf-8") + except Exception: + pass + + get_job_store().submit(request_id, _run()) + return {"request_id": request_id, "status": "pending"} + + # ---------------- v2：任务状态 ---------------- + @app.get("/runs/{request_id}/status", tags=["v2"]) + async def get_run_status(request_id: str): + """查询任务当前状态（pending / running / done / failed）。""" + info = get_job_store().get(request_id) + if info is None: + raise HTTPException(status_code=404, detail=f"任务 {request_id} 不存在") + ws_state = None + ws_path = Path("runs") / request_id / "workspace.json" + if ws_path.exists(): + import json + try: + ws_data = json.loads(ws_path.read_text(encoding="utf-8")) + ws_state = ws_data.get("meta", {}).get("state") or ws_data.get("status") + except Exception: + ws_state = "done" return { - "response": result.response, - "request_id": result.request_id, - "status": result.status, - "fast_path": result.fast_path, - "rounds_used": result.rounds_used, - "api_input_tokens": result.api_input_tokens, - "api_output_tokens": result.api_output_tokens, - "cost_est": result.cost_est, - "model_used": result.model_used, - "latency_ms": round(result.latency_ms, 2), - "route": result.route, - "workspace_path": result.workspace_path, - "error": result.error, + "request_id": request_id, + "status": info.state, # pending | running | done | failed + "ws_state": ws_state, + "started_at": info.started_at, + "finished_at": info.finished_at, + "error": info.error, + # PipelineResult 等效字段（扁平存储在 TaskInfo 中） + "response": info.response, + "pipeline_status": info.status, + "fast_path": info.fast_path, + "rounds_used": info.rounds_used, + "api_input_tokens": info.api_input_tokens, + "api_output_tokens": info.api_output_tokens, + "cost_est": info.cost_est, + "model_used": info.model_used, + "latency_ms": round(info.latency_ms, 2) if info.latency_ms else None, + "route": info.route, + "workspace_path": info.workspace_path, } + # ---------------- v2：SSE 实时流 ---------------- + @app.get("/runs/{request_id}/stream", tags=["v2"]) + async def stream_run(request_id: str): + """SSE 端点：实时推送 workspace.json 状态变化（供前端协作可视化）。""" + from fastapi.responses import StreamingResponse + + async def event_generator(): + import json # noqa: F401 + cancel_event = get_job_store().new_cancel(request_id) + try: + async for ev in get_job_store().watch_workspace(request_id, cancel_event): + payload = json.dumps(ev, ensure_ascii=False) + yield f"data: {payload}\n\n" + finally: + cancel_event.set() + + return StreamingResponse( + event_generator(), + media_type="text/event-stream", + headers={ + "Cache-Control": "no-cache", + "Connection": "keep-alive", + "X-Accel-Buffering": "no", + }, + ) + # ---------------- v1：/chat/legacy ---------------- @app.post("/chat/legacy", tags=["chat"]) async def chat_legacy(req: QueryRequest): @@ -251,6 +390,225 @@ try: raise HTTPException(status_code=404, detail=f"审核记录不存在或已审核: {review_id}") return {"ok": True, "review_id": review_id, "verdict": verdict} + # ---------------- 模型池（多价位异构模型） ---------------- + @app.get("/pool", tags=["pool"]) + async def pool_list(): + """读取模型池（角色指派 + 条目列表，api_key 打码）。""" + return get_pool().list() + + @app.post("/pool", tags=["pool"]) + async def pool_upsert(entry: dict): + """新增或更新池条目（按 id）。api_key 留空表示保留原值。""" + try: + get_pool().upsert(entry) + except ValueError as e: + raise HTTPException(status_code=400, detail=f"条目非法: {e}") + rebuild_pipeline() + return get_pool().list() + + @app.delete("/pool/{entry_id}", tags=["pool"]) + async def pool_delete(entry_id: str): + ok = get_pool().delete(entry_id) + if ok: + rebuild_pipeline() + return {"ok": ok, **get_pool().list()} + + @app.put("/pool/roles", tags=["pool"]) + async def pool_roles(roles: dict): + """指派角色：{"architect": "请构建前端：cd webapp && npm run build
") + + # ---------------- metrics API（JSON REST 端点）-------------------- + # 前端通过 axios GET /api/metrics 调用此端点获取指标数据 + @app.get("/metrics", include_in_schema=False) + async def metrics_page(): + return await metrics() + + @app.get("/api/metrics", tags=["system"]) async def metrics(): out = { "router": get_router().stats.summary(), diff --git a/gateway/model_pool.py b/gateway/model_pool.py new file mode 100644 index 0000000..9c73b75 --- /dev/null +++ b/gateway/model_pool.py @@ -0,0 +1,270 @@ +"""模型池（PoolStore）—— 多价位异构模型注册表。 + +设计（《实现方案_v4_模型池与工具智能体.md》D1）： +- 叙事从"端云分工"泛化为"按价位分工"：local（零边际成本，内置 llama.cpp）、 + budget（低价 API）、premium（高价 API）。位置只是价位的属性之一。 +- 池条目存"端点 + 凭据 + 模型名 + 价位 + 单价（$/1M tokens）"，不存模型权重。 +- roles 把池条目指派给三个角色：architect（决策/终审）、worker（实现/自验证）、 + agent（智能体工具循环）。角色留空 = 沿用经典单模型设置（向后兼容）。 +- 持久化到 config/model_pool.json（gitignore，与 settings.json 同级）。 +""" + +import json +import re +import threading +from pathlib import Path +from typing import Any, Dict, Optional + +_POOL_PATH = Path(__file__).resolve().parent.parent / "config" / "model_pool.json" + +# 合法取值 +TIERS = ("local", "budget", "premium") +BACKENDS = ("mock", "llama_server", "openai") +ROLES = ("architect", "worker", "agent") + +# 池条目允许的字段（其余字段拒绝写入） +ENTRY_FIELDS = { + "id", "name", "tier", "backend", "base_url", "model", "api_key", + "price_in", "price_out", "temperature", "max_tokens", "enabled", +} + +# 单价默认值（$/1M tokens）；local 档为 0 +PRICE_DEFAULTS = {"local": 0.0, "budget": 0.1, "premium": 1.0} + + +def _empty_pool() -> Dict[str, Any]: + return { + "roles": {"architect": "", "worker": "", "agent": ""}, + "entries": [], + } + + +class PoolError(ValueError): + """池条目/角色配置非法。""" + + +class PoolStore: + """模型池注册表（内存 + model_pool.json 持久化，线程安全）。""" + + def __init__(self, path: Optional[Path] = None): + self._path = Path(path) if path else _POOL_PATH + self._lock = threading.Lock() + self._data = _empty_pool() + self.load() + + # ---------- 持久化 ---------- + def load(self) -> None: + if self._path.exists(): + try: + raw = json.loads(self._path.read_text(encoding="utf-8")) + self._data = { + "roles": {**_empty_pool()["roles"], + **(raw.get("roles") or {})}, + "entries": list(raw.get("entries") or []), + } + except Exception: + self._data = _empty_pool() + else: + self._data = _empty_pool() + + def save(self) -> None: + self._path.parent.mkdir(parents=True, exist_ok=True) + self._path.write_text( + json.dumps(self._data, ensure_ascii=False, indent=2), encoding="utf-8") + + # ---------- 条目 CRUD ---------- + def list(self) -> Dict[str, Any]: + """返回完整池（api_key 打码）。""" + with self._lock: + return { + "roles": dict(self._data["roles"]), + "entries": [self._masked(e) for e in self._data["entries"]], + } + + def get(self, entry_id: str) -> Optional[Dict[str, Any]]: + with self._lock: + for e in self._data["entries"]: + if e.get("id") == entry_id: + return dict(e) + return None + + def upsert(self, entry: Dict[str, Any]) -> Dict[str, Any]: + """新增或更新条目（按 id）。返回打码后的条目。""" + clean = self._validate(entry) + with self._lock: + entries = self._data["entries"] + for i, e in enumerate(entries): + if e.get("id") == clean["id"]: + # 空 api_key 表示保留原值（前端不回传明文） + if not clean.get("api_key"): + clean["api_key"] = e.get("api_key", "") + entries[i] = clean + self.save() + return self._masked(clean) + entries.append(clean) + self.save() + return self._masked(clean) + + def delete(self, entry_id: str) -> bool: + with self._lock: + before = len(self._data["entries"]) + self._data["entries"] = [ + e for e in self._data["entries"] if e.get("id") != entry_id] + changed = len(self._data["entries"]) != before + if changed: + # 清空指向被删条目的角色指派 + for role, rid in self._data["roles"].items(): + if rid == entry_id: + self._data["roles"][role] = "" + self.save() + return changed + + # ---------- 角色指派 ---------- + def set_roles(self, roles: Dict[str, str]) -> Dict[str, str]: + """指派角色 -> 池条目 id（空串 = 沿用经典设置）。""" + with self._lock: + ids = {e.get("id") for e in self._data["entries"]} + for role, rid in roles.items(): + if role not in ROLES: + raise PoolError(f"未知角色: {role}") + if rid and rid not in ids: + raise PoolError(f"角色 {role} 指向不存在的模型条目: {rid}") + self._data["roles"][role] = rid or "" + self.save() + return dict(self._data["roles"]) + + def resolve(self, role: str) -> Optional[Dict[str, Any]]: + """解析角色当前生效的池条目（未指派/条目禁用时返回 None = 用经典设置）。""" + if role not in ROLES: + return None + with self._lock: + rid = self._data["roles"].get(role, "") + for e in self._data["entries"]: + if e.get("id") == rid: + return dict(e) if e.get("enabled", True) else None + return None + + def find_by_model(self, model: str) -> Optional[Dict[str, Any]]: + """按模型名找条目（用于按模型计价分账）。""" + with self._lock: + for e in self._data["entries"]: + if e.get("model") == model: + return dict(e) + return None + + # ---------- 校验与工具 ---------- + def _validate(self, entry: Dict[str, Any]) -> Dict[str, Any]: + if not isinstance(entry, dict): + raise PoolError("条目必须是对象") + unknown = set(entry) - ENTRY_FIELDS + if unknown: + raise PoolError(f"非法字段: {sorted(unknown)}") + eid = str(entry.get("id") or "").strip() + if not eid: + # 未提供 id 时按名称生成 slug + base = re.sub(r"[^a-zA-Z0-9_-]+", "-", + str(entry.get("name") or entry.get("model") or "model")).strip("-").lower() + eid = base or "model" + with self._lock: + exist = {e.get("id") for e in self._data["entries"]} + if eid in exist: + n = 2 + while f"{eid}-{n}" in exist: + n += 1 + eid = f"{eid}-{n}" + elif not re.fullmatch(r"[a-zA-Z0-9_-]{1,64}", eid): + raise PoolError("id 只允许字母/数字/-/_，长度 1-64") + backend = entry.get("backend", "openai") + if backend not in BACKENDS: + raise PoolError(f"backend 必须是 {BACKENDS} 之一") + tier = entry.get("tier", "budget") + if tier not in TIERS: + raise PoolError(f"tier 必须是 {TIERS} 之一") + if backend != "mock" and not str(entry.get("base_url") or "").strip(): + raise PoolError("非 mock 后端必须填写 base_url") + if backend != "mock" and not str(entry.get("model") or "").strip(): + raise PoolError("非 mock 后端必须填写 model") + try: + price_in = float(entry.get("price_in", PRICE_DEFAULTS[tier])) + price_out = float(entry.get("price_out", PRICE_DEFAULTS[tier])) + except (TypeError, ValueError): + raise PoolError("price_in/price_out 必须是数字") + if price_in < 0 or price_out < 0: + raise PoolError("单价不能为负") + try: + temperature = float(entry.get("temperature", 0.3)) + except (TypeError, ValueError): + temperature = 0.3 + try: + max_tokens = int(entry.get("max_tokens", 4096)) + except (TypeError, ValueError): + max_tokens = 4096 + return { + "id": eid, + "name": str(entry.get("name") or entry.get("model") or eid), + "tier": tier, + "backend": backend, + "base_url": str(entry.get("base_url") or "").strip(), + "model": str(entry.get("model") or "").strip(), + "api_key": str(entry.get("api_key") or ""), + "price_in": price_in, + "price_out": price_out, + "temperature": temperature, + "max_tokens": max_tokens, + "enabled": bool(entry.get("enabled", True)), + } + + @staticmethod + def _masked(entry: Dict[str, Any]) -> Dict[str, Any]: + out = dict(entry) + key = out.get("api_key") or "" + out["api_key_set"] = bool(key) + out["api_key"] = (key[:6] + "…") if key else "" + return out + + +def entry_to_architect_cfg(entry: Dict[str, Any]) -> Dict[str, Any]: + """池条目 -> build_architect 配置段。""" + cfg: Dict[str, Any] = { + "model": entry.get("model") or "local", + "base_url": entry.get("base_url") or "http://127.0.0.1:8901/v1", + "temperature": float(entry.get("temperature", 0.2)), + } + if entry.get("api_key"): + cfg["api_key"] = entry["api_key"] + if entry.get("max_tokens"): + cfg["max_tokens"] = int(entry["max_tokens"]) + return cfg + + +def entry_to_worker_cfg(entry: Dict[str, Any]) -> Dict[str, Any]: + """池条目 -> build_worker 配置段。""" + return { + "backend": entry.get("backend") or "openai", + "base_url": entry.get("base_url") or "", + "model": entry.get("model") or "", + "temperature": float(entry.get("temperature", 0.3)), + } + + +def compute_cost(entry: Dict[str, Any], input_tokens: int, output_tokens: int) -> float: + """按条目单价估算成本（USD）。price 单位：$/1M tokens。""" + return (input_tokens / 1e6) * float(entry.get("price_in", 0.0)) + \ + (output_tokens / 1e6) * float(entry.get("price_out", 0.0)) + + +# ---------- 全局单例 ---------- +_store: Optional[PoolStore] = None + + +def get_pool() -> PoolStore: + global _store + if _store is None: + _store = PoolStore() + return _store + + +def reset_pool() -> None: + """测试用：重置全局池单例。""" + global _store + _store = None diff --git a/gateway/settings.py b/gateway/settings.py index 9600af8..c836a38 100644 --- a/gateway/settings.py +++ b/gateway/settings.py @@ -22,12 +22,13 @@ DEFAULTS: Dict[str, Any] = { "port": 8901, # llama_server 端口 "temperature": 0.3, "max_fix_attempts": 2, + "per_step_timeout_s": 15, "code_timeout_s": 10, }, "architect": { - "model": "deepseek-chat", - "base_url": "https://api.deepseek.com/v1", - "api_key_env": "DEEPSEEK_API_KEY", + "model": "deepseek-v4-flash", + "base_url": "https://api.deepseek.com", + "api_key": "", }, "pipeline": { "fast_path": True, @@ -35,6 +36,11 @@ DEFAULTS: Dict[str, Any] = { "api_token_cap": 8000, "breach_policy": "architect_do", }, + "agent": { + "workspace_dir": "agent_workspace", # 智能体工作区根目录（越界即拒） + "max_rounds": 8, # 工具循环轮数上限 + "token_cap": 20000, # 单次智能体任务 token 熔断 + }, } diff --git a/router_system/v2stats.py b/router_system/v2stats.py index 6044192..7a6f832 100644 --- a/router_system/v2stats.py +++ b/router_system/v2stats.py @@ -22,6 +22,7 @@ class V2Stats: self._api_input_tokens = 0 self._api_output_tokens = 0 self._api_cost_usd = 0.0 + self._by_model: Dict[str, Dict[str, Any]] = {} self._recent: List[Dict[str, Any]] = [] def record(self, result) -> None: @@ -38,6 +39,17 @@ class V2Stats: self._api_input_tokens += getattr(result, "api_input_tokens", 0) self._api_output_tokens += getattr(result, "api_output_tokens", 0) self._api_cost_usd += getattr(result, "cost_est", 0.0) + # 按模型分账（token / 成本 / 次数） + model = getattr(result, "model_used", None) or "unknown" + in_tok = getattr(result, "api_input_tokens", 0) + out_tok = getattr(result, "api_output_tokens", 0) + bucket = self._by_model.setdefault(model, { + "requests": 0, "input_tokens": 0, "output_tokens": 0, "cost_est_usd": 0.0, + }) + bucket["requests"] += 1 + bucket["input_tokens"] += in_tok + bucket["output_tokens"] += out_tok + bucket["cost_est_usd"] = round(bucket["cost_est_usd"] + getattr(result, "cost_est", 0.0), 6) self._recent.append({ "request_id": getattr(result, "request_id", ""), "status": status, @@ -69,6 +81,10 @@ class V2Stats: "total": self._api_input_tokens + self._api_output_tokens, "cost_est_usd": round(self._api_cost_usd, 6), }, + "by_model": { + m: {**b, "cost_est_usd": round(b["cost_est_usd"], 6)} + for m, b in self._by_model.items() + }, } diff --git a/tests/test_agent_api.py b/tests/test_agent_api.py new file mode 100644 index 0000000..8a8c9a6 --- /dev/null +++ b/tests/test_agent_api.py @@ -0,0 +1,170 @@ +"""智能体端点测试：注入脚本化 chat_fn，不依赖真实模型/API key。""" +import json +import time + +import pytest + +pytest.importorskip("fastapi") +pytest.importorskip("httpx") + +from fastapi.testclient import TestClient + +import gateway.agent as ag +import gateway.api as ga +from gateway.model_pool import PoolStore +import gateway.model_pool as mp + + +@pytest.fixture() +def agent_env(tmp_path, monkeypatch): + """隔离：池/服务/工作区全部指向临时目录，chat_fn 用脚本替身。""" + mp.reset_pool() + mp._store = PoolStore(path=tmp_path / "pool.json") + ag.reset_agent_service() + service = ag.AgentService(run_dir=tmp_path / "agent_runs") + ag._service = service + # 快照用户真实设置，测试后原样恢复（settings.json 是活文件，不能污染） + store = ga.settings_store() + snapshot = json.loads(json.dumps(store._data, ensure_ascii=False)) + # 工作区指向临时目录 + 给经典回退一个假 key（防止 .env 缺失时 400） + store.update({"agent": {"workspace_dir": str(tmp_path / "ws")}, + "architect": {"api_key": "sk-fake-test"}}) + ga.rebuild_pipeline() + + script = [] + + def set_script(events): + script.clear() + script.extend(events) + + def fake_chat_factory(acfg): + async def chat_fn(messages, tools_spec): + if not script: + return {"content": "（脚本用尽）好的。", "tool_calls": [], "usage": {}} + return script.pop(0) + return chat_fn + + monkeypatch.setattr(ga, "build_agent_chat", fake_chat_factory) + yield {"service": service, "set_script": set_script, "ws": tmp_path / "ws"} + + store._data = snapshot + store.save() + mp.reset_pool() + ag.reset_agent_service() + ga.rebuild_pipeline() + + +@pytest.fixture() +def client(): + return TestClient(ga.app) + + +def _wait_done(service, rid, timeout=10.0): + t0 = time.time() + while time.time() - t0 < timeout: + info = service.get(rid) + if info and info.state in ("done", "failed"): + return info + time.sleep(0.05) + return service.get(rid) + + +def test_agent_full_flow(agent_env, client): + """写文件 -> 最终答复：验证事件、工作区落盘、状态终态。""" + agent_env["set_script"]([ + {"content": None, + "tool_calls": [{"id": "c1", "name": "write_file", + "arguments": {"path": "notes.md", "content": "# 笔记"}}], + "usage": {"prompt_tokens": 30, "completion_tokens": 6}}, + {"content": "已创建 notes.md，任务完成。", "tool_calls": [], + "usage": {"prompt_tokens": 40, "completion_tokens": 8}}, + ]) + r = client.post("/agent", json={"task": "帮我建一个 notes.md"}) + assert r.status_code == 200 + rid = r.json()["request_id"] + assert r.json()["status"] == "running" + + info = _wait_done(agent_env["service"], rid) + assert info.state == "done", info.error + assert "notes.md" in info.response + + # 工作区真实落盘 + assert (agent_env["ws"] / "notes.md").read_text(encoding="utf-8") == "# 笔记" + + # 事件序列 + events = client.get(f"/agent/{rid}/events").json() + kinds = [e["type"] for e in events] + assert "tool_call" in kinds and "tool_result" in kinds and "final" in kinds + assert events[-1]["reason"] == "answer" + + # 状态端点 + st = client.get(f"/agent/{rid}/status").json() + assert st["state"] == "done" + assert st["prompt_tokens"] == 70 and st["completion_tokens"] == 14 + + # 工作区浏览端点 + ls = client.get("/agent/workspace").json() + assert ls["ok"] is True + assert any(e["name"] == "notes.md" for e in ls["entries"]) + f = client.get("/agent/file", params={"path": "notes.md"}).json() + assert f["content"] == "# 笔记" + + +def test_agent_jail_via_api(agent_env, client): + """工具结果为 ok=False（越界被拒），循环仍能继续到最终答复。""" + agent_env["set_script"]([ + {"content": None, + "tool_calls": [{"id": "c1", "name": "read_file", + "arguments": {"path": "../../secret.txt"}}], + "usage": {"prompt_tokens": 10, "completion_tokens": 2}}, + {"content": "越界访问被拒绝。", "tool_calls": [], "usage": {}}, + ]) + r = client.post("/agent", json={"task": "读一下上级目录"}) + rid = r.json()["request_id"] + info = _wait_done(agent_env["service"], rid) + assert info.state == "done" + events = agent_env["service"].read_events(rid) + tool_result = next(e for e in events if e["type"] == "tool_result") + assert tool_result["ok"] is False + + # 文件读取 API 直接越界 -> 404/400 + r2 = client.get("/agent/file", params={"path": "../../x.txt"}) + assert r2.status_code in (400, 404) + + +def test_agent_model_from_pool(agent_env, client, monkeypatch): + """池 agent 角色（或显式 pool_id）应被采用；mock 池模型拒绝。""" + from gateway.agent import OpenAICompatChat + captured = {} + real_factory = None + + # 先放一个 openai 池条目并指派 agent 角色 + client.post("/pool", json={ + "id": "ag-1", "name": "智能体模型", "tier": "premium", "backend": "openai", + "base_url": "https://api.example.com", "model": "big-model-x", + "api_key": "sk-abc1234567", "enabled": True, + }) + client.put("/pool/roles", json={"agent": "ag-1"}) + + # /agent 不带 pool_id -> 用池 agent 角色 + r = client.post("/agent", json={"task": "hi"}) + assert r.status_code == 200 + assert r.json()["model"] == "big-model-x" + + # mock 条目 -> 400 + client.post("/pool", json={ + "id": "mk-1", "name": "mock", "tier": "local", "backend": "mock", + "model": "mock", "enabled": True, + }) + r2 = client.post("/agent", json={"task": "hi", "pool_id": "mk-1"}) + assert r2.status_code == 400 + + +def test_agent_task_validation(agent_env, client): + assert client.post("/agent", json={"task": ""}).status_code == 400 + assert client.post("/agent", json={}).status_code == 400 + + +def test_agent_404(agent_env, client): + assert client.get("/agent/ghost/status").status_code == 404 + assert client.get("/agent/ghost/events").json() == [] diff --git a/tests/test_model_pool.py b/tests/test_model_pool.py new file mode 100644 index 0000000..adf311c --- /dev/null +++ b/tests/test_model_pool.py @@ -0,0 +1,171 @@ +"""模型池（PoolStore）测试：条目校验/CRUD/角色指派/管线解析/成本分账。""" +import pytest + +pytest.importorskip("fastapi") + +from fastapi.testclient import TestClient + +import gateway.api as ga +import gateway.model_pool as mp +from gateway.model_pool import PoolStore, compute_cost, entry_to_architect_cfg, entry_to_worker_cfg + + +@pytest.fixture() +def pool(tmp_path): + """独立文件的全局池（不污染 config/model_pool.json）。""" + mp.reset_pool() + store = PoolStore(path=tmp_path / "model_pool.json") + mp._store = store + yield store + mp.reset_pool() + ga.rebuild_pipeline() + + +@pytest.fixture() +def client(): + return TestClient(ga.app) + + +def _entry(**over): + base = { + "id": "prem-1", "name": "旗舰模型", "tier": "premium", "backend": "openai", + "base_url": "https://api.deepseek.com", "model": "deepseek-v4-pro", + "api_key": "sk-test-1234567890", "price_in": 1.0, "price_out": 2.0, + "enabled": True, + } + base.update(over) + return base + + +# ---------------- PoolStore 单元 ---------------- + +def test_pool_upsert_and_mask(pool): + masked = pool.upsert(_entry()) + assert masked["api_key_set"] is True + assert "sk-test" not in masked["api_key"] # 明文不打回 + data = pool.list() + assert data["entries"][0]["model"] == "deepseek-v4-pro" + assert data["entries"][0]["api_key_set"] is True + + +def test_pool_upsert_keeps_key_when_blank(pool): + pool.upsert(_entry()) + pool.upsert(_entry(api_key="")) # 前端不回传明文 -> 保留 + assert pool.get("prem-1")["api_key"] == "sk-test-1234567890" + + +def test_pool_validation(pool): + with pytest.raises(ValueError): + pool.upsert(_entry(tier="超豪华")) + with pytest.raises(ValueError): + pool.upsert(_entry(backend="magic")) + with pytest.raises(ValueError): + pool.upsert(_entry(backend="openai", base_url="")) # 非 mock 缺端点 + with pytest.raises(ValueError): + pool.upsert(_entry(hack="x")) # 未知字段 + with pytest.raises(ValueError): + pool.upsert(_entry(price_in=-1)) + + +def test_pool_roles_and_resolve(pool): + pool.upsert(_entry()) + pool.upsert(_entry(id="local-1", tier="local", backend="llama_server", + base_url="http://127.0.0.1:8901/v1", model="qwen3.5-4b", + price_in=0, price_out=0)) + assert pool.resolve("architect") is None # 未指派 + pool.set_roles({"architect": "prem-1", "worker": "local-1"}) + assert pool.resolve("architect")["id"] == "prem-1" + assert pool.resolve("worker")["id"] == "local-1" + assert pool.resolve("agent") is None + # 指派不存在的条目 + with pytest.raises(ValueError): + pool.set_roles({"agent": "ghost"}) + # 删除条目 -> 角色自动清空 + pool.delete("prem-1") + assert pool.resolve("architect") is None + + +def test_pool_disabled_entry_not_resolved(pool): + pool.upsert(_entry(enabled=False)) + pool.set_roles({"architect": "prem-1"}) + assert pool.resolve("architect") is None # 禁用 -> 回退经典设置 + + +def test_entry_cfg_mapping(pool): + e = pool.get("prem-1") or _entry() + acfg = entry_to_architect_cfg(_entry()) + assert acfg["model"] == "deepseek-v4-pro" + assert acfg["api_key"] == "sk-test-1234567890" + wcfg = entry_to_worker_cfg(_entry()) + assert wcfg["backend"] == "openai" + + +def test_compute_cost(): + e = {"price_in": 1.0, "price_out": 2.0} + assert compute_cost(e, 1_000_000, 500_000) == pytest.approx(2.0) + assert compute_cost({"price_in": 0, "price_out": 0}, 999, 999) == 0.0 + + +# ---------------- API 端点 ---------------- + +def test_pool_api_crud(pool, client): + r = client.get("/pool") + assert r.status_code == 200 + assert r.json()["roles"]["architect"] == "" + r2 = client.post("/pool", json=_entry()) + assert r2.status_code == 200 + assert len(r2.json()["entries"]) == 1 + # 非法条目 -> 400 + r3 = client.post("/pool", json=_entry(tier="bad")) + assert r3.status_code == 400 + # 角色指派 + r4 = client.put("/pool/roles", json={"architect": "prem-1"}) + assert r4.status_code == 200 + assert r4.json()["roles"]["architect"] == "prem-1" + # 删除 + r5 = client.delete("/pool/prem-1") + assert r5.status_code == 200 + assert r5.json()["roles"]["architect"] == "" + + +def test_build_pipeline_uses_pool(pool, monkeypatch): + """池指派应覆盖经典设置，测试 override 最后生效。""" + pool.upsert(_entry()) + pool.set_roles({"architect": "prem-1"}) + ga.rebuild_pipeline() + + captured = {} + + def fake_build_architect(cfg): + captured["architect"] = dict(cfg) + from router_system.architect import ArchitectClient + return ArchitectClient(model=cfg.get("model", "m"), api_key="k") + + monkeypatch.setattr(ga, "build_architect", fake_build_architect) + pipe = ga.build_v2_pipeline(worker_cfg_override={"backend": "mock"}) + assert pipe is not None + assert captured["architect"]["model"] == "deepseek-v4-pro" # 池条目生效 + ga.rebuild_pipeline() + + +def test_v2stats_by_model(): + from router_system.v2stats import V2Stats + + class R: + request_id = "x" + fast_path = False + status = "done" + rounds_used = 1 + api_input_tokens = 1000 + api_output_tokens = 500 + cost_est = 0.002 + model_used = "deepseek-v4-pro" + route = [] + + s = V2Stats() + s.record(R()) + summary = s.summary() + bucket = summary["by_model"]["deepseek-v4-pro"] + assert bucket["requests"] == 1 + assert bucket["input_tokens"] == 1000 + assert bucket["cost_est_usd"] == pytest.approx(0.002)