Coverage for agentos/server/agent_api.py: 47%
122 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-05 20:52 +0800
1"""Agent API — FastAPI REST endpoint for ProductionAgent.
3Serves ProductionAgent as HTTP API with health check, metrics,
4task dispatch, and streaming support.
6v1.9.13: Initial — POST /run, /run/stream, GET /health, /stats.
7"""
9from __future__ import annotations
11import time
12import uuid
13from dataclasses import dataclass, field
14from typing import Any, AsyncIterator, Optional
16from fastapi import FastAPI, HTTPException
17from fastapi.responses import JSONResponse, StreamingResponse
18from pydantic import BaseModel, Field
20from agentos.agent.production import ProductionAgent, ProductionConfig
21from agentos.agent.tool_agent import ToolExecutor, AgentResult
22from agentos.llm.base import LLMProvider, Message, MessageRole
23from agentos.llm.smart_cache import SmartCache, CacheConfig
24from agentos.agent.model_router import ModelRouter
26__all__ = [
27 "AgentAPI",
28 "AgentAPIRequest",
29 "AgentAPIResponse",
30 "AgentAPIStats",
31 "create_agent_api",
32]
35# ── Pydantic Models ──────────────────────────────────────────────────
38class AgentAPIRequest(BaseModel):
39 task: str = Field(..., description="The task string to execute.")
40 session_id: str | None = Field(default=None, description="Session identifier for audit log grouping.")
41 budget_usd: float | None = Field(default=None, description="Override daily budget (default from config).")
42 enable_audit: bool = Field(default=True)
43 enable_cache: bool = Field(default=True)
46class AgentAPIResponse(BaseModel):
47 id: str = Field(default_factory=lambda: f"req-{uuid.uuid4().hex[:6]}")
48 success: bool
49 task: str
50 output: str
51 error: str | None = None
52 model: str | None = None
53 complexity: str | None = None
54 total_steps: int = 0
55 total_tokens: int = 0
56 cost_usd: float = 0.0
57 duration_ms: float = 0.0
58 cache_hit: bool = False
61class AgentAPIStats(BaseModel):
62 uptime_seconds: float
63 total_requests: int
64 success_count: int
65 failure_count: int
66 avg_latency_ms: float
67 total_cost_usd: float
68 cache_hit_rate: float
69 cache_savings_usd: float
70 budget_remaining_usd: float
73# ── Agent API Server ─────────────────────────────────────────────────
76@dataclass
77class AgentAPI:
78 """FastAPI app wrapping a ProductionAgent.
80 Usage:
81 from agentos.server.agent_api import create_agent_api
82 app = create_agent_api(provider, executor)
83 # uvicorn.run(app, host="0.0.0.0", port=8000)
85 Endpoints:
86 POST /agent/run — execute task synchronously
87 POST /agent/run/stream — execute task with SSE streaming
88 GET /agent/health — health check
89 GET /agent/stats — runtime statistics
90 GET /agent/budget — remaining budget info
91 """
93 provider: LLMProvider
94 executor: ToolExecutor
95 router: ModelRouter | None = None
96 cache: SmartCache | None = None
97 config: ProductionConfig = field(default_factory=ProductionConfig)
99 # runtime stats
100 _start_time: float = field(default_factory=time.time, init=False)
101 _total_requests: int = field(default=0, init=False)
102 _success_count: int = field(default=0, init=False)
103 _failure_count: int = field(default=0, init=False)
104 _total_latency_ms: float = field(default=0.0, init=False)
106 def build_app(self) -> FastAPI:
107 """Build and return a FastAPI app."""
108 app = FastAPI(title="AgentOS Agent API", version="1.9.13")
109 self_app = self
111 # agent factory — per request to isolate state
112 def _build_agent(
113 enable_audit: bool = True,
114 enable_cache: bool = True,
115 budget_usd: float | None = None,
116 session_id: str | None = None,
117 ) -> ProductionAgent:
118 cfg = ProductionConfig(
119 enable_audit=enable_audit,
120 enable_cache=enable_cache,
121 audit_log_dir=self_app.config.audit_log_dir,
122 budget_usd=budget_usd or self_app.config.budget_usd,
123 session_id=session_id or "",
124 )
125 return ProductionAgent(
126 provider=self_app.provider,
127 tool_executor=self_app.executor,
128 config=cfg,
129 router=self_app.router,
130 cache=self_app.cache,
131 )
133 @app.post("/agent/run", response_model=AgentAPIResponse)
134 async def agent_run(req: AgentAPIRequest):
135 self_app._total_requests += 1
136 t0 = time.time()
138 try:
139 agent = _build_agent(
140 enable_audit=req.enable_audit,
141 enable_cache=req.enable_cache,
142 budget_usd=req.budget_usd,
143 session_id=req.session_id,
144 )
145 result = agent.run(req.task)
146 elapsed_ms = (time.time() - t0) * 1000
147 self_app._total_latency_ms += elapsed_ms
149 if result.success:
150 self_app._success_count += 1
151 else:
152 self_app._failure_count += 1
154 last_text = ""
155 if result.final_answer:
156 last_text = str(result.final_answer)
157 elif result.steps:
158 last_text = str(result.steps[-1].final_answer
159 if hasattr(result.steps[-1], "final_answer")
160 else "")
162 cache_hit = (
163 agent.cache_hit_rate > 0 and result.success
164 )
166 return AgentAPIResponse(
167 success=result.success,
168 task=req.task,
169 output=last_text,
170 error=result.error,
171 model=getattr(agent.last_model, "name", None),
172 complexity=None,
173 total_steps=result.total_steps,
174 total_tokens=result.total_tokens,
175 cost_usd=round(result.total_cost_usd, 6),
176 duration_ms=round(elapsed_ms, 2),
177 cache_hit=cache_hit,
178 )
180 except Exception as e:
181 self_app._failure_count += 1
182 raise HTTPException(status_code=500, detail=str(e))
184 @app.post("/agent/run/stream")
185 async def agent_run_stream(req: AgentAPIRequest):
186 self_app._total_requests += 1
188 try:
189 agent = _build_agent(
190 enable_audit=req.enable_audit,
191 enable_cache=req.enable_cache,
192 budget_usd=req.budget_usd,
193 session_id=req.session_id,
194 )
196 async def event_stream() -> AsyncIterator[str]:
197 for step in agent.run_stream(req.task):
198 if hasattr(step, "to_dict"):
199 import json
200 yield f"data: {json.dumps(step.to_dict())}\n\n"
201 else:
202 yield f"data: {str(step)}\n\n"
203 yield "data: [DONE]\n\n"
205 return StreamingResponse(
206 event_stream(),
207 media_type="text/event-stream",
208 )
210 except Exception as e:
211 self_app._failure_count += 1
212 raise HTTPException(status_code=500, detail=str(e))
214 @app.get("/agent/health")
215 async def agent_health():
216 return {
217 "status": "ok",
218 "uptime_seconds": round(time.time() - self_app._start_time, 2),
219 "provider": self_app.provider.provider_name,
220 "model": self_app.provider.model_name,
221 "cache_enabled": self_app.cache is not None,
222 "router_enabled": self_app.router is not None,
223 }
225 @app.get("/agent/stats", response_model=AgentAPIStats)
226 async def agent_stats():
227 total = self_app._total_requests
228 avg_lat = (
229 self_app._total_latency_ms / total if total > 0 else 0.0
230 )
232 # build a temp agent to get latest budget cache stats
233 agent = _build_agent()
234 hit_rate = agent.cache_hit_rate
235 savings = agent.cache_savings
236 budget = agent._router.daily_budget_remaining if agent._router else 0.0
238 total_cost = agent._router._total_spent if hasattr(agent._router, "_total_spent") else 0.0
240 return AgentAPIStats(
241 uptime_seconds=round(time.time() - self_app._start_time, 2),
242 total_requests=total,
243 success_count=self_app._success_count,
244 failure_count=self_app._failure_count,
245 avg_latency_ms=round(avg_lat, 2),
246 total_cost_usd=round(total_cost, 6),
247 cache_hit_rate=round(hit_rate, 4),
248 cache_savings_usd=round(savings, 6),
249 budget_remaining_usd=round(budget, 4),
250 )
252 @app.get("/agent/budget")
253 async def agent_budget():
254 agent = _build_agent()
255 remaining = agent._router.daily_budget_remaining if agent._router else 0.0
256 total_spent = getattr(agent._router, "_total_spent", 0.0)
257 return {
258 "daily_budget_usd": self_app.config.budget_usd,
259 "remaining_usd": round(remaining, 4),
260 "spent_usd": round(total_spent, 6),
261 "usage_pct": round(
262 (total_spent / self_app.config.budget_usd * 100)
263 if self_app.config.budget_usd > 0
264 else 0,
265 2,
266 ),
267 }
269 return app
272# ── Factory ─────────────────────────────────────────────────────────
275def create_agent_api(
276 provider: LLMProvider,
277 executor: ToolExecutor,
278 *,
279 router: ModelRouter | None = None,
280 cache: SmartCache | None = None,
281 config: ProductionConfig | None = None,
282) -> FastAPI:
283 """Create a FastAPI app with ProductionAgent endpoints.
285 Args:
286 provider: LLM provider for agent inference.
287 executor: Tool executor with registered tools.
288 router: Model router (auto-created if None).
289 cache: SmartCache for response caching (optional).
290 config: Production configuration (uses defaults if None).
292 Returns:
293 FastAPI application ready to serve.
294 """
295 api = AgentAPI(
296 provider=provider,
297 executor=executor,
298 router=router,
299 cache=cache,
300 config=config or ProductionConfig(),
301 )
302 return api.build_app()