Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-agents/src/lexigram/ai/agents/executor/streaming.py: 10%

1"""Streaming support for AgentExecutorImpl. 

2 

3Provides the astream() method for streaming agent execution events. 

4""" 

5 

6from __future__ import annotations 

7 

8from typing import TYPE_CHECKING, Any 

9import uuid 

10 

11from lexigram.contracts.ai.agents import AgentEvent, AgentEventType 

12from lexigram.logging import ( 

13 get_logger, 

14) 

15 

16logger = get_logger(__name__) 

17 

18if TYPE_CHECKING: 

19 from lexigram.ai.agents.executor.executor import AgentExecutorImpl 

20 from lexigram.contracts.ai.agents import AgentProtocol 

21 from lexigram.contracts.ai.guards import GuardPipelineProtocol 

22 

23 

24async def astream( 

25 self: AgentExecutorImpl, 

26 agent: AgentProtocol, 

27 message: str, 

28 session_id: str | None = None, 

29 user_id: str | None = None, 

30 **kwargs: Any, 

31): 

32 """Stream agent execution events. 

33 

34 Yields AgentEvent objects as the agent executes, enabling 

35 real-time monitoring of thoughts, tool calls, and messages. 

36 

37 Args: 

38 agent: The agent to execute. 

39 message: User's input message. 

40 session_id: Session ID for multi-turn memory. 

41 user_id: User ID for governance tracking. 

42 **kwargs: Additional parameters passed to the strategy. 

43 

44 Yields: 

45 AgentEvent objects with type, data, and run_id. 

46 """ 

47 run_id = str(uuid.uuid4()) 

48 

49 yield AgentEvent( 

50 type=AgentEventType.STARTED, 

51 data={ 

52 "agent_name": agent.name, 

53 "session_id": session_id, 

54 "user_id": user_id, 

55 "message": message, 

56 }, 

57 run_id=run_id, 

58 ) 

59 

60 governance_ok = await _check_governance(self, agent, user_id) 

61 if not governance_ok: 

62 yield AgentEvent( 

63 type=AgentEventType.ERROR, 

64 data={ 

65 "error": "Governance denied request", 

66 "agent_name": agent.name, 

67 }, 

68 run_id=run_id, 

69 ) 

70 yield AgentEvent( 

71 type=AgentEventType.FINISHED, 

72 data={ 

73 "agent_name": agent.name, 

74 "success": False, 

75 "error": "Governance denied", 

76 }, 

77 run_id=run_id, 

78 ) 

79 return 

80 

81 # Agent-level pipeline (AgentBuilder.with_guard_pipeline) wins over the 

82 # DI-provided one; the DI pipeline remains the standard default. 

83 agent_pipeline = getattr(agent, "guard_pipeline", None) 

84 guard_pipeline = ( 

85 agent_pipeline if agent_pipeline is not None else self._guard_pipeline 

86 ) 

87 

88 guard_ok = await _check_guard_input( 

89 self, guard_pipeline, message, session_id, user_id 

90 ) 

91 if not guard_ok: 

92 yield AgentEvent( 

93 type=AgentEventType.ERROR, 

94 data={ 

95 "error": "Input blocked by guard pipeline", 

96 "agent_name": agent.name, 

97 }, 

98 run_id=run_id, 

99 ) 

100 yield AgentEvent( 

101 type=AgentEventType.FINISHED, 

102 data={ 

103 "agent_name": agent.name, 

104 "success": False, 

105 "error": "Input blocked", 

106 }, 

107 run_id=run_id, 

108 ) 

109 return 

110 

111 try: 

112 strategy = getattr(agent, "strategy", None) 

113 if strategy is None: 

114 from lexigram.ai.agents.strategies.react import ReActStrategy 

115 

116 strategy = ReActStrategy() 

117 

118 history = await _load_history(self, message, session_id) 

119 tools = list(agent.tools) if agent.tools else [] 

120 await _merge_skills(self, tools) 

121 

122 strategy_history = [{"role": m.role, "content": m.content} for m in history] 

123 

124 result = await strategy.execute( 

125 message=message, 

126 tools=tools, 

127 history=strategy_history, 

128 llm=self._llm, # type: ignore[arg-type] 

129 system_prompt=agent.system_prompt, 

130 temperature=getattr(agent, "temperature", 0.7), 

131 tool_registry=kwargs.get("tool_registry"), 

132 memory=getattr(agent, "memory", None), 

133 guard_pipeline=guard_pipeline, 

134 **kwargs, 

135 ) 

136 

137 if result.is_err(): 

138 error = result.unwrap_err() 

139 yield AgentEvent( 

140 type=AgentEventType.ERROR, 

141 data={ 

142 "error": str(error), 

143 "error_type": type(error).__name__, 

144 "agent_name": agent.name, 

145 }, 

146 run_id=run_id, 

147 ) 

148 yield AgentEvent( 

149 type=AgentEventType.FINISHED, 

150 data={ 

151 "agent_name": agent.name, 

152 "success": False, 

153 "error": str(error), 

154 }, 

155 run_id=run_id, 

156 ) 

157 return 

158 

159 response = result.unwrap() 

160 

161 guard_ok = await _check_guard_output( 

162 self, guard_pipeline, response.message, message, session_id, user_id 

163 ) 

164 if not guard_ok: 

165 yield AgentEvent( 

166 type=AgentEventType.ERROR, 

167 data={ 

168 "error": "Output blocked by guard pipeline", 

169 "agent_name": agent.name, 

170 }, 

171 run_id=run_id, 

172 ) 

173 yield AgentEvent( 

174 type=AgentEventType.FINISHED, 

175 data={ 

176 "agent_name": agent.name, 

177 "success": False, 

178 "error": "Output blocked", 

179 }, 

180 run_id=run_id, 

181 ) 

