Research that keeps going.
Submit a question, close the tab, or restart the worker. The job, lease, checkpoint, and result live in PostgreSQL.
import asyncio import json import os import time from contextlib import asynccontextmanager from uuid import UUID, uuid4 import psycopg from fastapi import FastAPI, HTTPException from fastapi.responses import HTMLResponse, StreamingResponse from pydantic import BaseModel, Field from psycopg.rows import dict_row from psycopg.types.json import Jsonb SEEDED_JOB_ID = UUID("00000000-0000-4000-8000-000000000017") SEEDED_RESULT = { "mode": "Bundled document demo", "summary": "In this sample evidence set, heat pumps trade more retrofit planning for efficient heating and cooling. Gas furnaces can be simpler when a home already has usable gas service and ducts. Climate, electrical capacity, insulation, current distribution, local prices, and backup plans decide the better fit.", "findings": [ { "source": "Climate and efficiency notes", "note": "Start with the building envelope and the heating load. Equipment labels alone do not answer how either system will perform in a specific home.", }, { "source": "Costs and retrofit notes", "note": "Compare the complete project, including electrical work, ducts, gas service, cooling needs, incentives, and expected operating costs.", }, { "source": "Cold weather and resilience notes", "note": "Cold climate design depends on correct sizing, low temperature capacity, defrost behavior, and a deliberate backup strategy.", }, ], "sources": [ "climate-and-efficiency.md", "costs-and-retrofit.md", "cold-weather-and-resilience.md", ], } def connect(): return psycopg.connect(os.environ["DATABASE_URL"], row_factory=dict_row) def seed_completed_job() -> None: with connect() as connection: connection.execute( """ INSERT INTO research_jobs (id, question, status, current_step, total_steps, checkpoint, result, attempts) VALUES (%s, %s, 'done', 4, 4, %s, %s, 1) ON CONFLICT (id) DO NOTHING """, ( SEEDED_JOB_ID, "Compare heat pumps and gas furnaces for an older home.", Jsonb({"seeded": True}), Jsonb(SEEDED_RESULT), ), ) def wait_for_database() -> None: delay = 1.0 for attempt in range(8): try: seed_completed_job() return except psycopg.Error: if attempt == 7: raise time.sleep(delay) delay = min(delay * 2, 8) @asynccontextmanager async def lifespan(_app: FastAPI): await asyncio.to_thread(wait_for_database) yield app = FastAPI(title="Research queue", lifespan=lifespan) class JobInput(BaseModel): question: str = Field(min_length=8, max_length=500) restart_demo: bool = False def serialize_job(row: dict) -> dict: return { key: value.isoformat() if hasattr(value, "isoformat") else value for key, value in row.items() } def fetch_job(job_id: UUID) -> dict | None: with connect() as connection: row = connection.execute( "SELECT * FROM research_jobs WHERE id = %s", (job_id,), ).fetchone() return serialize_job(row) if row else None @app.get("/health") def health(): return {"ok": True} @app.get("/api/jobs") def list_jobs(): with connect() as connection: rows = connection.execute( "SELECT * FROM research_jobs ORDER BY created_at DESC LIMIT 30" ).fetchall() return [serialize_job(row) for row in rows] @app.get("/api/jobs/{job_id}") def get_job(job_id: UUID): job = fetch_job(job_id) if not job: raise HTTPException(status_code=404, detail="Job not found") return job @app.post("/api/jobs", status_code=202) def create_job(job_input: JobInput): job_id = uuid4() with connect() as connection: row = connection.execute( """ INSERT INTO research_jobs (id, question, restart_demo) VALUES (%s, %s, %s) RETURNING * """, (job_id, job_input.question.strip(), job_input.restart_demo), ).fetchone() return serialize_job(row) @app.get("/api/jobs/{job_id}/events") async def job_events(job_id: UUID): async def events(): previous = None for _ in range(600): job = await asyncio.to_thread(fetch_job, job_id) if not job: yield f"data: {json.dumps({'error': 'Job not found'})}\n\n" return snapshot = json.dumps(job, sort_keys=True, default=str) if snapshot != previous: yield f"data: {json.dumps(job, default=str)}\n\n" previous = snapshot if job["status"] in {"done", "failed"}: return await asyncio.sleep(1) return StreamingResponse(events(), media_type="text/event-stream") @app.get("/", response_class=HTMLResponse) def home(): return """
Submit a question, close the tab, or restart the worker. The job, lease, checkpoint, and result live in PostgreSQL.