Coverage for agentos/server/agent_api.py: 47%

1"""Agent API — FastAPI REST endpoint for ProductionAgent. 

2 

3Serves ProductionAgent as HTTP API with health check, metrics, 

4task dispatch, and streaming support. 

5 

6v1.9.13: Initial — POST /run, /run/stream, GET /health, /stats. 

7""" 

8 

9from __future__ import annotations 

10 

11import time 

12import uuid 

13from collections.abc import AsyncIterator 

14from dataclasses import dataclass, field 

15 

16from fastapi import FastAPI, HTTPException 

17from fastapi.responses import StreamingResponse 

18from pydantic import BaseModel, Field 

19 

20from agentos.agent.model_router import ModelRouter 

21from agentos.agent.production import ProductionAgent, ProductionConfig 

22from agentos.agent.tool_agent import ToolExecutor 

23from agentos.llm.base import LLMProvider 

24from agentos.llm.smart_cache import SmartCache 

25 

26__all__ = [ 

27 "AgentAPI", 

28 "AgentAPIRequest", 

29 "AgentAPIResponse", 

30 "AgentAPIStats", 

31 "create_agent_api", 

32] 

33 

34 

35# ── Pydantic Models ────────────────────────────────────────────────── 

36 

37 

38class AgentAPIRequest(BaseModel): 

39 task: str = Field(..., description="The task string to execute.") 

40 session_id: str | None = Field( 

41 default=None, description="Session identifier for audit log grouping." 

42 ) 

43 budget_usd: float | None = Field( 

44 default=None, description="Override daily budget (default from config)." 

45 ) 

46 enable_audit: bool = Field(default=True) 

47 enable_cache: bool = Field(default=True) 

48 

49 

50class AgentAPIResponse(BaseModel): 

51 id: str = Field(default_factory=lambda: f"req-{uuid.uuid4().hex[:6]}") 

52 success: bool 

53 task: str 

54 output: str 

55 error: str | None = None 

56 model: str | None = None 

57 complexity: str | None = None 

58 total_steps: int = 0 

59 total_tokens: int = 0 

60 cost_usd: float = 0.0 

61 duration_ms: float = 0.0 

62 cache_hit: bool = False 

63 

64 

65class AgentAPIStats(BaseModel): 

66 uptime_seconds: float 

67 total_requests: int 

68 success_count: int 

69 failure_count: int 

70 avg_latency_ms: float 

71 total_cost_usd: float 

72 cache_hit_rate: float 

73 cache_savings_usd: float 

74 budget_remaining_usd: float 

75 

76 

77# ── Agent API Server ───────────────────────────────────────────────── 

78 

79 

80@dataclass 

81class AgentAPI: 

82 """FastAPI app wrapping a ProductionAgent. 

83 

84 Usage: 

85 from agentos.server.agent_api import create_agent_api 

86 app = create_agent_api(provider, executor) 

87 # uvicorn.run(app, host="0.0.0.0", port=8000) 

88 

89 Endpoints: 

90 POST /agent/run — execute task synchronously 

91 POST /agent/run/stream — execute task with SSE streaming 

92 GET /agent/health — health check 

93 GET /agent/stats — runtime statistics 

94 GET /agent/budget — remaining budget info 

95 """ 

96 

97 provider: LLMProvider 

98 executor: ToolExecutor 

99 router: ModelRouter | None = None 

100 cache: SmartCache | None = None 

101 config: ProductionConfig = field(default_factory=ProductionConfig) 

102 

103 # runtime stats 

104 _start_time: float = field(default_factory=time.time, init=False) 

105 _total_requests: int = field(default=0, init=False) 

106 _success_count: int = field(default=0, init=False) 

107 _failure_count: int = field(default=0, init=False) 

108 _total_latency_ms: float = field(default=0.0, init=False) 

109 

110 def build_app(self) -> FastAPI: 

111 """Build and return a FastAPI app.""" 

112 app = FastAPI(title="AgentOS Agent API", version="1.9.13") 

113 self_app = self 

114 

115 # agent factory — per request to isolate state 

116 def _build_agent( 

117 enable_audit: bool = True, 

118 enable_cache: bool = True, 

119 budget_usd: float | None = None, 

120 session_id: str | None = None, 

121 ) -> ProductionAgent: 

122 cfg = ProductionConfig( 

123 enable_audit=enable_audit, 

124 enable_cache=enable_cache, 

125 audit_log_dir=self_app.config.audit_log_dir, 

126 budget_usd=budget_usd or self_app.config.budget_usd, 

127 session_id=session_id or "", 

128 ) 

129 return ProductionAgent( 

130 provider=self_app.provider, 

131 tool_executor=self_app.executor, 

132 config=cfg, 

133 router=self_app.router, 

134 cache=self_app.cache, 

135 ) 

136 

137 @app.post("/agent/run", response_model=AgentAPIResponse) 

138 async def agent_run(req: AgentAPIRequest): 

139 self_app._total_requests += 1 

140 t0 = time.time() 

141 

142 try: 

143 agent = _build_agent( 

144 enable_audit=req.enable_audit, 

145 enable_cache=req.enable_cache, 

146 budget_usd=req.budget_usd, 

147 session_id=req.session_id, 

148 ) 

149 result = agent.run(req.task) 

150 elapsed_ms = (time.time() - t0) * 1000 

151 self_app._total_latency_ms += elapsed_ms 

152 

153 if result.success: 

154 self_app._success_count += 1 

155 else: 

156 self_app._failure_count += 1 

157 

158 last_text = "" 

159 if result.final_answer: 

160 last_text = str(result.final_answer) 

161 elif result.steps: 

162 last_text = str( 

163 result.steps[-1].final_answer 

164 if hasattr(result.steps[-1], "final_answer") 

165 else "" 

166 ) 

167 

168 cache_hit = agent.cache_hit_rate > 0 and result.success 

169 

170 return AgentAPIResponse( 

171 success=result.success, 

172 task=req.task, 

173 output=last_text, 

174 error=result.error, 

175 model=getattr(agent.last_model, "name", None), 

176 complexity=None, 

177 total_steps=result.total_steps, 

178 total_tokens=result.total_tokens, 

179 cost_usd=round(result.total_cost_usd, 6), 

180 duration_ms=round(elapsed_ms, 2), 

181 cache_hit=cache_hit, 

182 ) 

183 

184 except Exception as e: 

185 self_app._failure_count += 1 

186 raise HTTPException(status_code=500, detail=str(e)) 

187 

188 @app.post("/agent/run/stream") 

189 async def agent_run_stream(req: AgentAPIRequest): 

190 self_app._total_requests += 1 

191 

192 try: 

193 agent = _build_agent( 

194 enable_audit=req.enable_audit, 

195 enable_cache=req.enable_cache, 

196 budget_usd=req.budget_usd, 

197 session_id=req.session_id, 

198 ) 

199 

200 async def event_stream() -> AsyncIterator[str]: 

201 for step in agent.run_stream(req.task): 

202 if hasattr(step, "to_dict"): 

203 import json 

204 

205 yield f"data: {json.dumps(step.to_dict())}\n\n" 

206 else: 

207 yield f"data: {str(step)}\n\n" 

208 yield "data: [DONE]\n\n" 

209 

210 return StreamingResponse( 

211 event_stream(), 

212 media_type="text/event-stream", 

213 ) 

214 

215 except Exception as e: 

216 self_app._failure_count += 1 

217 raise HTTPException(status_code=500, detail=str(e)) 

218 

219 @app.get("/agent/health") 

220 async def agent_health(): 

221 return { 

222 "status": "ok", 

223 "uptime_seconds": round(time.time() - self_app._start_time, 2), 

224 "provider": self_app.provider.provider_name, 

225 "model": self_app.provider.model_name, 

226 "cache_enabled": self_app.cache is not None, 

227 "router_enabled": self_app.router is not None, 

228 } 

229 

230 @app.get("/agent/stats", response_model=AgentAPIStats) 

231 async def agent_stats(): 

232 total = self_app._total_requests 

233 avg_lat = self_app._total_latency_ms / total if total > 0 else 0.0 

234 

235 # build a temp agent to get latest budget cache stats 

236 agent = _build_agent() 

237 hit_rate = agent.cache_hit_rate 

238 savings = agent.cache_savings 

239 budget = agent._router.daily_budget_remaining if agent._router else 0.0 

240 

241 total_cost = ( 

242 agent._router._total_spent if hasattr(agent._router, "_total_spent") else 0.0 

243 ) 

244 

245 return AgentAPIStats( 

246 uptime_seconds=round(time.time() - self_app._start_time, 2), 

247 total_requests=total, 

248 success_count=self_app._success_count, 

249 failure_count=self_app._failure_count, 

250 avg_latency_ms=round(avg_lat, 2), 

251 total_cost_usd=round(total_cost, 6), 

252 cache_hit_rate=round(hit_rate, 4), 

253 cache_savings_usd=round(savings, 6), 

254 budget_remaining_usd=round(budget, 4), 

255 ) 

256 

257 @app.get("/agent/budget") 

258 async def agent_budget(): 

259 agent = _build_agent() 

260 remaining = agent._router.daily_budget_remaining if agent._router else 0.0 

261 total_spent = getattr(agent._router, "_total_spent", 0.0) 

262 return { 

263 "daily_budget_usd": self_app.config.budget_usd, 

264 "remaining_usd": round(remaining, 4), 

265 "spent_usd": round(total_spent, 6), 

266 "usage_pct": round( 

267 ( 

268 (total_spent / self_app.config.budget_usd * 100) 

269 if self_app.config.budget_usd > 0 

270 else 0 

271 ), 

272 2, 

273 ), 

274 } 

275 

276 return app 

277 

278 

279# ── Factory ───────────────────────────────────────────────────────── 

280 

281 

282def create_agent_api( 

283 provider: LLMProvider, 

284 executor: ToolExecutor, 

285 *, 

286 router: ModelRouter | None = None, 

287 cache: SmartCache | None = None, 

288 config: ProductionConfig | None = None, 

289) -> FastAPI: 

290 """Create a FastAPI app with ProductionAgent endpoints. 

291 

292 Args: 

293 provider: LLM provider for agent inference. 

294 executor: Tool executor with registered tools. 

295 router: Model router (auto-created if None). 

296 cache: SmartCache for response caching (optional). 

297 config: Production configuration (uses defaults if None). 

298 

299 Returns: 

300 FastAPI application ready to serve. 

301 """ 

302 api = AgentAPI( 

303 provider=provider, 

304 executor=executor, 

305 router=router, 

306 cache=cache, 

307 config=config or ProductionConfig(), 

308 ) 

309 return api.build_app()