feat(rag): align LightRAG native citations and MCP bridge

This commit is contained in:
lix-2026
2026-06-09 09:20:56 +08:00
parent 0e8b03daf8
commit 8e3be8b0b7
39 changed files with 2255 additions and 2732 deletions
-336
View File
@@ -402,51 +402,6 @@ struct MindmapOutlineToolPayload {
outline: Vec<MindmapOutlineItemPayload>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
struct DocsSearchDatasetPayload {
workspace_id: String,
documents: Vec<index_fts::SearchDocumentRecord>,
mindmaps: Vec<index_fts::SearchMindmapRecord>,
tables: Vec<index_fts::SearchTableRecord>,
table_rows: Vec<index_fts::SearchTableRowRecord>,
assets: Vec<index_fts::SearchAssetRecord>,
}
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
struct DocsSearchToolPayload {
query: String,
workspace_id: Option<String>,
limit: Option<u32>,
include_deleted: Option<bool>,
page_id: Option<String>,
title_only: Option<bool>,
exact: Option<bool>,
include_ocr: Option<bool>,
time_range: Option<String>,
time_field: Option<String>,
custom_range_from: Option<String>,
custom_range_to: Option<String>,
#[serde(default)]
datasets: Vec<DocsSearchDatasetPayload>,
}
#[derive(Debug, Clone, Deserialize)]
#[serde(rename_all = "camelCase")]
struct DocsReadToolPayload {
document_id: String,
max_chars: Option<u32>,
include_content: Option<bool>,
title: Option<String>,
workspace_id: Option<String>,
parent_id: Option<String>,
updated_at: Option<String>,
raw_text_length: Option<u64>,
raw_text: Option<String>,
content: Option<Value>,
}
#[derive(Debug, Serialize, PartialEq)]
#[serde(rename_all = "camelCase")]
pub struct RuntimeSuccess {
@@ -4128,126 +4083,6 @@ fn execute_tool_result(
"results": results,
}))
}
"docs_search" => {
let payload: DocsSearchToolPayload = parse_tool_input(&args, &data)?;
let normalized_query = payload.query.trim().to_string();
if normalized_query.is_empty() {
return Ok(json!({
"ok": true,
"source": infer_tool_source(&data),
"query": normalized_query,
"results": [],
"enqueueAssetIds": [],
}));
}
let limit = payload.limit.unwrap_or(12).clamp(1, 30) as usize;
let requested_workspace = payload
.workspace_id
.as_ref()
.map(|value| value.trim())
.filter(|value| !value.is_empty())
.map(str::to_string);
let mut enqueue_asset_ids = Vec::<String>::new();
let mut results = Vec::<Value>::new();
for dataset in payload.datasets {
if let Some(workspace_id) = requested_workspace.as_ref() {
if dataset.workspace_id != *workspace_id {
continue;
}
}
let evaluation = evaluate_search_documents(
&SearchDocumentsRequest {
query: normalized_query.clone(),
workspace_id: dataset.workspace_id.clone(),
page_id: payload.page_id.clone(),
limit,
title_only: payload.title_only.unwrap_or(false),
exact: payload.exact.unwrap_or(false),
include_ocr: payload.include_ocr.unwrap_or(false),
time_range: payload.time_range.clone().unwrap_or_else(|| "any".into()),
time_field: payload
.time_field
.clone()
.unwrap_or_else(|| "updated".into()),
custom_range_from: payload.custom_range_from.clone(),
custom_range_to: payload.custom_range_to.clone(),
},
&SearchDocumentsDataset {
documents: dataset.documents,
mindmaps: dataset.mindmaps,
tables: dataset.tables,
table_rows: dataset.table_rows,
assets: dataset.assets,
},
);
enqueue_asset_ids.extend(evaluation.enqueue_asset_ids);
for item in evaluation.results {
results.push(json!({
"id": item.id,
"title": item.title,
"snippet": item.snippet,
"updatedAt": item.updated_at,
"createdAt": item.created_at,
"matchField": item.match_field,
"hasOcr": item.has_ocr,
"publicPath": item.public_path,
"score": item.score,
"workspaceId": dataset.workspace_id,
}));
}
}
results.sort_by(|left, right| {
let left_score = left.get("score").and_then(Value::as_f64).unwrap_or(0.0);
let right_score = right.get("score").and_then(Value::as_f64).unwrap_or(0.0);
right_score
.partial_cmp(&left_score)
.unwrap_or(std::cmp::Ordering::Equal)
.then_with(|| {
let left_updated =
left.get("updatedAt").and_then(Value::as_str).unwrap_or("");
let right_updated =
right.get("updatedAt").and_then(Value::as_str).unwrap_or("");
right_updated.cmp(left_updated)
})
});
results.truncate(limit);
enqueue_asset_ids.sort();
enqueue_asset_ids.dedup();
Ok(json!({
"ok": true,
"source": infer_tool_source(&data),
"query": normalized_query,
"results": results,
"enqueueAssetIds": enqueue_asset_ids,
"includeDeleted": payload.include_deleted.unwrap_or(false),
}))
}
"docs_read" => {
let payload: DocsReadToolPayload = parse_tool_input(&args, &data)?;
let max_chars = payload.max_chars.unwrap_or(2500).clamp(200, 20_000) as usize;
let raw_text = payload.raw_text.unwrap_or_default();
let raw_text_length = payload
.raw_text_length
.unwrap_or_else(|| raw_text.chars().count() as u64);
let trimmed = trim_text_for_docs_read(&raw_text, max_chars);
let mut result = json!({
"ok": true,
"source": infer_tool_source(&data),
"documentId": payload.document_id,
"title": payload.title.unwrap_or_else(|| "".into()),
"workspaceId": payload.workspace_id.unwrap_or_else(|| "".into()),
"parentId": payload.parent_id,
"updatedAt": payload.updated_at,
"rawTextLength": raw_text_length,
"rawText": trimmed,
});
if payload.include_content.unwrap_or(false) {
if let Some(map) = result.as_object_mut() {
map.insert("content".into(), payload.content.unwrap_or(Value::Null));
}
}
Ok(result)
}
"doc_insert_blocks" => {
let blocks = normalize_blocks_from_value(&data);
let specs = parse_insert_specs(&args)?;
@@ -4604,44 +4439,6 @@ fn build_tool_plan_steps(
}]);
}
if invocation.tool == "docs_search" {
return Ok(vec![
RuntimeToolPlanStep {
kind: "transport".into(),
name: "docs_search.dataset".into(),
function_name: None,
description: "通过 transport 拉取文档、导图、表格与 OCR 搜索数据集,再交由 Rust runtime 统一排序。".into(),
args_json: tool_wire.args_json.clone(),
},
RuntimeToolPlanStep {
kind: "transform".into(),
name: "docs_search".into(),
function_name: None,
description: "在 Rust runtime 内复用统一搜索评估器,输出跨页文档搜索结果。".into(),
args_json: tool_wire.args_json.clone(),
},
]);
}
if invocation.tool == "docs_read" {
return Ok(vec![
RuntimeToolPlanStep {
kind: "transport".into(),
name: "docs_read.document".into(),
function_name: None,
description: "通过 transport 读取目标文档 meta/content,再交由 Rust runtime 统一裁剪正文与返回结构。".into(),
args_json: tool_wire.args_json.clone(),
},
RuntimeToolPlanStep {
kind: "transform".into(),
name: "docs_read".into(),
function_name: None,
description: "在 Rust runtime 内标准化跨页文档读取结果。".into(),
args_json: tool_wire.args_json.clone(),
},
]);
}
if invocation.tool.starts_with("onlyoffice_") {
let description = match invocation.tool.as_str() {
"onlyoffice_session_resolve" => {
@@ -7464,15 +7261,6 @@ fn strip_html_tags(value: &str) -> String {
result.trim().to_string()
}
fn trim_text_for_docs_read(value: &str, max_chars: usize) -> String {
let taken = value.chars().take(max_chars).collect::<String>();
if value.chars().count() > max_chars {
format!("{taken}")
} else {
taken
}
}
fn walk_mindmap_summaries(root: &MindmapTreeNode, max_nodes: usize) -> Vec<RuntimeMindmapSummary> {
let mut list = Vec::new();
let mut queue = VecDeque::from([(root, None::<String>, 0usize)]);
@@ -14509,130 +14297,6 @@ mod tests {
}
}
#[test]
fn docs_search_tool_plan_uses_transport_and_transform_steps() {
let plan = execute_runtime_input(RuntimeInput::Tool {
context: demo_context(),
tool: RuntimeToolInvocationWire {
tool: "docs_search".into(),
kind: "query".into(),
mode: Some("plan".into()),
args_json: json!({
"query": "rust",
}),
target: None,
reason: Some("跨页搜索".into()),
refs: vec!["task-050".into()],
},
data: None,
})
.expect("docs_search plan should build");
match plan {
RuntimeExecutionPlan::Tool(plan) => {
assert_eq!(plan.tool_name, "docs_search");
assert_eq!(plan.toolset_id, "toolset.docs_read");
assert_eq!(plan.steps.len(), 2);
assert_eq!(plan.steps[0].name, "docs_search.dataset");
assert_eq!(plan.steps[1].name, "docs_search");
}
_ => panic!("expected tool plan"),
}
}
#[test]
fn docs_search_tool_result_uses_rust_search_evaluation() {
let result = execute_runtime_query(RuntimeInput::Tool {
context: demo_context(),
tool: RuntimeToolInvocationWire {
tool: "docs_search".into(),
kind: "query".into(),
mode: Some("result".into()),
args_json: json!({
"query": "rust",
"limit": 5,
}),
target: None,
reason: Some("跨页搜索".into()),
refs: vec!["task-050".into()],
},
data: Some(json!({
"source": "convex",
"datasets": [
{
"workspaceId": "ws_1",
"documents": [
{
"id": "page_1",
"workspaceId": "ws_1",
"title": "Rust 文档",
"rawText": "这里记录 rust runtime 收口",
"createdAt": "2026-04-15T00:00:00Z",
"updatedAt": "2026-04-16T00:00:00Z"
}
],
"mindmaps": [],
"tables": [],
"tableRows": [],
"assets": []
}
]
})),
})
.expect("docs_search result should build");
assert_eq!(result.get("query").and_then(Value::as_str), Some("rust"));
assert_eq!(
result
.get("results")
.and_then(Value::as_array)
.map(Vec::len),
Some(1)
);
assert_eq!(result["results"][0]["id"], json!("page_1"));
}
#[test]
fn docs_read_tool_result_trims_text_and_keeps_content_optional() {
let long_text = "a".repeat(260);
let result = execute_runtime_query(RuntimeInput::Tool {
context: demo_context(),
tool: RuntimeToolInvocationWire {
tool: "docs_read".into(),
kind: "query".into(),
mode: Some("result".into()),
args_json: json!({
"documentId": "page_1",
"maxChars": 200,
"includeContent": true,
}),
target: None,
reason: Some("读取文档".into()),
refs: vec!["task-050".into()],
},
data: Some(json!({
"source": "convex",
"title": "示例页面",
"workspaceId": "ws_1",
"parentId": "parent_1",
"updatedAt": "2026-04-16T00:00:00Z",
"rawText": long_text,
"rawTextLength": 260,
"content": {"blocks": [{"id": "b1"}]}
})),
})
.expect("docs_read result should build");
assert_eq!(result["documentId"], json!("page_1"));
let trimmed = result["rawText"]
.as_str()
.expect("rawText should be string");
assert_eq!(trimmed.chars().count(), 201);
assert!(trimmed.ends_with('…'));
assert_eq!(result["rawTextLength"], json!(260));
assert_eq!(result["content"]["blocks"][0]["id"], json!("b1"));
}
#[test]
fn doc_get_tool_plan_uses_documents_content_query() {
let plan = execute_runtime_input(RuntimeInput::Tool {