LucasUn commited on
Commit
784a7d0
·
verified ·
1 Parent(s): 30c384e

Guard gateway questions to first sentence

Browse files
Files changed (1) hide show
  1. app.py +511 -501
app.py CHANGED
@@ -1,501 +1,511 @@
1
- import asyncio
2
- import base64
3
- import hashlib
4
- import hmac
5
- import json
6
- import logging
7
- import os
8
- import re
9
- import secrets
10
- import time
11
- from typing import Any, Literal
12
-
13
- import gradio as gr
14
- import uvicorn
15
- from fastapi import FastAPI, Header, HTTPException
16
- from fastapi.middleware.cors import CORSMiddleware
17
- from fastapi.responses import HTMLResponse
18
- from gradio_client import Client
19
- from pydantic import BaseModel, Field
20
-
21
-
22
- SESSION_SECRET = os.getenv("SESSION_SECRET", "")
23
- HF_TOKEN = os.getenv("HF_TOKEN", "")
24
- TARGET_SPACE_ID = os.getenv("TARGET_SPACE_ID", "ASQ2/ASQ")
25
- INITIAL_PROMPT = os.getenv("ASQ_INITIAL_PROMPT", "Hello")
26
- TARGET_CHAT_API_NAME = os.getenv("TARGET_CHAT_API_NAME", "/chat_flow")
27
- TARGET_INITIALIZE_API_NAME = os.getenv("TARGET_INITIALIZE_API_NAME", "/initialize_flow")
28
- SESSION_TTL_SECONDS = int(os.getenv("SESSION_TTL_SECONDS", "14400"))
29
- MAX_BACKGROUND_CHARS = int(os.getenv("MAX_BACKGROUND_CHARS", "20000"))
30
- PREPARE_TIMEOUT_SECONDS = int(os.getenv("PREPARE_TIMEOUT_SECONDS", "600"))
31
- PREPARE_RETRY_SECONDS = int(os.getenv("PREPARE_RETRY_SECONDS", "15"))
32
- PAIRING_TTL_SECONDS = int(os.getenv("PAIRING_TTL_SECONDS", "600"))
33
-
34
- logger = logging.getLogger("asq_gateway")
35
-
36
-
37
- class PrepareRequest(BaseModel):
38
- background: str = Field(min_length=20, max_length=MAX_BACKGROUND_CHARS)
39
-
40
-
41
- class JoinRequest(BaseModel):
42
- code: str = Field(min_length=6, max_length=6, pattern=r"^\d{6}$")
43
-
44
-
45
- class SuggestRequest(BaseModel):
46
- transcript: str = Field(min_length=1, max_length=4000)
47
-
48
-
49
- sessions: dict[str, dict[str, Any]] = {}
50
- interview_codes: dict[str, str] = {}
51
- pairing_codes: dict[str, str] = {}
52
- pairings_by_token: dict[str, dict[str, Any]] = {}
53
- preparation_tasks: set[asyncio.Task[Any]] = set()
54
-
55
-
56
- def _b64encode(data: bytes) -> str:
57
- return base64.urlsafe_b64encode(data).decode("ascii").rstrip("=")
58
-
59
-
60
- def _b64decode(data: str) -> bytes:
61
- return base64.urlsafe_b64decode(data + "=" * (-len(data) % 4))
62
-
63
-
64
- def _sign(payload: str) -> str:
65
- return _b64encode(
66
- hmac.new(SESSION_SECRET.encode("utf-8"), payload.encode("utf-8"), hashlib.sha256).digest()
67
- )
68
-
69
-
70
- def _cleanup_expired() -> None:
71
- now = int(time.time())
72
- expired_ids = [sid for sid, session in sessions.items() if session["expires_at"] <= now]
73
- for sid in expired_ids:
74
- session = sessions.pop(sid, None)
75
- if session and session.get("interview_code"):
76
- interview_codes.pop(session["interview_code"], None)
77
- expired_pair_tokens = [
78
- token for token, pairing in pairings_by_token.items() if pairing["expires_at"] <= now
79
- ]
80
- for pair_token in expired_pair_tokens:
81
- pairing = pairings_by_token.pop(pair_token, None)
82
- if pairing:
83
- pairing_codes.pop(pairing["code"], None)
84
-
85
-
86
- def _new_interview_code() -> str:
87
- _cleanup_expired()
88
- for _ in range(100):
89
- code = f"{secrets.randbelow(1_000_000):06d}"
90
- if code not in interview_codes:
91
- return code
92
- raise RuntimeError("Could not allocate an interview code")
93
-
94
-
95
- def _new_pairing_code() -> str:
96
- _cleanup_expired()
97
- for _ in range(100):
98
- code = f"{secrets.randbelow(1_000_000):06d}"
99
- if code not in pairing_codes and code not in interview_codes:
100
- return code
101
- raise RuntimeError("Could not allocate a pairing code")
102
-
103
-
104
- def create_token(session_id: str, role: Literal["owner", "glasses"], expires_at: int) -> str:
105
- payload = _b64encode(
106
- json.dumps(
107
- {"sid": session_id, "role": role, "exp": expires_at}, separators=(",", ":")
108
- ).encode("utf-8")
109
- )
110
- return f"{payload}.{_sign(payload)}"
111
-
112
-
113
- def verify_token(token: str, roles: set[str] | None = None) -> tuple[dict[str, Any], str]:
114
- if not SESSION_SECRET:
115
- raise HTTPException(status_code=503, detail="Gateway is not configured")
116
- try:
117
- payload, signature = token.split(".", 1)
118
- if not hmac.compare_digest(signature, _sign(payload)):
119
- raise ValueError("bad signature")
120
- data = json.loads(_b64decode(payload))
121
- if int(data["exp"]) <= int(time.time()):
122
- raise ValueError("expired")
123
- role = str(data["role"])
124
- session = sessions.get(str(data["sid"]))
125
- if not session:
126
- raise ValueError("unknown session")
127
- if roles is not None and role not in roles:
128
- raise HTTPException(status_code=403, detail="This session cannot perform that action")
129
- return session, role
130
- except HTTPException:
131
- raise
132
- except (ValueError, KeyError, json.JSONDecodeError):
133
- raise HTTPException(status_code=401, detail="Invalid or expired session")
134
-
135
-
136
- def _bearer(authorization: str | None, roles: set[str] | None = None) -> tuple[dict[str, Any], str]:
137
- if not authorization or not authorization.startswith("Bearer "):
138
- raise HTTPException(status_code=401, detail="Missing session token")
139
- return verify_token(authorization[7:], roles)
140
-
141
-
142
- def _extract_turn_count(value: Any) -> int:
143
- if isinstance(value, (int, float)):
144
- return int(value)
145
- match = re.search(r"\d+", str(value or ""))
146
- return int(match.group(0)) if match else 0
147
-
148
-
149
- def _message_text(message: Any) -> str:
150
- if isinstance(message, str):
151
- return message.strip()
152
- if isinstance(message, dict):
153
- content = message.get("content", message.get("text", ""))
154
- return _message_text(content)
155
- if isinstance(message, (list, tuple)):
156
- parts = [_message_text(item) for item in message]
157
- return "\n".join(part for part in parts if part).strip()
158
- return str(message or "").strip()
159
-
160
-
161
- def extract_initial_state(result: Any) -> tuple[str, list[Any], int]:
162
- # /initialize_flow returns:
163
- # turn markdown, interview history, hypotheses, summary, strategy, status.
164
- if not isinstance(result, (list, tuple)) or len(result) < 6:
165
- raise RuntimeError("Unexpected initialization response from the ASQ Space")
166
- hypotheses = str(result[2] or "").strip()
167
- history = result[1] if isinstance(result[1], list) else []
168
- turns = _extract_turn_count(result[0])
169
- if not hypotheses:
170
- raise RuntimeError("ASQ returned no initial hypotheses")
171
- return hypotheses, history, turns
172
-
173
-
174
- def extract_question(result: Any) -> tuple[str, list[Any], int, str | None]:
175
- # /chat_flow returns:
176
- # history, cleared input, turn markdown, hypotheses, summary, strategy, status.
177
- if not isinstance(result, (list, tuple)) or len(result) < 7:
178
- raise RuntimeError("Unexpected response from the ASQ Space")
179
- history = result[0]
180
- turns = _extract_turn_count(result[2])
181
- hypotheses = str(result[3]).strip() if result[3] is not None else None
182
- if not isinstance(history, list) or not history:
183
- raise RuntimeError("ASQ returned no conversation history")
184
- assistant_messages = [
185
- message
186
- for message in history
187
- if isinstance(message, dict) and message.get("role") == "assistant"
188
- ]
189
- question = _message_text((assistant_messages or history)[-1])
190
- if not question:
191
- raise RuntimeError("ASQ returned an empty question")
192
- return question, history, turns, hypotheses
193
-
194
-
195
- def initialize_asq(background: str) -> tuple[str, list[Any], int]:
196
- if not HF_TOKEN:
197
- raise RuntimeError("HF_TOKEN is not configured")
198
- client = Client(TARGET_SPACE_ID, hf_token=HF_TOKEN, verbose=False)
199
- result = client.predict(background, api_name=TARGET_INITIALIZE_API_NAME)
200
- return extract_initial_state(result)
201
-
202
-
203
- def call_asq(transcript: str, session: dict[str, Any]) -> tuple[str, list[Any], int, str | None]:
204
- if not HF_TOKEN:
205
- raise RuntimeError("HF_TOKEN is not configured")
206
- client = Client(TARGET_SPACE_ID, hf_token=HF_TOKEN, verbose=False)
207
- result = client.predict(
208
- transcript,
209
- session["history"],
210
- session["background"],
211
- api_name=TARGET_CHAT_API_NAME,
212
- )
213
- return extract_question(result)
214
-
215
-
216
- async def prepare_session(session_id: str, background: str) -> None:
217
- session = sessions.get(session_id)
218
- if not session:
219
- return
220
- deadline = time.monotonic() + PREPARE_TIMEOUT_SECONDS
221
- attempt = 0
222
- while True:
223
- attempt += 1
224
- try:
225
- hypotheses, history, turns = await asyncio.to_thread(initialize_asq, background)
226
- break
227
- except Exception as error:
228
- remaining = deadline - time.monotonic()
229
- logger.warning("ASQ preparation attempt %s failed: %s", attempt, type(error).__name__)
230
- if remaining <= 0:
231
- if session_id in sessions:
232
- session["status"] = "error"
233
- session["error"] = "Unable to start the model. Please try again."
234
- return
235
- await asyncio.sleep(min(PREPARE_RETRY_SECONDS, remaining))
236
- if session_id in sessions:
237
- session.update(
238
- status="ready",
239
- hypotheses=hypotheses,
240
- history=history,
241
- turns=turns,
242
- error=None,
243
- )
244
-
245
-
246
- def session_view(session: dict[str, Any], role: str) -> dict[str, Any]:
247
- response: dict[str, Any] = {
248
- "status": session["status"],
249
- "ready": session["status"] == "ready",
250
- "expiresAt": session["expires_at"],
251
- "question": session.get("current_question", ""),
252
- "questionRevision": session.get("question_revision", 0),
253
- "turns": session.get("turns", 0),
254
- }
255
- if session.get("error"):
256
- response["error"] = session["error"]
257
- if role == "owner":
258
- response["interviewCode"] = session.get("interview_code")
259
- response["hypotheses"] = session.get("hypotheses", "")
260
- return response
261
-
262
-
263
- api = FastAPI(title="ASQ Rokid Gateway", version="0.5.0")
264
- api.add_middleware(
265
- CORSMiddleware,
266
- allow_origins=["*"],
267
- allow_methods=["GET", "POST", "OPTIONS"],
268
- allow_headers=["Authorization", "Content-Type"],
269
- )
270
-
271
-
272
- @api.get("/", response_class=HTMLResponse)
273
- def mobile_app() -> str:
274
- return MOBILE_APP_HTML
275
-
276
-
277
- @api.get("/health")
278
- def health() -> dict[str, Any]:
279
- return {
280
- "ok": True,
281
- "configured": bool(SESSION_SECRET and HF_TOKEN),
282
- "target": TARGET_SPACE_ID,
283
- "version": "0.5.0",
284
- }
285
-
286
-
287
- @api.post("/auth")
288
- def create_owner_session() -> dict[str, Any]:
289
- if not SESSION_SECRET:
290
- raise HTTPException(status_code=503, detail="Gateway is not configured")
291
- _cleanup_expired()
292
- session_id = secrets.token_urlsafe(18)
293
- expires_at = int(time.time()) + SESSION_TTL_SECONDS
294
- sessions[session_id] = {
295
- "background": "",
296
- "hypotheses": "",
297
- "history": [],
298
- "turns": 0,
299
- "current_question": "",
300
- "question_revision": 0,
301
- "status": "draft",
302
- "error": None,
303
- "interview_code": None,
304
- "expires_at": expires_at,
305
- }
306
- return {
307
- "token": create_token(session_id, "owner", expires_at),
308
- "expiresIn": SESSION_TTL_SECONDS,
309
- }
310
-
311
-
312
- @api.post("/prepare", status_code=202)
313
- async def prepare(
314
- request: PrepareRequest,
315
- authorization: str | None = Header(default=None),
316
- ) -> dict[str, Any]:
317
- session, _ = _bearer(authorization, {"owner"})
318
- if session["status"] == "starting":
319
- return session_view(session, "owner")
320
- old_code = session.get("interview_code")
321
- if old_code:
322
- interview_codes.pop(old_code, None)
323
- code = _new_interview_code()
324
- session_id = next(sid for sid, value in sessions.items() if value is session)
325
- interview_codes[code] = session_id
326
- session.update(
327
- background=request.background.strip(),
328
- hypotheses="",
329
- history=[],
330
- turns=0,
331
- current_question="",
332
- question_revision=0,
333
- status="starting",
334
- error=None,
335
- interview_code=code,
336
- )
337
- task = asyncio.create_task(prepare_session(session_id, session["background"]))
338
- preparation_tasks.add(task)
339
- task.add_done_callback(preparation_tasks.discard)
340
- return session_view(session, "owner")
341
-
342
-
343
- @api.post("/pair")
344
- def create_pairing() -> dict[str, Any]:
345
- _cleanup_expired()
346
- code = _new_pairing_code()
347
- pair_token = secrets.token_urlsafe(24)
348
- expires_at = int(time.time()) + PAIRING_TTL_SECONDS
349
- pairings_by_token[pair_token] = {
350
- "code": code,
351
- "session_id": None,
352
- "expires_at": expires_at,
353
- }
354
- pairing_codes[code] = pair_token
355
- return {
356
- "pairingCode": code,
357
- "pairingToken": pair_token,
358
- "expiresIn": PAIRING_TTL_SECONDS,
359
- }
360
-
361
-
362
- @api.get("/pair/status")
363
- def pairing_status(
364
- x_pairing_token: str | None = Header(default=None, alias="X-Pairing-Token"),
365
- ) -> dict[str, Any]:
366
- _cleanup_expired()
367
- pairing = pairings_by_token.get(x_pairing_token or "")
368
- if not pairing:
369
- raise HTTPException(status_code=401, detail="Invalid or expired pairing")
370
- session_id = pairing.get("session_id")
371
- if not session_id:
372
- return {"paired": False, "expiresAt": pairing["expires_at"]}
373
- session = sessions.get(session_id)
374
- if not session:
375
- raise HTTPException(status_code=410, detail="Interview session expired")
376
- return {
377
- "paired": True,
378
- "token": create_token(session_id, "glasses", session["expires_at"]),
379
- **session_view(session, "glasses"),
380
- }
381
-
382
-
383
- @api.post("/pair/claim")
384
- def claim_pairing(
385
- request: JoinRequest,
386
- authorization: str | None = Header(default=None),
387
- ) -> dict[str, bool]:
388
- session, _ = _bearer(authorization, {"owner"})
389
- _cleanup_expired()
390
- pair_token = pairing_codes.get(request.code)
391
- pairing = pairings_by_token.get(pair_token or "")
392
- if not pairing or not hmac.compare_digest(request.code, pairing.get("code") or ""):
393
- raise HTTPException(status_code=401, detail="Incorrect or expired glasses pairing code")
394
- session_id = next(sid for sid, value in sessions.items() if value is session)
395
- pairing["session_id"] = session_id
396
- return {"paired": True}
397
-
398
-
399
- @api.post("/join")
400
- def join(request: JoinRequest) -> dict[str, Any]:
401
- _cleanup_expired()
402
- session_id = interview_codes.get(request.code)
403
- session = sessions.get(session_id or "")
404
- if not session or not hmac.compare_digest(request.code, session.get("interview_code") or ""):
405
- raise HTTPException(status_code=401, detail="Incorrect or expired interview code")
406
- return {
407
- "token": create_token(session_id, "glasses", session["expires_at"]),
408
- **session_view(session, "glasses"),
409
- }
410
-
411
-
412
- @api.get("/session")
413
- def get_session(authorization: str | None = Header(default=None)) -> dict[str, Any]:
414
- session, role = _bearer(authorization)
415
- return session_view(session, role)
416
-
417
-
418
- @api.post("/start")
419
- async def start(authorization: str | None = Header(default=None)) -> dict[str, Any]:
420
- session, role = _bearer(authorization, {"owner", "glasses"})
421
- if session["status"] != "ready":
422
- raise HTTPException(status_code=409, detail="Interview is not ready")
423
- if session.get("current_question"):
424
- return session_view(session, role)
425
- try:
426
- question, history, turns, hypotheses = await asyncio.to_thread(
427
- call_asq, INITIAL_PROMPT, session
428
- )
429
- except Exception as error:
430
- raise HTTPException(status_code=502, detail=f"ASQ inference failed: {error}")
431
- session["history"] = history
432
- session["turns"] = turns
433
- session["current_question"] = question
434
- session["question_revision"] = session.get("question_revision", 0) + 1
435
- if hypotheses:
436
- session["hypotheses"] = hypotheses
437
- return session_view(session, role)
438
-
439
-
440
- @api.post("/suggest")
441
- async def suggest(
442
- request: SuggestRequest,
443
- authorization: str | None = Header(default=None),
444
- ) -> dict[str, str]:
445
- session, _ = _bearer(authorization, {"owner", "glasses"})
446
- if session["status"] != "ready":
447
- raise HTTPException(status_code=409, detail="Interview is not ready")
448
- try:
449
- question, history, turns, hypotheses = await asyncio.to_thread(
450
- call_asq, request.transcript.strip(), session
451
- )
452
- except Exception as error:
453
- raise HTTPException(status_code=502, detail=f"ASQ inference failed: {error}")
454
- session["history"] = history
455
- session["turns"] = turns
456
- session["current_question"] = question
457
- session["question_revision"] = session.get("question_revision", 0) + 1
458
- if hypotheses:
459
- session["hypotheses"] = hypotheses
460
- return {"question": question}
461
-
462
-
463
- MOBILE_APP_HTML = """<!doctype html>
464
- <html lang="en"><head><meta charset="utf-8"><meta name="viewport" content="width=device-width,initial-scale=1">
465
- <title>ASQ Interview Setup</title>
466
- <style>
467
- :root{font-family:Inter,system-ui,sans-serif;color:#16201d;background:#eef5f1}*{box-sizing:border-box}
468
- 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}
469
- 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}
470
- 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}
471
- 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}
472
- .hypo{white-space:pre-wrap;line-height:1.5;background:#f7faf8;padding:14px;border-radius:12px;max-height:320px;overflow:auto}
473
- </style></head><body><main class="card">
474
- <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>
475
- <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>
476
- <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>
477
- </main><script>
478
- let token="", timer=null; const $=id=>document.getElementById(id);
479
- 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}
480
- function errorMessage(e){$("status").textContent=e.message;$("progress").classList.remove("hidden")}
481
- async function begin(){try{const data=await api("/auth",{method:"POST"});token=data.token;$("prepareBtn").disabled=false}catch(e){errorMessage(e)}}
482
- $("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}};
483
- async function poll(){try{show(await api("/session"))}catch(e){clearInterval(timer);errorMessage(e)}}
484
- 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…"}}
485
- begin();
486
- </script></body></html>"""
487
-
488
-
489
- with gr.Blocks(title="ASQ Rokid Gateway") as status_page:
490
- gr.Markdown(
491
- "# ASQ Rokid Gateway\n"
492
- "This service provides phone preparation and Rokid interview pairing APIs.\n\n"
493
- "Open the site root on a phone to prepare an interview. Use `GET /health` for service status."
494
- )
495
-
496
-
497
- api = gr.mount_gradio_app(api, status_page, path="/ui")
498
-
499
-
500
- if __name__ == "__main__":
501
- uvicorn.run(api, host="0.0.0.0", port=7860)
 
 
 
 
 
 
 
 
 
 
 