182 return 

183 

184 yield AgentEvent( 

185 type=AgentEventType.MESSAGE, 

186 data={ 

187 "message": response.message, 

188 "agent_name": agent.name, 

189 }, 

190 run_id=run_id, 

191 ) 

192 

193 yield AgentEvent( 

194 type=AgentEventType.FINISHED, 

195 data={ 

196 "agent_name": agent.name, 

197 "success": True, 

198 "message": response.message, 

199 "steps": response.step_count, 

200 "tool_calls": response.tool_call_count, 

201 "total_tokens": response.total_tokens, 

202 "duration_ms": response.duration_ms, 

203 }, 

204 run_id=run_id, 

205 ) 

206 

207 except Exception as e: 

208 logger.exception("astream_execution_failed", agent=agent.name) 

209 yield AgentEvent( 

210 type=AgentEventType.ERROR, 

211 data={ 

212 "error": str(e), 

213 "error_type": type(e).__name__, 

214 "agent_name": agent.name, 

215 }, 

216 run_id=run_id, 

217 ) 

218 yield AgentEvent( 

219 type=AgentEventType.FINISHED, 

220 data={ 

221 "agent_name": agent.name, 

222 "success": False, 

223 "error": str(e), 

224 }, 

225 run_id=run_id, 

226 ) 

227 

228 

229async def _check_governance( 

230 self: AgentExecutorImpl, 

231 agent: AgentProtocol, 

232 user_id: str | None, 

233) -> bool: 

234 """Check governance and emit events.""" 

235 if not self._governance: 

236 return True 

237 

238 try: 

239 model = getattr(self._llm, "model", "unknown") 

240 provider = getattr(self._llm, "provider", "unknown") 

241 allowed = await self._governance.check_request( 

242 model=model, 

243 provider=provider, 

244 user_id=user_id, 

245 ) 

246 if not allowed: 

247 logger.warning( 

248 "agent_governance_denied", 

249 agent=agent.name, 

250 user_id=user_id, 

251 ) 

252 return False 

253 except Exception as e: 

254 logger.warning("governance_check_failed", error=str(e)) 

255 

256 return True 

257 

258 

259async def _check_guard_input( 

260 self: AgentExecutorImpl, 

261 guard_pipeline: GuardPipelineProtocol | None, 

262 message: str, 

263 session_id: str | None, 

264 user_id: str | None, 

265) -> bool: 

266 """Check input guard pipeline.""" 

267 if not guard_pipeline: 

268 return True 

269 

270 try: 

271 guard_res = await guard_pipeline.check_input( 

272 content=message, 

273 metadata={"session_id": session_id, "user_id": user_id}, 

274 ) 

275 if not guard_res.is_ok(): 

276 return False 

277 

278 agg = guard_res.unwrap() 

279 blocked = bool(getattr(agg, "blocked", False)) 

280 if blocked: 

281 logger.warning("input_blocked", session_id=session_id) 

282 return False 

283 except (RuntimeError, OSError) as exc: 

284 logger.warning("guard_input_check_failed", error=str(exc)) 

285 return False 

286 return True 

287 

288 

289async def _check_guard_output( 

290 self: AgentExecutorImpl, 

291 guard_pipeline: GuardPipelineProtocol | None, 

292 output: str, 

293 original_input: str, 

294 session_id: str | None, 

295 user_id: str | None, 

296) -> bool: 

297 """Check output guard pipeline.""" 

298 if not guard_pipeline: 

299 return True 

300 

301 try: 

302 guard_res = await guard_pipeline.check_output( 

303 content=output, 

304 original_input=original_input, 

305 metadata={"session_id": session_id, "user_id": user_id}, 

306 ) 

307 if not guard_res.is_ok(): 

308 return False 

309 

310 agg = guard_res.unwrap() 

311 blocked = bool(getattr(agg, "blocked", False)) 

312 if blocked: 

313 logger.warning("output_blocked", session_id=session_id) 

314 return False 

315 except (RuntimeError, OSError) as exc: 

316 logger.warning("guard_output_check_failed", error=str(exc)) 

317 return False 

318 return True 

319 

320 

321async def _load_history( 

322 self: AgentExecutorImpl, 

323 message: str, 

324 session_id: str | None, 

325) -> list[Any]: 

326 """Load conversation history.""" 

327 from lexigram.contracts.ai.llm import ChatMessage, Role 

328 

329 history: list[ChatMessage] = [] 

330 

331 if self._working_memory: 

332 try: 

333 entries = await self._working_memory.assemble( 

334 query=message, 

335 token_budget=4096, 

336 owner_id=session_id or "anonymous", 

337 ) 

338 history = [ 

339 ChatMessage(role=Role(e.role), content=e.content) for e in entries 

340 ] 

341 except Exception as e: 

342 logger.warning("working_memory_assemble_failed", error=str(e)) 

343 elif self._memory and session_id: 

344 try: 

345 if hasattr(self._memory, "get_messages"): 

346 msgs = await self._memory.get_messages() 

347 history = [ 

348 ChatMessage(role=Role(m.role), content=m.content) for m in msgs 

349 ] 

350 except Exception as e: 

351 logger.warning("memory_load_failed", error=str(e)) 

352 

353 return history 

354 

355 

356async def _merge_skills( 

357 self: AgentExecutorImpl, 

358 tools: list[Any], 

359) -> None: 

360 """Merge skills as additional tools.""" 

361 if not self._skill_registry: 

362 return 

363 

364 try: 

365 skill_schemas = self._skill_registry.get_schemas() 

366 tools.extend(skill_schemas) 

367 except Exception as e: 

368 logger.warning("skill_registry_merge_failed", error=str(e))