Coverage for agentos/core/async_loop.py: 42%

1""" 

2Async agent execution loop with concurrency support. 

3 

4Provides async/await versions of the core agent loop for high-throughput 

5scenarios where multiple agents run concurrently. 

6""" 

7 

8from __future__ import annotations 

9 

10import asyncio 

11import time 

12from collections.abc import AsyncIterator, Awaitable, Callable 

13from dataclasses import dataclass 

14from typing import Any 

15 

16from agentos.core.context import AgentContext 

17from agentos.core.streaming import StreamChunk 

18 

19 

20@dataclass 

21class AsyncLoopConfig: 

22 """Configuration for async agent execution loop.""" 

23 

24 max_concurrency: int = 10 

25 """Max concurrent agent invocations.""" 

26 

27 timeout_seconds: float = 300.0 

28 """Per-invocation timeout.""" 

29 

30 retry_on_timeout: bool = True 

31 """Whether to retry timed-out invocations.""" 

32 

33 max_retries: int = 3 

34 """Max retries on transient failures.""" 

35 

36 collect_metrics: bool = True 

37 """Whether to collect timing metrics.""" 

38 

39 

40@dataclass 

41class AsyncInvocationResult: 

42 """Result of a single async agent invocation.""" 

43 

44 agent_id: str 

45 success: bool 

46 output: Any = None 

47 error: str | None = None 

48 latency_ms: float = 0.0 

49 retries: int = 0 

50 

51 

52class AsyncAgentLoop: 

53 """ 

54 Async execution loop for agent invocations. 

55 

56 Supports: 

57 - Concurrent multi-agent execution with semaphore-based throttling 

58 - Per-invocation timeouts via asyncio.wait_for 

59 - Automatic retry with exponential backoff 

60 - Streaming output via async generators 

61 - Metrics collection (p50/p95/p99 latency) 

62 

63 Example:: 

64 

65 loop = AsyncAgentLoop(config=AsyncLoopConfig(max_concurrency=5)) 

66 results = await loop.run_all([task1, task2, task3]) 

67 """ 

68 

69 def __init__(self, config: AsyncLoopConfig | None = None): 

70 self.config = config or AsyncLoopConfig() 

71 self._semaphore = asyncio.Semaphore(self.config.max_concurrency) 

72 self._metrics: list[float] = [] 

73 

74 async def run_single( 

75 self, 

76 agent_id: str, 

77 fn: Callable[..., Awaitable[Any]], 

78 *args: Any, 

79 **kwargs: Any, 

80 ) -> AsyncInvocationResult: 

81 """ 

82 Run a single agent invocation with timeout and retry. 

83 

84 Args: 

85 agent_id: Identifier for the agent invocation. 

86 fn: Async callable to execute. 

87 *args: Positional args for fn. 

88 **kwargs: Keyword args for fn. 

89 

90 Returns: 

91 AsyncInvocationResult with success/failure details. 

92 """ 

93 async with self._semaphore: 

94 return await self._execute_with_retry(agent_id, fn, args, kwargs) 

95 

96 async def _execute_with_retry( 

97 self, 

98 agent_id: str, 

99 fn: Callable[..., Awaitable[Any]], 

100 args: tuple, 

101 kwargs: dict, 

102 ) -> AsyncInvocationResult: 

103 last_error: str | None = None 

104 t0 = time.perf_counter() 

105 

106 for attempt in range(self.config.max_retries + 1): 

107 try: 

108 result = await asyncio.wait_for( 

109 fn(*args, **kwargs), 

110 timeout=self.config.timeout_seconds, 

111 ) 

112 latency = (time.perf_counter() - t0) * 1000 

113 if self.config.collect_metrics: 

114 self._metrics.append(latency) 

115 return AsyncInvocationResult( 

116 agent_id=agent_id, 

117 success=True, 

118 output=result, 

119 latency_ms=latency, 

120 retries=attempt, 

121 ) 

122 except TimeoutError: 

123 last_error = f"Timeout after {self.config.timeout_seconds}s" 