1
+ import asyncio
2
+ import base64
3
+ import hashlib
4
+ import hmac
5
+ import json
6
+ import logging
7
+ import os
8
+ import re
9
+ import secrets
10
+ import time
11
+ from typing import Any, Literal
12
+
13
+ import gradio as gr
14
+ import uvicorn
15
+ from fastapi import FastAPI, Header, HTTPException
16
+ from fastapi.middleware.cors import CORSMiddleware
17
+ from fastapi.responses import HTMLResponse
18
+ from gradio_client import Client
19
+ from pydantic import BaseModel, Field
20
+
21
+
22
+ SESSION_SECRET = os.getenv("SESSION_SECRET", "")
23
+ HF_TOKEN = os.getenv("HF_TOKEN", "")
24
+ TARGET_SPACE_ID = os.getenv("TARGET_SPACE_ID", "ASQ2/ASQ")
25
+ INITIAL_PROMPT = os.getenv("ASQ_INITIAL_PROMPT", "Hello")
26
+ TARGET_CHAT_API_NAME = os.getenv("TARGET_CHAT_API_NAME", "/chat_flow")
27
+ TARGET_INITIALIZE_API_NAME = os.getenv("TARGET_INITIALIZE_API_NAME", "/initialize_flow")
28
+ SESSION_TTL_SECONDS = int(os.getenv("SESSION_TTL_SECONDS", "14400"))
29
+ MAX_BACKGROUND_CHARS = int(os.getenv("MAX_BACKGROUND_CHARS", "20000"))
30
+ PREPARE_TIMEOUT_SECONDS = int(os.getenv("PREPARE_TIMEOUT_SECONDS", "600"))
31
+ PREPARE_RETRY_SECONDS = int(os.getenv("PREPARE_RETRY_SECONDS", "15"))
32
+ PAIRING_TTL_SECONDS = int(os.getenv("PAIRING_TTL_SECONDS", "600"))
33
+
34
+ logger = logging.getLogger("asq_gateway")
35
+
36
+
37
+ class PrepareRequest(BaseModel):
38
+ background: str = Field(min_length=20, max_length=MAX_BACKGROUND_CHARS)
39
+
40
+
41
+ class JoinRequest(BaseModel):
42
+ code: str = Field(min_length=6, max_length=6, pattern=r"^\d{6}$")
43
+
44
+
45
+ class SuggestRequest(BaseModel):
46
+ transcript: str = Field(min_length=1, max_length=4000)
47
+
48
+
49
+ sessions: dict[str, dict[str, Any]] = {}
50
+ interview_codes: dict[str, str] = {}
51
+ pairing_codes: dict[str, str] = {}
52
+ pairings_by_token: dict[str, dict[str, Any]] = {}
53
+ preparation_tasks: set[asyncio.Task[Any]] = set()
54
+
55
+
56
+ def _b64encode(data: bytes) -> str:
57
+ return base64.urlsafe_b64encode(data).decode("ascii").rstrip("=")
58
+
59
+
60
+ def _b64decode(data: str) -> bytes:
61
+ return base64.urlsafe_b64decode(data + "=" * (-len(data) % 4))
62
+
63
+
64
+ def _sign(payload: str) -> str:
65
+ return _b64encode(
66
+ hmac.new(SESSION_SECRET.encode("utf-8"), payload.encode("utf-8"), hashlib.sha256).digest()
67
+ )
68
+
69
+
70
+ def _cleanup_expired() -> None:
71
+ now = int(time.time())
72
+ expired_ids = [sid for sid, session in sessions.items() if session["expires_at"] <= now]
73
+ for sid in expired_ids:
74
+ session = sessions.pop(sid, None)
75
+ if session and session.get("interview_code"):
76
+ interview_codes.pop(session["interview_code"], None)
77
+ expired_pair_tokens = [
78
+ token for token, pairing in pairings_by_token.items() if pairing["expires_at"] <= now
79
+ ]
80
+ for pair_token in expired_pair_tokens:
81
+ pairing = pairings_by_token.pop(pair_token, None)
82
+ if pairing:
83
+ pairing_codes.pop(pairing["code"], None)
84
+
85
+
86
+ def _new_interview_code() -> str:
87
+ _cleanup_expired()
88
+ for _ in range(100):
89
+ code = f"{secrets.randbelow(1_000_000):06d}"
90
+ if code not in interview_codes:
91
+ return code
92
+ raise RuntimeError("Could not allocate an interview code")
93
+
94
+
95
+ def _new_pairing_code() -> str:
96
+ _cleanup_expired()
97
+ for _ in range(100):
98
+ code = f"{secrets.randbelow(1_000_000):06d}"
99
+ if code not in pairing_codes and code not in interview_codes:
100
+ return code
101
+ raise RuntimeError("Could not allocate a pairing code")
102
+
103
+
104
+ def create_token(session_id: str, role: Literal["owner", "glasses"], expires_at: int) -> str:
105
+ payload = _b64encode(
106
+ json.dumps(
107
+ {"sid": session_id, "role": role, "exp": expires_at}, separators=(",", ":")
108
+ ).encode("utf-8")
109
+ )
110
+ return f"{payload}.{_sign(payload)}"
111
+
112
+
113
+ def verify_token(token: str, roles: set[str] | None = None) -> tuple[dict[str, Any], str]:
114
+ if not SESSION_SECRET:
115
+ raise HTTPException(status_code=503, detail="Gateway is not configured")
116
+ try:
117
+ payload, signature = token.split(".", 1)
118
+ if not hmac.compare_digest(signature, _sign(payload)):
119
+ raise ValueError("bad signature")
120
+ data = json.loads(_b64decode(payload))
121
+ if int(data["exp"]) <= int(time.time()):
122
+ raise ValueError("expired")
123
+ role = str(data["role"])
124
+ session = sessions.get(str(data["sid"]))
125
+ if not session:
126
+ raise ValueError("unknown session")
127
+ if roles is not None and role not in roles:
128
+ raise HTTPException(status_code=403, detail="This session cannot perform that action")
129
+ return session, role
130
+ except HTTPException:
131
+ raise
132
+ except (ValueError, KeyError, json.JSONDecodeError):
133
+ raise HTTPException(status_code=401, detail="Invalid or expired session")
134
+
135
+
136
+ def _bearer(authorization: str | None, roles: set[str] | None = None) -> tuple[dict[str, Any], str]:
137
+ if not authorization or not authorization.startswith("Bearer "):
138
+ raise HTTPException(status_code=401, detail="Missing session token")
139
+ return verify_token(authorization[7:], roles)
140
+
141
+
142
+ def _extract_turn_count(value: Any) -> int:
143
+ if isinstance(value, (int, float)):
144
+ return int(value)
145
+ match = re.search(r"\d+", str(value or ""))
146
+ return int(match.group(0)) if match else 0
147
+
148
+
149
+ def _message_text(message: Any) -> str:
150
+ if isinstance(message, str):
151
+ return message.strip()
152
+ if isinstance(message, dict):
153
+ content = message.get("content", message.get("text", ""))
154
+ return _message_text(content)
155
+ if isinstance(message, (list, tuple)):
156
+ parts = [_message_text(item) for item in message]
157
+ return "\n".join(part for part in parts if part).strip()
158
+ return str(message or "").strip()
159
+
160
+
161
+ def _first_complete_sentence(text: str) -> str:
162
+ cleaned = text.strip()
163
+ sentence_end = re.search(r"[.!?。!?]", cleaned)
164
+ if sentence_end:
165
+ return cleaned[: sentence_end.end()].strip()
166
+ return cleaned.splitlines()[0].strip() if cleaned else ""
167
+
168
+
169
+ def extract_initial_state(result: Any) -> tuple[str, list[Any], int]:
170
+ # /initialize_flow returns:
171
+ # turn markdown, interview history, hypotheses, summary, strategy, status.
172
+ if not isinstance(result, (list, tuple)) or len(result) < 6:
173
+ raise RuntimeError("Unexpected initialization response from the ASQ Space")
174
+ hypotheses = str(result[2] or "").strip()
175
+ history = result[1] if isinstance(result[1], list) else []
176
+ turns = _extract_turn_count(result[0])
177
+ if not hypotheses:
178
+ raise RuntimeError("ASQ returned no initial hypotheses")
179
+ return hypotheses, history, turns
180
+
181
+
182
+ def extract_question(result: Any) -> tuple[str, list[Any], int, str | None]:
183
+ # /chat_flow returns:
184
+ # history, cleared input, turn markdown, hypotheses, summary, strategy, status.
185
+ if not isinstance(result, (list, tuple)) or len(result) < 7:
186
+ raise RuntimeError("Unexpected response from the ASQ Space")
187
+ history = result[0]
188
+ turns = _extract_turn_count(result[2])
189
+ hypotheses = str(result[3]).strip() if result[3] is not None else None
190
+ if not isinstance(history, list) or not history:
191
+ raise RuntimeError("ASQ returned no conversation history")
192
+ assistant_messages = [
193
+ message
194
+ for message in history
195
+ if isinstance(message, dict) and message.get("role") == "assistant"
196
+ ]
197
+ question = _first_complete_sentence(
198
+ _message_text((assistant_messages or history)[-1])
199
+ )
200
+ if not question:
201
+ raise RuntimeError("ASQ returned an empty question")
202
+ return question, history, turns, hypotheses
203
+
204
+
205
+ def initialize_asq(background: str) -> tuple[str, list[Any], int]:
206
+ if not HF_TOKEN:
207
+ raise RuntimeError("HF_TOKEN is not configured")
208
+ client = Client(TARGET_SPACE_ID, hf_token=HF_TOKEN, verbose=False)
209
+ result = client.predict(background, api_name=TARGET_INITIALIZE_API_NAME)
210
+ return extract_initial_state(result)
211
+
212
+
213
+ def call_asq(transcript: str, session: dict[str, Any]) -> tuple[str, list[Any], int, str | None]:
214
+ if not HF_TOKEN:
215
+ raise RuntimeError("HF_TOKEN is not configured")
216
+ client = Client(TARGET_SPACE_ID, hf_token=HF_TOKEN, verbose=False)
217
+ result = client.predict(
218
+ transcript,
219
+ session["history"],
220
+ session["background"],
221
+ api_name=TARGET_CHAT_API_NAME,
222
+ )
223
+ return extract_question(result)
224
+
225
+
226
+ async def prepare_session(session_id: str, background: str) -> None:
227
+ session = sessions.get(session_id)
228
+ if not session:
229
+ return
230
+ deadline = time.monotonic() + PREPARE_TIMEOUT_SECONDS
231
+ attempt = 0
232
+ while True:
233
+ attempt += 1
234
+ try:
235
+ hypotheses, history, turns = await asyncio.to_thread(initialize_asq, background)
236
+ break
237
+ except Exception as error:
238
+ remaining = deadline - time.monotonic()
239
+ logger.warning("ASQ preparation attempt %s failed: %s", attempt, type(error).__name__)
240
+ if remaining <= 0:
241
+ if session_id in sessions:
242
+ session["status"] = "error"
243
+ session["error"] = "Unable to start the model. Please try again."
244
+ return
245
+ await asyncio.sleep(min(PREPARE_RETRY_SECONDS, remaining))
246
+ if session_id in sessions:
247
+ session.update(
248
+ status="ready",
249
+ hypotheses=hypotheses,
250
+ history=history,
251
+ turns=turns,
252
+ error=None,
253
+ )
254
+
255
+
256
+ def session_view(session: dict[str, Any], role: str) -> dict[str, Any]:
257
+ response: dict[str, Any] = {
258
+ "status": session["status"],
259
+ "ready": session["status"] == "ready",
260
+ "expiresAt": session["expires_at"],
261
+ "question": session.get("current_question", ""),
262
+ "questionRevision": session.get("question_revision", 0),
263
+ "turns": session.get("turns", 0),
264
+ }
265
+ if session.get("error"):
266
+ response["error"] = session["error"]
267
+ if role == "owner":
268
+ response["interviewCode"] = session.get("interview_code")
269
+ response["hypotheses"] = session.get("hypotheses", "")
270
+ return response
271
+
272
+
273
+ api = FastAPI(title="ASQ Rokid Gateway", version="0.5.0")
274
+ api.add_middleware(
275
+ CORSMiddleware,
276
+ allow_origins=["*"],
277
+ allow_methods=["GET", "POST", "OPTIONS"],
278
+ allow_headers=["Authorization", "Content-Type"],
279
+ )
280
+
281
+
282
+ @api.get("/", response_class=HTMLResponse)
283
+ def mobile_app() -> str:
284
+ return MOBILE_APP_HTML
285
+
286
+
287
+ @api.get("/health")
288
+ def health() -> dict[str, Any]:
289
+ return {
290
+ "ok": True,
291
+ "configured": bool(SESSION_SECRET and HF_TOKEN),
292
+ "target": TARGET_SPACE_ID,
293
+ "version": "0.5.0",
294
+ }
295
+
296
+
297
+ @api.post("/auth")
298
+ def create_owner_session() -> dict[str, Any]:
299
+ if not SESSION_SECRET:
300
+ raise HTTPException(status_code=503, detail="Gateway is not configured")
301
+ _cleanup_expired()
302
+ session_id = secrets.token_urlsafe(18)
303
+ expires_at = int(time.time()) + SESSION_TTL_SECONDS
304
+ sessions[session_id] = {
305
+ "background": "",
306
+ "hypotheses": "",
307
+ "history": [],
308
+ "turns": 0,
309
+ "current_question": "",
310
+ "question_revision": 0,
311
+ "status": "draft",
312
+ "error": None,
313
+ "interview_code": None,
314
+ "expires_at": expires_at,
315
+ }
316
+ return {
317
+ "token": create_token(session_id, "owner", expires_at),
318
+ "expiresIn": SESSION_TTL_SECONDS,
319
+ }
320
+
321
+
322
+ @api.post("/prepare", status_code=202)
323
+ async def prepare(
324
+ request: PrepareRequest,
325
+ authorization: str | None = Header(default=None),
326
+ ) -> dict[str, Any]:
327
+ session, _ = _bearer(authorization, {"owner"})
328
+ if session["status"] == "starting":
329
+ return session_view(session, "owner")
330
+ old_code = session.get("interview_code")
331
+ if old_code:
332
+ interview_codes.pop(old_code, None)
333
+ code = _new_interview_code()
334
+ session_id = next(sid for sid, value in sessions.items() if value is session)
335
+ interview_codes[code] = session_id
336
+ session.update(
337
+ background=request.background.strip(),
338
+ hypotheses="",
339
+ history=[],
340
+ turns=0,
341
+ current_question="",
342
+ question_revision=0,
343
+ status="starting",
344
+ error=None,
345
+ interview_code=code,
346
+ )
347
+ task = asyncio.create_task(prepare_session(session_id, session["background"]))
348
+ preparation_tasks.add(task)
349
+ task.add_done_callback(preparation_tasks.discard)
350
+ return session_view(session, "owner")
351
+
352
+
353
+ @api.post("/pair")
354
+ def create_pairing() -> dict[str, Any]:
355
+ _cleanup_expired()
356
+ code = _new_pairing_code()
357
+ pair_token = secrets.token_urlsafe(24)
358
+ expires_at = int(time.time()) + PAIRING_TTL_SECONDS
359
+ pairings_by_token[pair_token] = {
360
+ "code": code,
361
+ "session_id": None,
362
+ "expires_at": expires_at,
363
+ }
364
+ pairing_codes[code] = pair_token
365
+ return {
366
+ "pairingCode": code,
367
+ "pairingToken": pair_token,
368
+ "expiresIn": PAIRING_TTL_SECONDS,
369
+ }
370
+
371
+
372
+ @api.get("/pair/status")
373
+ def pairing_status(
374
+ x_pairing_token: str | None = Header(default=None, alias="X-Pairing-Token"),
375
+ ) -> dict[str, Any]:
376
+ _cleanup_expired()
377
+ pairing = pairings_by_token.get(x_pairing_token or "")
378
+ if not pairing:
379
+ raise HTTPException(status_code=401, detail="Invalid or expired pairing")
380
+ session_id = pairing.get("session_id")
381
+ if not session_id:
382
+ return {"paired": False, "expiresAt": pairing["expires_at"]}
383
+ session = sessions.get(session_id)
384
+ if not session:
385
+ raise HTTPException(status_code=410, detail="Interview session expired")
386
+ return {
387
+ "paired": True,
388
+ "token": create_token(session_id, "glasses", session["expires_at"]),
389
+ **session_view(session, "glasses"),
390
+ }
391
+
392
+
393
+ @api.post("/pair/claim")
394
+ def claim_pairing(
395
+ request: JoinRequest,
396
+ authorization: str | None = Header(default=None),
397
+ ) -> dict[str, bool]:
398
+ session, _ = _bearer(authorization, {"owner"})
399
+ _cleanup_expired()
400
+ pair_token = pairing_codes.get(request.code)
401
+ pairing = pairings_by_token.get(pair_token or "")
402
+ if not pairing or not hmac.compare_digest(request.code, pairing.get("code") or ""):
403
+ raise HTTPException(status_code=401, detail="Incorrect or expired glasses pairing code")
404
+ session_id = next(sid for sid, value in sessions.items() if value is session)
405
+ pairing["session_id"] = session_id
406
+ return {"paired": True}
407
+
408
+
409
+ @api.post("/join")
410
+ def join(request: JoinRequest) -> dict[str, Any]:
411
+ _cleanup_expired()
412
+ session_id = interview_codes.get(request.code)
413
+ session = sessions.get(session_id or "")
414
+ if not session or not hmac.compare_digest(request.code, session.get("interview_code") or ""):
415
+ raise HTTPException(status_code=401, detail="Incorrect or expired interview code")
416
+ return {
417
+ "token": create_token(session_id, "glasses", session["expires_at"]),
418
+ **session_view(session, "glasses"),
419
+ }
420
+
421
+
422
+ @api.get("/session")
423
+ def get_session(authorization: str | None = Header(default=None)) -> dict[str, Any]:
424
+ session, role = _bearer(authorization)
425
+ return session_view(session, role)
426
+
427
+
428
+ @api.post("/start")
429
+ async def start(authorization: str | None = Header(default=None)) -> dict[str, Any]:
430
+ session, role = _bearer(authorization, {"owner", "glasses"})
431
+ if session["status"] != "ready":
432
+ raise HTTPException(status_code=409, detail="Interview is not ready")
433
+ if session.get("current_question"):
434
+ return session_view(session, role)
435
+ try:
436
+ question, history, turns, hypotheses = await asyncio.to_thread(
437
+ call_asq, INITIAL_PROMPT, session
438
+ )
439
+ except Exception as error:
440
+ raise HTTPException(status_code=502, detail=f"ASQ inference failed: {error}")
441
+ session["history"] = history
442
+ session["turns"] = turns
443
+ session["current_question"] = question
444
+ session["question_revision"] = session.get("question_revision", 0) + 1
445
+ if hypotheses:
446
+ session["hypotheses"] = hypotheses
447
+ return session_view(session, role)
448
+
449
+
450
+ @api.post("/suggest")
451
+ async def suggest(
452
+ request: SuggestRequest,
453
+ authorization: str | None = Header(default=None),
454
+ ) -> dict[str, str]:
455
+ session, _ = _bearer(authorization, {"owner", "glasses"})
456
+ if session["status"] != "ready":
457
+ raise HTTPException(status_code=409, detail="Interview is not ready")
458
+ try:
459
+ question, history, turns, hypotheses = await asyncio.to_thread(
460
+ call_asq, request.transcript.strip(), session
461
+ )
462
+ except Exception as error:
463
+ raise HTTPException(status_code=502, detail=f"ASQ inference failed: {error}")
464
+ session["history"] = history
465
+ session["turns"] = turns
466
+ session["current_question"] = question
467
+ session["question_revision"] = session.get("question_revision", 0) + 1
468
+ if hypotheses:
469
+ session["hypotheses"] = hypotheses
470
+ return {"question": question}
471
+
472
+
473
+ MOBILE_APP_HTML = """<!doctype html>
474
+ <html lang="en"><head><meta charset="utf-8"><meta name="viewport" content="width=device-width,initial-scale=1">
475
+ <title>ASQ Interview Setup</title>
476
+ <style>
477
+ :root{font-family:Inter,system-ui,sans-serif;color:#16201d;background:#eef5f1}*{box-sizing:border-box}
478
+ 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}
479
+ 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}
480
+ 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}
481
+ 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}
482
+ .hypo{white-space:pre-wrap;line-height:1.5;background:#f7faf8;padding:14px;border-radius:12px;max-height:320px;overflow:auto}
483
+ </style></head><body><main class="card">
484
+ <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>
485
+ <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>
486
+ <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>
487
+ </main><script>
488
+ let token="", timer=null; const $=id=>document.getElementById(id);
489
+ 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}
490
+ function errorMessage(e){$("status").textContent=e.message;$("progress").classList.remove("hidden")}
491
+ async function begin(){try{const data=await api("/auth",{method:"POST"});token=data.token;$("prepareBtn").disabled=false}catch(e){errorMessage(e)}}
492
+ $("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}};
493
+ async function poll(){try{show(await api("/session"))}catch(e){clearInterval(timer);errorMessage(e)}}
494
+ 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…"}}
495
+ begin();
496
+ </script></body></html>"""
497
+
498
+
499
+ with gr.Blocks(title="ASQ Rokid Gateway") as status_page:
500
+ gr.Markdown(
501
+ "# ASQ Rokid Gateway\n"
502
+ "This service provides phone preparation and Rokid interview pairing APIs.\n\n"
503
+ "Open the site root on a phone to prepare an interview. Use `GET /health` for service status."
504
+ )
505
+
506
+
507
+ api = gr.mount_gradio_app(api, status_page, path="/ui")
508
+
509
+
510
+ if __name__ == "__main__":
511
+ uvicorn.run(api, host="0.0.0.0", port=7860)