Move to builder

Author: EricLBuehler

Cargo.lock

@@ -23,6 +23,13 @@ Internally, we use a BERT model (Snowflake/snowflake-arctic-embed-l-v2.0)[https:
 - Python: `search_bert_model` in the Runner
 - Server: `search-bert-model` before the model type selector (`plain`/`vision-plain`)
 
+## Specifying a custom search callback
+
+By default, mistral.rs uses a DuckDuckGo-based search callback. To override this, you can provide your own search function:
+
+- Rust: use `.with_search_callback(...)` on the model builder with an `Arc<dyn Fn(&SearchFunctionParameters) -> anyhow::Result<Vec<SearchResult>> + Send + Sync>`.
+- Python: pass the `search_callback` keyword argument to `Runner`, which should be a function `def search_callback(query: str) -> List[Dict[str, str]]` returning a list of results with keys `"title"`, `"description"`, `"url"`, and `"content"`.
+
 ## HTTP server
 **Be sure to add `--enable-search`!**
 

docs/WEB_SEARCH.md

@@ -1,7 +1,6 @@
 use crate::{
     distributed,
     embedding::bert::BertPipeline,
-    search,
     pipeline::{
         llg::{constraint_from_llg_grammar, llg_grammar_from_constraint},
         text_models_inputs_processor::PagedAttentionMeta,
@@ -10,6 +9,7 @@ use crate::{
     prefix_cacher::PrefixCacheManagerV2,
     response::CompletionChoice,
     scheduler::{Scheduler, SchedulerOutput},
+    search,
     sequence::{SeqStepType, StopReason},
     CompletionResponse, SchedulerConfig, DEBUG,
 };

mistralrs-core/src/engine/mod.rs

@@ -104,6 +104,7 @@ pub use sampler::{
     CustomLogitsProcessor, DrySamplingParams, SamplingParams, StopTokens, TopLogprob,
 };
 pub use scheduler::{DefaultSchedulerMethod, SchedulerConfig};
+pub use search::{SearchCallback, SearchFunctionParameters, SearchResult};
 use serde::Serialize;
 pub use speech_models::{utils as speech_utils, SpeechGenerationConfig, SpeechLoaderType};
 use tokio::runtime::Runtime;
@@ -112,7 +113,6 @@ pub use tools::{
     CalledFunction, Function, Tool, ToolCallResponse, ToolCallType, ToolChoice, ToolType,
 };
 pub use topology::{LayerTopology, Topology};
-pub use search::SearchCallback;
 pub use utils::debug::initialize_logging;
 pub use utils::memory_usage::MemoryUsage;
 pub use utils::normal::{ModelDType, TryIntoDType};
@@ -202,12 +202,14 @@ pub struct MistralRsBuilder {
 }
 
 impl MistralRsBuilder {
+    /// Creates a new builder with the given pipeline, scheduler method, logging flag,
+    /// and optional embedding model for web search. To override the search callback,
+    /// use `.with_search_callback(...)` on the builder.
     pub fn new(
         pipeline: Arc<tokio::sync::Mutex<dyn Pipeline>>,
         method: SchedulerConfig,
         throughput_logging: bool,
         search_embedding_model: Option<BertEmbeddingModel>,
-        search_callback: Option<Arc<SearchCallback>>,
     ) -> Self {
         Self {
             pipeline,
@@ -220,7 +222,7 @@ impl MistralRsBuilder {
             disable_eos_stop: None,
             throughput_logging_enabled: throughput_logging,
             search_embedding_model,
-            search_callback,
+            search_callback: None,
         }
     }
     pub fn with_log(mut self, log: String) -> Self {

mistralrs-core/src/lib.rs

@@ -16,7 +16,8 @@ use crate::{Function, Tool, ToolType, WebSearchOptions, WebSearchUserLocation};
 
 /// Callback used to override how search results are gathered. The returned
 /// vector must be sorted in decreasing order of relevance.
-pub type SearchCallback = dyn Fn(&SearchFunctionParameters) -> Result<Vec<SearchResult>> + Send + Sync;
+pub type SearchCallback =
+    dyn Fn(&SearchFunctionParameters) -> Result<Vec<SearchResult>> + Send + Sync;
 
 pub(crate) fn search_tool_called(name: &str) -> bool {
     name == SEARCH_TOOL_NAME || name == EXTRACT_TOOL_NAME

mistralrs-core/src/search/mod.rs

@@ -32,20 +32,19 @@ use mistralrs_core::{
     SpeechLoader, StopTokens, TokenSource, TokenizationRequest, Tool, Topology,
     VisionLoaderBuilder, VisionSpecificConfig,
 };
-use mistralrs_core::{search::SearchFunctionParameters, search::SearchResult, SearchCallback};
+use mistralrs_core::{SearchCallback, SearchFunctionParameters, SearchResult};
 use pyo3::prelude::*;
-use pyo3::types::{PyType, PyList};
-use pyo3::PyObject;
+use pyo3::types::{PyList, PyType};
 use pyo3::Bound;
+use pyo3::PyObject;
 use std::fs::File;
 mod anymoe;
 mod requests;
 mod stream;
 mod util;
 mod which;
-use which::{Architecture, DiffusionArchitecture, SpeechLoaderType, VisionArchitecture, Which};
-// (keep imports minimal – if needed later, re-introduce)
 use mistralrs_core::ModelDType;
+use which::{Architecture, DiffusionArchitecture, SpeechLoaderType, VisionArchitecture, Which};
 
 static DEVICE: OnceLock<Result<Device>> = OnceLock::new();
 
@@ -95,18 +94,24 @@ static NEXT_REQUEST_ID: Mutex<RefCell<usize>> = Mutex::new(RefCell::new(0));
 fn wrap_search_callback(cb: PyObject) -> Arc<SearchCallback> {
     Arc::new(move |params: &SearchFunctionParameters| {
         Python::with_gil(|py| {
-            let list = cb.call1(py, (params.query.clone(),))?.downcast::<PyList>()?;
+            let obj = cb.call1(py, (params.query.clone(),))?;
+            let list = obj.downcast_bound::<PyList>(py)?;
             let mut results = Vec::new();
             for item in list.iter() {
                 let title: String = item.get_item("title")?.extract()?;
                 let description: String = item.get_item("description")?.extract()?;
                 let url: String = item.get_item("url")?.extract()?;
                 let content: String = item.get_item("content")?.extract()?;
-                results.push(SearchResult { title, description, url, content });
+                results.push(SearchResult {
+                    title,
+                    description,
+                    url,
+                    content,
+                });
             }
             Ok(results)
         })
-        .map_err(|e| anyhow::anyhow!(e.to_string()))
+        .map_err(|e: PyErr| anyhow::anyhow!(e.to_string()))
     })
 }
 
@@ -842,7 +847,11 @@ impl Runner {
             None
         };
         let cb = search_callback.map(wrap_search_callback);
-        let mistralrs = MistralRsBuilder::new(pipeline, scheduler_config, false, bert_model, cb)
+        let mut builder = MistralRsBuilder::new(pipeline, scheduler_config, false, bert_model);
+        if let Some(cb) = cb {
+            builder = builder.with_search_callback(cb);
+        }
+        let mistralrs = builder
             .with_no_kv_cache(no_kv_cache)
             .with_prefix_cache_n(prefix_cache_n)
             .build();

mistralrs-pyo3/src/lib.rs

@@ -1,6 +1,6 @@
 //! ## mistral.rs instance for server builder.
 
-use std::num::NonZeroUsize;
+use std::{num::NonZeroUsize, sync::Arc};
 
 use anyhow::{Context, Result};
 use candle_core::Device;
@@ -9,7 +9,7 @@ use mistralrs_core::{
     parse_isq_value, AutoDeviceMapParams, BertEmbeddingModel, DefaultSchedulerMethod,
     DeviceLayerMapMetadata, DeviceMapMetadata, DeviceMapSetting, Loader, LoaderBuilder,
     MemoryGpuConfig, MistralRsBuilder, ModelSelected, PagedAttentionConfig, SchedulerConfig,
-    TokenSource,
+    SearchCallback, TokenSource,
 };
 use tracing::info;
 
@@ -18,6 +18,9 @@ use crate::types::{LoadedPipeline, SharedMistralRsState};
 pub mod defaults {
     //! Provides the default values used for the mistral.rs instance for server.
     //! These defaults can be used for CLI argument fallbacks, config loading, or general initialization.
+
+    use std::sync::Arc;
+
     pub const DEVICE: Option<candle_core::Device> = None;
     pub const SEED: Option<u64> = None;
     pub const LOG: Option<String> = None;
@@ -42,6 +45,7 @@ pub mod defaults {
     pub const ENABLE_SEARCH: bool = false;
     pub const SEARCH_BERT_MODEL: Option<String> = None;
     pub const TOKEN_SOURCE: mistralrs_core::TokenSource = mistralrs_core::TokenSource::CacheToken;
+    pub const SEARCH_CALLBACK: Option<Arc<mistralrs_core::SearchCallback>> = None;
 }
 
 /// A builder for creating a mistral.rs instance with configured options for the mistral.rs server.
@@ -169,6 +173,9 @@ pub struct MistralRsForServerBuilder {
 
     /// Specify a Hugging Face model ID for a BERT model to assist web searching. Defaults to Snowflake Arctic Embed L.
     search_bert_model: Option<String>,
+
+    /// Optional override search callback
+    search_callback: Option<Arc<SearchCallback>>,
 }
 
 impl Default for MistralRsForServerBuilder {
@@ -199,6 +206,7 @@ impl Default for MistralRsForServerBuilder {
             cpu: defaults::CPU,
             enable_search: defaults::ENABLE_SEARCH,
             search_bert_model: defaults::SEARCH_BERT_MODEL,
+            search_callback: defaults::SEARCH_CALLBACK,
         }
     }
 }
@@ -460,6 +468,12 @@ impl MistralRsForServerBuilder {
         self
     }
 
+    /// Override the search function used when `web_search_options` is enabled.
+    pub fn with_search_callback(mut self, callback: Arc<SearchCallback>) -> Self {
+        self.search_callback = Some(callback);
+        self
+    }
+
     /// Builds the configured mistral.rs instance.
     ///
     /// ### Examples

mistralrs-server-core/src/mistralrs_for_server_builder.rs

@@ -1,11 +1,11 @@
 use anyhow::Result;
-use walkdir::WalkDir;
-use std::fs;
-use std::sync::Arc;
 use mistralrs::{
-    BertEmbeddingModel, IsqType, RequestBuilder, SearchCallback, SearchResult,
-    TextMessageRole, TextMessages, TextModelBuilder, WebSearchOptions,
+    BertEmbeddingModel, IsqType, RequestBuilder, SearchResult, TextMessageRole, TextMessages,
+    TextModelBuilder, WebSearchOptions,
 };
+use std::fs;
+use std::sync::Arc;
+use walkdir::WalkDir;
 
 fn local_search(query: &str) -> Result<Vec<SearchResult>> {
     let mut results = Vec::new();
@@ -40,16 +40,12 @@ async fn main() -> Result<()> {
         .build()
         .await?;
 
-    let messages = TextMessages::new().add_message(
-        TextMessageRole::User,
-        "Where is Cargo.toml in this repo?",
-    );
-    let messages = RequestBuilder::from(messages).with_web_search_options(
-        WebSearchOptions {
-            search_description: Some("Local filesystem search".to_string()),
-            ..Default::default()
-        },
-    );
+    let messages =
+        TextMessages::new().add_message(TextMessageRole::User, "Where is Cargo.toml in this repo?");
+    let messages = RequestBuilder::from(messages).with_web_search_options(WebSearchOptions {
+        search_description: Some("Local filesystem search".to_string()),
+        ..Default::default()
+    });
 
     let response = model.send_chat_request(messages).await?;
     println!("{}", response.choices[0].message.content.as_ref().unwrap());

mistralrs/examples/local_search/main.rs

@@ -116,10 +116,13 @@ impl AnyMoeModelBuilder {
             scheduler_method,
             self.base.throughput_logging,
             self.base.search_bert_model,
-            self.base.search_callback.clone(),
-        )
-        .with_no_kv_cache(self.base.no_kv_cache)
-        .with_no_prefix_cache(self.base.prefix_cache_n.is_none());
+        );
+        if let Some(cb) = self.base.search_callback.clone() {
+            runner = runner.with_search_callback(cb);
+        }
+        runner = runner
+            .with_no_kv_cache(self.base.no_kv_cache)
+            .with_no_prefix_cache(self.base.prefix_cache_n.is_none());
 
         if let Some(n) = self.base.prefix_cache_n {
             runner = runner.with_prefix_cache_n(n)

mistralrs/src/anymoe.rs

@@ -95,7 +95,7 @@ impl DiffusionModelBuilder {
             method: DefaultSchedulerMethod::Fixed(self.max_num_seqs.try_into()?),
         };
 
-        let runner = MistralRsBuilder::new(pipeline, scheduler_method, false, None, None);
+        let runner = MistralRsBuilder::new(pipeline, scheduler_method, false, None);
 
         Ok(Model::new(runner.build()))
     }