feat: add mistral and phi4 tool parser #2510
Conversation
WalkthroughAdds stricter JSON extraction and fallback handling in the tool-calling JSON parser, introduces two new model-specific parser configurations (mistral, phi4), and wires them into detection. Also adds a debug print for parsed JSON. Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant Detector as detect_and_parse_tool_call
participant Config as ToolCallConfig
participant JSONP as JsonParser
Client->>Detector: parse(input, parser="mistral"/"phi4"/...)
Detector->>Config: select config by parser key
Config-->>Detector: ToolCallConfig (start/end tokens, Json)
Detector->>JSONP: try_tool_call_parse_json(input, config)
JSONP->>JSONP: extract '{...}' segments, trim, validate
JSONP-->>Detector: parsed tool calls or stripped fallback
Detector-->>Client: result
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Poem
Tip 🔌 Remote MCP (Model Context Protocol) integration is now available!Pro plan users can now connect to remote MCP servers from the Integrations page. Connect with popular remote MCPs such as Notion and Linear to add more context to your reviews and chats. 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. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR/Issue comments)Type Other keywords and placeholders
Status, Documentation and Community
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (5)
lib/llm/src/postprocessor/tool_calling/json_parser.rs (1)
77-81: Clarify intent and consider array segments for single-token casesThe comment says “brace/bracket,” but the code only searches for '}'. If an array follows the start token (e.g., “[{...}, ...]”), this branch is skipped and we fall back. That’s OK for many cases, but we’ll miss multiple calls if there are several “single-token” segments each starting with '['.
- Update the comment for accuracy.
- Optional: also handle s.starts_with('[') by cutting at the last ']' and attempting to parse it as a JSON array.
Example minimal comment fix within this block:
- // Trim trailing non-JSON by cutting at the last closing brace/bracket + // Trim trailing non-JSON by cutting at the last closing braceIf you want to also handle arrays here (optional), I can provide a patch that safely parses '[' segments and merges/normalizes them. Let me know.
lib/llm/src/postprocessor/tool_calling/parsers.rs (4)
103-112: Mistral config looks good; consider adding a brief doc commentConfiguration is straightforward and consistent with how the single-token handler expects the start token. A short doc comment helps future maintainers.
+ /// Default configuration for Mistral tool calls. + /// Models may emit a `[TOOL_CALLS]` prefix followed by a JSON array. pub fn mistral() -> Self {
561-579: Test name nit: missing underscoreMinor nit: “with_start_tokenwith_new_lines” -> “with_start_token_with_new_lines”.
- fn test_mistralai_mistral_7b_instruct_v03_single_with_start_tokenwith_new_lines() { + fn test_mistralai_mistral_7b_instruct_v03_single_with_start_token_with_new_lines() {
841-857: Test name nit: says “multiple” but input contains a single callThe name “..._multiple_with_new_lines” is misleading here; the input has one call.
- fn test_detect_and_parse_tool_call_default_parser_llama3_json_without_python_tag_multiple_with_new_lines( + fn test_detect_and_parse_tool_call_default_parser_llama3_json_without_python_tag_with_new_lines(
157-181: Optional: avoid per-call HashMap allocation in detect_and_parse_tool_callConstructing the parser map on every call is small but avoidable. Consider a static Lazy or a direct match on the parser string.
Happy to draft a small refactor using once_cell::sync::Lazy if you want.
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (2)
lib/llm/src/postprocessor/tool_calling/json_parser.rs(2 hunks)lib/llm/src/postprocessor/tool_calling/parsers.rs(9 hunks)
🧰 Additional context used
🧬 Code Graph Analysis (1)
lib/llm/src/postprocessor/tool_calling/parsers.rs (1)
lib/llm/src/protocols/openai/chat_completions/aggregator.rs (1)
default(67-69)
🪛 GitHub Actions: Pre Merge Validation of (ai-dynamo/dynamo/refs/pull/2510/merge) by ayushag-nv.
lib/llm/src/postprocessor/tool_calling/parsers.rs
[error] 1-1: Pre-commit trailing-whitespace check failed; file modified by hook. Trailing whitespace detected and fixed by the 'trailing-whitespace' hook.
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (4)
- GitHub Check: Build and Test - dynamo
- GitHub Check: pre-merge-rust (lib/bindings/python)
- GitHub Check: pre-merge-rust (lib/runtime/examples)
- GitHub Check: pre-merge-rust (.)
🔇 Additional comments (14)
lib/llm/src/postprocessor/tool_calling/parsers.rs (14)
114-123: phi4 config is appropriate; relies on single-token array handlingUsing "functools" (without the '[') is a good choice because the array fallback needs the '[' to remain. With the proposed json_parser fallback fix (strip from first occurrence), this will also work when the token isn’t at position 0.
Consider adding a test where the response contains a short preface before “functools[...]” to exercise this path.
167-169: Wiring new parsers into detection map looks correctBoth "mistral" and "phi4" are added alongside existing parsers. Behavior remains backward compatible.
332-360: Great newline-tolerance coverage for nemotron_deciThis test strengthens confidence in the regex extractor with multi-line arrays inside tags.
413-438: Hermes multi-call with newlines: solid coverageConfirms the end-token with leading newline is matched and multiple tool_call blocks are handled.
483-507: Mistral simple with newlines: goodValidates that plain arrays (no start token present) still parse under the Mistral config.
509-524: Mistral multiple calls (no start token): goodConfirms vector parsing path.
526-547: Mistral multiple with newlines: goodExercises pretty-printed arrays.
597-622: Mistral with start token, multiple and newlines: goodValidates that a prefixed array parses as multiple calls.
635-648: Default llama3_json with newlines: goodCovers the plain JSON object path under the default-detected parser.
662-676: llama3_json with python tag + newlines: goodCovers the single-token without end-token flow for objects.
677-696: llama3_json multiple python tags + newlines: goodConfirms multiple single-token segments are aggregated.
822-840: Default parser with python tag + newlines: goodValidates detection default path handles the python tag.
870-924: phi4 tests are thorough; add a prefaced variant to harden furtherThe coverage for single, multiple, nested, and parameters-vs-arguments is great. Add one case with a natural-language preface before “functools[...]” to ensure robustness when the token isn’t at position 0. This will be fully supported after the json_parser fallback fix.
Would you like me to add a test like:
- input: "Tool calls:\nfunctools[ { ... }, { ... } ]"
- parser: Some("phi4")
- expectation: parse both calls
332-360: Note on trailing whitespace pre-commit failureCI flagged and auto-fixed trailing whitespace. Locally running pre-commit hooks before pushing will avoid future CI churn.
If you don’t have pre-commit set up locally:
- Install: pip install pre-commit
- Init: pre-commit install
- Run: pre-commit run -a
Also applies to: 413-438, 483-507, 526-547, 561-579, 597-622, 635-648, 662-676, 677-696, 822-840, 841-857
Signed-off-by: Hannah Zhang <hannahz@nvidia.com>
Overview:
Details:
Where should the reviewer start?
Related Issues: (use one of the action keywords Closes / Fixes / Resolves / Relates to)
Summary by CodeRabbit
New Features
Bug Fixes
Tests