Spaces:
Sleeping
Sleeping
File size: 3,369 Bytes
936905f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 | """
main.py — FastAPI application for AgentForge.
Exposes POST /run, GET /health, GET /benchmark.
"""
import asyncio
import logging
import time
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import memory
from config import get_settings
from orchestrator import run_pipeline
logging.basicConfig(level=logging.INFO, format="%(levelname)s %(name)s: %(message)s")
logger = logging.getLogger(__name__)
settings = get_settings()
app = FastAPI(title="AgentForge", version="1.0.0")
# ── CORS ───────────────────────────────────────────────────────────────────────
app.add_middleware(
CORSMiddleware,
allow_origins=settings.allowed_origins_list,
allow_origin_regex=r"https://.*\.vercel\.app",
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# ── Request / Response models ──────────────────────────────────────────────────
class RunRequest(BaseModel):
task: str
# ── Routes ────────────────────────────────────────────────────────────────────
@app.get("/health")
async def health():
return {
"status": "ok",
"redis": "connected" if memory.is_available() else "unavailable",
"model": settings.model_name,
}
@app.post("/run")
async def run(req: RunRequest):
if not req.task.strip():
raise HTTPException(status_code=400, detail="task cannot be empty")
logger.info("POST /run — task: %s…", req.task[:80])
result = await run_pipeline(req.task)
return result
BENCHMARK_TASK = "Compare PostgreSQL vs MongoDB for a startup in 2025"
@app.get("/benchmark")
async def benchmark():
"""
Run the same task twice. First call hits LLM; second hits Redis cache.
Returns timing comparison to demonstrate cache speedup.
"""
# Flush any existing cache for this task so benchmark is honest
from orchestrator import decompose_task
subtasks = await decompose_task(BENCHMARK_TASK)
for s in subtasks:
key = memory.cache_key(s["prompt"])
if memory.is_available() and memory._client:
try:
memory._client.delete(key)
except Exception:
pass
t0 = time.perf_counter()
run1 = await run_pipeline(BENCHMARK_TASK)
time1 = time.perf_counter() - t0
t0 = time.perf_counter()
run2 = await run_pipeline(BENCHMARK_TASK)
time2 = time.perf_counter() - t0
return {
"task": BENCHMARK_TASK,
"run1_cold": {
"total_time": round(time1, 2),
"parallel_time": run1["parallel_time"],
"sequential_time": run1["sequential_time"],
"speedup": run1["speedup"],
},
"run2_cached": {
"total_time": round(time2, 2),
"parallel_time": run2["parallel_time"],
"sequential_time": run2["sequential_time"],
},
"cache_speedup": round(time1 / time2, 1) if time2 > 0 else "∞",
}
|