vent_agent_2.0/agents/execution.py

"""聊天、审批恢复和定时执行共用的准入与运行装配。"""
import asyncio
from dataclasses import dataclass

from api import concurrency
from billing import guard
from db import chat_store


@dataclass
class Admission:
    session: concurrency.SessionLease
    slot: concurrency.TaskSlot

    def release(self):
        self.slot.release()
        self.session.release()


async def admit(user_name: str, session_id: str) -> Admission:
    session = concurrency.acquire_session(session_id)
    slot = None
    try:
        slot = concurrency.acquire_task_slot(user_name)
        if not await asyncio.to_thread(chat_store.get_session_for_user, session_id, user_name):
            raise concurrency.SessionMissingError("会话不存在")
        await asyncio.to_thread(guard.precheck, user_name)
        return Admission(session, slot)
    except BaseException:
        if slot:
            slot.release()
        session.release()
        raise


async def build_stream(admission: Admission, *, agent=None, **kwargs):
    from agents import agent_registry
    from api.sse_core import ChatStream, StreamRequest
    try:
        if agent is None:
            agent = await asyncio.to_thread(agent_registry.get_agent)
        graph_config = kwargs["config"]
        graph_config.setdefault("configurable", {}).update(
            workspace_user=kwargs["user_name"], workspace_session=kwargs["session_id"],
        )
        return ChatStream(StreamRequest(agent=agent, slot=admission.slot,
                         session_lease=admission.session, **kwargs))
    except BaseException:
        admission.release()
        raise