《AI Agent 工程实战》18.2 故障演练：模型超时/DB 宕机/队列积压

# 三个演练共用的公共定义（卷一的类型契约，见 2.3.2）
from dataclasses import dataclass, field
from typing import Any, Literal

Role = Literal["system", "user", "assistant", "tool"]

@dataclass
class ToolCall:
    id: str
    name: str
    arguments: dict[str, Any]

@dataclass
class Message:
    role: Role
    content: str = ""
    tool_calls: list[ToolCall] = field(default_factory=list)
    tool_call_id: str | None = None

@dataclass
class ModelResponse:
    message: Message
    stop_reason: str = "end_turn"
    usage: dict[str, int] = field(default_factory=dict)

# 接 18.2.1 的公共定义之后运行
import concurrent.futures as cf
import time


class SlowModel:
    """假模型：每次 complete 睡 delay 秒，模拟上游卡住。"""
    def __init__(self, name: str, delay: float) -> None:
        self.name, self.delay = name, delay

    def complete(self, messages, tools) -> ModelResponse:
        time.sleep(self.delay)
        return ModelResponse(message=Message(role="assistant", content=f"[{self.name}] 正常回答"))


class InstantModel:
    def __init__(self, name: str, text: str) -> None:
        self.name, self.text = name, text

    def complete(self, messages, tools) -> ModelResponse:
        return ModelResponse(message=Message(role="assistant", content=self.text))


_POOL = cf.ThreadPoolExecutor(max_workers=8)  # 常驻池：超时不阻塞后续调用


def call_with_timeout(model, messages, tools, timeout: float) -> ModelResponse:
    fut = _POOL.submit(model.complete, messages, tools)
    try:
        return fut.result(timeout=timeout)
    except cf.TimeoutError:
        raise TimeoutError(f"{model.name} 超过 {timeout:.1f}s 未返回") from None


class TimeoutRouter:
    """带单次调用超时的降级链：超时/失败即换下一个后端。"""
    def __init__(self, backends, timeout: float) -> None:
        self.backends, self.timeout = backends, timeout
        self.trace: list[tuple[str, str]] = []

    def complete(self, messages, tools) -> ModelResponse:
        last = None
        for b in self.backends:
            t0 = time.perf_counter()
            try:
                resp = call_with_timeout(b, messages, tools, self.timeout)
            except Exception as e:
                self.trace.append((b.name, f"STOP {type(e).__name__}  {time.perf_counter()-t0:.2f}s"))
                last = e
                continue
            self.trace.append((b.name, f"OK  {time.perf_counter()-t0:.2f}s"))
            return resp
        raise last


msgs = [Message(role="user", content="GIL 下 CPU 密集怎么并发？")]
print("== 无超时保护：直接调主模型 ==")
t0 = time.perf_counter()
r = SlowModel("primary", 5.0).complete(msgs, [])
print(f"返回: {r.message.content}  耗时 {time.perf_counter()-t0:.2f}s")

print("\n== 有超时 + 降级：timeout=1.0s ==")
router = TimeoutRouter(
    [SlowModel("primary", 5.0),
     InstantModel("fallback-haiku", "（降级）主模型超时，改由备用模型作答。")],
    timeout=1.0)
t0 = time.perf_counter()
resp = router.complete(msgs, [])
print("最终回答:", resp.message.content)
for name, outcome in router.trace:
    print(f"  {name:<16} {outcome}")
print(f"请求总耗时 {time.perf_counter()-t0:.2f}s  成功: {resp.message.content != ''}")

== 无超时保护：直接调主模型 ==
返回: [primary] 正常回答  耗时 5.00s

== 有超时 + 降级：timeout=1.0s ==
最终回答: （降级）主模型超时，改由备用模型作答。
  primary          STOP TimeoutError  1.01s
  fallback-haiku   OK  0.00s
请求总耗时 1.01s  成功: True

# 接 18.2.1 的公共定义之后运行（本演练不依赖 Message 等，可独立跑）
import time
import psycopg

GOOD = "postgresql://postgres:postgres@localhost:15435/research"
BAD = "postgresql://postgres:postgres@localhost:15436/research"  # 端口无人监听


