qzl
30 lines
1.0 KiB
Python
30 lines
1.0 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import re
|
|
from typing import Any
|
|
|
|
from ag_ui.core.events import BaseEvent
|
|
from ag_ui.encoder.encoder import EventEncoder
|
|
|
|
_EVENT_TYPE_RE = re.compile(r"^[A-Z0-9_]+$")
|
|
_ENCODER = EventEncoder()
|
|
|
|
|
|
def to_sse_event(stream_id: str, event: dict[str, Any]) -> str:
|
|
safe_stream_id = str(stream_id).replace("\r", "").replace("\n", "")
|
|
try:
|
|
event_model = BaseEvent.model_validate(event)
|
|
event_type = event_model.type.value
|
|
encoded_data = _ENCODER.encode(event_model)
|
|
return f"id: {safe_stream_id}\nevent: {event_type}\n{encoded_data}"
|
|
except Exception: # noqa: BLE001
|
|
raw_event_type = (
|
|
str(event.get("type", "MESSAGE")).replace("\r", "").replace("\n", "")
|
|
)
|
|
event_type = (
|
|
raw_event_type if _EVENT_TYPE_RE.fullmatch(raw_event_type) else "MESSAGE"
|
|
)
|
|
payload = json.dumps(event, ensure_ascii=True, separators=(",", ":"))
|
|
return f"id: {safe_stream_id}\nevent: {event_type}\ndata: {payload}\n\n"
|