zl-q
55 lines
1.3 KiB
Python
55 lines
1.3 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import TYPE_CHECKING, Any, Protocol, TypeVar
|
|
|
|
from ag_ui.core import BaseEvent
|
|
|
|
if TYPE_CHECKING:
|
|
pass
|
|
|
|
T = TypeVar("T", bound=BaseEvent)
|
|
|
|
|
|
class CodecLike(Protocol):
|
|
def to_wire(self, event: dict[str, Any]) -> dict[str, Any]: ...
|
|
|
|
|
|
class StoreLike(Protocol):
|
|
async def persist(self, event: dict[str, Any]) -> None: ...
|
|
|
|
|
|
class BusLike(Protocol):
|
|
async def publish(self, *, session_id: str, event: dict[str, Any]) -> str: ...
|
|
|
|
|
|
def is_base_event(event: Any) -> bool:
|
|
return isinstance(event, BaseEvent)
|
|
|
|
|
|
def to_dict(event: BaseEvent | dict[str, Any]) -> dict[str, Any]:
|
|
if isinstance(event, BaseEvent):
|
|
return event.model_dump(by_alias=True, exclude_none=True)
|
|
return event
|
|
|
|
|
|
class AgentScopeEventPipeline:
|
|
_codec: CodecLike
|
|
_store: StoreLike
|
|
_bus: BusLike
|
|
|
|
def __init__(self, *, codec: CodecLike, store: StoreLike, bus: BusLike) -> None:
|
|
self._codec = codec
|
|
self._store = store
|
|
self._bus = bus
|
|
|
|
async def emit(
|
|
self,
|
|
*,
|
|
session_id: str,
|
|
event: "BaseEvent | dict[str, Any]",
|
|
) -> str:
|
|
event_dict = to_dict(event)
|
|
wire_event = self._codec.to_wire(event_dict)
|
|
await self._store.persist(event_dict)
|
|
return await self._bus.publish(session_id=session_id, event=wire_event)
|