chore: harden lightrag knowledge rag runtime

2026-06-09 22:12:09 +08:00
parent 922965d30f
commit f698aad964
22 changed files with 899 additions and 101 deletions
@@ -3,11 +3,14 @@ use crate::context::RequestContext;
 use crate::error::WebError;
 use crate::routes::evidence::{citation_markdown_for_locator, citation_url_for_locator};
 use crate::routes::local_folder_source;
+use axum::body::Body;
 use axum::extract::{Extension, Json, Query, State};
-use axum::http::StatusCode;
+use axum::http::{header, StatusCode};
+use axum::response::Response;
 use core_protocol::evidence::{
    EvidenceBBox, EvidenceLocator, EvidenceOpenAction, EvidenceResourceKind,
 };
+use futures_util::TryStreamExt;
 use serde::{Deserialize, Serialize};
 use serde_json::{json, Value};
 use std::collections::hash_map::DefaultHasher;
@@ -327,6 +330,59 @@ pub async fn status(
    })))
 }

+pub async fn pipeline_events(
+    Extension(context): Extension<RequestContext>,
+) -> Result<Response, WebError> {
+    let endpoint = lightrag_endpoint();
+    let url = format!(
+        "{}/documents/pipeline_status/events",
+        endpoint.trim_end_matches('/')
+    );
+    let client = reqwest::Client::builder()
+        .connect_timeout(Duration::from_secs(10))
+        .build()
+        .map_err(|error| {
+            WebError::internal(format!("LightRAG SSE client 初始化失败: {error}"))
+                .with_context(&context)
+        })?;
+    let mut request = client.get(&url).header("accept", "text/event-stream");
+    if let Some(api_key) = lightrag_api_key() {
+        request = request.header("X-API-Key", api_key);
+    }
+    let upstream = request.send().await.map_err(|error| {
+        WebError::bad_gateway_code(
+            "knowledge_rag_lightrag_events_unreachable",
+            format!("无法连接 LightRAG pipeline 事件流: {error}"),
+        )
+        .with_context(&context)
+    })?;
+    let upstream_status = upstream.status();
+    if !upstream_status.is_success() {
+        let text = upstream.text().await.unwrap_or_default();
+        return Err(WebError::bad_gateway_code(
+            "knowledge_rag_lightrag_events_error",
+            format!("LightRAG pipeline 事件流返回 HTTP {upstream_status}: {text}"),
+        )
+        .with_context(&context));
+    }
+    let stream = upstream.bytes_stream().map_err(|error| {
+        std::io::Error::new(
+            std::io::ErrorKind::Other,
+            format!("LightRAG pipeline 事件流读取失败: {error}"),
+        )
+    });
+    Response::builder()
+        .status(StatusCode::OK)
+        .header(header::CONTENT_TYPE, "text/event-stream; charset=utf-8")
+        .header(header::CACHE_CONTROL, "no-cache, no-transform")
+        .header("x-accel-buffering", "no")
+        .body(Body::from_stream(stream))
+        .map_err(|error| {
+            WebError::internal(format!("LightRAG pipeline 事件响应构造失败: {error}"))
+                .with_context(&context)
+        })
+}
+
 pub(crate) async fn sync_registry_for_root(
    state: &AppState,
    context: &RequestContext,
@@ -1447,13 +1503,29 @@ fn lightrag_pipeline_status_summary(value: &Value) -> Value {
        .get("latest_message")
        .and_then(Value::as_str)
        .unwrap_or_default();
-    let progress = parse_lightrag_chunk_progress(latest_message);
+    let history_messages = lightrag_pipeline_history_messages(value);
+    let progress = parse_lightrag_chunk_progress(latest_message).or_else(|| {
+        history_messages
+            .iter()
+            .rev()
+            .find_map(|message| parse_lightrag_chunk_progress(message))
+    });
    json!({
        "ok": true,
        "busy": value.get("busy").and_then(Value::as_bool).unwrap_or(false),
+        "destructiveBusy": value.get("destructive_busy").and_then(Value::as_bool).unwrap_or(false),
        "scanning": value.get("scanning").and_then(Value::as_bool).unwrap_or(false),
+        "scanningExclusive": value.get("scanning_exclusive").and_then(Value::as_bool).unwrap_or(false),
+        "requestPending": value.get("request_pending").and_then(Value::as_bool).unwrap_or(false),
+        "pendingEnqueues": value.get("pending_enqueues").and_then(Value::as_u64).unwrap_or(0),
+        "pendingRequests": value.get("pending_requests").and_then(Value::as_bool).unwrap_or(false),
+        "docs": value.get("docs").and_then(Value::as_u64).unwrap_or(0),
+        "batches": value.get("batchs").and_then(Value::as_u64).unwrap_or(0),
+        "currentBatch": value.get("cur_batch").and_then(Value::as_u64).unwrap_or(0),
        "jobName": value.get("job_name").and_then(Value::as_str).unwrap_or_default(),
+        "jobStart": value.get("job_start").cloned().unwrap_or(Value::Null),
        "latestMessage": latest_message,
+        "historyMessages": history_messages,
        "cancellationRequested": value.get("cancellation_requested").and_then(Value::as_bool).unwrap_or(false),
        "cancellationReason": value.get("cancellation_reason").cloned().unwrap_or(Value::Null),
        "progress": progress.map(|progress| json!({
@@ -1464,6 +1536,25 @@ fn lightrag_pipeline_status_summary(value: &Value) -> Value {
    })
 }

+fn lightrag_pipeline_history_messages(value: &Value) -> Vec<String> {
+    value
+        .get("history_messages")
+        .and_then(Value::as_array)
+        .map(|messages| {
+            messages
+                .iter()
+                .filter_map(Value::as_str)
+                .rev()
+                .take(80)
+                .map(ToOwned::to_owned)
+                .collect::<Vec<_>>()
+                .into_iter()
+                .rev()
+                .collect()
+        })
+        .unwrap_or_default()
+}
+
 fn lightrag_rerank_status_summary(value: &Value) -> Value {
    let configuration = value.get("configuration").unwrap_or(&Value::Null);
    let queue = value.get("rerank_queue_status").unwrap_or(&Value::Null);
@@ -1510,13 +1601,15 @@ fn parse_lightrag_chunk_progress(message: &str) -> Option<LightRagChunkProgress>
    let current = current.trim().parse::<u64>().ok()?;
    let (total, _) = rest.split_once(' ')?;
    let total = total.trim().parse::<u64>().ok()?;
-    let doc_token = message
-        .split_whitespace()
-        .find(|token| token.starts_with("doc-") && token.contains("-chunk-"))?;
+    let doc_token = message.split_whitespace().find(|token| {
+        token.starts_with("doc-") && (token.contains("-chunk-") || token.contains("-mm-"))
+    })?;
    let doc_id = doc_token
        .rsplit_once("-chunk-")
+        .or_else(|| doc_token.rsplit_once("-mm-"))
        .map(|(doc_id, _)| doc_id)
        .unwrap_or(doc_token)
+        .trim_end_matches(|ch: char| !ch.is_ascii_alphanumeric())
        .to_string();
    if total == 0 || current == 0 || doc_id.is_empty() {
        return None;
@@ -3630,6 +3723,11 @@ mod tests {
        let status = json!({
            "busy": true,
            "scanning": true,
+            "destructive_busy": false,
+            "pending_enqueues": 1,
+            "docs": 1,
+            "batchs": 2,
+            "cur_batch": 1,
            "job_name": "book.docx",
            "latest_message": "Chunk 212 of 333 extracted 4 Ent + 0 Rel doc-ff0b60997a285a85e5704a114d7b3ffa-chunk-212"
        });
@@ -3638,6 +3736,10 @@ mod tests {

        assert_eq!(summary["busy"], true);
        assert_eq!(summary["scanning"], true);
+        assert_eq!(summary["pendingEnqueues"], 1);
+        assert_eq!(summary["docs"], 1);
+        assert_eq!(summary["batches"], 2);
+        assert_eq!(summary["currentBatch"], 1);
        assert_eq!(summary["progress"]["current"], 212);
        assert_eq!(summary["progress"]["total"], 333);
        assert_eq!(
@@ -3646,6 +3748,31 @@ mod tests {
        );
    }

+    #[test]
+    fn lightrag_pipeline_summary_extracts_chunk_progress_from_history() {
+        let status = json!({
+            "busy": true,
+            "scanning": false,
+            "job_name": "book.pdf",
+            "latest_message": "Merging stage 1/1: book.pdf",
+            "history_messages": [
+                "Analyzing multimodal: doc-e51400269a2c41a2fcb22f1071193fad",
+                "Chunk 5 of 21 extracted 2 Ent + 1 Rel doc-e51400269a2c41a2fcb22f1071193fad-mm-drawing-001",
+                "Merging stage 1/1: book.pdf"
+            ]
+        });
+
+        let summary = lightrag_pipeline_status_summary(&status);
+
+        assert_eq!(summary["progress"]["current"], 5);
+        assert_eq!(summary["progress"]["total"], 21);
+        assert_eq!(
+            summary["progress"]["docId"],
+            "doc-e51400269a2c41a2fcb22f1071193fad"
+        );
+        assert_eq!(summary["historyMessages"].as_array().unwrap().len(), 3);
+    }
+
    #[test]
    fn old_registry_json_defaults_indexed_roots() {
        let registry: KnowledgeRagSourceRegistry = serde_json::from_value(json!({