vent_agent_2.0/agents/agent_registry.py

"""agent 单例与失效重建。

为什么要单例：``create_deep_agent`` 要编译整张 LangGraph 图，耗时在百毫秒
到秒级；每次请求都编译会让首 token 时间不可接受。

为什么要失效重建：图一旦编译，模型实例、思考级别、middleware 栈、
compression 阈值都固化了。改了这些配置只能重建，不能原地改。

「单飞编译」（single-flight）：并发请求同时发现缓存为空时，只让第一个
真正去编译，其余等它完成——否则冷启动瞬间会同时编译 N 次。
"""

from __future__ import annotations

import asyncio
import threading
from typing import Any

from obs import get_logger

logger = get_logger(__name__)


class _AgentSlot:
    """持有编译好的 agent 及其代数（generation）。

    代数用于判断「我拿到的是不是已经被失效过的那一版」——失效后旧实例
    可能仍在服务在途请求，不能强行打断，只能让新请求拿新的。
    """

    def __init__(self) -> None:
        self._agent: Any | None = None
        self._generation = 0
        self._lock = threading.Lock()

    @property
    def generation(self) -> int:
        return self._generation

    def peek(self) -> Any | None:
        """快路径：已编译则直接返回，否则 None（不触发编译）。"""
        return self._agent

    def get(self) -> Any:
        """取 agent，必要时编译。并发调用只编译一次。"""
        agent = self._agent
        if agent is not None:
            return agent

        with self._lock:
            # 二次检查：可能在等锁期间已被别的线程编译好
            if self._agent is None:
                logger.debug("首次编译 agent（generation=%d）", self._generation + 1)
                self._agent = self._compile()
                self._generation += 1
                logger.info("agent 编译完成（generation=%d）", self._generation)
            return self._agent

    def _compile(self) -> Any:
        from agents.agent_factory import build_agent

        return build_agent()

    def invalidate(self) -> int:
        """丢弃当前实例，返回新代数。下次 ``get()`` 会重新编译。"""
        with self._lock:
            self._agent = None
            self._generation += 1
            logger.debug("agent 已失效，下次请求将重建（generation=%d）", self._generation)
            return self._generation


_slot = _AgentSlot()


def get_agent() -> Any:
    """取主 agent（必要时编译）。"""
    return _slot.get()


def peek_agent() -> Any | None:
    """快路径取 agent；未编译返回 None。热路径上先用它避免持锁。"""
    return _slot.peek()


def invalidate() -> int:
    """使当前 agent 失效。改模型/思考级别/压缩配置后调用。"""
    from agents import model_factory

    model_factory.clear_cache()
    return _slot.invalidate()


def generation() -> int:
    """当前代数。"""
    return _slot.generation


async def prewarm() -> bool:
    """后台预热：把耗时的图编译挪到线程池，不阻塞事件循环。

    启动时调用。失败不抛——预热失败应当在首次真实请求时暴露，
    而不是让整个服务起不来。
    """
    try:
        await asyncio.to_thread(get_agent)
        return True
    except Exception as exc:
        logger.exception("agent 预热失败，将在首次请求时重试", extra={"kv": {"event": "agent.prewarm_failed"}})
        return False