Skip to content

Commit 3b86af3

Browse files
committed
fix(rag): 限制聊天检索超时
Fixes #66
1 parent 648c2bc commit 3b86af3

6 files changed

Lines changed: 366 additions & 73 deletions

File tree

src-tauri/crates/core/src/rag.rs

Lines changed: 34 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,10 +8,22 @@ use async_trait::async_trait;
88
use sea_orm::DatabaseConnection;
99

1010
use crate::error::{AQBotError, Result};
11-
use crate::types::{RagContextResult, RagRetrievedItem, RagSourceResult};
11+
use crate::types::{RagContextResult, RagRetrievedItem, RagSourceError, RagSourceResult};
1212
use crate::vector_store::{EmbeddingRecord, VectorSearchResult, VectorStore};
1313
use crate::{document_parser, text_chunker};
1414

15+
fn format_rag_failure_message(reason: impl AsRef<str>) -> String {
16+
const PREFIX: &str = "检索失败";
17+
let reason = reason.as_ref().trim();
18+
if reason.is_empty() {
19+
return PREFIX.to_string();
20+
}
21+
if reason.starts_with(PREFIX) {
22+
return reason.to_string();
23+
}
24+
format!("{PREFIX}:{reason}")
25+
}
26+
1527
// ── Trait ────────────────────────────────────────────────────────────────────
1628

1729
/// A source of RAG content that can be searched and indexed.
@@ -314,6 +326,7 @@ pub async fn collect_rag_context(
314326

315327
let mut context_parts = Vec::new();
316328
let mut source_results = Vec::new();
329+
let mut errors = Vec::new();
317330

318331
for src_ref in &sources {
319332
let source = src_ref.source();
@@ -426,6 +439,16 @@ pub async fn collect_rag_context(
426439
src_ref.container_id,
427440
e
428441
);
442+
let source_type_str = match src_ref.source_type {
443+
RAGSourceType::Knowledge => "knowledge",
444+
RAGSourceType::Memory => "memory",
445+
};
446+
errors.push(RagSourceError {
447+
source_type: source_type_str.to_string(),
448+
container_id: src_ref.container_id.clone(),
449+
message: format_rag_failure_message(format!("重排序失败:{e}")),
450+
});
451+
continue;
429452
}
430453
}
431454
}
@@ -475,6 +498,15 @@ pub async fn collect_rag_context(
475498
src_ref.container_id,
476499
e
477500
);
501+
let source_type_str = match src_ref.source_type {
502+
RAGSourceType::Knowledge => "knowledge",
503+
RAGSourceType::Memory => "memory",
504+
};
505+
errors.push(RagSourceError {
506+
source_type: source_type_str.to_string(),
507+
container_id: src_ref.container_id.clone(),
508+
message: format_rag_failure_message(format!("向量检索失败:{e}")),
509+
});
478510
}
479511
}
480512
}
@@ -511,6 +543,7 @@ pub async fn collect_rag_context(
511543
RagContextResult {
512544
context_parts,
513545
source_results,
546+
errors,
514547
}
515548
}
516549

src-tauri/crates/core/src/types.rs

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1042,13 +1042,25 @@ pub struct RagSourceResult {
10421042
pub items: Vec<RagRetrievedItem>,
10431043
}
10441044

1045+
/// Retrieval failure for a single RAG source.
1046+
#[derive(Debug, Clone, Serialize, Deserialize)]
1047+
pub struct RagSourceError {
1048+
/// "knowledge" or "memory"
1049+
pub source_type: String,
1050+
pub container_id: String,
1051+
pub message: String,
1052+
}
1053+
10451054
/// Combined results of RAG context collection.
10461055
#[derive(Debug, Clone, Serialize, Deserialize)]
10471056
pub struct RagContextResult {
10481057
/// Formatted context parts for injection into system prompt.
10491058
pub context_parts: Vec<String>,
10501059
/// Structured results for frontend display.
10511060
pub source_results: Vec<RagSourceResult>,
1061+
/// Structured failures for frontend display.
1062+
#[serde(default, skip_serializing_if = "Vec::is_empty")]
1063+
pub errors: Vec<RagSourceError>,
10521064
}
10531065

10541066
/// Tauri event emitted after RAG context retrieval completes.
@@ -1057,6 +1069,8 @@ pub struct RagContextRetrievedEvent {
10571069
pub conversation_id: String,
10581070
pub message_id: Option<String>,
10591071
pub sources: Vec<RagSourceResult>,
1072+
#[serde(default, skip_serializing_if = "Vec::is_empty")]
1073+
pub errors: Vec<RagSourceError>,
10601074
}
10611075

10621076
// === Embedding Types ===

0 commit comments

Comments
 (0)