chore: enable tool call array parsing #2466
Conversation
WalkthroughPublic APIs for tool-call parsing now return Vec instead of Option across JSON parser, parsers, and tools modules. Logic shifts from single-item extraction to collecting multiple tool calls. Aggregator updated to consume and set multiple tool calls per choice, adjusting logging and finish_reason handling accordingly. Changes
Sequence Diagram(s)sequenceDiagram
participant Model as LLM Output
participant Parser as Parser(detector/json)
participant Tools as Tools Mapper
participant Aggregator as OpenAI Aggregator
Model->>Parser: detect_and_parse_tool_call(message)
Parser-->>Model: Vec<ToolCallResponse> (possibly empty)
alt Non-empty
Parser->>Tools: map to Vec<ChatCompletionMessageToolCall>
Tools-->>Parser: Vec<ChatCompletionMessageToolCall>
Parser->>Aggregator: tool_calls Vec
Aggregator->>Aggregator: set choice.tool_calls, clear text, finish_reason=ToolCalls
else Empty
Parser-->>Aggregator: []
Aggregator->>Aggregator: no changes
end
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 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: 0
🔭 Outside diff range comments (1)
lib/llm/src/postprocessor/tool_calling/json_parser.rs (1)
54-95: Doc comments are stale: update to reflect Vec-return semantics and multi-item handlingThe docs still describe Option-return and "last item only", which no longer matches the implementation and public API.
Apply this doc update:
/// # Return /// -/// - `Ok(Some(ToolCallResponse))` if parsing succeeds -/// - `Ok(None)` if input format is unrecognized or invalid JSON -/// - `Err(...)` if JSON is valid but deserialization or argument re-serialization fails +/// - `Ok(Vec<ToolCallResponse>)` containing zero or more parsed tool calls +/// - Empty vector if input is unrecognized/invalid for parsing +/// - `Err(...)` only if argument re-serialization fails /// /// # Note on List Handling /// -/// When the input contains a list of tool calls (either with `parameters` or `arguments`), -/// only the **last item** in the list is returned. This design choice assumes that the -/// most recent tool call in a list is the one to execute. +/// When the input contains a list of tool calls (either with `parameters` or `arguments`), +/// all valid items are returned in order. /// /// # Errors /// -/// Returns a `Result::Err` only if an inner `serde_json::to_string(...)` fails -/// (e.g., if the arguments are not serializable). +/// Returns a `Result::Err` only if an inner `serde_json::to_string(...)` fails +/// (e.g., if the arguments are not serializable). /// /// # Examples /// /// ```ignore /// let input = r#"<TOOLCALL>[{ "name": "search", "parameters": { "query": "rust" } }]</TOOLCALL>"#; -/// let result = try_tool_call_parse_json(input)?; -/// assert!(result.is_some()); +/// let result = try_tool_call_parse_json(input, &JsonParserConfig::default())?; +/// assert_eq!(result.len(), 1); /// ```
🧹 Nitpick comments (10)
lib/llm/src/postprocessor/tool_calling/json_parser.rs (4)
107-111: Avoid panicking on config mismatch; return an error insteadUsing assert! will panic at runtime. Prefer a controlled error to aid observability and avoid taking down the process.
- assert!( - tool_call_start_tokens.len() == tool_call_end_tokens.len(), - "Tool call start and end tokens must have the same length" - ); + if tool_call_start_tokens.len() != tool_call_end_tokens.len() { + anyhow::bail!("Tool call start and end tokens must have the same length"); + }
172-179: Simplify collection of list items with iterator collect()Current code builds a Vec and conditionally returns it. You can streamline and let collect() handle empty lists.
- let mut results = Vec::new(); - for item in list { - results.push(parse(item.name, item.parameters)?); - } - if !results.is_empty() { - return Ok(results); - } + return list + .into_iter() + .map(|item| parse(item.name, item.parameters)) + .collect();Also, the preceding comment still says “We pop the last item” which is now outdated. Consider updating it to “Collect all items”.
194-200: Mirror the iterator-based collection for arguments listsSame refactor opportunity as the parameters-variant; also update the stale “take the last item” comment above.
- let mut results = Vec::new(); - for item in list { - results.push(parse(item.name, item.arguments)?); - } - if !results.is_empty() { - return Ok(results); - } + return list + .into_iter() + .map(|item| parse(item.name, item.arguments)) + .collect();
112-126: Future enhancement: support multiple wrapper blocks, not just the last matchEven with array support, extract_tool_call_content() still takes the last match only (“TODO: Handle multiple tool calls”). If models emit multiple wrapper blocks (e.g., two <tool_call>...</tool_call>), earlier ones will be ignored.
I can propose an approach to return Vec<&str> from the extractor and iterate all blocks if/when you want to tackle this.
lib/llm/src/protocols/openai/chat_completions/aggregator.rs (2)
166-184: Handle multiple tool calls: LGTM; consider redacting/truncating arguments in logsThe multi-call handling (assign vector, clear text, set finish_reason) is correct. However, logging full arguments can leak sensitive data and blow up logs.
Truncate previews and log lengths instead of full payloads:
- for tool_call in &tool_calls { - tracing::debug!( - tool_call_id = %tool_call.id, - function_name = %tool_call.function.name, - arguments = %tool_call.function.arguments, - "Parsed structured tool call from aggregated content" - ); - } + for tool_call in &tool_calls { + let args = &tool_call.function.arguments; + let preview = if args.len() > 1024 { + format!("{}… (truncated, {} bytes)", &args[..1024], args.len()) + } else { + args.clone() + }; + tracing::debug!( + tool_call_id = %tool_call.id, + function_name = %tool_call.function.name, + arguments_preview = %preview, + arguments_len = args.len(), + "Parsed structured tool call from aggregated content" + ); + }
163-187: Add targeted tests for multi-call extraction to prevent regressionsTests currently validate text aggregation only. Add a case where
choice.textcontains a wrapper with two tool calls and assert:
choice.message.tool_callslength is 2choice.message.contentis Nonefinish_reason == ToolCallsI can provide a ready-to-run test if helpful.
lib/llm/src/postprocessor/tool_calling/tools.rs (2)
15-35: Unconditional map+collect simplifies control flowThe empty-vec branch is redundant; mapping an empty
parsedalready yields an empty vec.- let parsed = detect_and_parse_tool_call(message, parser_str)?; - if !parsed.is_empty() { - Ok(parsed - .into_iter() - .map( - |parsed| async_openai::types::ChatCompletionMessageToolCall { - id: parsed.id, - r#type: async_openai::types::ChatCompletionToolType::Function, - function: async_openai::types::FunctionCall { - name: parsed.function.name, - arguments: parsed.function.arguments, - }, - }, - ) - .collect()) - } else { - Ok(vec![]) - } + let parsed = detect_and_parse_tool_call(message, parser_str)?; + Ok(parsed + .into_iter() + .map(|parsed| async_openai::types::ChatCompletionMessageToolCall { + id: parsed.id, + r#type: async_openai::types::ChatCompletionToolType::Function, + function: async_openai::types::FunctionCall { + name: parsed.function.name, + arguments: parsed.function.arguments, + }, + }) + .collect())
42-66: Same simplification applies to streaming variantEnumerate + map is good; remove the conditional for cleaner code.
- let parsed = detect_and_parse_tool_call(message, parser_str)?; - if !parsed.is_empty() { - Ok(parsed - .into_iter() - .enumerate() - .map( - |(idx, parsed)| async_openai::types::ChatCompletionMessageToolCallChunk { - index: idx as u32, - id: Some(parsed.id), - r#type: Some(async_openai::types::ChatCompletionToolType::Function), - function: Some(async_openai::types::FunctionCallStream { - name: Some(parsed.function.name), - arguments: Some(parsed.function.arguments), - }), - // Add other fields as needed if required by the struct definition - }, - ) - .collect()) - } else { - Ok(vec![]) - } + let parsed = detect_and_parse_tool_call(message, parser_str)?; + Ok(parsed + .into_iter() + .enumerate() + .map(|(idx, parsed)| async_openai::types::ChatCompletionMessageToolCallChunk { + index: idx as u32, + id: Some(parsed.id), + r#type: Some(async_openai::types::ChatCompletionToolType::Function), + function: Some(async_openai::types::FunctionCallStream { + name: Some(parsed.function.name), + arguments: Some(parsed.function.arguments), + }), + }) + .collect())lib/llm/src/postprocessor/tool_calling/parsers.rs (2)
256-268: Rename tests to match Vec semanticsTest names still say “returns_none…”, but the contract is now “empty Vec”. Renaming avoids confusion.
- fn returns_none_on_invalid_input() { + fn returns_empty_vec_on_invalid_input() { @@ - fn returns_none_on_valid_json_wrong_shape() { + fn returns_empty_vec_on_valid_json_wrong_shape() {
430-453: Fix comments in error-handling test to reflect empty-Vec behaviorComments still mention Ok(None). Update to Ok(empty Vec) to avoid drift.
- // Known parser, but invalid input (not JSON) should return Ok(None) + // Known parser, but invalid input (not JSON) should return Ok(empty Vec) @@ - // Known parser, but valid JSON with wrong shape should return Ok(None) + // Known parser, but valid JSON with wrong shape should return Ok(empty Vec)
📜 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 (4)
lib/llm/src/postprocessor/tool_calling/json_parser.rs(4 hunks)lib/llm/src/postprocessor/tool_calling/parsers.rs(21 hunks)lib/llm/src/postprocessor/tool_calling/tools.rs(2 hunks)lib/llm/src/protocols/openai/chat_completions/aggregator.rs(1 hunks)
🧰 Additional context used
🧬 Code Graph Analysis (2)
lib/llm/src/postprocessor/tool_calling/tools.rs (1)
lib/llm/src/postprocessor/tool_calling/parsers.rs (1)
detect_and_parse_tool_call(136-157)
lib/llm/src/protocols/openai/chat_completions/aggregator.rs (1)
lib/llm/src/postprocessor/tool_calling/tools.rs (1)
try_tool_call_parse_aggregate(14-36)
⏰ 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/runtime/examples)
- GitHub Check: pre-merge-rust (lib/bindings/python)
- GitHub Check: pre-merge-rust (.)
🔇 Additional comments (5)
lib/llm/src/postprocessor/tool_calling/json_parser.rs (2)
95-101: Signature change to Result<Vec<…>> looks goodReturning a Vec aligns the JSON parser with the rest of the tool-calling stack and enables multi-call support. Logging the effective config is also helpful for diagnostics.
150-164: Single-object branches correctly wrap in a VecThe shift from Option to Vec with single-element vectors is consistent and avoids special-casing at call sites.
lib/llm/src/postprocessor/tool_calling/parsers.rs (3)
113-133: Vec-based return for try_tool_call_parse is consistent and future-proofDispatching to the JSON parser and returning Vec harmonizes the surface API and enables multiple tool calls. Clear error semantics for unimplemented formats look good.
136-157: Default parser fallback improves ergonomicsDefaulting to "default" when parser_str is None/empty is a nice usability improvement. Error message for unknown parsers is explicit.
170-176: Tests comprehensively cover single- and multi-call flows across configsGood coverage for parameters vs arguments fields, wrapper tags, and default parser behavior. The array-of-calls tests validate the primary PR objective.
Also applies to: 184-189, 195-206, 209-220, 222-232, 235-254, 270-306, 308-320, 405-427, 513-535
elyasmnvidian
left a comment
elyasmnvidian
left a comment
There was a problem hiding this comment.
looks good, please address coderabbit feedback
ayushag-nv
force-pushed
the
ayushag/multi-tool-call-parsing
branch
from
5385993 to
cfeb800
Compare
Signed-off-by: Hannah Zhang <hannahz@nvidia.com>
Overview:
This PR enables parsing of multiple tool calls present inside an array structure.
Details:
Currently for these kind of tool calls only last one was considered. This PR enables parsing of multiple tool calls present in an array structure and enables end to end.
Where should the reviewer start?
Related Issues: (use one of the action keywords Closes / Fixes / Resolves / Relates to)