Refactor RAG query architecture to async job queue to fix mobile background disconnects

app.py

@@ -819,10 +819,12 @@ def delete_document(source_name: str):
 
 # ── Query ─────────────────────────────────────────────────────────────────────
 
-@app.route("/api/query", methods=["POST"])
+_query_jobs = {}
+
+@app.route("/api/query/start", methods=["POST"])
 @limiter.limit("120 per minute")
-def query():
-    """RAG query — returns a streaming SSE response."""
+def query_start():
+    """Starts a RAG query job and returns a job_id."""
     data = request.get_json()
     q    = escape((data or {}).get("query", "").strip())
     top_k = (data or {}).get("top_k")
@@ -832,68 +834,82 @@ def query():
     use_bm25 = False   # Disabled per request
     use_gpu = bool((data or {}).get("use_gpu", False))
     cpu_threads = int((data or {}).get("cpu_threads", 2))
+    
     if not q:
         return jsonify({"error": "Empty query"}), 400
 
     token = get_session_token()
-
     chunk_count = vector_store.count()
     if chunk_count == 0:
         return jsonify({"error": "No documents ingested yet. Please upload documents first."}), 400
 
     log.info("Query received (%d chars) | vector store has %d chunks", len(q), chunk_count)
-
     remote_addr = request.remote_addr
 
-    def _generate(sess_token):
-        import queue
-        q_events = queue.Queue()
-
-        def _run():
-            config.current_session.set(sess_token)
-            try:
-                def cb(status):
-                    if isinstance(status, dict):
-                        q_events.put(status)
-                    else:
-                        q_events.put({"status": status})
-                ans, metrics = run_query_crew(q, top_k=top_k, max_tokens=max_tokens, use_vector=use_vector, use_graph=use_graph, use_bm25=use_bm25, session_token=token, status_callback=cb, use_gpu=use_gpu, cpu_threads=cpu_threads)
-                q_events.put({"done": True, "answer": ans, "metrics": metrics})
-                
-                gen_time = metrics.get("time_seconds", 0)
-                log_query_summary(token, remote_addr, q, top_k or 10, gen_time, True)
-            except Exception as e:
-                log.exception("Query pipeline error")
-                q_events.put({"error": str(e)})
-                log_query_summary(token, remote_addr, q, top_k or 10, 0, False, str(e))
+    job_id = uuid.uuid4().hex[:8]
+    _query_jobs[job_id] = {
+        "events": [],
+        "done": False,
+        "error": None
+    }
 
-        _llm_query_queue.put((_run, ()))
+    def _run():
+        config.current_session.set(token)
+        try:
+            def cb(status):
+                if isinstance(status, dict):
+                    _query_jobs[job_id]["events"].append(status)
+                else:
+                    _query_jobs[job_id]["events"].append({"status": status})
+                    
+            ans, metrics = run_query_crew(q, top_k=top_k, max_tokens=max_tokens, use_vector=use_vector, use_graph=use_graph, use_bm25=use_bm25, session_token=token, status_callback=cb, use_gpu=use_gpu, cpu_threads=cpu_threads)
+            
+            for i in range(0, len(ans), 80):
+                _query_jobs[job_id]["events"].append({"chunk": ans[i:i + 80]})
+            _query_jobs[job_id]["events"].append({"metrics": metrics})
+            _query_jobs[job_id]["events"].append({"done": True})
+            _query_jobs[job_id]["done"] = True
+            
+            gen_time = metrics.get("time_seconds", 0)
+            log_query_summary(token, remote_addr, q, top_k or 10, gen_time, True)
+        except Exception as e:
+            log.exception("Query pipeline error")
+            _query_jobs[job_id]["events"].append({"error": str(e)})
+            _query_jobs[job_id]["error"] = str(e)
+            _query_jobs[job_id]["done"] = True
+            log_query_summary(token, remote_addr, q, top_k or 10, 0, False, str(e))
+
+    _llm_query_queue.put((_run, ()))
+    return jsonify({"job_id": job_id})
+
+@app.route("/api/query/stream/<job_id>")
+def query_stream(job_id):
+    """Streams events for a specific query job starting from an offset."""
+    offset = int(request.args.get("offset", 0))
+    job = _query_jobs.get(job_id)
+    
+    if not job:
+        return jsonify({"error": "Job not found or expired"}), 404
 
+    def _generate():
+        import time
+        idx = offset
         while True:
-            try:
-                event = q_events.get(timeout=10)
-            except queue.Empty:
-                yield ": keep-alive\n\n"
-                continue
-
-            if "error" in event:
-                yield f"data: {json.dumps({'error': event['error']})}\n\n"
-                break
-            elif "status" in event:
-                yield f"data: {json.dumps({'status': event['status']})}\n\n"
-            elif "done" in event:
-                answer = event["answer"]
-                log.info("Query answered — %d chars", len(answer))
-                for i in range(0, len(answer), 80):
-                    chunk   = answer[i:i + 80]
-                    payload = json.dumps({"chunk": chunk})
-                    yield f"data: {payload}\n\n"
-                yield f"data: {json.dumps({'metrics': event['metrics']})}\n\n"
-                yield "data: {\"done\": true}\n\n"
+            while idx < len(job["events"]):
+                event = job["events"][idx]
+                yield f"data: {json.dumps(event)}\n\n"
+                if "error" in event or "done" in event:
+                    return
+                idx += 1
+
+            if job["error"] or job["done"]:
                 break
+                
+            time.sleep(0.5)
+            yield ": keep-alive\n\n"
 
     return Response(
-        stream_with_context(_generate(token)),
+        stream_with_context(_generate()),
         mimetype="text/event-stream",
         headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
     )

static/app.js

@@ -853,35 +853,56 @@ async function submitQuery() {
     if (!useVector) $(`ms-vector-${qId}`)?.classList.add('disabled');
     if (!useBm25) $(`ms-bm25-${qId}`)?.classList.add('disabled');
 
-    const resp = await fetch('/api/query', {
+    const startResp = await fetch('/api/query/start', {
       method:  'POST',
       headers: { 'Content-Type': 'application/json' },
       body:    JSON.stringify({ query: q, top_k: topK, max_tokens: maxTokens, use_vector: useVector, use_graph: useGraph, use_bm25: useBm25, use_gpu: useGpu, cpu_threads: cpuThreads }),
     });
 
-    if (!resp.ok) {
-      const errData = await resp.json();
-      throw new Error(errData.error || `HTTP ${resp.status}`);
+    if (!startResp.ok) {
+      const errData = await startResp.json();
+      throw new Error(errData.error || `HTTP ${startResp.status}`);
     }
 
-    const reader  = resp.body.getReader();
-    const decoder = new TextDecoder();
-    let   buffer  = '';
+    const startData = await startResp.json();
+    const jobId = startData.job_id;
 
-    while (true) {
-      const { done, value } = await reader.read();
-      if (done) break;
-      buffer += decoder.decode(value, { stream: true });
-      const lines = buffer.split('\n');
-      buffer = lines.pop();
+    let currentOffset = 0;
+    let jobDone = false;
 
-      for (const line of lines) {
-        if (!line.startsWith('data: ')) continue;
-        let payload;
-        try { payload = JSON.parse(line.slice(6)); }
-        catch(pe) { diag.warn('SSE parse error:', pe, line); continue; }
+    while (!jobDone) {
+      try {
+        const resp = await fetch(`/api/query/stream/${jobId}?offset=${currentOffset}`);
+        if (!resp.ok) {
+          const errData = await resp.json().catch(() => ({}));
+          if (resp.status === 404) throw new Error("Job not found or expired on server.");
+          throw new Error(errData.error || `HTTP ${resp.status}`);
+        }
 
-        if (payload.error) throw new Error(payload.error);
+        const reader  = resp.body.getReader();
+        const decoder = new TextDecoder();
+        let   buffer  = '';
+
+        while (!jobDone) {
+          const { done, value } = await reader.read();
+          if (done) break;
+          buffer += decoder.decode(value, { stream: true });
+          const lines = buffer.split('\n');
+          buffer = lines.pop();
+
+          for (const line of lines) {
+            if (!line.startsWith('data: ')) continue;
+            let payload;
+            try { payload = JSON.parse(line.slice(6)); }
+            catch(pe) { diag.warn('SSE parse error:', pe, line); continue; }
+
+            // Successfully received a valid JSON payload, increment offset for next reconnect
+            currentOffset++;
+
+            if (payload.error) {
+              jobDone = true;
+              throw new Error(payload.error);
+            }
 
         if (payload.status) {
           diag.info('SSE status:', payload.status);
@@ -979,7 +1000,19 @@ async function submitQuery() {
           }
         }
 
-        if (payload.done) { diag.info('SSE: done'); break; }
+        if (payload.done) { 
+          diag.info('SSE: done'); 
+          jobDone = true;
+          break; 
+        }
+      }
+    }
+      } catch (err) {
+        if (err.message.includes("Job not found") || jobDone) {
+          throw err;
+        }
+        diag.warn('Network dropped, reconnecting to job...', err);
+        await new Promise(r => setTimeout(r, 1500));
       }
     }
 
@@ -1006,11 +1039,7 @@ async function submitQuery() {
     diag.error('Query error:', err);
     stopTimer();
     
-    // Check if the browser recently suspended the tab (common on mobile)
     let errorMsg = err.message;
-    if (document.hidden || document.visibilityState === 'hidden') {
-      errorMsg = "Connection lost because the app was moved to the background. Please keep the app open during analysis to prevent the network from dropping.";
-    }
 
     if (currentMsId) {
       const el = $(currentMsId);