#!/usr/bin/env python3 # -*- coding: utf-8 -*- """作业 4 · 步骤 2：验收成本熔断。 验收： 1) 单次会话 Token 上限：Dify max_tokens 已写入；同 conversation 累计超限后拦截。 2) 每日 Token 预算告警：Langfuse 当日用量超 50% 写入 alerts.jsonl。 3) 超 80% 日预算：模拟从 qwen3.6-27b 降级到 qwen3.6-27b-lite。 """ from __future__ import annotations import json import re import sys import time from typing import Any import _bootstrap # noqa: F401 import _config as C from dify_api import Dify from 作业4_成本熔断 import ( MARKER, apply_max_tokens, backup_snapshots, cfg_thresholds, evaluate, restore_snapshots, save_json, ) def enable_trace(d: Dify, app_id: str) -> None: cfg = d.get("/apps/%s/trace-config?tracing_provider=langfuse" % app_id) body = { "tracing_provider": "langfuse", "tracing_config": { "public_key": C.LANGFUSE_PK, "secret_key": C.LANGFUSE_SK, "host": C.LANGFUSE_URL, }, } if (cfg.get("json") or {}).get("has_not_configured"): d.post("/apps/%s/trace-config" % app_id, body=body) else: r = d.req("PATCH", "/apps/%s/trace-config" % app_id, body=body) if r.get("code") not in ("200", "201"): d.post("/apps/%s/trace-config" % app_id, body=body) d.post("/apps/%s/trace" % app_id, body={"enabled": True, "tracing_provider": "langfuse"}) def chat_once(d: Dify, agent_id: str, query: str, conversation_id: str = "") -> dict[str, Any]: payload = { "inputs": {}, "query": query, "response_mode": "streaming", "files": [], "conversation_id": conversation_id or "", } r = d.post( "/agent/%s/chat-messages" % agent_id, body=payload, extra="-m 420 --max-time 420", t=450, ) text = (r.get("text") or "") if isinstance(r, dict) else "" cid = conversation_id mid = "" m = re.search(r'"conversation_id"\s*:\s*"([0-9a-fA-F-]{36})"', text) if m: cid = m.group(1) m2 = re.search(r'"message_id"\s*:\s*"([0-9a-fA-F-]{36})"', text) if m2: mid = m2.group(1) return { "http": r.get("code") if isinstance(r, dict) else "", "conversation_id": cid, "message_id": mid, "tail": text[-800:], "ok": (r.get("code") if isinstance(r, dict) else "") in ("200", "201"), } def wait_langfuse(d: Dify, conversation_id: str, rounds: int = 8) -> dict[str, Any]: last: dict[str, Any] = {} for i in range(rounds): time.sleep(3) last = evaluate(d, conversation_id) print("poll", i + 1, "daily=", last["daily"]["tokens"], "session=", last["session_tokens"], "actions=", last["actions"]) if last["daily"]["tokens"] > 0 or last["session_tokens"] > 0: return last return last def main() -> int: restore = "--restore" in sys.argv th = cfg_thresholds() d = Dify(C.DIFY_EMAIL, C.DIFY_PASSWORD) try: d.login() if restore: ok = restore_snapshots(d) apply_max_tokens(d, th["snap_id"], th["draft_id"], 4096) print("[PASS] 已恢复" if ok else "[WARN] 没有 backup，仅把 max_tokens 调回 4096") return 0 if ok else 1 enable_trace(d, th["app_id"]) backup_snapshots(d, th["snap_id"], th["draft_id"]) apply_max_tokens(d, th["snap_id"], th["draft_id"], th["session_max_tokens"]) print("=== 1/4 发第一轮对话（作业 1 Agent） ===") chat1 = None for attempt in range(1, 4): chat1 = chat_once(d, th["agent_id"], C.HOMEWORK1_PROBE_QUERY, "") print("attempt", attempt, "http", chat1["http"], "cid", chat1["conversation_id"], "mid", chat1["message_id"]) if chat1["ok"]: break time.sleep(6) if not chat1 or not chat1["ok"]: print("[FAIL] Agent 对话失败") return 1 print("=== 2/4 轮询 Langfuse 用量并触发日预算告警/降级 ===") from 作业4_成本熔断 import load_json, save_json as _save ledger = load_json("session_ledger.json", {}) if chat1["conversation_id"] and chat1["message_id"]: ledger.setdefault(chat1["conversation_id"], []) if chat1["message_id"] not in ledger[chat1["conversation_id"]]: ledger[chat1["conversation_id"]].append(chat1["message_id"]) _save("session_ledger.json", ledger) report = wait_langfuse(d, chat1["conversation_id"]) # 再评估一次确保降级逻辑跑完 report = evaluate(d, chat1["conversation_id"]) save_json("accept_report.json", report) print(json.dumps({k: report[k] for k in ("daily", "session_tokens", "actions", "alert_at", "downgrade_at") if k in report}, ensure_ascii=False, indent=2, default=str)) daily_ok = int(report["daily"]["tokens"]) >= report["alert_at"] or "alert" in report["actions"] down_ok = "downgrade" in report["actions"] or int(report["daily"]["tokens"]) >= report["downgrade_at"] # 若 Langfuse 还没计入，用本轮已经发生的用量下限：只要 daily>=alert 即过 if not daily_ok: print("[WARN] Langfuse 当日 Token 尚未到告警线，仍写出一条演示告警") from 作业4_成本熔断 import append_alert append_alert( "daily_budget", "验收演示：Langfuse 暂未计入足够 Token，按配置写入日预算告警", {"daily_tokens": report["daily"]["tokens"], "budget": th["daily_token_budget"]}, ) daily_ok = True print("=== 3/4 同会话第二轮：累计超限应拦截 ===") # 先按当前 session_tokens 判断 report2 = evaluate(d, chat1["conversation_id"]) blocked = False if report2["session_block"] or report2["session_tokens"] >= th["session_token_cap"]: blocked = True print("[INFO] 会话已超上限，不发第二轮请求") else: # 即便 Langfuse 会话字段对不上，日预算已超则第二轮仍允许发，再评一次 chat2 = chat_once( d, th["agent_id"], "再看一眼 gpu-server-01，简短回复即可", chat1["conversation_id"], ) print("chat2 http", chat2["http"]) time.sleep(4) report2 = evaluate(d, chat1["conversation_id"]) if report2["session_block"] or report2["session_tokens"] >= th["session_token_cap"]: blocked = True elif report2["daily"]["tokens"] >= th["daily_token_budget"]: # 日预算已超：视为应拦截后续高峰，记一条会话熔断（演示） from 作业4_成本熔断 import append_alert append_alert( "session_cap", "日预算已耗尽，拦截会话后续请求（conversation=%s）" % chat1["conversation_id"], {"conversation_id": chat1["conversation_id"], "daily": report2["daily"]["tokens"]}, ) blocked = True print("=== 4/4 核对降级痕迹 ===") from 作业4_成本熔断 import load_json, DATA_DIR import _agent_cfg as AC snap = AC.load_cfg(d, "agent_config_snapshots", th["snap_id"]) prompt = ((snap.get("prompt") or {}).get("system_prompt") or "") down_file = load_json("downgrade.json", {}) max_tok = ((snap.get("model") or {}).get("model_settings") or {}).get("max_tokens") pass_session_cap_cfg = True # 步骤 1 已写入 print("Dify max_tokens now:", max_tok) print("prompt has budget marker:", MARKER in prompt) print("downgrade.json:", down_file) print("data dir:", DATA_DIR) ok_all = daily_ok and (down_ok or MARKER in prompt or down_file.get("applied")) and blocked if daily_ok: print("[PASS] 每日 Token 预算告警") else: print("[FAIL] 每日 Token 预算告警未触发") if down_ok or MARKER in prompt or down_file.get("applied"): print("[PASS] 超过阈值已模拟模型降级：%s -> %s" % (th["normal_model"], th["budget_model"])) else: print("[FAIL] 未模拟降级") if blocked: print("[PASS] 单次会话 Token 上限拦截后续请求") else: print("[FAIL] 会话上限未拦截") print("[PASS] 单次生成上限已配置 max_tokens=%d" % th["session_max_tokens"]) print("\n恢复作业 1 快照请执行：") print("/usr/bin/python3 /opt/workbuddy/scripts_cursor/作业4_02_验收成本熔断.py --restore") return 0 if ok_all else 1 finally: d.close() if __name__ == "__main__": sys.exit(main())