"""DataInventor: open-source training data creation and curation. Bring any OpenAI-compatible endpoint.""" from __future__ import annotations import base64 import csv import io import json import os import secrets import threading import urllib.parse import urllib.request import time import uuid import sys from pathlib import Path sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) # the datainventor package sits next to app/ from fastapi import FastAPI, HTTPException, Request from fastapi.responses import FileResponse, JSONResponse, RedirectResponse, Response from fastapi.staticfiles import StaticFiles from pydantic import BaseModel from publish import LICENSES, publish as hub_publish, slug from engine import HF_BASE, HF_MODEL, Endpoint, EndpointError, Job, make_blueprint, run_job, stats ROOT = Path(__file__).resolve().parent.parent STATIC = ROOT / "static" MAX_ROWS = int(os.environ.get("MAX_ROWS", "200")) MAX_RUNNING = int(os.environ.get("MAX_RUNNING", "2")) SERVER_KEY = os.environ.get("HF_TOKEN", "") app = FastAPI(title="DataInventor") JOBS: dict[str, Job] = {} ENDPOINTS: dict[str, Endpoint] = {} # kept apart from Job so a key never lands in a response GUARD = threading.Lock() class SpecIn(BaseModel): format: str = "instruction" request: str domains: list[str] = [] subdomains: list[str] = [] languages: list[str] = ["English"] constraints: str = "" rows: int = 100 blueprint: dict | None = None def make_endpoint() -> Endpoint: """The endpoint is fixed by the host (environment variables). Visitors never supply a URL or a key.""" key = os.environ.get("DI_API_KEY") or SERVER_KEY if not key: raise HTTPException(503, "This host has no model configured. Set DI_API_KEY (or HF_TOKEN) and restart.") return Endpoint(base_url=os.environ.get("DI_BASE_URL", HF_BASE).rstrip("/"), api_key=key, model=os.environ.get("DI_MODEL", HF_MODEL), disable_thinking=os.environ.get("DI_DISABLE_THINKING", "0") != "0", workers=max(1, min(int(os.environ.get("DI_WORKERS", "4")), 8)), min_max_tokens=int(os.environ.get("DI_MIN_MAX_TOKENS", "0"))) @app.get("/api/defaults") def defaults(): return {"model": os.environ.get("DI_MODEL", HF_MODEL), "configured": bool(os.environ.get("DI_API_KEY") or SERVER_KEY), "max_rows": MAX_ROWS} @app.get("/api/domains") def domains(): return JSONResponse(json.loads((STATIC / "domains.json").read_text())) @app.post("/api/blueprint") def blueprint(s: SpecIn): ep = make_endpoint() if not s.request.strip(): raise HTTPException(400, "Describe the dataset first.") try: return make_blueprint(ep, s.request, s.domains, s.subdomains, s.languages, s.constraints, s.rows) except EndpointError as err: raise HTTPException(502, str(err)) @app.post("/api/jobs") def start(s: SpecIn): ep = make_endpoint() if not s.request.strip(): raise HTTPException(400, "Describe the dataset first.") if not 1 <= s.rows <= MAX_ROWS: raise HTTPException(400, f"Rows must be between 1 and {MAX_ROWS} on this host.") with GUARD: running = sum(1 for j in JOBS.values() if j.state in ("queued", "planning", "generating")) if running >= MAX_RUNNING: raise HTTPException(429, "Two datasets are already being written. Wait for one to finish.") jid = uuid.uuid4().hex[:10] cfg = s.model_dump(exclude={"engine"}) job = Job(id=jid, config=cfg) JOBS[jid] = job ENDPOINTS[jid] = ep for old in sorted(JOBS, key=lambda k: JOBS[k].started)[:-20]: JOBS.pop(old, None); ENDPOINTS.pop(old, None) threading.Thread(target=run_job, args=(job, ep), daemon=True).start() return {"id": jid} def get_job(jid: str) -> Job: if jid not in JOBS: raise HTTPException(404, "That dataset is no longer here. Start a new run.") return JOBS[jid] @app.get("/api/jobs/{jid}") def status(jid: str, since: int = 0): j = get_job(jid) with j.lock: rows = list(j.rows) out = {"id": j.id, "state": j.state, "message": j.message, "target": j.target, "count": len(rows), "attempts": j.attempts, "rejected": dict(j.rejected), "config": {k: j.config.get(k) for k in ("format", "request", "domains", "languages", "constraints", "rows")}, "slug": slug(j.config.get("request", "")), "elapsed": round((j.finished or time.time()) - j.started, 1), "new_rows": rows[since:since + 50]} if j.state in ("done", "stopped", "failed"): out["stats"] = stats(rows) out["blueprint"] = j.blueprint return out @app.post("/api/jobs/{jid}/stop") def stop(jid: str): get_job(jid).stop.set() return {"ok": True} @app.get("/api/jobs/{jid}/download") def download(jid: str, format: str = "csv"): j = get_job(jid) rows = list(j.rows) if not rows: raise HTTPException(404, "There are no rows to download yet.") pref = "chosen" in rows[0] cols = ["prompt", "chosen", "rejected"] if pref else ["prompt", "completion"] meta = ["area", "subtopic", "topic", "persona", "shape", "difficulty", "language"] name = f"datainventor-{jid}" if format == "jsonl": body = "\n".join(json.dumps({k: r.get(k) for k in cols + meta}, ensure_ascii=False) for r in rows) return Response(body, media_type="application/jsonl", headers={"Content-Disposition": f'attachment; filename="{name}.jsonl"'}) buf = io.StringIO() w = csv.writer(buf) w.writerow(cols + meta) for r in rows: w.writerow([r.get(k, "") for k in cols + meta]) return Response(buf.getvalue().encode("utf-8-sig"), media_type="text/csv", headers={"Content-Disposition": f'attachment; filename="{name}.csv"'}) # ---------------------------------------------------------------- sign in with Hugging Face (publishing) OAUTH_ID = os.environ.get("OAUTH_CLIENT_ID", "") OAUTH_SECRET = os.environ.get("OAUTH_CLIENT_SECRET", "") OAUTH_URL = os.environ.get("OPENID_PROVIDER_URL", "https://huggingface.co").rstrip("/") SPACE_HOST = os.environ.get("SPACE_HOST", "") DEV_TOKEN = os.environ.get("DI_DEV_USER_TOKEN", "") # local testing only SESSIONS: dict[str, dict] = {} PUBLISHES: dict[str, list] = {} SESSION_SECONDS = 2 * 3600 def current_session(request: Request) -> dict | None: sid = request.cookies.get("di_sid") s = SESSIONS.get(sid or "") if s and s["exp"] > time.time(): return s if DEV_TOKEN: from huggingface_hub import HfApi return {"user": HfApi(token=DEV_TOKEN).whoami()["name"], "token": DEV_TOKEN, "exp": time.time() + 600, "sid": "dev"} return None def redirect_uri() -> str: return f"https://{SPACE_HOST}/oauth/callback" @app.get("/api/me") def me(request: Request): s = current_session(request) return {"user": s["user"] if s else None, "can_sign_in": bool(OAUTH_ID and SPACE_HOST) or bool(DEV_TOKEN), "host": SPACE_HOST} @app.get("/oauth/login") def login(job: str = ""): if not (OAUTH_ID and SPACE_HOST): raise HTTPException(404, "Sign-in is not available on this host.") state = secrets.token_urlsafe(24) q = urllib.parse.urlencode({"client_id": OAUTH_ID, "redirect_uri": redirect_uri(), "response_type": "code", "scope": "openid profile contribute-repos", "state": state}) resp = RedirectResponse(f"{OAUTH_URL}/oauth/authorize?{q}") resp.set_cookie("di_state", state + "|" + (job if job.isalnum() else ""), max_age=600, httponly=True, secure=True, samesite="lax") return resp @app.get("/oauth/callback") def callback(request: Request, code: str = "", state: str = "", error: str = ""): saved = request.cookies.get("di_state", "") want, _, job = saved.partition("|") if error or not code or not want or not secrets.compare_digest(want, state): return RedirectResponse("/?signin=failed") basic = base64.b64encode(f"{OAUTH_ID}:{OAUTH_SECRET}".encode()).decode() try: data = urllib.parse.urlencode({"grant_type": "authorization_code", "code": code, "redirect_uri": redirect_uri(), "client_id": OAUTH_ID}).encode() tok = json.load(urllib.request.urlopen(urllib.request.Request(f"{OAUTH_URL}/oauth/token", data, {"Authorization": "Basic " + basic, "Content-Type": "application/x-www-form-urlencoded"}), timeout=30)) info = json.load(urllib.request.urlopen(urllib.request.Request(f"{OAUTH_URL}/oauth/userinfo", headers={"Authorization": "Bearer " + tok["access_token"]}), timeout=30)) except Exception: return RedirectResponse("/?signin=failed") sid = secrets.token_urlsafe(24) SESSIONS[sid] = {"user": info.get("preferred_username") or info.get("name"), "token": tok["access_token"], "exp": time.time() + SESSION_SECONDS, "sid": sid} for k in [k for k, v in SESSIONS.items() if v["exp"] < time.time()]: SESSIONS.pop(k, None) resp = RedirectResponse("/?signin=ok" + (f"&job={job}" if job else "")) resp.set_cookie("di_sid", sid, max_age=SESSION_SECONDS, httponly=True, secure=True, samesite="lax") resp.delete_cookie("di_state") return resp @app.post("/api/signout") def signout(request: Request): SESSIONS.pop(request.cookies.get("di_sid", ""), None) resp = JSONResponse({"ok": True}) resp.delete_cookie("di_sid") return resp class PublishIn(BaseModel): name: str private: bool = True license: str = "apache-2.0" confirm: bool = False @app.post("/api/jobs/{jid}/publish") def publish(jid: str, body: PublishIn, request: Request): s = current_session(request) if not s: raise HTTPException(401, "Sign in with Hugging Face to publish.") if not body.confirm: raise HTTPException(400, "Confirm that you have the right to publish this data.") if body.license not in LICENSES: raise HTTPException(400, "Choose one of the offered licenses.") j = get_job(jid) if j.state not in ("done", "stopped") or not j.rows: raise HTTPException(400, "Wait for the run to finish, or stop it, before publishing.") now = time.time() recent = [t for t in PUBLISHES.get(s["user"], []) if now - t < 3600] if len(recent) >= 5: raise HTTPException(429, "You have published five datasets in the last hour. Try again later.") try: url = hub_publish(s["token"], s["user"], body.name.strip(), body.private, body.license, j.config, list(j.rows), stats(list(j.rows)), ENDPOINTS[jid].model if jid in ENDPOINTS else "") except ValueError as e: raise HTTPException(400, str(e)) except Exception as e: msg = str(e) if "409" in msg or "already" in msg.lower(): raise HTTPException(409, "You already have a dataset with that name. Pick another name.") if "403" in msg or "401" in msg: raise HTTPException(403, "Hugging Face refused the request. Sign in again and allow DataInventor to create datasets.") raise HTTPException(502, "Publishing failed. Try again in a moment.") PUBLISHES[s["user"]] = recent + [now] return {"url": url} from curate_api import router as curate_router # noqa: E402 (imports make_endpoint from this module) app.include_router(curate_router) from visualize_api import router as vis_router # noqa: E402 app.include_router(vis_router) def _asset_version() -> str: import hashlib h = hashlib.sha1() for f in ("app.js", "curate.js", "visualize.js", "app.css"): h.update((STATIC / f).read_bytes()) return h.hexdigest()[:10] INDEX_HTML = (STATIC / "index.html").read_text().replace(".js\"", f".js?v={_asset_version()}\"").replace(".css\"", f".css?v={_asset_version()}\"") @app.get("/") def index(): # never cache the page itself (embeds kept showing an old build); assets carry a content hash instead return Response(INDEX_HTML, media_type="text/html", headers={"Cache-Control": "no-cache"}) app.mount("/static", StaticFiles(directory=STATIC), name="static")