codex/xianren_studio
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: 新增知识库选项卡与完整操作界面——库管理、多格式文档上传（文本/PDF）、自动分块、FTS 全文检索、预览/重切/删除、结果复制

Browse Source
This commit is contained in:
xianrenge
2026-08-17 14:13:09 +08:00
parent 0f942182be
commit ab734478e8
15 changed files with 1686 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/core/src/knowledge_base.rs
+474
View File
@@ -0,0 +1,474 @@
use crate::error::Result;
use rusqlite::{params, Connection};
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct KnowledgeBase {
pub id: String,
pub name: String,
pub description: String,
pub chunk_size: i64,
pub chunk_overlap: i64,
pub doc_count: i64,
pub chunk_count: i64,
pub created_at: String,
pub updated_at: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct KbDocument {
pub id: String,
pub kb_id: String,
pub name: String,
pub file_type: String,
pub file_size: i64,
pub char_count: i64,
pub chunk_count: i64,
pub created_at: String,
pub updated_at: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct KbDocumentDetail {
#[serde(flatten)]
pub doc: KbDocument,
pub content: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct KbSearchHit {
pub chunk_id: String,
pub document_id: String,
pub document_name: String,
pub seq: i64,
pub content: String,
}
const KB_COLUMNS: &str = "id, name, description, chunk_size, chunk_overlap, created_at, updated_at";
const DOC_COLUMNS: &str = "id, kb_id, name, file_type, file_size, char_count, chunk_count, created_at, updated_at";
// ---------- 知识库 ----------
pub fn list_knowledge_bases(db: &Connection) -> Result<Vec<KnowledgeBase>> {
let mut stmt = db.prepare(&format!(
"SELECT {KB_COLUMNS},
(SELECT COUNT(*) FROM kb_documents d WHERE d.kb_id = kb.id),
(SELECT COUNT(*) FROM kb_chunks c WHERE c.kb_id = kb.id)
FROM knowledge_bases kb ORDER BY kb.updated_at DESC, kb.created_at DESC"
))?;
let rows = stmt.query_map([], row_to_kb)?;
let mut out = Vec::new();
for row in rows {
out.push(row?);
}
Ok(out)
}
pub fn get_knowledge_base(db: &Connection, id: &str) -> Result<Option<KnowledgeBase>> {
let mut stmt = db.prepare(&format!(
"SELECT {KB_COLUMNS},
(SELECT COUNT(*) FROM kb_documents d WHERE d.kb_id = kb.id),
(SELECT COUNT(*) FROM kb_chunks c WHERE c.kb_id = kb.id)
FROM knowledge_bases kb WHERE kb.id = ?1"
))?;
let mut rows = stmt.query_map(params![id], row_to_kb)?;
match rows.next() {
Some(row) => Ok(Some(row?)),
None => Ok(None),
}
}
pub fn insert_knowledge_base(db: &Connection, kb: &KnowledgeBase) -> Result<()> {
db.execute(
"INSERT INTO knowledge_bases (id, name, description, chunk_size, chunk_overlap)
VALUES (?1, ?2, ?3, ?4, ?5)",
params![
kb.id,
kb.name,
kb.description,
kb.chunk_size.clamp(100, 10000),
kb.chunk_overlap.clamp(0, 10000),
],
)?;
Ok(())
}
pub fn update_knowledge_base(db: &Connection, kb: &KnowledgeBase) -> Result<()> {
db.execute(
"UPDATE knowledge_bases SET name = ?1, description = ?2,
chunk_size = ?3, chunk_overlap = ?4, updated_at = datetime('now')
WHERE id = ?5",
params![
kb.name,
kb.description,
kb.chunk_size.clamp(100, 10000),
kb.chunk_overlap.clamp(0, 10000),
kb.id,
],
)?;
Ok(())
}
pub fn delete_knowledge_base(db: &Connection, id: &str) -> Result<()> {
db.execute(
"DELETE FROM knowledge_bases WHERE id = ?1",
params![id],
)?;
Ok(())
}
// ---------- 文档 ----------
pub fn list_documents(db: &Connection, kb_id: &str) -> Result<Vec<KbDocument>> {
let mut stmt = db.prepare(&format!(
"SELECT {DOC_COLUMNS} FROM kb_documents
WHERE kb_id = ?1 ORDER BY updated_at DESC, created_at DESC"
))?;
let rows = stmt.query_map(params![kb_id], row_to_doc)?;
let mut out = Vec::new();
for row in rows {
out.push(row?);
}
Ok(out)
}
pub fn get_document(db: &Connection, id: &str) -> Result<Option<KbDocument>> {
let mut stmt = db.prepare(&format!(
"SELECT {DOC_COLUMNS} FROM kb_documents WHERE id = ?1"
))?;
let mut rows = stmt.query_map(params![id], row_to_doc)?;
match rows.next() {
Some(row) => Ok(Some(row?)),
None => Ok(None),
}
}
pub fn get_document_detail(db: &Connection, id: &str) -> Result<Option<KbDocumentDetail>> {
let mut stmt = db.prepare(&format!(
"SELECT {DOC_COLUMNS}, content FROM kb_documents WHERE id = ?1"
))?;
let mut rows = stmt.query_map(params![id], |row| {
let doc = row_to_doc(row)?;
let content: String = row.get(9)?;
Ok(KbDocumentDetail { doc, content })
})?;
match rows.next() {
Some(row) => Ok(Some(row?)),
None => Ok(None),
}
}
/// 插入文档及其分块（同一事务内完成，成功后自动写入 FTS 索引）。
pub fn insert_document_with_chunks(
db: &Connection,
doc: &KbDocument,
content: &str,
chunks: &[String],
) -> Result<()> {
let tx = db.unchecked_transaction()?;
tx.execute(
"INSERT INTO kb_documents
(id, kb_id, name, file_type, file_size, char_count, chunk_count, content)
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8)",
params![
doc.id,
doc.kb_id,
doc.name,
doc.file_type,
doc.file_size,
doc.char_count,
chunks.len() as i64,
content,
],
)?;
for (i, chunk) in chunks.iter().enumerate() {
tx.execute(
"INSERT INTO kb_chunks (id, kb_id, document_id, seq, content)
VALUES (?1, ?2, ?3, ?4, ?5)",
params![
uuid::Uuid::new_v4().to_string(),
doc.kb_id,
doc.id,
(i + 1) as i64,
chunk,
],
)?;
}
tx.commit()?;
Ok(())
}
pub fn delete_document(db: &Connection, id: &str) -> Result<()> {
db.execute("DELETE FROM kb_documents WHERE id = ?1", params![id])?;
Ok(())
}
/// 重新切分文档：删除旧分块后按新参数重分并重建索引。
pub fn rechunk_document(
db: &Connection,
id: &str,
chunk_size: i64,
chunk_overlap: i64,
) -> Result<usize> {
let detail = get_document_detail(db, id)?
.ok_or_else(|| crate::error::CoreError::Other("document not found".into()))?;
let chunks = chunk_text(&detail.content, chunk_size as usize, chunk_overlap as usize);
let tx = db.unchecked_transaction()?;
tx.execute("DELETE FROM kb_chunks WHERE document_id = ?1", params![id])?;
for (i, chunk) in chunks.iter().enumerate() {
tx.execute(
"INSERT INTO kb_chunks (id, kb_id, document_id, seq, content)
VALUES (?1, ?2, ?3, ?4, ?5)",
params![
uuid::Uuid::new_v4().to_string(),
detail.doc.kb_id,
id,
(i + 1) as i64,
chunk,
],
)?;
}
tx.execute(
"UPDATE kb_documents SET chunk_count = ?1, updated_at = datetime('now') WHERE id = ?2",
params![chunks.len() as i64, id],
)?;
tx.commit()?;
Ok(chunks.len())
}
/// 把文本按字符数切块，相邻块之间保留 overlap 字符的重叠。
pub fn chunk_text(text: &str, chunk_size: usize, chunk_overlap: usize) -> Vec<String> {
let size = chunk_size.max(100);
let overlap = chunk_overlap.min(size.saturating_sub(1));
let chars: Vec<char> = text.chars().collect();
let mut chunks = Vec::new();
let mut start = 0usize;
while start < chars.len() {
let end = (start + size).min(chars.len());
let piece: String = chars[start..end].iter().collect();
if !piece.trim().is_empty() {
chunks.push(piece);
}
if end >= chars.len() {
break;
}
start = end.saturating_sub(overlap);
}
chunks
}
// ---------- 检索 ----------
/// 在知识库中检索：优先 FTS5（trigram）BM25 排序，短查询或解析失败回退 LIKE。
pub fn search(
db: &Connection,
kb_id: Option<&str>,
query: &str,
limit: usize,
) -> Result<Vec<KbSearchHit>> {
let q = query.trim();
if q.is_empty() {
return Ok(Vec::new());
}
let limit = limit.clamp(1, 100) as i64;
if q.chars().count() >= 3 {
let fts_q = format!("\"{}\"", q.replace('"', "\"\""));
let result = db.prepare(
"SELECT c.id, c.document_id, d.name, c.seq, c.content
FROM kb_chunks c
JOIN kb_documents d ON d.id = c.document_id
JOIN kb_chunks_fts ON kb_chunks_fts.rowid = c.rowid
WHERE (?1 IS NULL OR c.kb_id = ?1) AND kb_chunks_fts MATCH ?2
ORDER BY bm25(kb_chunks_fts) LIMIT ?3",
)
.and_then(|mut stmt| {
let rows = stmt.query_map(params![kb_id, fts_q, limit], row_to_hit)?;
rows.collect::<std::result::Result<Vec<_>, _>>()
});
if let Ok(hits) = result {
return Ok(hits);
}
}
// 回退：LIKE 子串匹配
let mut stmt = db.prepare(
"SELECT c.id, c.document_id, d.name, c.seq, c.content
FROM kb_chunks c
JOIN kb_documents d ON d.id = c.document_id
WHERE (?1 IS NULL OR c.kb_id = ?1) AND c.content LIKE '%' || ?2 || '%'
ORDER BY c.rowid LIMIT ?3",
)?;
let rows = stmt.query_map(params![kb_id, q, limit], row_to_hit)?;
let mut out = Vec::new();
for row in rows {
out.push(row?);
}
Ok(out)
}
fn row_to_kb(row: &rusqlite::Row<'_>) -> rusqlite::Result<KnowledgeBase> {
Ok(KnowledgeBase {
id: row.get(0)?,
name: row.get(1)?,
description: row.get(2)?,
chunk_size: row.get(3)?,
chunk_overlap: row.get(4)?,
doc_count: row.get(7)?,
chunk_count: row.get(8)?,
created_at: row.get(5)?,
updated_at: row.get(6)?,
})
}
fn row_to_doc(row: &rusqlite::Row<'_>) -> rusqlite::Result<KbDocument> {
Ok(KbDocument {
id: row.get(0)?,
kb_id: row.get(1)?,
name: row.get(2)?,
file_type: row.get(3)?,
file_size: row.get(4)?,
char_count: row.get(5)?,
chunk_count: row.get(6)?,
created_at: row.get(7)?,
updated_at: row.get(8)?,
})
}
fn row_to_hit(row: &rusqlite::Row<'_>) -> rusqlite::Result<KbSearchHit> {
Ok(KbSearchHit {
chunk_id: row.get(0)?,
document_id: row.get(1)?,
document_name: row.get(2)?,
seq: row.get(3)?,
content: row.get(4)?,
})
}
#[cfg(test)]
mod tests {
use super::*;
use rusqlite::Connection;
fn test_conn() -> Connection {
let conn = Connection::open_in_memory().unwrap();
conn.execute_batch(include_str!("schema.sql")).unwrap();
conn
}
fn insert_doc(conn: &Connection, id: &str, content: &str) {
let doc = KbDocument {
id: id.to_string(),
kb_id: "kb1".to_string(),
name: format!("{id}.md"),
file_type: "md".to_string(),
file_size: content.len() as i64,
char_count: content.chars().count() as i64,
chunk_count: 0,
created_at: String::new(),
updated_at: String::new(),
};
let chunks = chunk_text(content, 200, 20);
insert_document_with_chunks(conn, &doc, content, &chunks).unwrap();
}
#[test]
fn chunk_and_search_roundtrip() {
let conn = test_conn();
let kb = KnowledgeBase {
id: "kb1".to_string(),
name: "测试库".to_string(),
description: String::new(),
chunk_size: 200,
chunk_overlap: 20,
doc_count: 0,
chunk_count: 0,
created_at: String::new(),
updated_at: String::new(),
};
insert_knowledge_base(&conn, &kb).unwrap();
insert_doc(&conn, "d1", "仙人工作室是一个本地大模型桌面应用，支持聊天、智能体与工作流。");
let kbs = list_knowledge_bases(&conn).unwrap();
assert_eq!(kbs.len(), 1);
assert!(kbs[0].chunk_count >= 1);
let hits = search(&conn, Some("kb1"), "智能体", 10).unwrap();
assert_eq!(hits.len(), 1);
assert!(hits[0].content.contains("智能体"));
}
#[test]
fn short_query_falls_back_to_like() {
let conn = test_conn();
insert_knowledge_base(
&conn,
&KnowledgeBase {
id: "kb1".to_string(),
name: "测试库".to_string(),
description: String::new(),
chunk_size: 200,
chunk_overlap: 20,
doc_count: 0,
chunk_count: 0,
created_at: String::new(),
updated_at: String::new(),
},
)
.unwrap();
insert_doc(&conn, "d1", "今天是星期一，天气晴朗，适合出去散步。");
let hits = search(&conn, Some("kb1"), "散步", 10).unwrap();
assert_eq!(hits.len(), 1);
}
#[test]
fn delete_document_removes_chunks() {
let conn = test_conn();
insert_knowledge_base(
&conn,
&KnowledgeBase {
id: "kb1".to_string(),
name: "测试库".to_string(),
description: String::new(),
chunk_size: 200,
chunk_overlap: 20,
doc_count: 0,
chunk_count: 0,
created_at: String::new(),
updated_at: String::new(),
},
)
.unwrap();
insert_doc(&conn, "d1", "知识库测试内容，用于验证删除级联。");
delete_document(&conn, "d1").unwrap();
assert!(get_document(&conn, "d1").unwrap().is_none());
let hits = search(&conn, Some("kb1"), "知识库", 10).unwrap();
assert!(hits.is_empty());
}
#[test]
fn rechunk_changes_chunk_count() {
let conn = test_conn();
insert_knowledge_base(
&conn,
&KnowledgeBase {
id: "kb1".to_string(),
name: "测试库".to_string(),
description: String::new(),
chunk_size: 200,
chunk_overlap: 20,
doc_count: 0,
chunk_count: 0,
created_at: String::new(),
updated_at: String::new(),
},
)
.unwrap();
let long = "这是一段很长很长的知识库测试文本。".repeat(100);
insert_doc(&conn, "d1", &long);
let before = get_document(&conn, "d1").unwrap().unwrap().chunk_count;
let n = rechunk_document(&conn, "d1", 500, 50).unwrap();
let after = get_document(&conn, "d1").unwrap().unwrap().chunk_count;
assert_eq!(n as i64, after);
assert!(after < before);
assert!(!search(&conn, Some("kb1"), "测试文本", 5).unwrap().is_empty());
}
}
crates/core/src/lib.rs
+2
View File
@@ -1,6 +1,7 @@
pub mod agents;
pub mod app;
pub mod error;
pub mod knowledge_base;
pub mod mcp_servers;
pub mod models;
pub mod scheduled_tasks;
@@ -11,6 +12,7 @@ pub mod workflows;
pub use app::CoreApp;
pub use error::{CoreError, Result};
pub use knowledge_base::{KbDocument, KbDocumentDetail, KbSearchHit, KnowledgeBase};
pub use models::ModelInfo;
pub use agents::Agent;
pub use scheduled_tasks::ScheduledTask;
crates/core/src/schema.sql
+48
View File
@@ -127,3 +127,51 @@ CREATE TABLE IF NOT EXISTS scheduled_tasks (
created_at TEXT NOT NULL DEFAULT (datetime('now')),
updated_at TEXT NOT NULL DEFAULT (datetime('now'))
);
CREATE TABLE IF NOT EXISTS knowledge_bases (
id TEXT PRIMARY KEY,
name TEXT NOT NULL,
description TEXT NOT NULL DEFAULT '',
chunk_size INTEGER NOT NULL DEFAULT 500,
chunk_overlap INTEGER NOT NULL DEFAULT 50,
created_at TEXT NOT NULL DEFAULT (datetime('now')),
updated_at TEXT NOT NULL DEFAULT (datetime('now'))
);
CREATE TABLE IF NOT EXISTS kb_documents (
id TEXT PRIMARY KEY,
kb_id TEXT NOT NULL REFERENCES knowledge_bases(id) ON DELETE CASCADE,
name TEXT NOT NULL,
file_type TEXT NOT NULL DEFAULT '',
file_size INTEGER NOT NULL DEFAULT 0,
char_count INTEGER NOT NULL DEFAULT 0,
chunk_count INTEGER NOT NULL DEFAULT 0,
content TEXT NOT NULL DEFAULT '',
created_at TEXT NOT NULL DEFAULT (datetime('now')),
updated_at TEXT NOT NULL DEFAULT (datetime('now'))
);
CREATE TABLE IF NOT EXISTS kb_chunks (
id TEXT PRIMARY KEY,
kb_id TEXT NOT NULL REFERENCES knowledge_bases(id) ON DELETE CASCADE,
document_id TEXT NOT NULL REFERENCES kb_documents(id) ON DELETE CASCADE,
seq INTEGER NOT NULL DEFAULT 1,
content TEXT NOT NULL,
created_at TEXT NOT NULL DEFAULT (datetime('now'))
);
-- 全文检索索引（trigram 分词，对中文/短文本更友好；content='' 为 contentless-delete 模式）
CREATE VIRTUAL TABLE IF NOT EXISTS kb_chunks_fts USING fts5(content, content = '', tokenize = 'trigram');
CREATE TRIGGER IF NOT EXISTS kb_chunks_ai AFTER INSERT ON kb_chunks BEGIN
INSERT INTO kb_chunks_fts(rowid, content) VALUES (new.rowid, new.content);
END;
CREATE TRIGGER IF NOT EXISTS kb_chunks_ad AFTER DELETE ON kb_chunks BEGIN
INSERT INTO kb_chunks_fts(kb_chunks_fts, rowid) VALUES ('delete', old.rowid);
END;
CREATE TRIGGER IF NOT EXISTS kb_chunks_au AFTER UPDATE ON kb_chunks BEGIN
INSERT INTO kb_chunks_fts(kb_chunks_fts, rowid) VALUES ('delete', old.rowid);
INSERT INTO kb_chunks_fts(rowid, content) VALUES (new.rowid, new.content);
END;