From 3988a937be79347d41af9b54eb0862c55fc44876 Mon Sep 17 00:00:00 2001 From: Prins Kumar Date: Thu, 9 Jul 2026 15:04:46 +0530 Subject: [PATCH 1/3] fix: surface rebuild stage in UI during long ECC phases - add progress callbacks to embed, extract, summarize_communities; cache last ECC status to fix unknown time on timeout --- ecc/app/graphrag/graph_rag.py | 27 ++++++++++++++++++++++----- graphrag/app/routers/ui.py | 18 ++++++++++++++++-- 2 files changed, 38 insertions(+), 7 deletions(-) diff --git a/ecc/app/graphrag/graph_rag.py b/ecc/app/graphrag/graph_rag.py index 6dd5b0e..bbbbed9 100644 --- a/ecc/app/graphrag/graph_rag.py +++ b/ecc/app/graphrag/graph_rag.py @@ -333,7 +333,7 @@ async def load(conn: AsyncTigerGraphConnection): async def embed( - embed_chan: Channel, embedding_store: EmbeddingStore, graphname: str + embed_chan: Channel, embedding_store: EmbeddingStore, graphname: str, progress=None ): """ Creates and starts one worker for each embed job @@ -366,6 +366,11 @@ async def embed( n_embed += 1 if n_embed % 100 == 0: logger.info(f"Embedding Processing: {n_embed} embedded so far") + if progress is not None: + try: + progress(f"Embedding documents ({n_embed} embedded so far)") + except Exception: + pass except ChannelClosed: break except Exception: @@ -383,6 +388,7 @@ async def extract( extractor: BaseExtractor, conn: AsyncTigerGraphConnection, num_senders: int, + progress=None, ): """ Creates and starts one worker for each extract job @@ -412,6 +418,11 @@ async def extract( logger.info( f"Entity Extraction: {n_chunks} chunks dispatched" ) + if progress is not None: + try: + progress(f"Extracting entities ({n_chunks} chunks processed)") + except Exception: + pass except ChannelClosed: break except Exception: @@ -536,6 +547,7 @@ async def summarize_communities( comm_process_chan: Channel, upsert_chan: Channel, embed_chan: Channel, + progress=None, ): logger.info("Community summarization started") n_comm = 0 @@ -552,6 +564,11 @@ async def summarize_communities( logger.info( f"Community summarization: {n_comm} dispatched" ) + if progress is not None: + try: + progress(f"Summarizing communities ({n_comm} done so far)") + except Exception: + pass except ChannelClosed: break except Exception: @@ -634,9 +651,9 @@ async def new_doc_pipeline(): grp.create_task(upsert(upsert_chan)) grp.create_task(load(conn)) - grp.create_task(embed(embed_chan, embedding_store, graphname)) + grp.create_task(embed(embed_chan, embedding_store, graphname, progress)) grp.create_task( - extract(extract_chan, upsert_chan, embed_chan, extractor, conn, num_chunk_senders) + extract(extract_chan, upsert_chan, embed_chan, extractor, conn, num_chunk_senders, progress) ) logger.info("Join docs_chan") await docs_chan.join() @@ -687,11 +704,11 @@ async def new_doc_pipeline(): grp.create_task(communities(conn, comm_process_chan)) # summarize each community grp.create_task( - summarize_communities(conn, comm_process_chan, upsert_chan, embed_chan) + summarize_communities(conn, comm_process_chan, upsert_chan, embed_chan, progress) ) grp.create_task(upsert(upsert_chan)) grp.create_task(load(conn)) - grp.create_task(embed(embed_chan, embedding_store, graphname)) + grp.create_task(embed(embed_chan, embedding_store, graphname, progress)) logger.info("Join comm_process_chan") await comm_process_chan.join() logger.info("Join embed_chan") diff --git a/graphrag/app/routers/ui.py b/graphrag/app/routers/ui.py index 1b2682e..3e942ca 100644 --- a/graphrag/app/routers/ui.py +++ b/graphrag/app/routers/ui.py @@ -62,6 +62,10 @@ from common.metrics.prometheus_metrics import metrics as pmetrics from common.metrics.tg_proxy import TigerGraphConnectionProxy from common.utils.graph_locks import acquire_graph_lock, release_graph_lock, acquire_rebuild_lock, release_rebuild_lock, get_rebuilding_graph, get_current_operation + +# Cache the last successful ECC status response per graph so the UI +# still sees started_at and stage when ECC is too busy to respond. +_last_ecc_status_cache: dict = {} from supportai import supportai from common.py_schemas.schemas import ( AgentProgess, @@ -2258,7 +2262,7 @@ async def rebuild_and_monitor(): elapsed = 0 while elapsed < max_wait_time: - await asyncio.sleep(poll_interval) # Non-blocking sleep + await asyncio.sleep(poll_interval) elapsed += poll_interval try: @@ -2272,6 +2276,11 @@ async def rebuild_and_monitor(): status_data = status_response.json() is_running = status_data.get("is_running", False) status = status_data.get("status", "unknown") + + # Cache last known good response so the UI endpoint + # can fall back to it when ECC is too busy to respond. + if is_running and status_data.get("started_at"): + _last_ecc_status_cache[graphname] = status_data # Log every minute to avoid spam if elapsed % 60 == 0: @@ -2280,6 +2289,7 @@ async def rebuild_and_monitor(): # Check if ALL stages are complete if not is_running and status in ["completed", "failed", "idle"]: LogWriter.info(f"ECC rebuild finished for {graphname} with status: {status} after {elapsed}s") + _last_ecc_status_cache.pop(graphname, None) break else: LogWriter.warning(f"ECC status check returned {status_response.status_code} for {graphname}") @@ -2290,6 +2300,7 @@ async def rebuild_and_monitor(): if elapsed >= max_wait_time: LogWriter.error(f"ECC rebuild monitoring timed out for {graphname} after {max_wait_time}s") + _last_ecc_status_cache.pop(graphname, None) except Exception as e: LogWriter.error(f"Error during ECC rebuild monitoring for {graphname}: {e}") @@ -2342,9 +2353,12 @@ def get_rebuild_status( "error": f"ECC service returned status {response.status_code}" } except httpx.TimeoutException as e: - # ECC is busy (heavy processing) - assume rebuild is still running + # ECC is busy (heavy processing) - assume rebuild is still running. + # Return the last cached status so the UI keeps started_at and stage. LogWriter.warning(f"ECC status check timed out (ECC may be busy): {str(e)}") + cached = _last_ecc_status_cache.get(graphname, {}) return { + **cached, "graphname": graphname, "is_running": True, "status": "unknown", From 0d6aedea7e8e6ce95de4e8c3a51c400c7c3b0dcf Mon Sep 17 00:00:00 2001 From: Prins Kumar Date: Mon, 13 Jul 2026 14:59:17 +0530 Subject: [PATCH 2/3] fix: clear ECC status cache in finally and preserve cached status on timeout --- graphrag/app/routers/ui.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/graphrag/app/routers/ui.py b/graphrag/app/routers/ui.py index 3e942ca..8cfb891 100644 --- a/graphrag/app/routers/ui.py +++ b/graphrag/app/routers/ui.py @@ -2307,6 +2307,10 @@ async def rebuild_and_monitor(): import traceback LogWriter.error(traceback.format_exc()) finally: + # Always drop cached status when monitoring ends (success, + # timeout, or unexpected failure) so timeout fallbacks do + # not keep reporting a stale rebuild as active. + _last_ecc_status_cache.pop(graphname, None) # Release lock only after ALL stages complete (or timeout/error) release_rebuild_lock(graphname) LogWriter.info(f"Released global rebuild lock for {graphname}") @@ -2361,7 +2365,7 @@ def get_rebuild_status( **cached, "graphname": graphname, "is_running": True, - "status": "unknown", + "status": cached.get("status", "unknown"), "error": "ECC is busy processing, status check timed out. Rebuild likely still in progress." } except Exception as e: From 943501c48313cd39944c5911ed2dd66f17f755eb Mon Sep 17 00:00:00 2001 From: Prins Kumar Date: Wed, 22 Jul 2026 21:11:36 +0530 Subject: [PATCH 3/3] feat: show document-level chunking progress during rebuild Expose progress_current/total/pct from ECC and render a bar in KGAdmin so the UI shows real progress while chunking documents, without wiping the bar when stream_docs finishes early. --- ecc/app/graphrag/graph_rag.py | 156 +++++++++++++++++------- ecc/app/main.py | 42 ++++++- graphrag-ui/src/pages/setup/KGAdmin.tsx | 47 +++++++ 3 files changed, 197 insertions(+), 48 deletions(-) diff --git a/ecc/app/graphrag/graph_rag.py b/ecc/app/graphrag/graph_rag.py index bbbbed9..2ccdad3 100644 --- a/ecc/app/graphrag/graph_rag.py +++ b/ecc/app/graphrag/graph_rag.py @@ -57,10 +57,9 @@ async def stream_docs( Streams the document contents into the docs_chan, over ``ttl_batches`` partitions (each StreamIds call claims and returns one partition). - *progress* (optional) is a callable invoked once when document - streaming completes — runtime hands the rebuild status forward - from "Chunking documents" to "Extracting entities and - relationships" at that boundary. + *progress* is accepted for API compatibility with callers but is + not used here — the UI stage advances after chunk_docs finishes, + not when streaming completes. """ logger.info(f"streaming docs ({ttl_batches} batches)") n_docs = 0 @@ -87,11 +86,9 @@ async def stream_docs( # close the docs chan -- this function is the only sender logger.info("closing docs chan") docs_chan.close() - if progress is not None: - try: - progress("Extracting entities and relationships") - except Exception: - pass + # Do NOT advance the UI stage here. Streaming finishing does not + # mean chunking is done; advancing here used to clear the chunking + # progress bar before chunk_docs finished (and before the UI polled). async def stream_chunks( conn: AsyncTigerGraphConnection, @@ -163,23 +160,45 @@ async def chunk_docs( embed_chan: Channel, upsert_chan: Channel, extract_chan: Channel, + progress=None, + total_docs: int = 0, ): """ Creates and starts one worker for each document in the docs channel. + + When *total_docs* is known (unprocessed Document count), reports + document-level chunking progress via *progress* so the UI can show + a percentage bar. """ logger.info("Chunk Processing Start") - doc_tasks = [] n_docs = 0 + + async def _chunk_one(content): + nonlocal n_docs + await workers.chunk_doc(conn, content, upsert_chan, embed_chan, extract_chan) + n_docs += 1 + if progress is not None and total_docs > 0: + pct = min(100, int(100 * n_docs / total_docs)) + try: + progress( + f"Chunking documents ({n_docs}/{total_docs} — {pct}%)", + current=n_docs, + total=total_docs, + ) + except TypeError: + try: + progress(f"Chunking documents ({n_docs}/{total_docs} — {pct}%)") + except Exception: + pass + except Exception: + pass + async with asyncio.TaskGroup() as grp: while True: try: content = await docs_chan.get() - n_docs += 1 - task = grp.create_task( - workers.chunk_doc(conn, content, upsert_chan, embed_chan, extract_chan) - ) - doc_tasks.append(task) + grp.create_task(_chunk_one(content)) except ChannelClosed: break except Exception: @@ -187,6 +206,35 @@ async def chunk_docs( logger.info(f"Chunk Processing End: {n_docs} document(s) processed") + # Chunking finished — mark 100% then move stage forward and clear bar. + if progress is not None: + if total_docs > 0: + try: + progress( + f"Chunking documents ({n_docs}/{total_docs} — 100%)", + current=total_docs, + total=total_docs, + ) + except TypeError: + try: + progress(f"Chunking documents ({n_docs}/{total_docs} — 100%)") + except Exception: + pass + except Exception: + pass + try: + progress( + "Extracting entities and relationships", + clear_progress=True, + ) + except TypeError: + try: + progress("Extracting entities and relationships") + except Exception: + pass + except Exception: + pass + logger.info("closing extract_chan") await extract_chan.put(None) @@ -333,7 +381,7 @@ async def load(conn: AsyncTigerGraphConnection): async def embed( - embed_chan: Channel, embedding_store: EmbeddingStore, graphname: str, progress=None + embed_chan: Channel, embedding_store: EmbeddingStore, graphname: str ): """ Creates and starts one worker for each embed job @@ -366,11 +414,6 @@ async def embed( n_embed += 1 if n_embed % 100 == 0: logger.info(f"Embedding Processing: {n_embed} embedded so far") - if progress is not None: - try: - progress(f"Embedding documents ({n_embed} embedded so far)") - except Exception: - pass except ChannelClosed: break except Exception: @@ -388,7 +431,6 @@ async def extract( extractor: BaseExtractor, conn: AsyncTigerGraphConnection, num_senders: int, - progress=None, ): """ Creates and starts one worker for each extract job @@ -418,11 +460,6 @@ async def extract( logger.info( f"Entity Extraction: {n_chunks} chunks dispatched" ) - if progress is not None: - try: - progress(f"Extracting entities ({n_chunks} chunks processed)") - except Exception: - pass except ChannelClosed: break except Exception: @@ -547,7 +584,6 @@ async def summarize_communities( comm_process_chan: Channel, upsert_chan: Channel, embed_chan: Channel, - progress=None, ): logger.info("Community summarization started") n_comm = 0 @@ -564,11 +600,6 @@ async def summarize_communities( logger.info( f"Community summarization: {n_comm} dispatched" ) - if progress is not None: - try: - progress(f"Summarizing communities ({n_comm} done so far)") - except Exception: - pass except ChannelClosed: break except Exception: @@ -599,11 +630,21 @@ async def run(graphname: str, conn: AsyncTigerGraphConnection, progress=None): where the job is. """ - def _report(msg: str) -> None: + def _report(msg: str, current=None, total=None, clear_progress=False) -> None: if progress is None: return try: - progress(msg) + progress( + msg, + current=current, + total=total, + clear_progress=clear_progress, + ) + except TypeError: + try: + progress(msg) + except Exception: + pass except Exception: pass @@ -640,20 +681,51 @@ def _report(msg: str) -> None: # remaining producer racing chunk_docs' load_q writes. async def new_doc_pipeline(): await sc_task + # Count unprocessed Documents before StreamIds claims them, + # so chunking can report document-level % progress. + total_docs = 0 + try: + count_result = conn.getVertexCount( + "Document", where="epoch_processed=0" + ) + if asyncio.iscoroutine(count_result): + count_result = await count_result + total_docs = int(count_result or 0) + except Exception as e: + logger.warning( + f"Could not count unprocessed Documents for progress: {e}" + ) + if total_docs > 0: + _report( + f"Chunking documents (0/{total_docs} — 0%)", + current=0, + total=total_docs, + ) + logger.info( + f"Chunking progress denominator: {total_docs} unprocessed Document(s)" + ) async with asyncio.TaskGroup() as inner: inner.create_task( stream_docs(conn, docs_chan, 100, progress=progress) ) inner.create_task( - chunk_docs(conn, docs_chan, embed_chan, upsert_chan, extract_chan) + chunk_docs( + conn, + docs_chan, + embed_chan, + upsert_chan, + extract_chan, + progress=progress, + total_docs=total_docs, + ) ) grp.create_task(new_doc_pipeline()) grp.create_task(upsert(upsert_chan)) grp.create_task(load(conn)) - grp.create_task(embed(embed_chan, embedding_store, graphname, progress)) + grp.create_task(embed(embed_chan, embedding_store, graphname)) grp.create_task( - extract(extract_chan, upsert_chan, embed_chan, extractor, conn, num_chunk_senders, progress) + extract(extract_chan, upsert_chan, embed_chan, extractor, conn, num_chunk_senders) ) logger.info("Join docs_chan") await docs_chan.join() @@ -678,7 +750,7 @@ async def new_doc_pipeline(): # schema edits could still leave queries missing. community_start = time.perf_counter() if community_detection_switch: - _report("Detecting communities") + _report("Detecting communities", clear_progress=True) await install_queries(COMMUNITY_QUERIES, conn) logger.info("Community Processing Start") @@ -704,11 +776,11 @@ async def new_doc_pipeline(): grp.create_task(communities(conn, comm_process_chan)) # summarize each community grp.create_task( - summarize_communities(conn, comm_process_chan, upsert_chan, embed_chan, progress) + summarize_communities(conn, comm_process_chan, upsert_chan, embed_chan) ) grp.create_task(upsert(upsert_chan)) grp.create_task(load(conn)) - grp.create_task(embed(embed_chan, embedding_store, graphname, progress)) + grp.create_task(embed(embed_chan, embedding_store, graphname)) logger.info("Join comm_process_chan") await comm_process_chan.join() logger.info("Join embed_chan") @@ -738,7 +810,7 @@ async def new_doc_pipeline(): f"IN_COMMUNITY pair missing on schema" ) if mirrorable: - _report("Updating domain types") + _report("Updating domain types", clear_progress=True) await graphrag_mirror_communities(conn, mirrorable) community_end = time.perf_counter() logger.info("Community Processing End") diff --git a/ecc/app/main.py b/ecc/app/main.py index 9889dc0..c4828d2 100644 --- a/ecc/app/main.py +++ b/ecc/app/main.py @@ -226,6 +226,9 @@ def rebuild_status( "is_running": task_info.get("status") == "running", "status": task_info.get("status"), "stage": task_info.get("stage"), + "progress_current": task_info.get("progress_current"), + "progress_total": task_info.get("progress_total"), + "progress_pct": task_info.get("progress_pct"), "started_at": task_info.get("started_at"), "completed_at": task_info.get("completed_at"), "failed_at": task_info.get("failed_at"), @@ -240,14 +243,33 @@ def rebuild_status( } -def _set_stage(task_key: str, msg: str) -> None: +def _set_stage( + task_key: str, + msg: str, + current=None, + total=None, + clear_progress: bool = False, +) -> None: """Update the human-readable stage label for an in-flight task. - Pulled out so individual stage transitions don't have to know - about the ``running_tasks`` schema. + + Optional *current*/*total* populate a progress bar. Progress fields + are left unchanged unless new values are provided or + *clear_progress* is True — otherwise a later stage string (e.g. + from stream_docs finishing early) would wipe the chunking bar + before the UI can poll it. """ info = running_tasks.get(task_key) - if info is not None: - info["stage"] = msg + if info is None: + return + info["stage"] = msg + if current is not None and total is not None and total > 0: + info["progress_current"] = int(current) + info["progress_total"] = int(total) + info["progress_pct"] = min(100, int(100 * current / total)) + elif clear_progress: + info.pop("progress_current", None) + info.pop("progress_total", None) + info.pop("progress_pct", None) async def run_with_tracking(task_key: str, run_func, graphname: str, conn): @@ -304,7 +326,15 @@ async def run_with_tracking(task_key: str, run_func, graphname: str, conn): # ``run_func`` may ignore the kwarg (the supportai legacy path # does); the call falls back to the no-progress signature on # ``TypeError``. - progress_cb = lambda msg: _set_stage(task_key, msg) + def progress_cb(msg, current=None, total=None, clear_progress=False): + _set_stage( + task_key, + msg, + current=current, + total=total, + clear_progress=clear_progress, + ) + try: await run_func(graphname, conn, progress=progress_cb) except TypeError: diff --git a/graphrag-ui/src/pages/setup/KGAdmin.tsx b/graphrag-ui/src/pages/setup/KGAdmin.tsx index a2691a7..62c7ba1 100644 --- a/graphrag-ui/src/pages/setup/KGAdmin.tsx +++ b/graphrag-ui/src/pages/setup/KGAdmin.tsx @@ -234,6 +234,7 @@ const KGAdmin = () => { if (!open) { setRefreshMessage(""); setPollingActive(false); + setRebuildProgress(null); } }; @@ -471,6 +472,12 @@ const KGAdmin = () => { const isRebuildRunningRef = useRef(false); const [isCheckingStatus, setIsCheckingStatus] = useState(false); const [pollingActive, setPollingActive] = useState(false); + // Document-level chunking progress from ECC rebuild_status (optional). + const [rebuildProgress, setRebuildProgress] = useState<{ + pct: number; + current: number; + total: number; + } | null>(null); // Load available graphs. First seed from sessionStorage so the // dropdown shows something immediately, then refresh from @@ -1124,20 +1131,43 @@ const KGAdmin = () => { setRefreshMessage( `⚠️ A rebuild is already in progress for "${graphName}" (started at ${startTime})${stage}. Please wait for it to complete.` ); + if ( + typeof statusData.progress_pct === "number" && + typeof statusData.progress_total === "number" && + statusData.progress_total > 0 + ) { + setRebuildProgress({ + pct: statusData.progress_pct, + current: statusData.progress_current ?? 0, + total: statusData.progress_total, + }); + } else if ( + !(typeof statusData.stage === "string" && + statusData.stage.includes("Chunking documents")) + ) { + // Keep last chunking bar if a poll lands between updates; + // clear only once we've left the chunking stage. + setRebuildProgress(null); + } } else if (wasRunning && statusData.status === "completed") { setRefreshMessage(`✅ Rebuild completed successfully for "${graphName}".`); setPollingActive(false); + setRebuildProgress(null); } else if (statusData.status === "failed") { setRefreshMessage(`❌ Previous rebuild failed: ${statusData.error || "Unknown error"}`); setPollingActive(false); + setRebuildProgress(null); } else if (statusData.status === "error") { setRefreshMessage(`❌ Failed to check rebuild status: ${statusData.error || "Unknown error"}`); setPollingActive(false); + setRebuildProgress(null); } else if (statusData.status === "unknown") { setRefreshMessage(`⚠️ ECC service returned unknown status. It may be unavailable.`); setPollingActive(false); + setRebuildProgress(null); } else { setRefreshMessage(""); + setRebuildProgress(null); } } else { setRefreshMessage(`❌ Failed to check rebuild status (HTTP ${statusResponse.status}).`); @@ -2673,6 +2703,23 @@ const KGAdmin = () => { {refreshMessage} )} + + {isRebuildRunning && rebuildProgress && ( +
+
+ Chunking documents + + {rebuildProgress.current}/{rebuildProgress.total} ({rebuildProgress.pct}%) + +
+
+
+
+
+ )}