llm_manager/backend/
web_context.rs

1use tracing::info;
2
3use crate::backend::web_search;
4
5const WEB_SEARCH_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(15);
6
7/// Result of building an injected prompt with web search context.
8pub struct InjectedPrompt {
9    /// The modified message content with web context prepended.
10    pub content: String,
11    /// Whether web search was actually performed.
12    pub performed: bool,
13}
14
15fn log(cb: &std::sync::Mutex<Option<Box<dyn Fn(String) + Send + Sync>>>, msg: String) {
16    if let Some(c) = cb.lock().unwrap().as_ref() {
17        c(msg);
18    }
19}
20
21/// Build the full prompt to send to llama-server, including web search context injection.
22/// Returns the original request if no web search is needed or the preset doesn't match.
23pub async fn build_injected_prompt(
24    preset_name: &str,
25    messages: &serde_json::Value,
26    web_search_enabled: bool,
27    web_search_engine: &str,
28    web_search_engine_url: &str,
29    web_search_api_key: &str,
30    log_callback: &std::sync::Mutex<Option<Box<dyn Fn(String) + Send + Sync>>>,
31) -> InjectedPrompt {
32    log(log_callback, format!("Web search: preset='{}', enabled={}", preset_name, web_search_enabled));
33    if !web_search_enabled {
34        log(log_callback, "Web search: disabled in config, skipping".into());
35        return InjectedPrompt {
36            content: String::new(),
37            performed: false,
38        };
39    }
40
41    let messages_array = match messages.get("messages").and_then(|m| m.as_array()) {
42        Some(m) => {
43            info!("Web search: found {} messages", m.len());
44            log(log_callback, format!("Web search: found {} messages", m.len()));
45            m
46        }
47        None => {
48            info!("Web search: no messages array in request");
49            log(log_callback, "Web search: no messages array in request".into());
50            return InjectedPrompt {
51                content: String::new(),
52                performed: false,
53            };
54        }
55    };
56
57    if messages_array.is_empty() {
58        log(log_callback, "Web search: empty messages array".into());
59        return InjectedPrompt {
60            content: String::new(),
61            performed: false,
62        };
63    }
64
65    let last_msg = messages_array.last().unwrap();
66    let user_content = last_msg.get("content");
67    let content = match user_content {
68        Some(serde_json::Value::String(s)) => {
69            info!("Web search: content is String ({} chars)", s.len());
70            log(log_callback, format!("Web search: content is String ({} chars)", s.len()));
71            s.clone()
72        }
73        Some(serde_json::Value::Array(parts)) => {
74            info!("Web search: content is Array with {} parts", parts.len());
75            let text_parts: Vec<&str> = parts
76                .iter()
77                .filter_map(|p| p.get("text").and_then(|t| t.as_str()))
78                .collect();
79            if text_parts.is_empty() {
80                info!("Web search: no text parts in array");
81                log(log_callback, "Web search: no text parts in array".into());
82                return InjectedPrompt {
83                    content: String::new(),
84                    performed: false,
85                };
86            }
87            let joined = text_parts.join(" ");
88            info!("Web search: joined content ({} chars)", joined.len());
89            log(log_callback, format!("Web search: joined content ({} chars)", joined.len()));
90            joined
91        }
92        _ => {
93            info!("Web search: content type is {:?}", user_content.map(|v| {
94                match v {
95                    serde_json::Value::Null => "null",
96                    serde_json::Value::Bool(_) => "bool",
97                    serde_json::Value::Number(_) => "number",
98                    serde_json::Value::String(_) => "string",
99                    serde_json::Value::Array(_) => "array",
100                    serde_json::Value::Object(_) => "object",
101                }
102            }));
103            log(log_callback, "Web search: unsupported content type".into());
104            return InjectedPrompt {
105                content: String::new(),
106                performed: false,
107            };
108        }
109    };
110
111    let needs = web_search::needs_search(&content);
112    info!("Web search: needs_search={} for '{}'", needs, &content[..content.char_indices().nth(80).map(|(i, _)| i).unwrap_or(content.len())]);
113    log(log_callback, format!("Web search: needs_search={} for '{}'", needs, &content[..content.char_indices().nth(80).map(|(i, _)| i).unwrap_or(content.len())]));
114    if !needs {
115        log(log_callback, "Web search: no search keywords found, skipping".into());
116        return InjectedPrompt {
117            content: String::new(),
118            performed: false,
119        };
120    }
121
122    info!("Web search: triggering for message: {}", &content[..content.char_indices().nth(100).map(|(i, _)| i).unwrap_or(content.len())]);
123    log(log_callback, format!("Web search: triggering for: '{}'", &content[..content.char_indices().nth(100).map(|(i, _)| i).unwrap_or(content.len())]));
124
125    let query = content.to_string();
126    let engine = web_search_engine.to_string();
127    let engine_url = web_search_engine_url.to_string();
128    let api_key = web_search_api_key.to_string();
129    log(log_callback, format!("Web search: engine={}, url={}", engine, engine_url));
130    let search_handle = tokio::spawn(async move {
131        web_search::gather_search_context(&query, &engine, &engine_url, &api_key).await
132    });
133
134    let search_result = match tokio::time::timeout(WEB_SEARCH_TIMEOUT, search_handle).await {
135        Ok(Ok(Ok((ctx, sources)))) => {
136            info!("Web search: gathered context ({} chars)", ctx.len());
137            log(log_callback, format!("Web search: gathered {} chars, {} sources", ctx.len(), sources.len()));
138            (ctx, sources)
139        }
140        Ok(Ok(Err(e))) => {
141            info!("Web search failed: {}", e);
142            log(log_callback, format!("Web search failed: {}", e));
143            return InjectedPrompt {
144                content: String::new(),
145                performed: false,
146            };
147        }
148        Ok(Err(e)) => {
149            info!("Web search task panicked: {}", e);
150            log(log_callback, format!("Web search task panicked: {}", e));
151            return InjectedPrompt {
152                content: String::new(),
153                performed: false,
154            };
155        }
156        Err(_) => {
157            info!("Web search timed out");
158            log(log_callback, "Web search timed out".into());
159            return InjectedPrompt {
160                content: String::new(),
161                performed: false,
162            };
163        }
164    };
165
166    let (search_context, sources) = search_result;
167
168    // Build sources list
169    let sources_section = if sources.is_empty() {
170        String::new()
171    } else {
172        let sources_list: String = sources
173            .iter()
174            .enumerate()
175            .map(|(i, url)| format!("{}. {}", i + 1, url))
176            .collect::<Vec<_>>()
177            .join("\n");
178        format!("\n\n---\n\n**Sources:**\n{}\n\n**When using information from these sources, display the original URL as a reference.**", sources_list)
179    };
180
181    info!("Web search: gathered context ({} chars)", search_context.len());
182
183    let new_content = format!(
184        "[WEB CONTEXT]\nINSTRUCTION: Cite sources using inline markdown links in your answer. Format: [source name](URL). Place links directly after the facts they support. If you find PDF link, add them to the list with brief description. Do NOT include claims you cannot verify.\n\n{}\n[END WEB CONTEXT]\n\n{}\n\n---\n\n{}",
185        search_context, sources_section, content
186    );
187
188    #[allow(unused_variables)]
189    if let Some(cb) = log_callback.lock().unwrap().as_ref() {
190        // cb(format!("Web search: results injected ({} chars)", search_context.len()));
191    }
192
193    InjectedPrompt {
194        content: new_content,
195        performed: true,
196    }
197}
llm_manager/backend/web_context.rs

llm_manager/backend/
web_context.rs
