50 lines
1.6 KiB
Python
50 lines
1.6 KiB
Python
|
import asyncio
|
||
|
|
|
||
|
|
import httpx
|
||
|
|
|
||
|
|
from app.config import settings
|
||
|
|
|
||
|
|
HEX_API_BASE = "https://app.hex.tech/api/v1"
|
||
|
|
HEADERS = {
|
||
|
|
"Authorization": f"Bearer {settings.HEX_API_TOKEN}",
|
||
|
|
"Content-Type": "application/json",
|
||
|
|
}
|
||
|
|
|
||
|
|
NOTEBOOKS = {
|
||
|
|
"distraction_patterns": settings.HEX_NB_DISTRACTIONS,
|
||
|
|
"focus_trends": settings.HEX_NB_FOCUS_TRENDS,
|
||
|
|
"weekly_report": settings.HEX_NB_WEEKLY_REPORT,
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
async def run_notebook(notebook_key: str, user_id: str) -> dict:
|
||
|
|
project_id = NOTEBOOKS.get(notebook_key)
|
||
|
|
if not project_id:
|
||
|
|
raise ValueError(f"Unknown notebook: {notebook_key}")
|
||
|
|
|
||
|
|
async with httpx.AsyncClient(timeout=60) as http:
|
||
|
|
# Trigger run — POST /projects/{projectId}/runs
|
||
|
|
resp = await http.post(
|
||
|
|
f"{HEX_API_BASE}/projects/{project_id}/runs",
|
||
|
|
headers=HEADERS,
|
||
|
|
json={"inputParams": {"user_id": user_id}},
|
||
|
|
)
|
||
|
|
resp.raise_for_status()
|
||
|
|
run_id = resp.json()["runId"]
|
||
|
|
|
||
|
|
# Poll for completion — GET /projects/{projectId}/runs/{runId}
|
||
|
|
for _ in range(30):
|
||
|
|
status_resp = await http.get(
|
||
|
|
f"{HEX_API_BASE}/projects/{project_id}/runs/{run_id}",
|
||
|
|
headers=HEADERS,
|
||
|
|
)
|
||
|
|
status_resp.raise_for_status()
|
||
|
|
data = status_resp.json()
|
||
|
|
if data["status"] == "COMPLETED":
|
||
|
|
return {"status": "COMPLETED", "elapsed": data.get("elapsedTime")}
|
||
|
|
if data["status"] in ("ERRORED", "KILLED", "UNABLE_TO_ALLOCATE_KERNEL"):
|
||
|
|
raise RuntimeError(f"Hex run failed: {data['status']}")
|
||
|
|
await asyncio.sleep(2)
|
||
|
|
|
||
|
|
raise TimeoutError("Hex notebook run timed out")
|