1"""

2 Streaming SSE (Server-Sent Events) endpoint for agent interactions.

3

4 Provides real-time streaming of agent outputs via HTTP SSE, enabling

5 browser-based chat UIs and real-time monitoring dashboards.

6"""

7

8 from __future__ import annotations

9

10 import asyncio

11 import json

12 import time

13 from collections import defaultdict

14 from dataclasses import dataclass, field

15 from typing import Any, AsyncIterator, Optional

16

17

18@dataclass

19 class StreamEvent:

20 """Single SSE event emitted by the stream."""

21

22 event: str

23 """Event type: 'chunk', 'tool_call', 'tool_result', 'done', 'error'."""

24

25 data: dict[str, Any]

26 """Event payload as JSON-serializable dict."""

27

28 id: Optional[str] = None

29 """Optional event ID for resume support."""

30

31 retry: Optional[int] = None

32 """Reconnection retry interval in milliseconds."""

33

34 def to_sse(self) -> str:

35 """Format as SSE wire format."""

36 lines: list[str] = []

37 if self.id:

38 lines.append(f"id: {self.id}")

39 if self.event:

40 lines.append(f"event: {self.event}")

41 lines.append(f"data: {json.dumps(self.data, ensure_ascii=False)}")

42 if self.retry:

43 lines.append(f"retry: {self.retry}")

44 lines.append("") # blank line terminates event

45 return "\n".join(lines)

46

47

48@dataclass

49 class StreamSession:

50 """Track an active streaming session."""

51

52 session_id: str

53 started_at: float = field(default_factory=time.time)

54 events_emitted: int = 0

55 last_event_at: float = 0.0

56 metadata: dict[str, Any] = field(default_factory=dict)

57

58

59 class StreamingAgent:

60 """

61 Agent that emits Server-Sent Events for real-time streaming.

62

63 Example (FastAPI integration)::

64

65 streaming = StreamingAgent(agent_loop)

66

67 @app.get("/agent/stream")

68 async def stream():

69 return StreamingResponse(

70 streaming.stream_chat("What is quantum computing?", "session-1"),

71 media_type="text/event-stream"

72 )

73 """

74

75 def __init__(

76 self,

77 agent_loop: Any = None,

78 heartbeat_interval: float = 15.0,

79 ):

80 """

81 Args:

82 agent_loop: The underlying agent loop (sync or async).

83 heartbeat_interval: Seconds between heartbeat keepalive events.

84 """

85 self._loop = agent_loop

86 self._heartbeat = heartbeat_interval

87 self._sessions: dict[str, StreamSession] = defaultdict(StreamSession)

88

89 async def stream_chat(

90 self,

91 message: str,

92 session_id: str = "default",

93 ) -> AsyncIterator[str]:

94 """

95 Stream a chat interaction as SSE events.

96

97 Yields:

98 SSE-formatted strings suitable for HTTP response body.

99 """

100 session = self._sessions[session_id]

101 session.session_id = session_id

102 t_start = time.time()

103

104 # Emit start event

105 yield StreamEvent(

106 event="start",

107 data={"session_id": session_id, "message": message},

108 ).to_sse()

109 session.events_emitted += 1

110

111 # Simulate streaming chunks (integrate with real agent loop)

112 chunks = self._generate_chunks(message)

113 heartbeat_task = asyncio.create_task(

114 self._heartbeat_loop(session_id)

115 )

116

117 try:

118 async for chunk in chunks:

119 yield StreamEvent(

120 event="chunk",

121 data={"content": chunk, "session_id": session_id},

122 ).to_sse()

123 session.events_emitted += 1

124 session.last_event_at = time.time()

125 finally:

126 heartbeat_task.cancel()

127 try:

128 await heartbeat_task

129 except asyncio.CancelledError:

130 pass

131

132 # Emit done event

133 total_ms = (time.time() - t_start) * 1000

134 yield StreamEvent(

135 event="done",

136 data={

137 "session_id": session_id,

138 "total_latency_ms": total_ms,

139 "events_emitted": session.events_emitted,

140 },

141 ).to_sse()

142

143 def stream_chat_sync(self, message: str, session_id: str = "default"):

144 """Synchronous wrapper for stream_chat."""

145 loop = asyncio.get_event_loop()

146 return _SyncSSEWrapper(

147 loop.run_until_complete(

148 self._collect_events(message, session_id)

149 )

150 )

151

152 async def _collect_events(

153 self, message: str, session_id: str

154 ) -> list[str]:

155 events: list[str] = []

156 async for sse in self.stream_chat(message, session_id):

157 events.append(sse)

158 return events

159

160 async def _generate_chunks(self, message: str) -> AsyncIterator[str]:

161 """Generate streaming text chunks. Override with real LLM integration."""

162 if self._loop and hasattr(self._loop, "run"):

163 # Integrate with actual agent loop

164 result = self._loop.run(message)

165 text = str(result.output) if hasattr(result, "output") else str(result)

166 words = text.split()

167 for i, word in enumerate(words):

168 chunk = word + (" " if i < len(words) - 1 else "")

169 yield chunk

170 await asyncio.sleep(0.02) # simulate streaming

171 else:

172 # Fallback: simulate streaming

173 words = message.split()

174 yield f"Processing: {message}\n"

175 await asyncio.sleep(0.3)

176 for i in range(3):

177 yield f"Agent step {i + 1}: analyzing...\n"

178 await asyncio.sleep(0.5)

179 yield f"Complete. Response for: {message}"

180

181 async def _heartbeat_loop(self, session_id: str) -> None:

182 """Send periodic heartbeat comments to keep connection alive."""

183 while True:

184 await asyncio.sleep(self._heartbeat)

185

186 def emit_tool_call(self, session_id: str, tool_name: str, args: dict) -> str:

187 """Emit a tool_call SSE event (non-streaming helper)."""

188 return StreamEvent(

189 event="tool_call",

190 data={

191 "session_id": session_id,

192 "tool": tool_name,

193 "arguments": args,

194 },

195 ).to_sse()

196

197 def emit_tool_result(

198 self, session_id: str, tool_name: str, result: Any

199 ) -> str:

200 """Emit a tool_result SSE event."""

201 return StreamEvent(

202 event="tool_result",

203 data={

204 "session_id": session_id,

205 "tool": tool_name,

206 "result": result,

207 },

208 ).to_sse()

209

210 def emit_error(self, session_id: str, error: str) -> str:

211 """Emit an error SSE event."""

212 return StreamEvent(

213 event="error",

214 data={"session_id": session_id, "error": error},

215 ).to_sse()

216

217 def get_session(self, session_id: str) -> Optional[StreamSession]:

218 return self._sessions.get(session_id)

219

220 def list_sessions(self) -> dict[str, StreamSession]:

221 return dict(self._sessions)

222

223

224 class _SyncSSEWrapper:

225 """Makes a list of SSE strings iterable for sync streaming."""

226

227 def __init__(self, events: list[str]):

228 self._events = events

229

230 def __iter__(self):

231 return iter(self._events)