custom_search/
main.rs

1use anyhow::Result;
2use mistralrs::{
3    IsqType, RequestBuilder, SearchResult, TextMessageRole, TextMessages, TextModelBuilder,
4    WebSearchOptions,
5};
6
7use std::fs;
8use std::sync::Arc;
9use walkdir::WalkDir;
10
11/// Very small helper that searches the current repository for files that match the given query.
12/// The file name is matched against the query (case-sensitive contains check) and the whole file
13/// content is returned so the model can decide which parts are useful.
14fn local_search(query: &str) -> Result<Vec<SearchResult>> {
15    let mut results = Vec::new();
16
17    for entry in WalkDir::new(".") {
18        let entry = entry?;
19        if entry.file_type().is_file() {
20            let name = entry.file_name().to_string_lossy();
21            if name.contains(query) {
22                let path = entry.path().display().to_string();
23                let content = fs::read_to_string(entry.path()).unwrap_or_default();
24
25                results.push(SearchResult {
26                    title: name.into_owned(),
27                    description: path.clone(),
28                    url: path,
29                    content,
30                });
31            }
32        }
33    }
34
35    // Simple ordering – longest titles (usually most specific) first.
36    results.sort_by_key(|r| r.title.clone());
37    results.reverse();
38    Ok(results)
39}
40
41#[tokio::main]
42async fn main() -> Result<()> {
43    // Build the model enabling web-search support.  We supply a custom search callback that is
44    // used **instead** of the default remote web search when we set `web_search_options` later
45    // in the request.
46
47    // The EmbeddingGemma reranker is **not** required even when using a custom callback – it is
48    // used inside the retrieval pipeline to cluster / rank the results that our callback returns.
49    let model = TextModelBuilder::new("NousResearch/Hermes-3-Llama-3.1-8B")
50        .with_isq(IsqType::Q4K)
51        .with_logging()
52        .with_search_callback(Arc::new(|params: &mistralrs::SearchFunctionParameters| {
53            // In a real application there could be network or database calls here – but for the
54            // sake of demonstration we simply perform a local filesystem search.
55            local_search(&params.query)
56        }))
57        .build()
58        .await?;
59
60    let messages =
61        TextMessages::new().add_message(TextMessageRole::User, "Where is Cargo.toml in this repo?");
62
63    // Enable searching for this request.  Because we provided a custom search callback above, the
64    // model will call **our** function instead of performing an online web search.
65    let messages =
66        RequestBuilder::from(messages).with_web_search_options(WebSearchOptions::default());
67
68    let response = model.send_chat_request(messages).await?;
69
70    println!("{}", response.choices[0].message.content.as_ref().unwrap());
71
72    Ok(())
73}
custom_search/main.rs

custom_search/
main.rs
