Spaces:
Sleeping
Sleeping
| import asyncio | |
| import base64 | |
| import hashlib | |
| import hmac | |
| import json | |
| import logging | |
| import os | |
| import re | |
| import secrets | |
| import time | |
| from typing import Any, Literal | |
| import gradio as gr | |
| import uvicorn | |
| from fastapi import FastAPI, Header, HTTPException | |
| from fastapi.middleware.cors import CORSMiddleware | |
| from fastapi.responses import HTMLResponse | |
| from gradio_client import Client | |
| from pydantic import BaseModel, Field | |
| SESSION_SECRET = os.getenv("SESSION_SECRET", "") | |
| HF_TOKEN = os.getenv("HF_TOKEN", "") | |
| TARGET_SPACE_ID = os.getenv("TARGET_SPACE_ID", "ASQ2/ASQ") | |
| INITIAL_PROMPT = os.getenv("ASQ_INITIAL_PROMPT", "Hello") | |
| TARGET_CHAT_API_NAME = os.getenv("TARGET_CHAT_API_NAME", "/chat_flow") | |
| TARGET_INITIALIZE_API_NAME = os.getenv("TARGET_INITIALIZE_API_NAME", "/initialize_flow") | |
| SESSION_TTL_SECONDS = int(os.getenv("SESSION_TTL_SECONDS", "14400")) | |
| MAX_BACKGROUND_CHARS = int(os.getenv("MAX_BACKGROUND_CHARS", "20000")) | |
| PREPARE_TIMEOUT_SECONDS = int(os.getenv("PREPARE_TIMEOUT_SECONDS", "600")) | |
| PREPARE_RETRY_SECONDS = int(os.getenv("PREPARE_RETRY_SECONDS", "15")) | |
| PAIRING_TTL_SECONDS = int(os.getenv("PAIRING_TTL_SECONDS", "600")) | |
| logger = logging.getLogger("asq_gateway") | |
| class PrepareRequest(BaseModel): | |
| background: str = Field(min_length=20, max_length=MAX_BACKGROUND_CHARS) | |
| class JoinRequest(BaseModel): | |
| code: str = Field(min_length=6, max_length=6, pattern=r"^\d{6}$") | |
| class SuggestRequest(BaseModel): | |
| transcript: str = Field(min_length=1, max_length=4000) | |
| sessions: dict[str, dict[str, Any]] = {} | |
| interview_codes: dict[str, str] = {} | |
| pairing_codes: dict[str, str] = {} | |
| pairings_by_token: dict[str, dict[str, Any]] = {} | |
| preparation_tasks: set[asyncio.Task[Any]] = set() | |
| def _b64encode(data: bytes) -> str: | |
| return base64.urlsafe_b64encode(data).decode("ascii").rstrip("=") | |
| def _b64decode(data: str) -> bytes: | |
| return base64.urlsafe_b64decode(data + "=" * (-len(data) % 4)) | |
| def _sign(payload: str) -> str: | |
| return _b64encode( | |
| hmac.new(SESSION_SECRET.encode("utf-8"), payload.encode("utf-8"), hashlib.sha256).digest() | |
| ) | |
| def _cleanup_expired() -> None: | |
| now = int(time.time()) | |
| expired_ids = [sid for sid, session in sessions.items() if session["expires_at"] <= now] | |
| for sid in expired_ids: | |
| session = sessions.pop(sid, None) | |
| if session and session.get("interview_code"): | |
| interview_codes.pop(session["interview_code"], None) | |
| expired_pair_tokens = [ | |
| token for token, pairing in pairings_by_token.items() if pairing["expires_at"] <= now | |
| ] | |
| for pair_token in expired_pair_tokens: | |
| pairing = pairings_by_token.pop(pair_token, None) | |
| if pairing: | |
| pairing_codes.pop(pairing["code"], None) | |
| def _new_interview_code() -> str: | |
| _cleanup_expired() | |
| for _ in range(100): | |
| code = f"{secrets.randbelow(1_000_000):06d}" | |
| if code not in interview_codes: | |
| return code | |
| raise RuntimeError("Could not allocate an interview code") | |
| def _new_pairing_code() -> str: | |
| _cleanup_expired() | |
| for _ in range(100): | |
| code = f"{secrets.randbelow(1_000_000):06d}" | |
| if code not in pairing_codes and code not in interview_codes: | |
| return code | |
| raise RuntimeError("Could not allocate a pairing code") | |
| def create_token(session_id: str, role: Literal["owner", "glasses"], expires_at: int) -> str: | |
| payload = _b64encode( | |
| json.dumps( | |
| {"sid": session_id, "role": role, "exp": expires_at}, separators=(",", ":") | |
| ).encode("utf-8") | |
| ) | |
| return f"{payload}.{_sign(payload)}" | |
| def verify_token(token: str, roles: set[str] | None = None) -> tuple[dict[str, Any], str]: | |
| if not SESSION_SECRET: | |
| raise HTTPException(status_code=503, detail="Gateway is not configured") | |
| try: | |
| payload, signature = token.split(".", 1) | |
| if not hmac.compare_digest(signature, _sign(payload)): | |
| raise ValueError("bad signature") | |
| data = json.loads(_b64decode(payload)) | |
| if int(data["exp"]) <= int(time.time()): | |
| raise ValueError("expired") | |
| role = str(data["role"]) | |
| session = sessions.get(str(data["sid"])) | |
| if not session: | |
| raise ValueError("unknown session") | |
| if roles is not None and role not in roles: | |
| raise HTTPException(status_code=403, detail="This session cannot perform that action") | |
| return session, role | |
| except HTTPException: | |
| raise | |
| except (ValueError, KeyError, json.JSONDecodeError): | |
| raise HTTPException(status_code=401, detail="Invalid or expired session") | |
| def _bearer(authorization: str | None, roles: set[str] | None = None) -> tuple[dict[str, Any], str]: | |
| if not authorization or not authorization.startswith("Bearer "): | |
| raise HTTPException(status_code=401, detail="Missing session token") | |
| return verify_token(authorization[7:], roles) | |
| def _extract_turn_count(value: Any) -> int: | |
| if isinstance(value, (int, float)): | |
| return int(value) | |
| match = re.search(r"\d+", str(value or "")) | |
| return int(match.group(0)) if match else 0 | |
| def _message_text(message: Any) -> str: | |
| if isinstance(message, str): | |
| return message.strip() | |
| if isinstance(message, dict): | |
| content = message.get("content", message.get("text", "")) | |
| return _message_text(content) | |
| if isinstance(message, (list, tuple)): | |
| parts = [_message_text(item) for item in message] | |
| return "\n".join(part for part in parts if part).strip() | |
| return str(message or "").strip() | |
| def _first_complete_sentence(text: str) -> str: | |
| cleaned = text.strip() | |
| sentence_end = re.search(r"[.!?。！？]", cleaned) | |
| if sentence_end: | |
| return cleaned[: sentence_end.end()].strip() | |
| return cleaned.splitlines()[0].strip() if cleaned else "" | |
| def extract_initial_state(result: Any) -> tuple[str, list[Any], int]: | |
| # /initialize_flow returns: | |
| # turn markdown, interview history, hypotheses, summary, strategy, status. | |
| if not isinstance(result, (list, tuple)) or len(result) < 6: | |
| raise RuntimeError("Unexpected initialization response from the ASQ Space") | |
| hypotheses = str(result[2] or "").strip() | |
| history = result[1] if isinstance(result[1], list) else [] | |
| turns = _extract_turn_count(result[0]) | |
| if not hypotheses: | |
| raise RuntimeError("ASQ returned no initial hypotheses") | |
| return hypotheses, history, turns | |
| def extract_question(result: Any) -> tuple[str, list[Any], int, str | None]: | |
| # /chat_flow returns: | |
| # history, cleared input, turn markdown, hypotheses, summary, strategy, status. | |
| if not isinstance(result, (list, tuple)) or len(result) < 7: | |
| raise RuntimeError("Unexpected response from the ASQ Space") | |
| history = result[0] | |
| turns = _extract_turn_count(result[2]) | |
| hypotheses = str(result[3]).strip() if result[3] is not None else None | |
| if not isinstance(history, list) or not history: | |
| raise RuntimeError("ASQ returned no conversation history") | |
| assistant_messages = [ | |
| message | |
| for message in history | |
| if isinstance(message, dict) and message.get("role") == "assistant" | |
| ] | |
| question = _first_complete_sentence( | |
| _message_text((assistant_messages or history)[-1]) | |
| ) | |
| if not question: | |
| raise RuntimeError("ASQ returned an empty question") | |
| return question, history, turns, hypotheses | |
| def initialize_asq(background: str) -> tuple[str, list[Any], int]: | |
| if not HF_TOKEN: | |
| raise RuntimeError("HF_TOKEN is not configured") | |
| client = Client(TARGET_SPACE_ID, hf_token=HF_TOKEN, verbose=False) | |
| result = client.predict(background, api_name=TARGET_INITIALIZE_API_NAME) | |
| return extract_initial_state(result) | |
| def call_asq(transcript: str, session: dict[str, Any]) -> tuple[str, list[Any], int, str | None]: | |
| if not HF_TOKEN: | |
| raise RuntimeError("HF_TOKEN is not configured") | |
| client = Client(TARGET_SPACE_ID, hf_token=HF_TOKEN, verbose=False) | |
| result = client.predict( | |
| transcript, | |
| session["history"], | |
| session["background"], | |
| api_name=TARGET_CHAT_API_NAME, | |
| ) | |
| return extract_question(result) | |
| async def prepare_session(session_id: str, background: str) -> None: | |
| session = sessions.get(session_id) | |
| if not session: | |
| return | |
| deadline = time.monotonic() + PREPARE_TIMEOUT_SECONDS | |
| attempt = 0 | |
| while True: | |
| attempt += 1 | |
| try: | |
| hypotheses, history, turns = await asyncio.to_thread(initialize_asq, background) | |
| break | |
| except Exception as error: | |
| remaining = deadline - time.monotonic() | |
| logger.warning("ASQ preparation attempt %s failed: %s", attempt, type(error).__name__) | |
| if remaining <= 0: | |
| if session_id in sessions: | |
| session["status"] = "error" | |
| session["error"] = "Unable to start the model. Please try again." | |
| return | |
| await asyncio.sleep(min(PREPARE_RETRY_SECONDS, remaining)) | |
| if session_id in sessions: | |
| session.update( | |
| status="ready", | |
| hypotheses=hypotheses, | |
| history=history, | |
| turns=turns, | |
| error=None, | |
| ) | |
| def session_view(session: dict[str, Any], role: str) -> dict[str, Any]: | |
| response: dict[str, Any] = { | |
| "status": session["status"], | |
| "ready": session["status"] == "ready", | |
| "expiresAt": session["expires_at"], | |
| "question": session.get("current_question", ""), | |
| "questionRevision": session.get("question_revision", 0), | |
| "turns": session.get("turns", 0), | |
| } | |
| if session.get("error"): | |
| response["error"] = session["error"] | |
| if role == "owner": | |
| response["interviewCode"] = session.get("interview_code") | |
| response["hypotheses"] = session.get("hypotheses", "") | |
| return response | |
| api = FastAPI(title="ASQ Rokid Gateway", version="0.5.0") | |
| api.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=["*"], | |
| allow_methods=["GET", "POST", "OPTIONS"], | |
| allow_headers=["Authorization", "Content-Type"], | |
| ) | |
| def mobile_app() -> str: | |
| return MOBILE_APP_HTML | |
| def health() -> dict[str, Any]: | |
| return { | |
| "ok": True, | |
| "configured": bool(SESSION_SECRET and HF_TOKEN), | |
| "target": TARGET_SPACE_ID, | |
| "version": "0.5.0", | |
| } | |
| def create_owner_session() -> dict[str, Any]: | |
| if not SESSION_SECRET: | |
| raise HTTPException(status_code=503, detail="Gateway is not configured") | |
| _cleanup_expired() | |
| session_id = secrets.token_urlsafe(18) | |
| expires_at = int(time.time()) + SESSION_TTL_SECONDS | |
| sessions[session_id] = { | |
| "background": "", | |
| "hypotheses": "", | |
| "history": [], | |
| "turns": 0, | |
| "current_question": "", | |
| "question_revision": 0, | |
| "status": "draft", | |
| "error": None, | |
| "interview_code": None, | |
| "expires_at": expires_at, | |
| } | |
| return { | |
| "token": create_token(session_id, "owner", expires_at), | |
| "expiresIn": SESSION_TTL_SECONDS, | |
| } | |
| async def prepare( | |
| request: PrepareRequest, | |
| authorization: str | None = Header(default=None), | |
| ) -> dict[str, Any]: | |
| session, _ = _bearer(authorization, {"owner"}) | |
| if session["status"] == "starting": | |
| return session_view(session, "owner") | |
| old_code = session.get("interview_code") | |
| if old_code: | |
| interview_codes.pop(old_code, None) | |
| code = _new_interview_code() | |
| session_id = next(sid for sid, value in sessions.items() if value is session) | |
| interview_codes[code] = session_id | |
| session.update( | |
| background=request.background.strip(), | |
| hypotheses="", | |
| history=[], | |
| turns=0, | |
| current_question="", | |
| question_revision=0, | |
| status="starting", | |
| error=None, | |
| interview_code=code, | |
| ) | |
| task = asyncio.create_task(prepare_session(session_id, session["background"])) | |
| preparation_tasks.add(task) | |
| task.add_done_callback(preparation_tasks.discard) | |
| return session_view(session, "owner") | |
| def create_pairing() -> dict[str, Any]: | |
| _cleanup_expired() | |
| code = _new_pairing_code() | |
| pair_token = secrets.token_urlsafe(24) | |
| expires_at = int(time.time()) + PAIRING_TTL_SECONDS | |
| pairings_by_token[pair_token] = { | |
| "code": code, | |
| "session_id": None, | |
| "expires_at": expires_at, | |
| } | |
| pairing_codes[code] = pair_token | |
| return { | |
| "pairingCode": code, | |
| "pairingToken": pair_token, | |
| "expiresIn": PAIRING_TTL_SECONDS, | |
| } | |
| def pairing_status( | |
| x_pairing_token: str | None = Header(default=None, alias="X-Pairing-Token"), | |
| ) -> dict[str, Any]: | |
| _cleanup_expired() | |
| pairing = pairings_by_token.get(x_pairing_token or "") | |
| if not pairing: | |
| raise HTTPException(status_code=401, detail="Invalid or expired pairing") | |
| session_id = pairing.get("session_id") | |
| if not session_id: | |
| return {"paired": False, "expiresAt": pairing["expires_at"]} | |
| session = sessions.get(session_id) | |
| if not session: | |
| raise HTTPException(status_code=410, detail="Interview session expired") | |
| return { | |
| "paired": True, | |
| "token": create_token(session_id, "glasses", session["expires_at"]), | |
| **session_view(session, "glasses"), | |
| } | |
| def claim_pairing( | |
| request: JoinRequest, | |
| authorization: str | None = Header(default=None), | |
| ) -> dict[str, bool]: | |
| session, _ = _bearer(authorization, {"owner"}) | |
| _cleanup_expired() | |
| pair_token = pairing_codes.get(request.code) | |
| pairing = pairings_by_token.get(pair_token or "") | |
| if not pairing or not hmac.compare_digest(request.code, pairing.get("code") or ""): | |
| raise HTTPException(status_code=401, detail="Incorrect or expired glasses pairing code") | |
| session_id = next(sid for sid, value in sessions.items() if value is session) | |
| pairing["session_id"] = session_id | |
| return {"paired": True} | |
| def join(request: JoinRequest) -> dict[str, Any]: | |
| _cleanup_expired() | |
| session_id = interview_codes.get(request.code) | |
| session = sessions.get(session_id or "") | |
| if not session or not hmac.compare_digest(request.code, session.get("interview_code") or ""): | |
| raise HTTPException(status_code=401, detail="Incorrect or expired interview code") | |
| return { | |
| "token": create_token(session_id, "glasses", session["expires_at"]), | |
| **session_view(session, "glasses"), | |
| } | |
| def get_session(authorization: str | None = Header(default=None)) -> dict[str, Any]: | |
| session, role = _bearer(authorization) | |
| return session_view(session, role) | |
| async def start(authorization: str | None = Header(default=None)) -> dict[str, Any]: | |
| session, role = _bearer(authorization, {"owner", "glasses"}) | |
| if session["status"] != "ready": | |
| raise HTTPException(status_code=409, detail="Interview is not ready") | |
| if session.get("current_question"): | |
| return session_view(session, role) | |
| try: | |
| question, history, turns, hypotheses = await asyncio.to_thread( | |
| call_asq, INITIAL_PROMPT, session | |
| ) | |
| except Exception as error: | |
| raise HTTPException(status_code=502, detail=f"ASQ inference failed: {error}") | |
| session["history"] = history | |
| session["turns"] = turns | |
| session["current_question"] = question | |
| session["question_revision"] = session.get("question_revision", 0) + 1 | |
| if hypotheses: | |
| session["hypotheses"] = hypotheses | |
| return session_view(session, role) | |
| async def suggest( | |
| request: SuggestRequest, | |
| authorization: str | None = Header(default=None), | |
| ) -> dict[str, str]: | |
| session, _ = _bearer(authorization, {"owner", "glasses"}) | |
| if session["status"] != "ready": | |
| raise HTTPException(status_code=409, detail="Interview is not ready") | |
| try: | |
| question, history, turns, hypotheses = await asyncio.to_thread( | |
| call_asq, request.transcript.strip(), session | |
| ) | |
| except Exception as error: | |
| raise HTTPException(status_code=502, detail=f"ASQ inference failed: {error}") | |
| session["history"] = history | |
| session["turns"] = turns | |
| session["current_question"] = question | |
| session["question_revision"] = session.get("question_revision", 0) + 1 | |
| if hypotheses: | |
| session["hypotheses"] = hypotheses | |
| return {"question": question} | |
| MOBILE_APP_HTML = """<!doctype html> | |
| <html lang="en"><head><meta charset="utf-8"><meta name="viewport" content="width=device-width,initial-scale=1"> | |
| <title>ASQ Interview Setup</title> | |
| <style> | |
| :root{font-family:Inter,system-ui,sans-serif;color:#16201d;background:#eef5f1}*{box-sizing:border-box} | |
| body{margin:0;padding:24px 16px}.card{max-width:760px;margin:auto;background:#fff;border-radius:20px;padding:24px;box-shadow:0 12px 36px #173b2b1a} | |
| h1{margin:0 0 6px;font-size:1.55rem}.muted{color:#65716d;margin:0 0 24px}.step{border-top:1px solid #e4ebe7;padding-top:20px;margin-top:20px} | |
| label{display:block;font-weight:650;margin-bottom:8px}input,textarea,button{font:inherit;width:100%;border-radius:12px}input,textarea{border:1px solid #b9c9c1;padding:13px}textarea{min-height:220px;resize:vertical} | |
| button{border:0;background:#087a54;color:#fff;padding:14px;font-weight:700;margin-top:12px}button:disabled{background:#9cada5}.hidden{display:none}.status{padding:14px;border-radius:12px;background:#edf7f2;margin-top:14px;white-space:pre-wrap} | |
| .hypo{white-space:pre-wrap;line-height:1.5;background:#f7faf8;padding:14px;border-radius:12px;max-height:320px;overflow:auto} | |
| </style></head><body><main class="card"> | |
| <h1>ASQ Interview Setup</h1><p class="muted">Open ASQ on the Rokid glasses first. The glasses will display a six-digit pairing code; enter that code on this phone.</p> | |
| <section id="prepare"><label for="pairingCode">Code shown on the glasses</label><input id="pairingCode" inputmode="numeric" maxlength="6" placeholder="Six digits"><label for="background" style="margin-top:16px">Case and interview background</label><textarea id="background" placeholder="Enter the case facts, context, existing evidence, interview goals, and any constraints..."></textarea><button id="prepareBtn" disabled>Connect glasses and prepare interview</button></section> | |
| <section id="progress" class="step hidden"><div id="status" class="status">Connecting the glasses…</div><div id="hypoWrap" class="hidden"><h2>Initial hypotheses</h2><div id="hypotheses" class="hypo"></div></div></section> | |
| </main><script> | |
| let token="", timer=null; const $=id=>document.getElementById(id); | |
| async function api(path, options={}){const headers={"Content-Type":"application/json",...(options.headers||{})};if(token)headers.Authorization=`Bearer ${token}`;const r=await fetch(path,{...options,headers});let data={};try{data=await r.json()}catch{}if(!r.ok)throw new Error(data.detail||`Request failed (${r.status})`);return data} | |
| function errorMessage(e){$("status").textContent=e.message;$("progress").classList.remove("hidden")} | |
| async function begin(){try{const data=await api("/auth",{method:"POST"});token=data.token;$("prepareBtn").disabled=false}catch(e){errorMessage(e)}} | |
| $("prepareBtn").onclick=async()=>{const pairingCode=$("pairingCode").value.trim();const background=$("background").value.trim();if(!/^\d{6}$/.test(pairingCode)){errorMessage("Enter the six-digit code shown on the glasses.");return}if(background.length<20){errorMessage("Please provide at least 20 characters of background.");return}$("prepareBtn").disabled=true;$("progress").classList.remove("hidden");$("status").textContent="Connecting the glasses and generating initial hypotheses…";try{await api("/pair/claim",{method:"POST",body:JSON.stringify({code:pairingCode})});const data=await api("/prepare",{method:"POST",body:JSON.stringify({background})});show(data);timer=setInterval(poll,3000)}catch(e){errorMessage(e);$("prepareBtn").disabled=false}}; | |
| async function poll(){try{show(await api("/session"))}catch(e){clearInterval(timer);errorMessage(e)}} | |
| function show(data){if(data.status==="ready"){clearInterval(timer);$("status").textContent="Ready. The glasses are connected and can begin the interview.";$("hypotheses").textContent=data.hypotheses||"No hypotheses returned.";$("hypoWrap").classList.remove("hidden")}else if(data.status==="error"){clearInterval(timer);$("status").textContent=data.error||"Model startup failed.";$("prepareBtn").disabled=false}else{$("status").textContent="Glasses connected. Starting the model and generating initial hypotheses…"}} | |
| begin(); | |
| </script></body></html>""" | |
| with gr.Blocks(title="ASQ Rokid Gateway") as status_page: | |
| gr.Markdown( | |
| "# ASQ Rokid Gateway\n" | |
| "This service provides phone preparation and Rokid interview pairing APIs.\n\n" | |
| "Open the site root on a phone to prepare an interview. Use `GET /health` for service status." | |
| ) | |
| api = gr.mount_gradio_app(api, status_page, path="/ui") | |
| if __name__ == "__main__": | |
| uvicorn.run(api, host="0.0.0.0", port=7860) | |