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, name: str = "main") -> None:
        self.name = name
        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() if self.name == "main" else build_agent(self.name)

    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()
_slots: dict[str, _AgentSlot] = {"main": _slot}
_slots_lock = threading.Lock()


def _named_slot(name: str) -> _AgentSlot:
    if name != "main":
        from extensions.registry import load_extensions
        load_extensions().business(name)
    with _slots_lock:
        return _slots.setdefault(name, _AgentSlot(name))


def get_agent(name: str = "main") -> Any:
    """按名称取主 Agent 或业务 Agent，分别单飞编译和缓存。"""
    return _named_slot(name).get()


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


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

    if name is not None:
        return _named_slot(name).invalidate()
    model_factory.clear_cache()
    with _slots_lock:
        slots = list(_slots.values())
    for slot in slots:
        slot.invalidate()
    return _slot.generation


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