vent_agent_2.0/runtime/business.py

"""业务级准入和停止感知，独立执行与委派共用。"""
import asyncio
from contextlib import contextmanager

from db import business_store
from extensions.registry import load_extensions
from langchain.agents.middleware.types import AgentMiddleware


class BusinessError(ValueError):
    pass


class BusinessStopped(BusinessError):
    pass


class BusinessBusy(BusinessError):
    pass


_counts: dict[str, int] = {}
_active: dict[str, dict[asyncio.Task, int]] = {}


def ensure_enabled(name):
    load_extensions().business(name)
    state = business_store.state(name)
    if not state["local_enabled"]:
        raise BusinessStopped(f"业务智能体已停止：{name}")


@contextmanager
def track(name):
    task = asyncio.current_task()
    group = _active.setdefault(name, {})
    group[task] = group.get(task, 0) + 1
    try:
        yield
    finally:
        if group.get(task, 0) <= 1:
            group.pop(task, None)
        else:
            group[task] -= 1


class BusinessLease:
    def __init__(self, name):
        ensure_enabled(name)
        limit = load_extensions().business(name).max_concurrency
        if _counts.get(name, 0) >= limit:
            raise BusinessBusy(f"业务智能体并发已达上限：{name}")
        self.name, self.released = name, False
        self._tracking = None
        _counts[name] = _counts.get(name, 0) + 1

    def bind(self):
        if self._tracking is None:
            self._tracking = track(self.name)
            self._tracking.__enter__()

    def release(self):
        if not self.released:
            if self._tracking is not None:
                self._tracking.__exit__(None, None, None)
            _counts[self.name] = max(0, _counts.get(self.name, 0) - 1)
            self.released = True


async def stop(name):
    load_extensions().business(name)
    await asyncio.to_thread(business_store.set_state, name, local_enabled=False)
    current = asyncio.current_task()
    targets = [task for task in list(_active.get(name, {})) if task is not current and not task.done()]
    for task in targets:
        task.cancel()
    if targets:
        await asyncio.gather(*targets, return_exceptions=True)


async def start(name):
    load_extensions().business(name)
    await asyncio.to_thread(business_store.set_state, name, local_enabled=True)


class BusinessGateMiddleware(AgentMiddleware):
    def __init__(self, agent_name):
        self.agent_name = agent_name

    @property
    def name(self):
        return f"BusinessGate_{self.agent_name}"

    async def abefore_agent(self, state, runtime):
        await asyncio.to_thread(ensure_enabled, self.agent_name)

    async def awrap_model_call(self, request, handler):
        await asyncio.to_thread(ensure_enabled, self.agent_name)
        with track(self.agent_name):
            return await handler(request)

    async def awrap_tool_call(self, request, handler):
        await asyncio.to_thread(ensure_enabled, self.agent_name)
        with track(self.agent_name):
            return await handler(request)