def probe(conninfo: str, timeout: int = 2) -> tuple[bool, str | None, float]:
    t0 = time.perf_counter()
    try:
        with psycopg.connect(conninfo, connect_timeout=timeout) as c:
            with c.cursor() as cur:
                cur.execute("SELECT 1")
                cur.fetchone()
        return True, None, time.perf_counter() - t0
    except Exception as e:
        head = f"{type(e).__name__}: {str(e).splitlines()[0]}"
        return False, head, time.perf_counter() - t0


def handle_request(ready: bool, ready_err: str | None) -> tuple[int, dict]:
    """就绪检查失败时立即返回 503，而不是挂死。"""
    if not ready:
        return 503, {"error": "service_unavailable",
                     "detail": "数据库暂不可达，请稍后重试",
                     "cause": ready_err.split(":")[0]}
    return 200, {"status": "ok"}


print("== 健康态：连真实容器（15435）==")
ok, err, dt = probe(GOOD)
print(f"ready={ok}  error={err}  elapsed={dt:.2f}s")

print("\n== 故障态：连不存在的端口（15436）==")
ok, err, dt = probe(BAD)
print(f"ready={ok}  elapsed={dt:.2f}s")
print(f"error: {err}")

print("\n== 探针摘流量后的请求行为 ==")
for label, (ok, err, _) in (("DB 正常", probe(GOOD)), ("DB 宕机", probe(BAD))):
    code, body = handle_request(ok, err)
    print(f"{label:<8} -> HTTP {code}  {body}")

== 健康态：连真实容器（15435）==
ready=True  error=None  elapsed=0.06s

== 故障态：连不存在的端口（15436）==
ready=False  elapsed=0.00s
error: OperationalError: connection failed: connection to server at "127.0.0.1", port 15436 failed: could not receive data from server: Connection refused

== 探针摘流量后的请求行为 ==
DB 正常    -> HTTP 200  {'status': 'ok'}
DB 宕机    -> HTTP 503  {'error': 'service_unavailable', 'detail': '数据库暂不可达，请稍后重试', 'cause': 'OperationalError'}

# 接 18.2.1 的公共定义之后运行（本演练用 Redis + asyncio，可独立跑）
import asyncio
import time
import redis

r = redis.Redis.from_url("redis://localhost:16380/0", decode_responses=True)
QKEY = "ab:drill18:queue"
ALERT_THRESHOLD, PRODUCE, BATCH, BATCH_LATENCY = 100, 500, 50, 0.05

r.delete(QKEY)
for i in range(PRODUCE):                       # ① 生产者：一口气塞 500 个
    r.rpush(QKEY, f"run-{i}")
depth = r.llen(QKEY)
print(f"入队 {PRODUCE} 个任务，队列深度 = {depth}")
print(f"告警阈值 = {ALERT_THRESHOLD}  触发告警 = {depth > ALERT_THRESHOLD}")

consumed, t0 = 0, time.perf_counter()          # ② 消费者：每批 50 个，慢于生产
while consumed < PRODUCE:
    for _ in range(BATCH):
        if r.lpop(QKEY) is None:
            break
        consumed += 1
    time.sleep(BATCH_LATENCY)
print(f"消费完 {consumed} 个任务耗时 {time.perf_counter()-t0:.2f}s，剩余深度 = {r.llen(QKEY)}")


async def bounded_demo() -> None:              # ③ 进程内背压：有界队列
    q: asyncio.Queue = asyncio.Queue(maxsize=20)
    peak = 0

    async def producer():
        nonlocal peak
        for i in range(200):
            await q.put(i)                     # 队列满则阻塞 —— 背压
            peak = max(peak, q.qsize())
        await q.put(None)

    async def consumer():
        while True:
            item = await q.get()
            if item is None:
                break
            await asyncio.sleep(0.002)

    await asyncio.gather(producer(), consumer())
    print(f"有界队列 maxsize=20，生产 200 个，峰值深度 = {peak}（被钉死在 maxsize）")


print()
asyncio.run(bounded_demo())
r.delete(QKEY)                                 # ④ 清理
print(f"清理完成，队列键存在 = {bool(r.exists(QKEY))}")

入队 500 个任务，队列深度 = 500
告警阈值 = 100  触发告警 = True
消费完 500 个任务耗时 2.37s，剩余深度 = 0

有界队列 maxsize=20，生产 200 个，峰值深度 = 20（被钉死在 maxsize）
清理完成，队列键存在 = False