124 if not self.config.retry_on_timeout: 

125 break 

126 except Exception as exc: 

127 last_error = f"{type(exc).__name__}: {exc}" 

128 if attempt >= self.config.max_retries: 

129 break 

130 

131 latency = (time.perf_counter() - t0) * 1000 

132 if self.config.collect_metrics: 

133 self._metrics.append(latency) 

134 return AsyncInvocationResult( 

135 agent_id=agent_id, 

136 success=False, 

137 error=last_error, 

138 latency_ms=latency, 

139 retries=attempt, 

140 ) 

141 

142 async def run_all( 

143 self, 

144 tasks: list[tuple[str, Callable[..., Awaitable[Any]], tuple, dict]], 

145 ) -> list[AsyncInvocationResult]: 

146 """ 

147 Run multiple agent invocations concurrently. 

148 

149 Args: 

150 tasks: List of (agent_id, async_fn, args, kwargs) tuples. 

151 

152 Returns: 

153 List of results in the same order as input tasks. 

154 """ 

155 coros = [ 

156 self.run_single(agent_id, fn, *args, **kwargs) for agent_id, fn, args, kwargs in tasks 

157 ] 

158 return list(await asyncio.gather(*coros)) 

159 

160 async def run_streaming( 

161 self, 

162 agent_id: str, 

163 stream_fn: Callable[[], AsyncIterator[StreamChunk]], 

164 ) -> AsyncIterator[StreamChunk]: 

165 """ 

166 Run an agent and yield streaming output chunks. 

167 

168 Args: 

169 agent_id: Identifier for the agent. 

170 stream_fn: Async generator yielding StreamChunk objects. 

171 

172 Yields: 

173 StreamChunk as they become available. 

174 """ 

175 async with self._semaphore: 

176 t0 = time.perf_counter() 

177 chunk_count = 0 

178 async for chunk in stream_fn(): 

179 chunk_count += 1 

180 yield chunk 

181 latency = (time.perf_counter() - t0) * 1000 

182 if self.config.collect_metrics: 

183 self._metrics.append(latency) 

184 

185 def get_latency_stats(self) -> dict[str, float]: 

186 """ 

187 Compute p50/p95/p99 latency from collected metrics. 

188 

189 Returns: 

190 Dict with keys p50_ms, p95_ms, p99_ms, mean_ms, count. 

191 """ 

192 if not self._metrics: 

193 return {"p50_ms": 0, "p95_ms": 0, "p99_ms": 0, "mean_ms": 0, "count": 0} 

194 sorted_ms = sorted(self._metrics) 

195 n = len(sorted_ms) 

196 

197 def percentile(p: float) -> float: 

198 idx = int(n * p / 100) 

199 return sorted_ms[min(idx, n - 1)] 

200 

201 return { 

202 "p50_ms": percentile(50), 

203 "p95_ms": percentile(95), 

204 "p99_ms": percentile(99), 

205 "mean_ms": sum(sorted_ms) / n, 

206 "count": n, 

207 } 

208 

209 def reset_metrics(self) -> None: 

210 """Clear accumulated latency metrics.""" 

211 self._metrics.clear() 

212 

213 

214class AsyncContextManager: 

215 """ 

216 Async-safe context manager for agent sessions. 

217 

218 Manages async context propagation across concurrent agent invocations. 

219 """ 

220 

221 def __init__(self, context: AgentContext): 

222 self._context = context 

223 self._store: dict[str, Any] = {} 

224 self._lock = asyncio.Lock() 

225 

226 async def get(self, key: str, default: Any = None) -> Any: 

227 async with self._lock: 

228 return self._store.get(key, default) 

229 

230 async def set(self, key: str, value: Any) -> None: 

231 async with self._lock: 

232 self._store[key] = value 

233 

234 async def update(self, mapping: dict[str, Any]) -> None: 

235 async with self._lock: 

236 self._store.update(mapping) 

237 

238 async def snapshot(self) -> dict[str, Any]: 

239 async with self._lock: 

240 return dict(self._store)