llm_wiki_lib/
query.rs

1//! Query Pipeline — User question → Read wiki → LLM answer.
2//!
3//! 1. Read all (or relevant) wiki .md files
4//! 2. Call LLM with question + wiki context
5//! 3. Return synthesized answer
6
7use anyhow::Result;
8use regex;
9use std::path::PathBuf;
10
11use crate::llm::Message;
12use crate::wiki::Wiki;
13
14/// Result of a query.
15#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
16pub struct QueryResult {
17    pub ok: bool,
18    pub question: String,
19    pub answer: String,
20    /// Paths of wiki files that were read to answer the question.
21    pub sources: Vec<String>,
22    /// Whether the wiki was empty.
23    #[serde(default)]
24    pub no_wiki_content: bool,
25}
26
27impl QueryResult {
28    pub fn ok(question: String, answer: String, sources: Vec<PathBuf>) -> Self {
29        Self {
30            ok: true,
31            question,
32            answer,
33            sources: sources
34                .into_iter()
35                .map(|p| p.to_string_lossy().to_string())
36                .collect(),
37            no_wiki_content: false,
38        }
39    }
40
41    pub fn empty(question: String) -> Self {
42        Self {
43            ok: true,
44            question,
45            answer: String::new(),
46            sources: vec![],
47            no_wiki_content: true,
48        }
49    }
50}
51
52/// Run the query pipeline.
53pub async fn run(wiki: &Wiki, question: &str) -> Result<QueryResult> {
54    let files = wiki.read_wiki_files()?;
55
56    if files.is_empty() {
57        return Ok(QueryResult::empty(question.to_string()));
58    }
59
60    // Combine all wiki content for context
61    let context = build_context(&files);
62    let system_md = wiki.config().system_md_content()?;
63    let prompt = build_query_prompt(&system_md, question, &context);
64
65    let llm = wiki.llm();
66    let messages = &[Message::system(&prompt)];
67    let answer = llm.chat(messages).await?;
68    let answer = strip_think_tags(&answer);
69
70    let sources: Vec<PathBuf> = files.iter().map(|f| f.path.clone()).collect();
71    Ok(QueryResult::ok(question.to_string(), answer, sources))
72}
73
74/// Strip LLM thinking tags (<think>/</think>) from model responses.
75fn strip_think_tags(s: &str) -> String {
76    let re = regex::Regex::new(r"(?s)<think>.*?</think>").unwrap();
77    re.replace_all(s, "").to_string()
78}
79
80/// Build a combined context string from wiki files.
81fn build_context(files: &[crate::wiki::WikiFile]) -> String {
82    if files.is_empty() {
83        return String::from("(知识库为空)");
84    }
85
86    let mut ctx = String::new();
87
88    for file in files {
89        let _filename = file
90            .path
91            .file_name()
92            .and_then(|n| n.to_str())
93            .unwrap_or("unknown");
94
95        ctx.push_str("## 文件: \n\n");
96        ctx.push_str(&file.content);
97        ctx.push_str("\n\n---\n\n");
98    }
99
100    ctx
101}
102
103/// Build the LLM prompt for querying.
104fn build_query_prompt(system_md: &str, question: &str, context: &str) -> String {
105    format!(
106        r#"# System Prompt
107
108{system_md}
109
110---
111
112# Task
113
114你是一个知识库助手。请根据以下知识库内容回答用户的问题。
115
116## 知识库内容
117{context}
118
119---
120
121## 用户问题
122
123{question}
124
125---
126
127## 要求
128
1291. 优先使用知识库中的信息回答
1302. 如果知识库中没有相关信息，直接说明"知识库中没有相关内容"
1313. 用中文回答
1324. 如有参考，提及来源文件名
1335. 简洁准确，不得编造信息
134"#
135    )
136}
llm_wiki_lib/query.rs

llm_wiki_lib/
query.rs
