
mistral.rs Python SDK 自定义搜索回调实战用本地文件系统搜索替换默认 Web 搜索【免费下载链接】mistral.rsFast, flexible LLM inference项目地址: https://gitcode.com/GitHub_Trending/mi/mistral.rs导读本文基于 mistral.rs 官方 Python SDK 示例custom_search源码见 examples/python/custom_search.py对应文档见 custom-search.md讲解如何在推理引擎中注入自定义搜索回调search_callback让模型在 Agent 式问答中通过你提供的检索函数获取上下文而非走内置的 DuckDuckGo 在线搜索。读完本文你将掌握Runner的enable_search/search_callback参数用法、回调返回的数据契约以及搜索工具在引擎内部的分派与重排机制并能举一反三地接入数据库、API 或内部知识库检索。一、为什么需要自定义搜索回调mistral.rs 内置了一套与 OpenAIweb_search_options兼容的 Web 搜索能力当模型需要最新信息时会调用名为mistralrs_search_the_web的搜索工具引擎默认通过 DuckDuckGo HTML 接口执行搜索并并发抓取结果页正文实现在 mistralrs-core/src/search/mod.rs 的run_search_tool中。但默认实现在以下场景并不适用需要检索内部知识库、私有文档、数据库或本地代码仓库而它们不在公网上需要自定义检索策略例如权限过滤、向量相似度检索、专属排序规则需要完全离线运行避免任何外部网络请求。为此mistral.rs 提供了search_callback机制你注册一个自定义回调函数引擎在工具分派阶段会优先调用它来获取搜索结果取代默认的在线搜索。下面的示例正是把搜索目标替换为当前目录下的文件系统。二、完整示例本地文件系统搜索以下代码完整来自官方示例 examples/python/custom_search.py在仓库根目录下运行即可复现from mistralrs import ( Runner, Which, ChatCompletionRequest, Architecture, WebSearchOptions, ) import os def local_search(query: str): results [] for root, _, files in os.walk(.): for f in files: if query in f: path os.path.join(root, f) try: content open(path).read() except Exception: content results.append( { title: f, description: path, url: path, content: content, } ) results.sort(keylambda r: r[title], reverseTrue) return results runner Runner( whichWhich.Plain( model_idNousResearch/Hermes-3-Llama-3.1-8B, archArchitecture.Llama, ), enable_searchTrue, search_callbacklocal_search, ) res runner.send_chat_completion_request( ChatCompletionRequest( modeldefault, messages[{role: user, content: Where is Cargo.toml in this repo?}], max_tokens64, web_search_optionsWebSearchOptions( search_descriptionLocal filesystem search ), ) ) print(res.choices[0].message.content)运行后模型会向用户提问 Where is Cargo.toml in this repo?local_search遍历当前目录树把文件名包含Cargo.toml的文件连同其完整内容返回给模型模型据此给出带路径的答案。三、逐段拆解三个关键契约3.1 回调函数入参与返回结构local_search(query: str)接收一个字符串查询即模型生成的搜索工具参数query返回一个字典列表每个字典必须包含以下四个字段字段含义在示例中的取值title结果标题文件名fdescription结果描述/摘要文件路径pathurl结果地址文件路径pathcontent完整正文内容文件内容读取失败则为空字符串这四个字段与引擎内部SearchResult结构体mistralrs-core/src/search/mod.rs一一对应是回调与引擎之间的数据契约。content是模型最终用于作答的关键信息因此示例里直接把整个文件内容塞进去——文件越大、token 越多引擎会在后续阶段按预算截断。3.2 Runner 配置enable_search 与 search_callbackrunner Runner( whichWhich.Plain( model_idNousResearch/Hermes-3-Llama-3.1-8B, archArchitecture.Llama, ), enable_searchTrue, search_callbacklocal_search, )enable_searchTrue开启搜索能力。从 Python 绑定源码mistralrs-pyo3/src/lib.rs看它决定是否加载用于结果重排的 embedding 模型search_embedding_model未显式指定时使用默认值embedding_gemma即google/embeddinggemma-300m见 mistralrs-core/src/engine/mod.rs 的SearchEmbeddingModel枚举。search_callbacklocal_search注册自定义回调。注意传入的是函数对象本身不调用引擎内部会将其包装为 Rust 侧的类型安全的回调。3.3 请求级配置WebSearchOptionsweb_search_optionsWebSearchOptions( search_descriptionLocal filesystem search )WebSearchOptions是请求级ChatCompletionRequest的web_search_options参数的搜索选项其完整字段定义在 mistralrs-core/src/request.rssearch_context_size搜索结果上下文大小user_location用户位置信息可用于搜索工具描述中附带位置filters域名过滤allowed_domains/blocked_domainsexternal_web_access是否允许外部访问return_token_budget返回结果的 token 预算search_content_types内容类型文本/图片image_settings图片搜索设置search_description覆盖搜索工具的描述文本——本例中用它把默认的 search the web 描述替换为 Local filesystem search引导模型理解这个工具检索的是本地文件系统extract_description覆盖内容提取工具的描述文本。在引擎构造搜索工具时mistralrs-core/src/search/mod.rs 的get_search_toolssearch_description若被指定就会替代内置的默认描述SEARCH_DESCRIPTION注入工具定义未指定时使用默认描述。四、源码级原理从 Python 回调到引擎分派4.1 Rust 侧的类型定义自定义搜索回调在 Rust 侧的类型是pub type SearchCallback dyn Fn(SearchFunctionParameters) - ResultVecSearchResult Send Sync;mistralrs-core/src/search/mod.rs其中SearchFunctionParameters只包含一个query: String字段同文件 L148-L151。也就是说回调本质上是一个查询 → 结果列表的纯函数这为接入任意检索后端文件系统、数据库、向量库、HTTP API留出了极大的自由度。4.2 Python 回调如何被包装Python 绑定通过wrap_search_callbackmistralrs-pyo3/src/lib.rs把 Python 函数包装成 Rust 的SearchCallback每次被调用时在 GIL 下调用 Python 函数传入params.query把返回值向下转型为PyList逐个元素提取title、description、url、content四个字段组装成SearchResult。这解释了为什么 Python 回调必须返回包含这四个键的字典列表——字段缺失或类型不符会直接导致回调报错。4.3 引擎分派回调优先于默认搜索在 Agent 工具分派流程中mistralrs-core/src/engine/tool_dispatch.rs引擎的逻辑是let base: VecSearchResult if let Some(cb) engine.search_callback { // 有自定义回调同步调用用户回调block_in_place match tokio::task::block_in_place(|| cb(params)) { ... } } else { // 无自定义回调走默认 DuckDuckGo 在线搜索 match search::run_search_tool(params).await { ... } };即只要配置了search_callback就完全跳过默认的在线搜索这正是本文示例实现本地文件系统搜索的关键。4.4 分派之后的检索增强管线拿到原始结果后引擎并不会直接把它丢给模型而是继续执行同为 tool_dispatch.rs 中的流程过滤apply_search_filters应用WebSearchFilters中配置的允许/屏蔽域名截断cap_content_lensearch/mod.rs用当前模型的 tokenizer 把每条结果的content截断到 token 预算以内重排若enable_searchTrue加载了 embedding 重排模型SearchPipeline则调用rank_document_chunks按与查询的相关性重排tool_dispatch.rs若未加载 embedding 模型则回退到 BM25 词法排序同文件 L209-L229预算控制按 token 长度从小到大累计直到填满max_toks预算保证送入模型的上下文不会超限。因此自定义回调只需保证返回结果按相关性降序排列即可源码注释明确要求 sorted in decreasing order of relevance剩余的处理截断、重排、预算由引擎完成。五、Rust API 对照ModelBuilder 方式同样的能力在 Rust 侧通过ModelBuilder暴露完整示例见 mistralrs/examples/advanced/search_callback/main.rslet model ModelBuilder::new(Qwen/Qwen3-4B) .with_auto_isq(IsqBits::Four) .with_logging() .with_search_callback(Arc::new(|params: mistralrs::SearchFunctionParameters| { local_search(params.query) })) .build() .await?;与 Python 版对应关系如下能力Python SDKRust SDK注册回调Runner(..., search_callbacklocal_search)ModelBuilder::with_search_callback(ArcSearchCallback)请求级开启ChatCompletionRequest(web_search_options...)RequestBuilder::from(messages).with_web_search_options(WebSearchOptions::default())返回结构含title/description/url/content的字典列表VecSearchResultRust 侧的with_search_callback定义于 mistralrs/src/builder_macros.rs并同样在 text_model.rs、auto_model.rs、gguf.rs 等各类模型构建器中生效。六、扩展思路与注意事项基于这套回调机制你可以把local_search替换为任意检索实现代码库语义检索把文件名匹配升级为基于 embedding 的向量相似度检索content填充命中的代码片段内部知识库问答回调内查询 SQLite/PostgreSQL/向量数据库把url设为记录主键HTTP API 代理回调内调用内部搜索服务将结果透传给模型多数据源聚合合并多个后端结果后统一排序记得按相关性降序返回。使用时有几点需要留意enable_searchTrue与search_callback的关系enable_search控制是否加载重排用的 embedding 模型而回调是否生效取决于是否传入search_callback。从源码看二者相互独立但要让搜索功能完整可用含重排建议两者同时配置。回调是同步阻塞的引擎通过block_in_place调用回调因此回调内部避免做长时间阻塞操作如需慢速检索应考虑在回调内做缓存或预先计算。模型选择示例选用NousResearch/Hermes-3-Llama-3.1-8BArchitecture.Llama这是一款具备较强工具调用能力的模型实际上只要模型支持工具调用均可搭配自定义搜索回调使用。search_description的作用把工具描述改成本地检索语义能显著提升模型对工具的调用意愿与参数质量建议针对你的检索场景自定义描述。综上search_callback是 mistral.rs 把外部检索与模型推理解耦的核心扩展点模型只负责提出查询、消化结果检索逻辑完全由你掌控从而让 Agent 在不依赖公网搜索的前提下获得高质量、可审计的上下文来源。【免费下载链接】mistral.rsFast, flexible LLM inference项目地址: https://gitcode.com/GitHub_Trending/mi/mistral.rs创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考