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..8cfb891 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,12 +2300,17 @@ 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}") 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}") @@ -2342,12 +2357,15 @@ 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", + "status": cached.get("status", "unknown"), "error": "ECC is busy processing, status check timed out. Rebuild likely still in progress." } except Exception as e: