Skip to content

Add web search hook callbacks #1426

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged

Conversation

EricLBuehler
Copy link
Owner

@EricLBuehler EricLBuehler commented Jun 4, 2025

Summary

  • implement SearchCallback trait and description overrides
  • allow engines to use custom search callbacks
  • plumb callback through builders for text, vision, gguf models
  • expose callback in Python API
  • add rust and python examples demonstrating local filesystem search

Testing

  • cargo fmt (fails: rustfmt not installed)
  • cargo clippy --workspace --all-targets -- -D warnings (fails: clippy not installed)

https://chatgpt.com/codex/tasks/task_e_683fa6e5c3488322b812986a1f087436

Summary by CodeRabbit

  • New Features

    • Added support for custom search callbacks, enabling user-defined search functions in Rust and Python APIs.
    • Introduced new example scripts demonstrating local file search integration with chat models.
    • Enhanced search tool descriptions with customizable options for greater flexibility.
  • Documentation

    • Added guidance on specifying and using custom search callbacks in Rust and Python.

Copy link

coderabbitai bot commented Jun 4, 2025

Walkthrough

This update introduces a customizable search callback mechanism across the codebase, allowing users to override the default web search with their own search logic in both Rust and Python APIs. The change propagates the callback through model builders, engine construction, and exposes it in the public API. Example scripts and documentation are added to demonstrate local filesystem search integration.

Changes

File(s) Change Summary
Cargo.toml, mistralrs/Cargo.toml Added walkdir crate as a workspace dependency.
docs/WEB_SEARCH.md Documented custom search callback usage in Rust and Python APIs.
examples/python/local_search.py, mistralrs/examples/local_search/main.rs Added Python and Rust examples demonstrating local filesystem search integration with chat models.
mistralrs-core/src/engine/mod.rs, mistralrs-core/src/engine/search_request.rs Added optional search_callback field to Engine; updated search logic to use callback if provided.
mistralrs-core/src/lib.rs Propagated search_callback through MistralRsBuilder, MistralRs, and related structs; re-exported callback types.
mistralrs-core/src/request.rs Added optional tool description override fields to WebSearchOptions.
mistralrs-core/src/search/mod.rs Added public SearchCallback type alias; updated tool descriptions to use custom overrides if present.
mistralrs-pyo3/src/lib.rs Enabled Python search callback integration in Runner via callback wrapping and builder update.
mistralrs-pyo3/mistralrs.pyi Added enable_search and search_callback parameters to Runner constructor with documentation.
mistralrs-server-core/src/mistralrs_for_server_builder.rs Added optional search_callback to server builder and exposed builder method for configuration.
mistralrs/src/lib.rs Publicly re-exported SearchCallback and SearchResult from core.
mistralrs/src/anymoe.rs, mistralrs/src/gguf_lora_model.rs, mistralrs/src/gguf_xlora_model.rs, mistralrs/src/lora_model.rs, mistralrs/src/speculative.rs, mistralrs/src/xlora_model.rs Modified model builder logic to conditionally attach search_callback during runner construction.
mistralrs/src/gguf.rs, mistralrs/src/text_model.rs, mistralrs/src/vision_model.rs Added search_callback field and builder method to respective model builders.
mistralrs-quant/kernels/marlin/marlin_kernel.cu Removed extraneous spaces in template argument brackets in lop3 calls.
scripts/generate_uqff_card.py Changed filename quoting style from single to double quotes around double quotes.

Sequence Diagram(s)

sequenceDiagram
    participant User
    participant ModelBuilder
    participant Engine
    participant SearchCallback
    participant SearchTool

    User->>ModelBuilder: Configure with .with_search_callback(cb)
    ModelBuilder->>Engine: Passes search_callback to Engine::new
    Engine->>SearchCallback: (if set) Invoke cb with query
    SearchCallback-->>Engine: Return search results
    Engine->>SearchTool: (else) Use default search tool
    SearchTool-->>Engine: Return search results
    Engine-->>User: Respond with search-augmented chat completion
Loading
sequenceDiagram
    participant PythonUser
    participant Runner
    participant PySearchCallback
    participant Engine

    PythonUser->>Runner: Runner(search_callback=py_cb)
    Runner->>Engine: Passes wrapped search_callback
    Engine->>PySearchCallback: Calls Python callback with query
    PySearchCallback-->>Engine: Returns list of file info dicts
    Engine-->>PythonUser: Responds with chat completion using local search results
Loading

Poem

In burrows deep, a rabbit found
A clever way to search around.
With custom code, both Rust and Py,
The search is yours—just give a try!
From local files to web afar,
Now models know just where you are.
🐇✨

✨ Finishing Touches
  • 📝 Generate Docstrings

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share
🪧 Tips

Chat

There are 3 ways to chat with CodeRabbit:

  • Review comments: Directly reply to a review comment made by CodeRabbit. Example:
    • I pushed a fix in commit <commit_id>, please review it.
    • Explain this complex logic.
    • Open a follow-up GitHub issue for this discussion.
  • Files and specific lines of code (under the "Files changed" tab): Tag @coderabbitai in a new review comment at the desired location with your query. Examples:
    • @coderabbitai explain this code block.
    • @coderabbitai modularize this function.
  • PR comments: Tag @coderabbitai in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
    • @coderabbitai gather interesting stats about this repository and render them as a table. Additionally, render a pie chart showing the language distribution in the codebase.
    • @coderabbitai read src/utils.ts and explain its main purpose.
    • @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.
    • @coderabbitai help me debug CodeRabbit configuration file.

Support

Need help? Create a ticket on our support page for assistance with any issues or questions.

Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments.

CodeRabbit Commands (Invoked using PR comments)

  • @coderabbitai pause to pause the reviews on a PR.
  • @coderabbitai resume to resume the paused reviews.
  • @coderabbitai review to trigger an incremental review. This is useful when automatic reviews are disabled for the repository.
  • @coderabbitai full review to do a full review from scratch and review all the files again.
  • @coderabbitai summary to regenerate the summary of the PR.
  • @coderabbitai generate docstrings to generate docstrings for this PR.
  • @coderabbitai generate sequence diagram to generate a sequence diagram of the changes in this PR.
  • @coderabbitai resolve resolve all the CodeRabbit review comments.
  • @coderabbitai configuration to show the current CodeRabbit configuration for the repository.
  • @coderabbitai help to get help.

Other keywords and placeholders

  • Add @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.
  • Add @coderabbitai summary to generate the high-level summary at a specific location in the PR description.
  • Add @coderabbitai anywhere in the PR title to generate the title automatically.

CodeRabbit Configuration File (.coderabbit.yaml)

  • You can programmatically configure CodeRabbit by adding a .coderabbit.yaml file to the root of your repository.
  • Please see the configuration documentation for more information.
  • If your editor has YAML language server enabled, you can add the path at the top of this file to enable auto-completion and validation: # yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json

Documentation and Community

  • Visit our Documentation for detailed information on how to use CodeRabbit.
  • Join our Discord Community to get help, request features, and share feedback.
  • Follow us on X/Twitter for updates and announcements.

Copy link

github-actions bot commented Jun 4, 2025

Code Metrics Report
===============================================================================
 Language            Files        Lines         Code     Comments       Blanks
===============================================================================
 C Header                3           62           53            0            9
 CSS                     1          473          408           14           51
 Dockerfile              1           42           23           10            9
 HTML                    1           73           61            4            8
 JavaScript              7         1248          936          174          138
 JSON                   14          123          122            0            1
 Makefile                1            6            5            0            1
 Python                 86         4045         3410          161          474
 Shell                   1           63           26           18           19
 Plain Text              3         3723            0         2413         1310
 TOML                   21          693          632           10           51
 YAML                    2           21           19            2            0
-------------------------------------------------------------------------------
 Jupyter Notebooks       3            0            0            0            0
 |- Markdown             2           77           32           31           14
 |- Python               2          205          178            1           26
 (Total)                            282          210           32           40
-------------------------------------------------------------------------------
 Markdown               59         5044            0         3844         1200
 |- BASH                10          111          105            2            4
 |- JSON                 2           42           42            0            0
 |- Python               7          121          109            0           12
 |- Rust                22          757          634            1          122
 |- TOML                 2           75           63            0           12
 (Total)                           6150          953         3847         1350
-------------------------------------------------------------------------------
 Rust                  375       132168       117614         2894        11660
 |- Markdown           174         2989           29         2649          311
 (Total)                         135157       117643         5543        11971
===============================================================================
 Total                 578       147784       123309         9544        14931
===============================================================================

Copy link

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 5

🧹 Nitpick comments (4)
mistralrs/src/speculative.rs (1)

102-102: LGTM! Consider the design implications for speculative decoding.

The integration correctly uses the target model's search callback for speculative decoding. This is logical since the target model is the primary decision maker.

Consider documenting why only the target model's search callback is used in speculative decoding, as it may not be immediately obvious to future maintainers that the draft model's callback (if any) is intentionally ignored.

examples/python/local_search.py (2)

8-9: Fix import order to follow PEP 8 conventions.

Standard library imports should come before third-party imports.

Apply this diff to fix the import order:

+import os
+
 from mistralrs import (
     Runner,
     Which,
     ChatCompletionRequest,
     Architecture,
     WebSearchOptions,
 )
-import os
🧰 Tools
🪛 Pylint (3.3.7)

[convention] 8-8: standard import "os" should be placed before first party import "mistralrs.Runner"

(C0411)


1-9: Add module and function docstrings for better documentation.

The example would benefit from documentation explaining its purpose and usage.

Add docstrings:

+"""
+Example demonstrating local filesystem search integration with mistral.rs.
+
+This script shows how to use a custom search callback to search local files
+and integrate the results with a chat completion model.
+"""
+
 from mistralrs import (
     Runner,
     Which,
     ChatCompletionRequest,
     Architecture,
     WebSearchOptions,
 )
 import os

+def local_search(query: str):
+    """
+    Search for files containing the query string in their names.
+    
+    Args:
+        query: Search string to match against filenames
+        
+    Returns:
+        List of search results sorted by title in descending order
+    """

Also applies to: 10-27

🧰 Tools
🪛 Pylint (3.3.7)

[convention] 1-1: Missing module docstring

(C0114)


[error] 1-7: Unable to import 'mistralrs'

(E0401)


[convention] 8-8: standard import "os" should be placed before first party import "mistralrs.Runner"

(C0411)

mistralrs/examples/local_search/main.rs (1)

47-52: Clarify search description to match implementation.

The search description claims "Local filesystem search" but the implementation only searches filenames, not file contents. This could mislead users about the search capabilities.

     let messages = RequestBuilder::from(messages).with_web_search_options(
         WebSearchOptions {
-            search_description: Some("Local filesystem search".to_string()),
+            search_description: Some("Local filename search".to_string()),
             ..Default::default()
         },
     );
📜 Review details

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between c53d346 and 1185b57.

📒 Files selected for processing (22)
  • Cargo.toml (1 hunks)
  • examples/python/local_search.py (1 hunks)
  • mistralrs-core/src/engine/mod.rs (4 hunks)
  • mistralrs-core/src/engine/search_request.rs (1 hunks)
  • mistralrs-core/src/lib.rs (11 hunks)
  • mistralrs-core/src/request.rs (1 hunks)
  • mistralrs-core/src/search/mod.rs (3 hunks)
  • mistralrs-pyo3/src/lib.rs (6 hunks)
  • mistralrs/Cargo.toml (1 hunks)
  • mistralrs/examples/local_search/main.rs (1 hunks)
  • mistralrs/src/anymoe.rs (1 hunks)
  • mistralrs/src/diffusion_model.rs (1 hunks)
  • mistralrs/src/gguf.rs (5 hunks)
  • mistralrs/src/gguf_lora_model.rs (1 hunks)
  • mistralrs/src/gguf_xlora_model.rs (1 hunks)
  • mistralrs/src/lib.rs (1 hunks)
  • mistralrs/src/lora_model.rs (1 hunks)
  • mistralrs/src/speculative.rs (1 hunks)
  • mistralrs/src/speech_model.rs (1 hunks)
  • mistralrs/src/text_model.rs (5 hunks)
  • mistralrs/src/vision_model.rs (5 hunks)
  • mistralrs/src/xlora_model.rs (1 hunks)
🧰 Additional context used
🧬 Code Graph Analysis (4)
mistralrs-core/src/engine/search_request.rs (1)
mistralrs-core/src/search/mod.rs (1)
  • run_search_tool (190-267)
mistralrs/src/speech_model.rs (2)
mistralrs-core/src/lib.rs (2)
  • new (205-225)
  • new (276-456)
mistralrs-pyo3/src/lib.rs (1)
  • new (527-851)
mistralrs/src/gguf.rs (4)
mistralrs/src/model.rs (1)
  • best_device (11-23)
mistralrs-core/src/lib.rs (1)
  • with_search_callback (256-259)
mistralrs/src/vision_model.rs (1)
  • with_search_callback (93-96)
mistralrs/src/text_model.rs (1)
  • with_search_callback (132-135)
mistralrs-core/src/lib.rs (3)
mistralrs/src/gguf.rs (1)
  • with_search_callback (79-82)
mistralrs/src/vision_model.rs (1)
  • with_search_callback (93-96)
mistralrs/src/text_model.rs (1)
  • with_search_callback (132-135)
🪛 GitHub Check: Clippy
mistralrs/src/lib.rs

[failure] 126-126:
module search is private

mistralrs-pyo3/src/lib.rs

[failure] 35-35:
module search is private


[failure] 35-35:
module search is private

🪛 GitHub Check: Check (ubuntu-latest, stable)
mistralrs/src/lib.rs

[failure] 126-126:
module search is private

mistralrs-pyo3/src/lib.rs

[failure] 35-35:
module search is private


[failure] 35-35:
module search is private


[failure] 98-98:
no method named downcast found for struct pyo3::Py in the current scope

🪛 GitHub Check: Docs
mistralrs/src/lib.rs

[failure] 126-126:
module search is private

🪛 Ruff (0.11.9)
examples/python/local_search.py

17-17: Use a context manager for opening files

(SIM115)

🪛 Pylint (3.3.7)
examples/python/local_search.py

[convention] 1-1: Missing module docstring

(C0114)


[error] 1-7: Unable to import 'mistralrs'

(E0401)


[convention] 10-10: Missing function or method docstring

(C0116)


[warning] 18-18: Catching too general exception Exception

(W0718)


[refactor] 17-17: Consider using 'with' for resource-allocating operations

(R1732)


[warning] 17-17: Using open without explicitly specifying an encoding

(W1514)


[convention] 8-8: standard import "os" should be placed before first party import "mistralrs.Runner"

(C0411)

⏰ Context from checks skipped due to timeout of 90000ms (3)
  • GitHub Check: Test Suite (ubuntu-latest, stable)
  • GitHub Check: Test Suite (macOS-latest, stable)
  • GitHub Check: Test Suite (windows-latest, stable)
🔇 Additional comments (34)
Cargo.toml (1)

81-81:

✅ Verification successful

LGTM! Verify the walkdir version is current and secure.

The addition of the walkdir dependency for recursive directory traversal in search examples is appropriate. Ensure version 2.5.0 is current and free from known vulnerabilities.


🏁 Script executed:

#!/bin/bash
# Check for the latest version of walkdir and any security advisories
echo "Latest walkdir version:"
curl -s https://crates.io/api/v1/crates/walkdir | jq '.crate.max_stable_version'

echo "Checking for security advisories:"
# Use cargo audit if available, or check RustSec database
if command -v cargo-audit &> /dev/null; then
    echo "walkdir" | cargo audit --json 2>/dev/null | jq '.vulnerabilities[]' || echo "No vulnerabilities found"
else
    echo "cargo-audit not available, manual check recommended"
fi

Length of output: 360


Dependency Addition Approved: walkdir v2.5.0 Verified

Confirmed that the workspace now depends on walkdir v2.5.0, which is the latest stable release on crates.io. No known security advisories were flagged (manual audit recommended if deeper review is desired). No further changes required.

mistralrs/Cargo.toml (1)

29-29: LGTM! Proper workspace dependency usage.

Correctly references the walkdir workspace dependency, ensuring version consistency across the workspace.

mistralrs/src/gguf_lora_model.rs (1)

87-87:

❓ Verification inconclusive

LGTM! Verify consistent search callback integration across all model builders.

The search callback integration follows the expected pattern of cloning from the underlying model builder. Ensure this pattern is consistently applied across all model builder implementations.


🏁 Script executed:

#!/bin/bash
# Check for consistent search_callback integration across model builders
echo "Searching for MistralRsBuilder::new calls to verify consistent search_callback usage:"
rg -A 5 -B 2 "MistralRsBuilder::new" --type rust

echo -e "\nSearching for search_callback field usage in model builders:"
ast-grep --pattern 'search_callback.clone()'

Length of output: 7399


LGTM — consistent search callback integration

All model builders that support search are now cloning self.*.search_callback. Models without search (speech_model.rs, diffusion_model.rs) correctly pass None.

  • Please verify the call in mistralrs-server-core/src/mistralrs_for_server_builder.rs:
    The MistralRsBuilder::new invocation there omits the search-callback argument. Confirm this is intentional or update it to include the callback clone.
mistralrs-core/src/request.rs (1)

113-116: LGTM! Well-documented optional fields.

The new search_description and extract_description fields are properly implemented with clear documentation and maintain backward compatibility through optional types.

mistralrs/src/gguf_xlora_model.rs (1)

99-99: LGTM! Consistent search callback forwarding.

The search callback is properly forwarded from the nested gguf_model to the MistralRsBuilder, maintaining consistency with the broader search callback integration.

mistralrs/src/anymoe.rs (1)

119-119: LGTM! Proper search callback integration.

The search callback is correctly forwarded from the base TextModelBuilder to the MistralRsBuilder, following the established pattern across model builders.

mistralrs/src/xlora_model.rs (1)

104-104: LGTM: Consistent callback parameter forwarding.

The change correctly forwards the cloned search_callback from the embedded text_model to the MistralRsBuilder::new constructor, maintaining consistency with the updated API across model builders.

mistralrs/src/diffusion_model.rs (1)

98-98: LGTM: Appropriate search callback parameter for diffusion models.

Correctly adds None as the search callback parameter to match the updated MistralRsBuilder::new signature. This is appropriate since diffusion models typically don't require search functionality.

mistralrs/src/lora_model.rs (1)

92-92: LGTM: Consistent callback forwarding pattern.

The change correctly forwards the cloned search_callback from the embedded text_model, maintaining consistency with the search callback integration pattern across model builders.

mistralrs-core/src/engine/search_request.rs (1)

75-80: LGTM: Core search callback functionality implemented correctly.

The conditional logic properly enables custom search callbacks while maintaining backward compatibility. The implementation:

  • Uses the custom callback cb(&tool_call_params) when available
  • Falls back to the default search::run_search_tool(&tool_call_params) otherwise
  • Maintains consistent error handling with .unwrap() matching the existing pattern

This change enables the extensible search functionality described in the PR objectives.

mistralrs/src/speech_model.rs (1)

114-114: LGTM: Correctly updated MistralRsBuilder::new call to match new signature.

The addition of the second None parameter properly aligns with the updated constructor that now accepts optional search embedding model and search callback parameters.

mistralrs-core/src/engine/mod.rs (4)

4-4: LGTM: Properly added search module import.

The import is necessary for the SearchCallback type used in the Engine struct.


76-76: LGTM: Correctly added search_callback field to Engine struct.

The field type Option<Arc<search::SearchCallback>> is appropriate for optional, thread-safe callback functionality.


110-110: LGTM: Properly added search_callback parameter to constructor.

The parameter follows the existing pattern and maintains consistency with other optional parameters.


133-133: LGTM: Correct field initialization.

Simple assignment of the search_callback parameter to the struct field is appropriate.

mistralrs/src/text_model.rs (6)

3-3: LGTM: Necessary import for SearchCallback type.

The import is required for the search_callback field and builder method.


8-8: LGTM: Added Arc import for thread-safe callback handling.

The Arc import is necessary for the Arc<SearchCallback> type used in the search_callback field.


30-30: LGTM: Properly added search_callback field to TextModelBuilder.

The field type Option<Arc<SearchCallback>> is consistent with the pattern used in other model builders.


120-120: LGTM: Correct default initialization.

Initializing search_callback to None is appropriate as it makes the feature optional by default.


131-135: LGTM: Well-implemented builder method.

The with_search_callback method follows the standard builder pattern and includes appropriate documentation.


375-375: LGTM: Correctly passing search_callback to MistralRsBuilder.

Using clone() on the Arc is the correct approach for sharing the callback across components.

mistralrs/src/vision_model.rs (6)

3-3: LGTM: Necessary import for SearchCallback type.

The import is required for the search_callback field and builder method.


8-8: LGTM: Added Arc import for thread-safe callback handling.

The Arc import is necessary for the Arc<SearchCallback> type used in the search_callback field.


31-31: LGTM: Properly added search_callback field to VisionModelBuilder.

The field type and placement are consistent with the TextModelBuilder implementation.


81-81: LGTM: Correct default initialization.

Initializing search_callback to None maintains consistency with other model builders.


92-96: LGTM: Consistent builder method implementation.

The with_search_callback method matches the implementation in TextModelBuilder and follows established patterns.


318-318: LGTM: Correctly passing search_callback to MistralRsBuilder.

Using clone() on the Arc maintains consistency with the text model implementation.

mistralrs-core/src/search/mod.rs (2)

17-19: Well-designed callback interface!

The SearchCallback type alias is properly defined with thread safety constraints (Send + Sync) and a clear contract. The function signature requiring sorted results in decreasing relevance order is well-documented.


147-154: Excellent backward-compatible customization feature.

The implementation properly allows custom descriptions while maintaining backward compatibility through fallbacks to the default constants. The approach is consistent for both search and extract tools.

Also applies to: 173-180

mistralrs/src/gguf.rs (1)

3-3: Consistent and well-integrated search callback support.

The implementation correctly follows the established builder pattern, maintains consistency with other model builders, and properly integrates the search callback throughout the build process. The field initialization and method signature match the patterns used in other model types.

Also applies to: 7-7, 22-22, 67-67, 78-82, 258-258

examples/python/local_search.py (1)

10-27: Functional implementation demonstrates the search callback feature well.

The local search function correctly implements the expected interface and provides a useful example of filesystem-based search integration.

🧰 Tools
🪛 Ruff (0.11.9)

17-17: Use a context manager for opening files

(SIM115)

🪛 Pylint (3.3.7)

[convention] 10-10: Missing function or method docstring

(C0116)


[warning] 18-18: Catching too general exception Exception

(W0718)


[refactor] 17-17: Consider using 'with' for resource-allocating operations

(R1732)


[warning] 17-17: Using open without explicitly specifying an encoding

(W1514)

mistralrs-core/src/lib.rs (1)

115-115: Comprehensive and well-integrated search callback support.

The implementation properly integrates the search callback throughout the entire MistralRs system:

  • Public re-export for external usage
  • Consistent propagation through builder pattern
  • Proper handling in engine creation and reboot scenarios
  • Thread-safe Arc wrapping maintained throughout

The integration is thorough and maintains consistency with existing patterns.

Also applies to: 163-163, 201-201, 210-210, 223-223, 255-259, 288-288, 312-312, 344-344, 366-366, 491-491

mistralrs/examples/local_search/main.rs (1)

33-41: LGTM! Good example structure.

The model setup demonstrates proper usage of the new search callback functionality with clear configuration.

mistralrs-pyo3/src/lib.rs (1)

525-525: Verify parameter integration consistency.

The new search_callback parameter is properly integrated into the function signatures and passed through to the builder. The implementation looks consistent.

Also applies to: 550-550, 844-844, 905-905

@@ -123,6 +123,7 @@
pub use model::{best_device, Model};
pub use speculative::TextSpeculativeBuilder;
pub use speech_model::SpeechModelBuilder;
pub use mistralrs_core::{SearchCallback, search::SearchResult};
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Verification agent

🧩 Analysis chain

Fix compilation error: cannot access private module.

The re-export of search::SearchResult will fail to compile because the search module is private in mistralrs_core.

Apply this fix to resolve the compilation error:

-pub use mistralrs_core::{SearchCallback, search::SearchResult};
+pub use mistralrs_core::SearchCallback;

Run the following script to verify the module visibility:


🏁 Script executed:

#!/bin/bash
# Description: Check the visibility of the search module in mistralrs_core
# Expected: Confirm that search module is private and SearchResult is not accessible

# Look for the search module declaration in mistralrs_core
rg -A 5 "mod search" mistralrs-core/src/

# Check if SearchResult is publicly exported in mistralrs_core
rg "pub.*SearchResult" mistralrs-core/src/

Length of output: 1002


Expose search module in mistralrs_core to fix private-module errors

The search module in mistralrs_core is declared as mod search;, so its contents (including SearchResult and SearchCallback) aren’t publicly accessible. You can resolve this by making the module public and updating your re-export paths.

• mistralrs-core/src/lib.rs

-mod search;
+pub mod search;

• mistralrs/src/lib.rs (line 126)

-pub use mistralrs_core::{SearchCallback, search::SearchResult};
+pub use mistralrs_core::search::{SearchCallback, SearchResult};
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
pub use mistralrs_core::{SearchCallback, search::SearchResult};
- pub use mistralrs_core::{SearchCallback, search::SearchResult};
+ pub use mistralrs_core::search::{SearchCallback, SearchResult};
🧰 Tools
🪛 GitHub Check: Clippy

[failure] 126-126:
module search is private

🪛 GitHub Check: Check (ubuntu-latest, stable)

[failure] 126-126:
module search is private

🪛 GitHub Check: Docs

[failure] 126-126:
module search is private

🤖 Prompt for AI Agents
In mistralrs-core/src/lib.rs around line 126, the search module is private,
causing a compilation error when re-exporting SearchResult from mistralrs_core
in mistralrs/src/lib.rs. To fix this, make the search module public by changing
its declaration from "mod search;" to "pub mod search;" in
mistralrs-core/src/lib.rs. Then update the re-export in mistralrs/src/lib.rs to
use the correct public path for SearchResult and SearchCallback.

Comment on lines +17 to +19
content = open(path).read()
except Exception:
content = ""
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion

Improve file handling with proper resource management.

The current file opening approach lacks proper resource management and error handling specificity.

Apply this diff to use a context manager and improve error handling:

-                try:
-                    content = open(path).read()
-                except Exception:
-                    content = ""
+                try:
+                    with open(path, 'r', encoding='utf-8') as f:
+                        content = f.read()
+                except (OSError, UnicodeDecodeError):
+                    content = ""
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
content = open(path).read()
except Exception:
content = ""
try:
with open(path, 'r', encoding='utf-8') as f:
content = f.read()
except (OSError, UnicodeDecodeError):
content = ""
🧰 Tools
🪛 Ruff (0.11.9)

17-17: Use a context manager for opening files

(SIM115)

🪛 Pylint (3.3.7)

[warning] 18-18: Catching too general exception Exception

(W0718)


[refactor] 17-17: Consider using 'with' for resource-allocating operations

(R1732)


[warning] 17-17: Using open without explicitly specifying an encoding

(W1514)

🤖 Prompt for AI Agents
In examples/python/local_search.py around lines 17 to 19, the file is opened
without using a context manager, which can lead to resource leaks. Replace the
open call with a with statement to ensure the file is properly closed after
reading. Also, catch specific exceptions related to file operations instead of a
general Exception to improve error handling clarity.

Comment on lines +10 to +31
fn local_search(query: &str) -> Result<Vec<SearchResult>> {
let mut results = Vec::new();
for entry in WalkDir::new(".") {
let entry = entry?;
if entry.file_type().is_file() {
let name = entry.file_name().to_string_lossy();
if name.contains(query) {
let path = entry.path().display().to_string();
let content = fs::read_to_string(entry.path()).unwrap_or_default();
results.push(SearchResult {
title: name.into_owned(),
description: path.clone(),
url: path,
content,
});
}
}
}
results.sort_by_key(|r| r.title.clone());
results.reverse();
Ok(results)
}
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion

Improve search logic and error handling.

The current implementation has several areas for improvement:

  1. Misleading search scope: The function only searches filenames but reads full file contents, which may confuse users about what's being searched.
  2. Questionable sorting: Sorting by title then reversing gives descending alphabetical order, which may not be the most useful for search relevance.
  3. Silent error handling: Using unwrap_or_default() for file reading silently ignores permission errors or other IO issues.

Consider this improved implementation:

 fn local_search(query: &str) -> Result<Vec<SearchResult>> {
     let mut results = Vec::new();
     for entry in WalkDir::new(".") {
         let entry = entry?;
         if entry.file_type().is_file() {
             let name = entry.file_name().to_string_lossy();
             if name.contains(query) {
                 let path = entry.path().display().to_string();
-                let content = fs::read_to_string(entry.path()).unwrap_or_default();
+                let content = fs::read_to_string(entry.path())
+                    .unwrap_or_else(|_| format!("Could not read file: {}", path));
                 results.push(SearchResult {
                     title: name.into_owned(),
                     description: path.clone(),
                     url: path,
                     content,
                 });
             }
         }
     }
-    results.sort_by_key(|r| r.title.clone());
-    results.reverse();
+    // Sort by relevance (exact matches first, then alphabetical)
+    results.sort_by(|a, b| {
+        let a_exact = a.title.eq_ignore_ascii_case(query);
+        let b_exact = b.title.eq_ignore_ascii_case(query);
+        match (a_exact, b_exact) {
+            (true, false) => std::cmp::Ordering::Less,
+            (false, true) => std::cmp::Ordering::Greater,
+            _ => a.title.cmp(&b.title),
+        }
+    });
     Ok(results)
 }
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
fn local_search(query: &str) -> Result<Vec<SearchResult>> {
let mut results = Vec::new();
for entry in WalkDir::new(".") {
let entry = entry?;
if entry.file_type().is_file() {
let name = entry.file_name().to_string_lossy();
if name.contains(query) {
let path = entry.path().display().to_string();
let content = fs::read_to_string(entry.path()).unwrap_or_default();
results.push(SearchResult {
title: name.into_owned(),
description: path.clone(),
url: path,
content,
});
}
}
}
results.sort_by_key(|r| r.title.clone());
results.reverse();
Ok(results)
}
fn local_search(query: &str) -> Result<Vec<SearchResult>> {
let mut results = Vec::new();
for entry in WalkDir::new(".") {
let entry = entry?;
if entry.file_type().is_file() {
let name = entry.file_name().to_string_lossy();
if name.contains(query) {
let path = entry.path().display().to_string();
let content = fs::read_to_string(entry.path())
.unwrap_or_else(|_| format!("Could not read file: {}", path));
results.push(SearchResult {
title: name.into_owned(),
description: path.clone(),
url: path,
content,
});
}
}
}
// Sort by relevance (exact matches first, then alphabetical)
results.sort_by(|a, b| {
let a_exact = a.title.eq_ignore_ascii_case(query);
let b_exact = b.title.eq_ignore_ascii_case(query);
match (a_exact, b_exact) {
(true, false) => std::cmp::Ordering::Less,
(false, true) => std::cmp::Ordering::Greater,
_ => a.title.cmp(&b.title),
}
});
Ok(results)
}
🤖 Prompt for AI Agents
In mistralrs/examples/local_search/main.rs around lines 10 to 31, improve the
local_search function by expanding the search to include file contents, not just
filenames, to better match user expectations. Replace the current sorting by
title with a relevance-based sorting, such as by the number of query occurrences
in content or filename. Change error handling for file reading to properly
propagate or log errors instead of silently ignoring them with
unwrap_or_default, ensuring IO issues are visible and handled appropriately.

@EricLBuehler EricLBuehler merged commit 39673eb into master Jun 4, 2025
12 of 13 checks passed
Copy link

@coderabbitai coderabbitai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 0

♻️ Duplicate comments (1)
examples/python/local_search.py (1)

17-20: 🛠️ Refactor suggestion

Improve file handling with proper resource management.

The file handling approach has the same issues identified in previous reviews.

The current implementation lacks proper resource management and error handling specificity. Apply the previously suggested fix:

-                try:
-                    content = open(path).read()
-                except Exception:
-                    content = ""
+                try:
+                    with open(path, 'r', encoding='utf-8') as f:
+                        content = f.read()
+                except (OSError, UnicodeDecodeError):
+                    content = ""
🧰 Tools
🪛 Ruff (0.11.9)

18-18: Use a context manager for opening files

(SIM115)

🪛 Pylint (3.3.7)

[warning] 19-19: Catching too general exception Exception

(W0718)


[refactor] 18-18: Consider using 'with' for resource-allocating operations

(R1732)


[warning] 18-18: Using open without explicitly specifying an encoding

(W1514)

🧹 Nitpick comments (2)
examples/python/local_search.py (2)

1-8: Address import organization and add module documentation.

The import organization doesn't follow PEP8 conventions, and the module lacks documentation.

Apply these improvements:

+"""
+Example demonstrating local filesystem search integration with mistralrs.
+
+This script shows how to use a custom search callback to search local files
+and integrate the results with a chat completion model.
+"""
+import os
+
 from mistralrs import (
     Runner,
     Which,
     ChatCompletionRequest,
     Architecture,
     WebSearchOptions,
 )
-import os
🧰 Tools
🪛 Pylint (3.3.7)

[convention] 1-1: Missing module docstring

(C0114)


[error] 1-7: Unable to import 'mistralrs'

(E0401)


[convention] 8-8: standard import "os" should be placed before first party import "mistralrs.Runner"

(C0411)


11-30: Add function documentation and improve implementation.

The function lacks documentation and could benefit from better structure.

Apply this improvement:

 def local_search(query: str):
+    """
+    Search for files containing the query string in their filename.
+    
+    Args:
+        query: String to search for in filenames
+        
+    Returns:
+        List of dictionaries containing file information
+    """
     results = []
🧰 Tools
🪛 Ruff (0.11.9)

18-18: Use a context manager for opening files

(SIM115)

🪛 Pylint (3.3.7)

[convention] 11-11: Missing function or method docstring

(C0116)


[warning] 19-19: Catching too general exception Exception

(W0718)


[refactor] 18-18: Consider using 'with' for resource-allocating operations

(R1732)


[warning] 18-18: Using open without explicitly specifying an encoding

(W1514)

📜 Review details

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 231d9f8 and d92d2b5.

📒 Files selected for processing (5)
  • docs/WEB_SEARCH.md (2 hunks)
  • examples/python/local_search.py (1 hunks)
  • mistralrs-pyo3/mistralrs.pyi (3 hunks)
  • mistralrs-quant/kernels/marlin/marlin_kernel.cu (4 hunks)
  • scripts/generate_uqff_card.py (1 hunks)
✅ Files skipped from review due to trivial changes (1)
  • mistralrs-pyo3/mistralrs.pyi
🚧 Files skipped from review as they are similar to previous changes (1)
  • docs/WEB_SEARCH.md
🧰 Additional context used
🪛 Ruff (0.11.9)
examples/python/local_search.py

18-18: Use a context manager for opening files

(SIM115)

🪛 Pylint (3.3.7)
examples/python/local_search.py

[convention] 1-1: Missing module docstring

(C0114)


[error] 1-7: Unable to import 'mistralrs'

(E0401)


[convention] 11-11: Missing function or method docstring

(C0116)


[warning] 19-19: Catching too general exception Exception

(W0718)


[refactor] 18-18: Consider using 'with' for resource-allocating operations

(R1732)


[warning] 18-18: Using open without explicitly specifying an encoding

(W1514)


[convention] 8-8: standard import "os" should be placed before first party import "mistralrs.Runner"

(C0411)

scripts/generate_uqff_card.py

[convention] 57-57: Constant name "file" doesn't conform to UPPER_CASE naming style

(C0103)

⏰ Context from checks skipped due to timeout of 90000ms (4)
  • GitHub Check: Check (windows-latest, stable)
  • GitHub Check: Test Suite (macOS-latest, stable)
  • GitHub Check: Test Suite (windows-latest, stable)
  • GitHub Check: Test Suite (ubuntu-latest, stable)
🔇 Additional comments (7)
mistralrs-quant/kernels/marlin/marlin_kernel.cu (4)

89-90: LGTM! Template syntax formatting improvement.

The removal of spaces around template argument brackets follows standard C++ conventions and improves code consistency.


113-115: LGTM! Consistent template formatting improvement.

The spacing corrections maintain consistency with the previous template specialization changes.


138-139: LGTM! Formatting consistency maintained.

The template bracket corrections continue the consistent pattern of improvements across all dequant specializations.


161-163: LGTM! Formatting improvements completed consistently.

This finalizes the systematic template syntax corrections across all dequant template specializations, ensuring uniform code style throughout the file.

scripts/generate_uqff_card.py (1)

57-57: LGTM! Stylistic improvement in string formatting.

The change from double quotes around the f-string to single quotes is a minor stylistic improvement that maintains the same functionality.

🧰 Tools
🪛 Pylint (3.3.7)

[convention] 57-57: Constant name "file" doesn't conform to UPPER_CASE naming style

(C0103)

examples/python/local_search.py (2)

33-40: LGTM! Proper integration of search callback feature.

The Runner configuration correctly demonstrates the new search callback functionality with appropriate parameters.


42-52: LGTM! Well-structured example usage.

The chat completion request effectively demonstrates the integration of local search with the model, using appropriate web search options to describe the local filesystem search capability.

@EricLBuehler
Copy link
Owner Author

@EricLBuehler EricLBuehler deleted the codex/implement-hook-for-custom-search-result-handling branch June 4, 2025 11:42
@coderabbitai coderabbitai bot mentioned this pull request Jun 4, 2025
Jeadie added a commit to spiceai/mistral.rs that referenced this pull request Jul 14, 2025
* Fix handling of Metal fused attn head dims (EricLBuehler#1234)

* Fix handling of metal attn head dims

* Fix handling of gemma3 1b when images

* Tweak default for paged attn builder

* Support paged attn for vision model rust api (EricLBuehler#1235)

* [Breaking] Support setting HF cache path (EricLBuehler#1237)

* Add it internally

* Add the apis

* Support tool calling for DeepSeek models (EricLBuehler#1239)

* Support tool calling for deepseek models

* Format

* Fix deepseek

* Server image processing refactor and fixes (EricLBuehler#1244)

* Fix strict gemma3 case

* Accept multiple images in the content array

* Fix multiple images in one array ct

* Add it to the python api

* Typos

* Optimized CUDA RoPE kernels (EricLBuehler#1247)

* Add the kernels

* It works

* Works

* Buulds

* Typo fix (add_speial_tokens to add_special_tokens) (EricLBuehler#1246)

* Fix typo

* Update mistralrs.pyi

* Fixes for UQFF + distributed layers (EricLBuehler#1250)

* Fixes for uqff + distributed layers

* Typo

* Automatic agentic search integration (`web_search_options`) (EricLBuehler#1243)

* Add the tool

* Actually search

* Clippy

* Sort of works

* Remove some debuggers

* tweak

* Add some rules

* Works great

* Tweak 'system' prompt

* Update mistralrs-core/src/search/mod.rs

Co-authored-by: Copilot <[email protected]>

* Typo

* Add it to all the apis

* Add bert model for similarity reranking

* Typos

* Early detection of tools

* Alias max_tokens -> max_completion_tokens too

* Customizable bert model

* Flip the enabler around

* Add docs

* Update readme

* Typo

---------

Co-authored-by: Copilot <[email protected]>

* Format kernels (EricLBuehler#1251)

* Update readme

* Update readme

* Remove test

* Add quantize guards for uqff deserialize (EricLBuehler#1252)

* Refactor cuBLASlt-related code (EricLBuehler#1253)

* Centralize cublaslt into mistralrs-quant

* Use cublaslt in unquant layer

* Use beautiful trait constants for simpler code

* Move tests

* Dispatch to unquant for cublaslt

* Dispatch to unquant for cublaslt

* Fix feature

* Add convert_to_gptq script

* Update deps, bump pyo3 version (EricLBuehler#1259)

* Faster cuda FP8 performance (EricLBuehler#1257)

* Avoid fp8 sync

* Fix dtype

* Rust 1.86 clippy (EricLBuehler#1260)

* Rust 1.86 clippy

* Clippy

* Refactor engine arch (EricLBuehler#1262)

* Refactor engine add_request

* Don't recompile regex

* Clippy

* Revamped LoRA support - removing the Ordering system! (EricLBuehler#1263)

* Play with varbuilder lifetimes

* Merge lora weights

* Clippy

* Lora works

* Support multiple loras

* Cleanup, remove adapter activation

* Complete merge

* Fast Metal-specific quantization method: AFQ (EricLBuehler#1264)

* Add mlx quantized kernels

* Add mlx quantized kernels

* Kernel launcher

* Add AFQ isq quant and dequant

* Some quantmethod things

* Begin to implement the qmm caller

* Clippy

* Much faster

* Cache kernels

* Docs

* Clippy

* Add it to uqff

* Support prequantized models from MLX (EricLBuehler#1265)

* Refactor quantizedconfig

* Support AFQ prequantized

* Update docs

* Update docs

* Automatic ISQ to select fastest & most accurate method (EricLBuehler#1266)

* Automatic isq

* typo

* Doc

* Improved usage metrics (EricLBuehler#1267)

* Fix cuda

* Bump tokio from 1.44.1 to 1.44.2 (EricLBuehler#1270)

Bumps [tokio](https://github.com/tokio-rs/tokio) from 1.44.1 to 1.44.2.
- [Release notes](https://github.com/tokio-rs/tokio/releases)
- [Commits](tokio-rs/tokio@tokio-1.44.1...tokio-1.44.2)

---
updated-dependencies:
- dependency-name: tokio
  dependency-version: 1.44.2
  dependency-type: direct:production
...

Signed-off-by: dependabot[bot] <[email protected]>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>

* Gather MM ops in mistralrs-quant (EricLBuehler#1272)

* Update the caller

* Wire things up

* Broadcase for afq gathermm

* Broadcase for afq gathermm

* Clippy

* Improve performance of deepseek models

* Typo fix

* BincountOp not used

* Implement Llama 4! (EricLBuehler#1268)

* Implement Llama 4

* Implement the main changes for the text model

* Make chunked mask

* Wire things up

* Add some EP

* Initial sketch of inputs processor

* Runs

* Progress

* all reduce moes

* It works!

* Some cleanup

* Faster moe block

* Add device map

* Make chunked matrix

* Fully working now!

* Reactivate cublaslt

* Fix shared mlp cublaslt

* Refactor to packed experts

* Complete merge

* It is a normal model now

* Fixes

* Set device for moe

* ISQ fixes

* Much faster sort kernel

* Faster loading!

* Faster loading!

* Fp8 cpu copy ops in candle backend

* Add the vision model

* Add mmproj layer

* Actually merge the inputs

* Sketch most of the image processor

* Add the rest of the image processor

* Implement the whole processor

* Add the loader

* Some fixes

* A batch of fixes

* Some fixes

* tmp

* Actually support isq

* Ok it works a bit

* Fix norm device

* It works

* A bit cleaner

* Support residul tensors

* Remove text loader

* Implement the device mapping system

* Fix auto device map

* Add examples

* Add model card

* Typo

* Remove superflous logging

* Fixes for Llama 4 UQFF loading (EricLBuehler#1275)

* Support sharding for UQFF (EricLBuehler#1276)

* Serialize sharded uqff files

* Loading

* Fix base64

* Fix bug for group-topk (group_limited_greedy) in deepseek models (EricLBuehler#1278)

* Support the DeepCoder model (EricLBuehler#1279)

* Add faq for metal not found

* Improved PagedAttn scheduling accuracy (EricLBuehler#1282)

* Scheduler ops by reference

* Ensure scheduler gets correct prompts

* Fix cuda build for copy_blocks

* Fixes for scheduling image seqs with pagedattn (EricLBuehler#1283)

* update to llguidance 0.7.16 (EricLBuehler#1284)

* update llguidance to 0.7.16 from crates.io; use ParserFactory

* add lark_llg.py example

* use new llguidance::Matcher APIs

* rework spec-decoding with llg

* more work on spec sampling

* check for parser stop

* fix clippy

* remove unneeded rollback

* update build_llg_factory to return Result

* Update dependencies (EricLBuehler#1286)

* Much faster image inputs processing (EricLBuehler#1289)

* Add more SDPA head dims for much faster SigLIP (EricLBuehler#1290)

* More sdpa head dims, faster vision models

* Move nonzero to above for faster metal synch

* Doc

* Update valid head dims

* Show throughput in interactive mode (EricLBuehler#1291)

* Update interactive mode throughput stats

* Accurate prompt t/s

* Accurate prompt t/s for usage

* Unify bitwise operations (EricLBuehler#1288)

* Unify bitwise ops

* Tests pass

* Fix cuda build

* Clippy

* Multimodal prefix caching support! (EricLBuehler#1298)

* Initial progress

* Support vision prefix caching

* Update docs

* Add multimodal data abstraction

* Interactive mode improvements (EricLBuehler#1299)

* More ergonomic image url parsing

* Add option to clear

* Add the Qwen 3 and Qwen 3 MoE models! (EricLBuehler#1285)

* Add qwen3 model

* Add enable_thinking

* Add initial qwen3 moe

* Add the moe model

* Format

* Fix order of norm

* Fix expert shapes

* Fix reverse

* Fix norm device for isq

* Fix nonzero when no nonzero

* Moe model runs

* Working qwen3 moe

* Add metal fp8 blockwise dequant

* Clean

* Typo

* Enable tool calling

* Streamlined ux

* Add some examples

* Add docs

* Fix dead link

* Remove interactive mode max_len

* Update QWEN3.md

* Hotfix for vision mode clear

* Revamped and streaming web search support (EricLBuehler#1301)

* Streaming web search

* Refactor a bit

* More refactoring

* Add some logging, parallelize some things

* Allow url

* Suppress warning, allow multi-turn searching

* Batch compute_similarities

* Cap content len

* Typos

* Doc

* Handle vision messages or different tool call prefixes (EricLBuehler#1302)

* Fix cuda

* Tune web search budget

* Simplify prefix cacher (EricLBuehler#1305)

* Use rustyline to handle non-ascii in interactive mode (EricLBuehler#1306)

The io::stdin().read_line() cannot handle non-ascii input, which caused
crash when use backspace to delete non-ascii characters.

Introduce rustyline to the interactive mode to solve the problem. Plus
it can bring more editing features in the future.

Close EricLBuehler#1140

* Add more tools for automatic search (EricLBuehler#1307)

* Add interactive mode history

* Add a website extraction tool

* Pass toks by reference

* Optimize prompt chunking

* Fix CPU hogging in interactive mode (EricLBuehler#1309)

The log enabler should be checked after the sleep instead of a busy
loop checking.

Since the interactive mode always disables the token speed logger, 100%
CPU was taken by this loop always.

* Add Metal precompilation support  (EricLBuehler#1311)

* Add metal precompilation for paged attn

* Add for mistralrs-quant

* Better constructor

* Dont always build

* Fix name for paged attn rebuild

* Reduce thrashing of Metal autorelease (EricLBuehler#1313)

* Reduce calls to autorelease

* Optimize clone_in_cache

* Refactor float8

* make `AdapterPaths` and `LoraAdapterPaths` public (EricLBuehler#1314)

Make `AdapterPaths` and `LoraAdapterPaths` public so `LocalModelPaths`
can be constructed outside of `mistralrs-core`.

* Refactor KV cache manager (EricLBuehler#1315)

* Refactor kv cache

* Refactor caches

* Fix some overflows

* Add `Audio` and `Speech` model categories (EricLBuehler#1317)

* add `Audio` to `ModelCategory`

* add `Speech` to `ModelCategory`

* fix to go back to PartialEq having an exhaustiveness check

* Remove has_conv2d from vision model API (EricLBuehler#1318)

* Unified/automatic flash attention enabler (EricLBuehler#1319)

* Remove from sdpa params

* Fix errors

* No warnings

* Log

* Clippy

* Fix cublaslt 4d mask (EricLBuehler#1320)

* Fix cublaslt 4d mask

* Clippy

* Keep caches on gpu

* Qwen VL models fixes (EricLBuehler#1322)

* Add some defaults

* Fix

* Fix one thing

* 2.5 vl works

* Use caching again

* Fix v2

* Move index inside loop

* Offset in ropeidx

* Default support for vision prefix caching is false

* Fixes for all vision models (EricLBuehler#1323)

* Fix phi input processor?

* Fix phi input processor

* Handle no_prefix_cache from pipeline

* Phi models confirmed 👍

* Fixed for phi inputs processors

* Fixed for phi4

* Llama 3 confirmed 😀

* Mistral 3 confirmed 😃

* Idefics 2/3 fixes

* Some fixes

* Remove unsafety

* Improved+faster LRU prefix cacher (EricLBuehler#1321)

* Show TTFT

* Use LRU prefix cacher

* Faster prefix cacher

* Inplace ISQ support and default to mmap (EricLBuehler#1277)

* Initial impl of immediate isq

* Immediate isq -> !loading_isq

* Varbuiler utils always using mmap!

* Log

* Add for packed experts

* Afq without copy

* Clarify

* Clippy

* Apple immediate isq

* Better logic for loading_isq

* Support showing ttft

* Rename

* Shared quantize guard

* Parallel progress bar

* Parallel loading for progress bars

* Actual ISQ support

* Conditional parallelism for NiceProgressBar

* Use conditional iterator

* Warn once

* Predicate for applying immediate isq

* Allow parallel

* Remove debug print

* Remove debug print

* Remove debug print

* Fix typos (EricLBuehler#1329)

* Fix Idefics 3 arch chat templating (EricLBuehler#1330)

* Update inputs merger

* Fix

* Better warning

* Better warning

* Better warning

* Nonzero ahead of time

* No f32

* Clippy

* Optimize get_logprobs

* Fix packed experts

* Update masking

* Use Sdpa in idefics3

* QuantMethod in idefics3 vision

* Remove a .contiguous

* Remove two space from PR comment (EricLBuehler#1331)

* Add automatic vision loader type (EricLBuehler#1332)

* Add automatic vision loader

* Remove references to --arch

* Update examples

* Add the Dia 1.6b TTS model! (EricLBuehler#1304)

* Add loading

* Add rope, mlp, most of attn

* Add encoder + encoder layer, decoder layer forwards

* Add decoder forwards

* Add prepare_audio_prompt

* prepare_generation mostly done

* Add a proper dia kvcache

* Add most of decoder_step

* Add the sampler

* Add the generation loop

* Wire things up

* Add speech pipeline

* Fixes

* Loads

* Some fixes

* f32

* Some progress

* Ok it runs upto dac decoding

* Add dac part loading

* Loads and runs at least

* Remove encodec

* Debugging

* Debugging

* Huh

* Complete merge

* Interactive

* Confirmed dac works at least

* Looks like encoder works

* Much progress

* Hmm

* Sampling

* Almost there

* Sampler

* Sampler

* Bf16 support

* Response

* Use it in interactive mode

* Fix oneshot

* Add openai api

* Add openai api

* Refactor loading

* Use naive sdpa for inplace

* Factor out

* Clippy

* Clippy

* Config

* Refactor config

* Metal clippy

* Fix t/s

* ISQ support

* Some fixes, nits

* Fix cuda

* Clippy

* Inhibit cublaslt for cuda

* Add server example

* Add python example

* Add rust api

* Add docs

* Update config.toml

* Fix .pyi

* Update readme

* config.toml tweak

* config.toml tweak

* config.toml tweak

* config.toml tweak

* config.toml tweak

* config.toml tweak

* config.toml tweak

* config.toml tweak

* config.toml tweak

* update `llguidance` to `0.7.20` (EricLBuehler#1334)

Update `llguidance` from `0.7.16` to `0.7.20` so that it has guidance-ai/llguidance#172 which is a fix for building on GCC 15.

* Add model category <> messages check (EricLBuehler#1335)

* Verify model category matches the messages

* Add vision chat

* Fixes

* Add element-wise normalization check (EricLBuehler#1340)

* Fix streaming example print statement (EricLBuehler#1339)

* Fix normalization formula in comment (EricLBuehler#1338)

* Fix image_to_pixels to handle non-RGB images (EricLBuehler#1337)

* Fix typo in expect messages (EricLBuehler#1342)

* Don't use mmap on cuda (EricLBuehler#1336)

* No mmap on cuda

* Simplify streaming tool call logic

* Remove debug

* Support AWQ format models (EricLBuehler#1350)

* Support AWQ format models

* Clippy fix

* Fix uqff dummy layer ISQ application (EricLBuehler#1351)

* Disable immediate isq if write_uqff (EricLBuehler#1352)

* Fixes for UQFF loading on CUDA, ISQ pack factor (EricLBuehler#1354)

* Fix logic for uqff on cuda

* Updated pack_factor

* Refactor Option references for model paths (EricLBuehler#1347)

* refactor: use Option refs in model path helpers

* Format

* Add a script for server benchmarking (EricLBuehler#1355)

* Serde alias

* Fix

* Update for tie_word_embeddings

* Print running/waiting

* 30 users

* Update num_users

* Update dummy paged attn

* Optimized Metal qmv_fast path (EricLBuehler#1356)

* Compile with lto

* Tweak profiles

* New, fast sampler for Metal! (EricLBuehler#1327)

* Show TTFT

* Use LRU prefix cacher

* Faster prefix cacher

* A bit of gpu sampling

* Minp but cpu for now

* Metal fast cumsum impl

* Sampling with fast topp kernel

* Hmm not perfect

* Add metal sort kernels

* Tmp

* Add single block sort

* Add most of multi block sort, just need copy op

* Add copy kernels

* Expose kernels

* Add a test

* Ok it works

* Structure things

* Add caching

* Rename

* Cpu is default

* CUDA case

* Topk

* Refactor Option references for model paths (EricLBuehler#1347)

* refactor: use Option refs in model path helpers

* Format

* Add a script for server benchmarking (EricLBuehler#1355)

* Serde alias

* Fix

* Update for tie_word_embeddings

* Print running/waiting

* 30 users

* Update num_users

* Update dummy paged attn

* Optimized Metal qmv_fast path (EricLBuehler#1356)

* Compile with lto

* Tweak profiles

* Fix topk

* Penalties

* Add logits processor, clippy fixes

* Fix chat port

* Remove warning

* Fix chat port

* Fix metal parallel sampling (EricLBuehler#1357)

* Cpu if parallel for now

* Tweak bench script

* Add immediate isq predicates for qwen3 (EricLBuehler#1358)

* Add immediate isq predicates for qwen3

* Fix parsing of "parse_isq_value" depedent of device

* Typo

* Fix gemma3 logging

* Regressions fixes (EricLBuehler#1359)

* Fix regression for mmap

* Revert EricLBuehler#1321

* Refactored matching_cache impl

* Clippy

* Revamped and smaller readme (EricLBuehler#1360)

* Expandable detail sections

* Refactor using derivative model

* Tweak quick examples

* Update llama

* Update llama

* Supported accelerators is a table

* Update installation guides

* Tweak apis

* Remove --port in quick examples

* Add demo gif

* Add gif in readme

* Update demo gif

* Update demo gif

* Update demo gif

* Add gif in readme

* Add gif in readme

* Add a web chat app! (EricLBuehler#1362)

* Initial

* Markdown

* Copy code

* Add model loading sidebar

* Support vision models

* Tweak isq

* Links go to another page

* Clear when switch model

* Fix html tags

* Add image support!

* More then one images

* Fix

* Improved textarea

* Tab for switching between vision and text

* No paged attn for now

* Prettier format

* Multiple models at once

* Better switching, clearing ability

* Mobile support

* Inline markdown parser

* Update examples

* Typos

* Support specifying isq

* Fix mobile

* Fixes

* Fix button on mobile

* Image height is capped

* Thumbnail

* Fix rotating kv cache edge case

* Add drag and drop for images

* Small things

* Sidebar is frozen now

* Better listner

* Add readme

* Tweak readme

* Add chat history support to web chat app (EricLBuehler#1363)

* Add chat history

* Support renaming

* Start immediately with new chat

* Add timestamp

* Prettier chat list

* Style

* Delete chat

* Fix copy button

* Fix markdown rendering

* Store things in cache

* Store things in cache

* Refactor web chat, fix multichat image restore (EricLBuehler#1364)

* Fix multichat image restoration.

* Clippy

* Refactor

* Refactor frontent

* Fix repeated immediate isq init (EricLBuehler#1365)

* Add images_ref

* Add debug impl

* Fix the bug

* Tweak style of buttons

* Add a spinner

* Move spinner

* Tweak emoji

* Add gif

* Tweak initial gif

* Include vision tower tensors in Mistral3 UQFF (EricLBuehler#1366)

* Fix mistral 3 uqff resitdual tensors for vision

* Rolling shard creation for uqff files (EricLBuehler#1367)

* Fix occasional unstability during isq of afq (EricLBuehler#1368)

* Fix unstability during isq of afq

* Clippy

* Fix web chat installation

* Support web chat file uploading (EricLBuehler#1370)

* Web chat fixes

* Fix thumbnail in message, reuse blank chat

* Add file uploading support

* Fix scroll

* Allowed extensions

* Preserve files as literals

* Support multiple clients

* Add a stop button

* New cache dir

* New cache dir

* Fix

* Refactor

* Update readme

* Tweak drag-and-drop css

* Add speech generation support to the web chat! (EricLBuehler#1373)

* Initial speech gen support for web chat

* Tweak ui

* Update docs

* Prefix caching for PagedAttention! (EricLBuehler#1369)

* Exposing some things for logical token blocks

* Prefix cache manager has the scheduler

* Refactor

* Get logical and physical blocks into the prefix cacher

* Hash and cache

* Pass physical block prefill

* Allocation of prefilled block tables

* Temp

* Dont always use 2

* Hmm

* Hmm

* It mostly works

* Increment refcount

* Support images!

* Add to dummy paged attn

* Fix some clippy

* Clippy

* More checks

* Include EricLBuehler#1371, closes EricLBuehler#1371

* Typos

* Update docs

* Metal PagedAttention accuracy improvements (EricLBuehler#1374)

* Fix subtle bug

* Fix half sum bug

* Format metal paged attention

* Handle images in paged attn scheduler (EricLBuehler#1375)

* Include schemas needed for chatcompletions endpoint (EricLBuehler#1353)

* EricLBuehler#1326: WIP include schemas needed for chat completions endpoint

 Conflicts:
	Cargo.lock
	mistralrs-server/src/main.rs

* EricLBuehler#1326: WIP define utoipa as a workspace dep since core and server both need it

* EricLBuehler#1326: first draft of handling schemas that use Either

* EricLBuehler#1326: first draft of handling schema for Grammar

* EricLBuehler#1326: Add in other endpoints to API docs.

* EricLBuehler#1326: Adjust code comments

* EricLBuehler#1326: Implement coderabbitai suggestions

- EricLBuehler#1353 (review)
- EricLBuehler#1353 (comment)

* Fix constraints with metal sampler

* Revert EricLBuehler#1375

* Fix case where prefix cacher returns no toks (EricLBuehler#1377)

* Fix AFQ UQFF serialization

* Faster UQFF serialization (EricLBuehler#1379)

* Faster UQFF serialization

* Fix uqff gemma3

* Improve gemma3 auto loader names

* UQFF creation for AFQ on CPU support (EricLBuehler#1380)

* Add afq cpu quantize/dequantize

* Clippy

* Improved device for afq quantize

* Improved dtype handling for cpu afq (de)quantize

* Improved generate_uqff_card

* Add fused CPU attention kernel! (EricLBuehler#1382)

* Working

* Fix warnings

* Allow mask

* Support bf16, f16

* Handle striding

* Parallelized

* Add initial vector flash attn

* Avoid repeated allocations

* Tiled kv

* Apply some clippy

* Some small fixes

* Chunked vec_dot

* Clipy

* Use T::zero

* Refactor attention backends (EricLBuehler#1384)

* Refactor attention code

* Refactor attention code

* Move into backends

* Set macOS thread affinity for CPU attn (EricLBuehler#1385)

* Use lazylock

* Format

* Fix metal warn build

* Faster Qwen 3 MoE support on Metal (EricLBuehler#1387)

* Fix load

* Use afq gather qmm

* Well it runs

* It works

* Polish

* Fast and slow options

* Remove quantized.rs

* Polish some more

* Refactor

* Add isq

* Update load in parallel

* Support fp8

* Refactor for FusedExperts

* Clippy

* Handle pack factor when loading prequantized models

* Use f32 only in moe

* Avoid using f32 so much

* Avoid using f32 so much

* Fix PagedAttention block leaks (EricLBuehler#1388)

* Warn and ignore if ignored

* Fix a block allocation leak

* Update bench.py

* Fix double free in block engine

* Do not apply ISQ if loading a prequantized model

* Fix cuda build again (EricLBuehler#1389)

* Fix cuda build

* Fix

* Format

* Fixes for cuda docker

* Update dockerfiles

* Bump version to 0.6.0 (EricLBuehler#1390)

* Bump version to 0.6.0

* Remove lower_level api

* Make a static dir

* Update deps

* Fix routing for static handler in web chat

* Fewer .contiguous calls for qwen3 moe (EricLBuehler#1391)

* Allow speech models to accept batched inputs (EricLBuehler#1393)

* Allow speech models to accept batched inputs

* Clippy

* Ring distributed backend for heterogeneous TP (EricLBuehler#1238)

* Begin work on ring distributed backend for Metal

* Add the actual ring functionality

* It loads and kind of runs

* It works

* Optimize buffer allocation

* Avoid copy

* It works

* Add allgather

* Fix load

* Ping-pong

* Small things

* Add config json

* Allow different ip address

* Read config once

* Read config when appropriate

* Replicate requests

* Small fix

* Fix small compat with openai

* Clippy

* Update docs

* Add deepseek tool calling chat template

* Add auto loader for vision/text detection! (EricLBuehler#1402)

* Add auto loader for vision/text detection

* Build fixes

* Add model loader

* Update docs

* Format

* Create Mistral.rs Server Core Lib: `mistralrs-server-core` (EricLBuehler#1346)

* First draft of exposing mistral server routes as lib

* make arg struct fields pub

* Take base path so utoipa swagger route can properly redirect

* Expose swagger routes and make it configurable

* Add base path option for swagger docs

* More work on modularizing mistralrs server

* Sync fork (+1 squashed commit)
Squashed commits:
[169ae9e] Sync fork

* Adjust fn params to use refs / individual params instead of args

* Start breaking down controller actions into smaller pieces

* Continue refactoring

* Make mods pub so they can be used outside crate

* Allow chat completion streamer to take a callback so that you can get the complete response when finished

WIP (+3 squashed commits)
Squashed commits:
[0061d87] WIP
[c484d56] WIP
[16f8a60] WIP

* Sync fork

* Adjust callback type

* Remove throughput_log arg that was removed in 26afcc3

* Implement defaults for Args (and use for Clap)

* Small code formatting tweaks

* Rename callback to match SSE event and code clean up

* Sync fork

* WIP: first very rough draft of server core builder. Doesn't meet parity with old functional approach yet (slower / unstable?).

* Clean up (+4 squashed commits)
Squashed commits:
[e1cff387] Sync fork
[d8301025] WIP debugging
[1ea9f8c8] Sync fork
[4fe28cf5] WIP: debug function

* WIP server core builders

* Code clean up

* Add on_chunk callback

* Code clean up

* First draft of creating version of mistral-server that uses server-core

Code clean up (+1 squashed commit)
Squashed commits:
[adea1693]

* Sync fork

* Add helper methods to builder to make optional args more ergonomic (since .build validates params)

* Start adding docs

* Start cleaning up crates deps

* Example commit of mistral-server with implementing server-core

* Start addressing CodeRabbit feedback

* Fix comment typo

* Tweak doc blocks

* - Update type alias naming for clarity (MistralRs instead of Mistral)
- CodeRabbit, don't use eprintln for lib (use trace)
- Allow buffer size to be passed in and default to Constant
- Allow router body limit to be passed in and default to Constant
- Update doc examples

* Typo

* Address CoderRabbitAI feedback

* Support linear rope for llama3 (EricLBuehler#1408)

* Hotfix for loading

* Fix vllama4 uqff loading (EricLBuehler#1409)

* Fix vllama4 uqff loading

* Fix regex

* Fix regex

* Maybe a fix

* Gracefully handle receiver disconnects (EricLBuehler#1410)

* Handle receiver disconnects

* Format

* Fix Qwen3 MoE device mapping irregularities (EricLBuehler#1411)

* Fix bias

* Fix lm_head packing case

* Account for gate

* Fix head dim

* Fix interactive mode URL parsing (EricLBuehler#1412)

* fix url regex in vision interactive mode

* Fix regex

* Clippy

* Refactor auto device map (EricLBuehler#1413)

* Refactor auto device map

* Refactor a bit more

* Clippy

* Enable runtime sampling tweaks in interactive mode (EricLBuehler#1414)

* Document runtime sampling commands

* Fix readme

* Tweak

* Bounds checking

* Tweak temp bounds

* Send streaming tokens every time

* Gumbel sampling for fast sampler (EricLBuehler#1416)

* Improved handling for initialize_logging

* Improved CPU flash attention accuracy & performance (EricLBuehler#1417)

* Downcast correctly

* Operate internally in f32

* Avoid some casts and striding

* Prefetch

* Provide chat_templates to container users (EricLBuehler#1419)

Models often come without chat templates requiring mapping them
from the source repository into a container for access by the
mistralrs-server.

Copy the templates from the build tree into the root of the image
to permit use via `--chat-template /chat_templates/something.json`

TODO:
  With the increase in quantized models and support for other
formats, the initial benchmark run during model load can be used
to qualify/select existing chat templates embedded into the binary
for models which do not come with any (to include output of the
functional failures in each test allowing users to modify the
ones already provided correctly to suit the model being loaded).

Co-authored-by: RageLtMan <rageltman [at] sempervictus>

* Faster cpu flash attn (EricLBuehler#1418)

* Faster cpu flash attn

* Prefetch

* Clippy

* Add some tests

* Add softcap tests

* Fix test_parse_image_url test

* Update tests

* Update tests

* Web search improvements (bm25, web chat) (EricLBuehler#1420)

* Fix web search blocking case

* Web search support in web chat

* Tweak ui

* Support fallback to bm25

* Clippy

* Reinject descriptions

* Propely handle consecutive searches (EricLBuehler#1421)

* Update extraction tool reinjection

* Looped

* Update docs (EricLBuehler#1422)

- lib.rs: clean up example var names and match logging change from EricLBuehler@201d6be
- server_builder: fix typo
- READMEs: link to crate docs

* Better tool call detection logic (EricLBuehler#1424)

* Add web search hook callbacks (EricLBuehler#1426)

* feat: add customizable search hook

* Move to builder

* Update docs

* Fix CUDA context switching, bind thread on CudaStorage drop (EricLBuehler#1428)

* Add CUDA context helper and use in Llama forward

* No flashparams?

* working

* Tweak

* Update to use dep

* conditionally build flash attention inputs (EricLBuehler#1429)

* Add AGENTS.md (EricLBuehler#1430)

* Support Qwen3 GGUF model (EricLBuehler#1432)

* Support QWen3 GGUF model

* Clippy fix

* cargo fmt

* Improved paged attn prefix caching (EricLBuehler#1434)

* Improved paged attn prefix caching

* Disable

* Clippy

* Temporary fix for qwen3 gguf tokenizer (EricLBuehler#1433)

* Temporary fix for qwen3 gguf tokenizer

* Typo fix

* Add tool callback support (EricLBuehler#1427)

* Add tool callback support

* Fixes

* Support named tool callbacks

* Update examples

* Update docs

* Clippy

* Centralize crate dependencies (EricLBuehler#1438)

* chore: centralize dependencies

* Format

* Fix bug in tokenizer created with gguf metadata (EricLBuehler#1440)

* Fix bug in tokenizer created with gguf metadata

* Clippy fix

* Update deps (EricLBuehler#1441)

* Small things

* Update deps

* Update deps

* Update breaking changes

* Doc fixes (EricLBuehler#1442)

* Mention uqff_maker

* Downgrade rustyline 16.0.0 -> 15.0.0 (EricLBuehler#1444)

* Add max_completion_tokens alias for server (EricLBuehler#1451)

* Audio input support (Phi 4 multimodal) (EricLBuehler#1448)

* Deps

* Add conformer

* Nemo loading

* Position embeds

* Load t5 attn bias

* Attn and feed forward

* Add conv module and glu pointwise

* Implement relative attn bias

* Add the forward methods

* Add encoder embedding

* Fix oproj

* Some loading

* Conformer loads!

* Fully loading speech stack

* Merger

* Dont need that

* First pass at audio processing

* Read samples

* Optional

* Small loading fix

* Runs but not correct yet

* Improved audio processing?

* Works with this

* Fix t5 attn bias

* It works!

* Comment

* Use some other crates

* Clippy

* Allow bf16 on metal

* Add prefix_audio

* Remove unused

* Typo

* User specified

* Add audio url parsing

* AudioProjectionMode -> InputMode

* Audio prefix caching

* Fix bug in audio prefix caching

* Support both at the same time!

* Tweak logging

* Support stereo

* Add mistralrs-audio

* Support batching

* Add server and rust api example

* Add python api

* Fix add_multimodal_message

* Fix unfold for conformer

* Streaming example

* Add web chat support

* Add modalities registry

* Fix offline cache issue for gguf models (EricLBuehler#1452)

* Add MCP server endpoints (EricLBuehler#1453)

* feat(server): add MCP server support

* Add mcp docs

* Add handle_list_tools_request

* Better launch, tool handling

* Tmp state

* Ok works

* Handle modalities

* Update docs

* Add ping

* Tweak temperature bounds, args

* MCP documentation pass (EricLBuehler#1455)

* Fix table

* Update mcp docs

* Improve readme header

* Improve readme header

* Integrate an MCP client (EricLBuehler#1456)

* Add builtin mcp client

* Use async loader

* Add headers

* Handle sse

* More flexible search request

* Add tool callbacks with tools, for mcp

* Add bearer token support

* Add websocket support

* Update docs

* Add python api

* Clippy

* Add http api, docs

* Tests pass

* Make these configs actually work

* Add docs

* Make mistralrs-mcp

* Refactor examples

* Update examples

* Add defaults

* Add defaults

* Add defaults

* Update docs

* Improved docs

* Add -y to npx usages

* Even better examples

* Update generate_wheels

* Update generate_wheels

* Update generate_wheels

* Fix Dockerfile.cuda-all

* Improve automatic tool call (EricLBuehler#1460)

* Improved auto tool call

* Add logging

* chore: `Dockerfile.cuda-all` configurable threads (EricLBuehler#1458)

* chore: `Dockerfile.cuda-all` - Merge `RUN` for `apt-get install` (EricLBuehler#1459)

* Add fallback definition for isnan (EricLBuehler#1463)

* chore: `Dockerfile` - Drop runtime rayon thread ENV (EricLBuehler#1465)

* chore: Dockerfile - Remove rayon threads env

* chore: Dockerfile - Improve formatting for `apt-get`

* Remove duplicate calls for api_dir_list (EricLBuehler#1474)

* Remove duplicate calls for api_dir_list

* Support local cache for api_dir_list

* Fix home folder for metal

* Capitalized

* Fix transient pyo3 dep (EricLBuehler#1478)

Co-authored-by: Eric Buehler <[email protected]>

* Fix objc dep with non macos (EricLBuehler#1480)

* Fix phi 3/4 + nccl issue (EricLBuehler#1481)

* Fix log

* Fix n kv heads

* Fix phi3.5 moe (EricLBuehler#1482)

* Fix phi3.5 moe accum device

* Fix again

* Fix again

* Support GLM4 model! (EricLBuehler#1437)

* Support GLM4 model

* Mention GLM4 model in ReadMe

* glm4 type hint

* Typo fix

* Fix unsupported chat_template function

* Clippy fix

* Refactor distributed backend (EricLBuehler#1484)

* Refactor distributed backend, check power of 2

* Fix compilation

* Cap metal paged attn kv allocation (EricLBuehler#1485)

* Better paged attn metal cap (EricLBuehler#1486)

* Better paged attn metal cap

* Small fix

* Comment

* Small fix

* Refactor

* Server core: consolidate and unify route handlers and API surface (EricLBuehler#1423)

* Start working on consolidating completion and chat_completion underlying implementations

* Move response channel to util mod for now (since it's used with streaming and non streaming)

* More work on consolidating completions and chat completions

* More WIP consolidation of server core handlers

* More WIP consolidation of server core handlers

* More WIP consolidation of server core handlers

* Update docs and restrict completion core visibility

* CodeRabbit feedback: remove logprobs warn from route handler since parse request also checks this

* Use consistent var name for completions mod

* Make route handler modules public API consistent (same fn names, etc.) and provide proxy fn that wrap core fns so core mod doesn't have to be pub
Make lib.rs example compile checked and update example

* Code formatting

* Typo

* Sync fork

* Sync fork

* Docs example fix

* Support qwen3 gguf (EricLBuehler#1488)

* Add qwen3 gguf

* Template fixup

* Make bos/eos token IDs optional (EricLBuehler#1493)

* Remove python deps from CUDA dockerfiles (EricLBuehler#1487)

* Handle noncontiguous v in naive_sdpa (EricLBuehler#1499)

Co-authored-by: Eric Buehler <[email protected]>

* Server Core: refactor Paged Attention configuration (EricLBuehler#1500)

* Use StorageModePrivate for Metal PA kv cache (EricLBuehler#1506)

* Fix OpenAI stream: emit field in tool-call deltas for schema compliance (EricLBuehler#1507)

* FP8 KV-cache quantization for PagedAttention (EricLBuehler#1400)

* Add most of paged attn kv quant

* It builds a bit

* All the functionality at least

* Small fix

* Add a scale

* Fix bf16 usage

* Make k_v_scale optional

* Collector

* Tweak collection

* Refactor

* Add to apis

* Add cuda impl

* Fix compilation

* Fixes

* Handle ENABLE_FP8

* Format

* Tweak

* Fix scaled_convert usage

* Fix cache_t size

* Fixed scale collection

* Actual fix

* Fix fp8 for CC<8

* Fix the usual String != &str bit (EricLBuehler#1483)

Co-authored-by: RageLtMan <rageltman [at] sempervictus>

* chore: `Dockerfile` - Drop runtime rayon thread ENV (EricLBuehler#1465)

* chore: Dockerfile - Remove rayon threads env

* chore: Dockerfile - Improve formatting for `apt-get`

* Remove duplicate calls for api_dir_list (EricLBuehler#1474)

* Remove duplicate calls for api_dir_list

* Support local cache for api_dir_list

* Fix home folder for metal

* Capitalized

* Fix transient pyo3 dep (EricLBuehler#1478)

Co-authored-by: Eric Buehler <[email protected]>

* Fix objc dep with non macos (EricLBuehler#1480)

* Fix phi 3/4 + nccl issue (EricLBuehler#1481)

* Fix log

* Fix n kv heads

* Fix phi3.5 moe (EricLBuehler#1482)

* Fix phi3.5 moe accum device

* Fix again

* Fix again

* Support GLM4 model! (EricLBuehler#1437)

* Support GLM4 model

* Mention GLM4 model in ReadMe

* glm4 type hint

* Typo fix

* Fix unsupported chat_template function

* Clippy fix

* Refactor distributed backend (EricLBuehler#1484)

* Refactor distributed backend, check power of 2

* Fix compilation

* Cap metal paged attn kv allocation (EricLBuehler#1485)

* Better paged attn metal cap (EricLBuehler#1486)

* Better paged attn metal cap

* Small fix

* Comment

* Small fix

* Refactor

* Server core: consolidate and unify route handlers and API surface (EricLBuehler#1423)

* Start working on consolidating completion and chat_completion underlying implementations

* Move response channel to util mod for now (since it's used with streaming and non streaming)

* More work on consolidating completions and chat completions

* More WIP consolidation of server core handlers

* More WIP consolidation of server core handlers

* More WIP consolidation of server core handlers

* Update docs and restrict completion core visibility

* CodeRabbit feedback: remove logprobs warn from route handler since parse request also checks this

* Use consistent var name for completions mod

* Make route handler modules public API consistent (same fn names, etc.) and provide proxy fn that wrap core fns so core mod doesn't have to be pub
Make lib.rs example compile checked and update example

* Code formatting

* Typo

* Sync fork

* Sync fork

* Docs example fix

* Support qwen3 gguf (EricLBuehler#1488)

* Add qwen3 gguf

* Template fixup

* Make bos/eos token IDs optional (EricLBuehler#1493)

* Remove python deps from CUDA dockerfiles (EricLBuehler#1487)

* Handle USE_FP8 for cuda

* Fix cuda warn

* Add readme

* Saturating sub in sequence state

---------

Co-authored-by: Eric Buehler <[email protected]>
Co-authored-by: RageLtMan <[email protected]>
Co-authored-by: Brennan Kinney <[email protected]>
Co-authored-by: Guoqing Bao <[email protected]>
Co-authored-by: Matthew Haynes <[email protected]>

* Validate model name in OpenAI API (EricLBuehler#1509)

* Validate model name in openai api

* Add docs, allow 'ignore'

* Updated examples for EricLBuehler#1509

* Fix mcp import in doc string (EricLBuehler#1510)

* Add multi-model support! (EricLBuehler#1512)

* Refactor MistralRs

* Working multi-model!

* Add mutli-model docs initially

* Update mistralrs-pyo3, mistralrs-bench, mistralrs

* Update apis for consistency

* API tweaks

* Logging tweaks

* Add examples, tweak cli

* Clearer pipeline id

* Fix config key semantics

* Format and clippy

* Tweak logging, fix example

* Clippy refactor

* Update examples

* Remove unused multi model docs

* Replace 'ignore' with 'default'

* Update docs

* Add stars label to readme (EricLBuehler#1513)

* Add CLAUDE.md

* Handle base_model.model case in lora (EricLBuehler#1514)

* Add thread_local! for engine-specific const/static (EricLBuehler#1517)

* Fix MCP doc test (EricLBuehler#1511)

* Allow disabling metal precompilation (EricLBuehler#1518)

* Allow disabling metal precompilation

* Simple preprocessor

* Simple docs

---------

Co-authored-by: Eric Buehler <[email protected]>

* Rust 1.88 clippy (EricLBuehler#1522)

* Rust 1.88 clippy

* Format

* Fix cuda warnings (EricLBuehler#1526)

* Avoid panic decoding tokens on error (EricLBuehler#1527)

* Split Marlin and Paged Attention kernels for faster build (EricLBuehler#1525)

* Split Marlin and Paged Attention kernels for faster build

* Typo fix

* chore: update llguidance (EricLBuehler#1535)

* chore: update llguidance

* chore: remove unused import

* Add the SmolLM3 model! (EricLBuehler#1501)

* Add model

* Update loader

* Fix llama config usage

* Docs

* Fix config no_rope_layers

* Fix tie_word_embeddings default

* Add chat template

* Embed the chat templates

* Fix embedding template

* enable_thinking default true

* Update examples

* XML tools for smollm3

* Add smollm3 docs

* Fix openai examples

* Clippy

---------

Co-authored-by: Eric Buehler <[email protected]>

* Add full Gemma 3n support! (EricLBuehler#1519)

* Add initial

* Loading for text model

* Add ple embeddings

* Add altup, laurel block

* Update rmsnorm

* Add mlp

* Update attn norm application

* Currently no kv shared

* Wire it up

* It runs

* Fix bf16

* Fix scaled embd

* Fixes for mean

* tmp

* Attn confirmed

* Fix target_magnitude

* Add shared kv

* Ok it works

* Remove npy

* Fix streaming

* Remove warnings

* Remove paged attn

* Refactor rope

* Add immediate isq

* Add vision & mproj

* Update image processor

* Vision merge runs, not correct

* Remove

* Add mobilenet v5

* Add multimodal vision embedding

* Fix load

* runs

* Fix gamma

* Works but just not vision tower

* It works!!

* Tweak

* Fix warnings

* Move vision tower

* Fix warn

* Update cache manager things

* Refactor

* Add audio model, it loads

* Add audio processing

* It runs at least

* tmp

* A bit better

* Audio works!!!!

* Fused attn in vision

* Clippy

* Update audio runner

* Optimized audio model

* Remove unused things

* Fix inputs processor bug

* Remove comments

* Clippy

* Small optimizations

* Format

* Correctly register modalities

* Add docs

* Update readme

* Runs there

* Fixed padding from Blaizzy/mlx-vlm#410

* Add better checks

* Fix sdpa n_kv_groups

* Vision encoder works!

* Rotate image

* Clippy

* Fix cuda loading

* Updated device mapper

* Fix overflow

* Fix dtype errors

* Refactor image/audio embeddings

* Fix metal

* Fix dtype mismatch

* Audio processing fixes

* Audio processing fixes

* Works

* Audio is good

* Fix boi/eoi too

* Embed the chat templates

* Better embedding accuracy in non f32

* More f32

* Support bf16 on metal

* Add more ISQ

* Fixed device map

* Clippy

* Gemma3n no paged attn

* Fix saturating sub

* Faster rmsnorm

* Use sdpa for vision model

* Fix ple bug

* Fix name

* Fix multiaudio

* Add matformer config loading

* Add docs

* Add support for matformer in auto device mapper

* Update docs

* Typos

* Tweak

* Tweak

* Fix multidevice

* Fix gemma3n text model auto device map

* Fix dims3

* Fix auto devic emap vision

* Non-metal keeps PLE on cpu

* Complete merge

* Vision dtype f16 -> f32

* Fix metal nm device

* Fix uqff

* Typos

* Reference uqff

* Fix tests

* Fix sequence length check (EricLBuehler#1546)

* update candle version (EricLBuehler#1545)

Co-authored-by: AlpineVibrations <[email protected]>

* add ios target to metal deps (EricLBuehler#1548)

---------

Signed-off-by: dependabot[bot] <[email protected]>
Co-authored-by: Eric Buehler <[email protected]>
Co-authored-by: Eric Buehler <[email protected]>
Co-authored-by: edwko <[email protected]>
Co-authored-by: Copilot <[email protected]>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Co-authored-by: Guoqing Bao <[email protected]>
Co-authored-by: Michał Moskal <[email protected]>
Co-authored-by: Chen Mulong <[email protected]>
Co-authored-by: Steph Wolski <[email protected]>
Co-authored-by: omahs <[email protected]>
Co-authored-by: Viktor Szépe <[email protected]>
Co-authored-by: Matthew Haynes <[email protected]>
Co-authored-by: RageLtMan <[email protected]>
Co-authored-by: Brennan Kinney <[email protected]>
Co-authored-by: Eric Buehler <[email protected]>
Co-authored-by: Sbargaoui <[email protected]>
Co-authored-by: Gaétan Lepage <[email protected]>
Co-authored-by: Ammar Elsabe <[email protected]>
Co-authored-by: luke <[email protected]>
Co-authored-by: AlpineVibrations <[email protected]>
Co-authored-by: Michael Tissen <[email protected]>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant