diff --git a/crates/gbnf-rs/AGENTS.md b/crates/gbnf-rs/AGENTS.md index 291e227..067a705 100644 --- a/crates/gbnf-rs/AGENTS.md +++ b/crates/gbnf-rs/AGENTS.md @@ -21,4 +21,9 @@ Generate GBNF grammar ASTs from `schemars::Schema` inputs. ## Constraints - Generated object rules include all properties - required properties appear first - - optional properties may follow in any order + - optional properties follow in fixed, deterministic order + - each optional property appears at most once (`?` semantics) + - when requireds exist, each optional is prefixed by a comma if present + - when no requireds exist, the first present optional has no leading comma + - when no requireds exist, optional properties can start at any position and may skip middle properties while preserving order and uniqueness + - implemented via recursive choices over downstream chains (O(n^2) size) diff --git a/crates/gbnf-rs/src/lib.rs b/crates/gbnf-rs/src/lib.rs index 4c80332..ce0a59b 100644 --- a/crates/gbnf-rs/src/lib.rs +++ b/crates/gbnf-rs/src/lib.rs @@ -258,35 +258,85 @@ impl Generator { optional_props.sort_by(|a, b| a.0.cmp(b.0)); if !optional_props.is_empty() { - let mut choices = Vec::new(); - for (name, subschema) in optional_props { - let subschema: Schema = subschema.clone().try_into().unwrap_or_default(); - let expr = self.expr_from_schema(&subschema, defs, cache, grammar); - let expr = self.maybe_rule(expr, grammar); - choices.push(Expr::Seq(vec![ + // Build optional property expressions in fixed order. + // When there are required fields already emitted (first == false), + // each optional field, if present, must be preceded by a comma. + // When no required fields exist (first == true), we construct a + // nested optional chain like: opt1 ("," opt2 ("," opt3)?)? + // and wrap the whole chain in Optional(...), allowing an empty object. + + // Precompute value expressions for each optional prop so we reuse + // the same rule/reference across multiple chains and avoid duplicates. + let prepared: Vec<(String, Expr)> = optional_props + .iter() + .map(|(name, subschema)| { + let subschema: Schema = (*subschema).clone().try_into().unwrap_or_default(); + let expr = self.expr_from_schema(&subschema, defs, cache, grammar); + let expr = self.maybe_rule(expr, grammar); + ((*name).clone(), expr) + }) + .collect(); + + // Helper to build a single property expr (without any leading comma) + let mk_prop_idx = |idx: usize| { + let (ref name, ref expr) = prepared[idx]; + Expr::Seq(vec![ Expr::Literal(format!("\"{}\"", name)), Expr::Ref("ws".into()), Expr::Literal(":".into()), Expr::Ref("ws".into()), - expr, - ])); - } + expr.clone(), + ]) + }; if first { - seq.push(Expr::Optional(Box::new(Expr::Seq(vec![ - Expr::Choice(choices.clone()), - Expr::Repeat(Box::new(Expr::Seq(vec![ + // No requireds: allow starting at any optional and skipping + // middles while preserving order and uniqueness. + // Build chain_i for each i: prop(i) followed by optional + // comma + choice of any chain_j where j > i. + let n = prepared.len(); + if n > 0 { + let mut chains: Vec = vec![Expr::Literal(String::new()); n]; + // Build from the end to the front so later chains exist. + for i in (0..n).rev() { + let prop = mk_prop_idx(i); + let mut choices: Vec = Vec::new(); + for j in (i + 1)..n { + choices.push(chains[j].clone()); + } + let chain_i = if choices.is_empty() { + prop + } else { + Expr::Seq(vec![ + prop, + Expr::Optional(Box::new(Expr::Seq(vec![ + Expr::Ref("ws".into()), + Expr::Literal(",".into()), + Expr::Ref("ws".into()), + Expr::Choice(choices), + ]))), + ]) + }; + chains[i] = chain_i; + } + let top_choice = if n == 1 { + chains.remove(0) + } else { + Expr::Choice(chains) + }; + seq.push(Expr::Optional(Box::new(top_choice))); + } + } else { + // Requireds present: add each optional as an independent ("," prop)? + for (idx, _) in prepared.iter().enumerate() { + let prop = mk_prop_idx(idx); + seq.push(Expr::Optional(Box::new(Expr::Seq(vec![ + Expr::Ref("ws".into()), Expr::Literal(",".into()), Expr::Ref("ws".into()), - Expr::Choice(choices), - ]))), - ])))); - } else { - seq.push(Expr::Repeat(Box::new(Expr::Seq(vec![ - Expr::Literal(",".into()), - Expr::Ref("ws".into()), - Expr::Choice(choices), - ])))); + prop, + ])))); + } } } } @@ -346,7 +396,7 @@ mod tests { let generator = Generator::new(); let grammar = generator.generate("params", &schema); insta::assert_snapshot!(grammar.to_string(), @r###" -params-r0 ::= "{" ws "\"file_path\"" ws ":" ws string "," ws "\"new_string\"" ws ":" ws string "," ws "\"old_string\"" ws ":" ws string ("," ws "\"expected_replacements\"" ws ":" ws number)* ws "}" +params-r0 ::= "{" ws "\"file_path\"" ws ":" ws string "," ws "\"new_string\"" ws ":" ws string "," ws "\"old_string\"" ws ":" ws string (ws "," ws "\"expected_replacements\"" ws ":" ws number)? ws "}" "###); } @@ -364,13 +414,32 @@ params-r0 ::= "{" ws "\"file_path\"" ws ":" ws string "," ws "\"new_string\"" ws let generator = Generator::new(); let grammar = generator.generate("params", &schema); insta::assert_snapshot!(grammar.to_string(), @r###" -params-r0 ::= "{" ws "\"path\"" ws ":" ws string ("," ws "\"ignore\"" ws ":" ws r1 | "\"include\"" ws ":" ws r2 | "\"include_hidden\"" ws ":" ws r3)* ws "}" +params-r0 ::= "{" ws "\"path\"" ws ":" ws string (ws "," ws "\"ignore\"" ws ":" ws r1)? (ws "," ws "\"include\"" ws ":" ws r2)? (ws "," ws "\"include_hidden\"" ws ":" ws r3)? ws "}" r1 ::= "[" ws (string ("," ws string)*)? ws "]" r2 ::= "[" ws (string ("," ws string)*)? ws "]" r3 ::= "true" | "false" "###); } + #[derive(JsonSchema)] + struct AllOptional { + a: Option, + b: Option, + c: Option>, + } + + #[test] + fn all_optional_object_grammar() { + let schema = schemars::schema_for!(AllOptional); + let generator = Generator::new(); + let grammar = generator.generate("params", &schema); + insta::assert_snapshot!(grammar.to_string(), @r###" +params-r0 ::= "{" ws ("\"a\"" ws ":" ws string (ws "," ws "\"b\"" ws ":" ws r1 (ws "," ws "\"c\"" ws ":" ws r2)? | "\"c\"" ws ":" ws r2)? | "\"b\"" ws ":" ws r1 (ws "," ws "\"c\"" ws ":" ws r2)? | "\"c\"" ws ":" ws r2)? ws "}" +r1 ::= "true" | "false" +r2 ::= "[" ws (string ("," ws string)*)? ws "]" +"###); + } + #[derive(JsonSchema)] struct Item { value: String, diff --git a/crates/llm/AGENTS.md b/crates/llm/AGENTS.md index bfc5afb..ec8cdce 100644 --- a/crates/llm/AGENTS.md +++ b/crates/llm/AGENTS.md @@ -91,10 +91,10 @@ Trait-based LLM client implementations for multiple providers. - implements `ToolExecutor` for MCP calls - unrecognized or unprefixed tool names return "{name} is not a valid tool name" - tool execution errors with `is_error` set propagate as `Err` - - tool call chunks insert assistant messages immediately before execution - - accumulated assistant content is flushed as a separate assistant message - - the tool-call assistant message carries empty `content` with `tool_calls` populated - - accumulated streamed content is appended as an assistant message after the stream completes + - tool orchestration message shape + - streamed assistant parts (text/thinking) and tool calls from a single response are combined into one assistant message in original order + - the tool result is appended as a separate tool-role message with the same call id + - follow-up assistant responses are appended after tool execution completes - Test utilities - `TestProvider` implements `LlmClient` - captures `ChatMessageRequest`s for assertions diff --git a/crates/llm/src/tools.rs b/crates/llm/src/tools.rs index a053c56..9caf951 100644 --- a/crates/llm/src/tools.rs +++ b/crates/llm/src/tools.rs @@ -14,7 +14,7 @@ use tokio_stream::{Stream, StreamExt}; use crate::{ AssistantMessage, AssistantPart, ChatMessage, ChatMessageRequest, JsonResult, LlmClient, - ResponseChunk, ToolCall, + ResponseChunk, }; #[async_trait] @@ -69,100 +69,84 @@ pub async fn run_tool_loop( tx.send(ToolEvent::RequestStarted).ok(); let mut handles: JoinSet<(String, String, Result>)> = JoinSet::new(); - let mut assistant_content: Option = None; - let mut assistant_thinking: Option = None; + let mut parts: Vec = Vec::new(); + let mut current_part: Option = None; while let Some(chunk) = stream.next().await { let chunk = chunk?; let mut done = false; - let mut tool_calls: Vec = Vec::new(); match &chunk { ResponseChunk::Content(content) => { - if !content.is_empty() { - assistant_content - .get_or_insert_with(String::new) - .push_str(content); + if let Some(AssistantPart::Text { text }) = current_part.as_mut() { + text.push_str(content); + } else { + if let Some(part) = current_part.take() { + parts.push(part); + } + current_part = Some(AssistantPart::Text { + text: content.into(), + }); } } ResponseChunk::Thinking(thinking) => { - if !thinking.is_empty() { - assistant_thinking - .get_or_insert_with(String::new) - .push_str(thinking); + if let Some(AssistantPart::Thinking { text }) = current_part.as_mut() { + text.push_str(thinking); + } else { + if let Some(part) = current_part.take() { + parts.push(part); + } + current_part = Some(AssistantPart::Thinking { + text: thinking.into(), + }); } } ResponseChunk::ToolCall(tc) => { - tool_calls.push(tc.clone()); + if let Some(part) = current_part.take() { + parts.push(part); + } + current_part = Some(AssistantPart::ToolCall(tc.clone())); + tx.send(ToolEvent::ToolStarted { + call_id: tc.id.clone(), + name: tc.name.clone(), + args: tc.arguments.clone(), + }) + .ok(); + let executor = tool_executor.clone(); + let name = tc.name.clone(); + let args = tc.arguments.clone(); + let call_id = tc.id.clone(); + handles.spawn(async move { + match args { + JsonResult::Content { content } => { + let res = executor.call(&name, content).await; + (call_id, name, res) + } + JsonResult::Error { .. } => ( + call_id, + name, + Err::>(Box::new( + std::io::Error::new( + std::io::ErrorKind::Other, + "Could not parse arguments as JSON", + ), + )), + ), + } + }); } ResponseChunk::Usage { .. } => {} ResponseChunk::Done => { done = true; } } - if !tool_calls.is_empty() { - if let Some(content) = assistant_content.take() { - chat_history - .lock() - .unwrap() - .push(ChatMessage::Assistant(AssistantMessage { - content: vec![AssistantPart::Text { text: content }], - })); - } - let mut parts: Vec = tool_calls - .clone() - .into_iter() - .map(AssistantPart::ToolCall) - .collect(); - if let Some(thinking) = assistant_thinking.take() { - parts.insert(0, AssistantPart::Thinking { text: thinking }); - } - chat_history - .lock() - .unwrap() - .push(ChatMessage::Assistant(AssistantMessage { content: parts })); - } tx.send(ToolEvent::Chunk(chunk)).ok(); - for call in tool_calls { - tx.send(ToolEvent::ToolStarted { - call_id: call.id.clone(), - name: call.name.clone(), - args: call.arguments.clone(), - }) - .ok(); - let executor = tool_executor.clone(); - let name = call.name.clone(); - let args = call.arguments.clone(); - let call_id = call.id.clone(); - handles.spawn(async move { - match args { - JsonResult::Content { content } => { - let res = executor.call(&name, content).await; - (call_id, name, res) - } - JsonResult::Error { .. } => ( - call_id, - name, - Err::>(Box::new( - std::io::Error::new( - std::io::ErrorKind::Other, - "Could not parse arguments as JSON", - ), - )), - ), - } - }); - } if done { break; } } - if assistant_content.is_some() || assistant_thinking.is_some() { - let mut parts: Vec = Vec::new(); - if let Some(thinking) = assistant_thinking.take() { - parts.push(AssistantPart::Thinking { text: thinking }); - } - if let Some(content) = assistant_content.take() { - parts.push(AssistantPart::Text { text: content }); - } + if let Some(part) = current_part.take() { + parts.push(part); + } + if !parts.is_empty() { chat_history .lock() .unwrap() @@ -280,30 +264,25 @@ mod tests { .await .unwrap(); let updated = history.lock().unwrap().clone(); - // Behavior: assistant preamble content, then separate tool-call message, tool result, and final assistant response - assert_eq!(updated.len(), 5); - // First assistant message should contain the preamble content with no tool calls - let preamble_msg = &updated[1]; - if let ChatMessage::Assistant(a) = preamble_msg { - assert_eq!(a.content.len(), 1); + // Behavior: assistant content and tool call are combined in one assistant message, + // followed by the tool result message and the final assistant response + assert_eq!(updated.len(), 4); + // First assistant message should contain the preamble content followed by the tool call + let first_assistant = &updated[1]; + if let ChatMessage::Assistant(a) = first_assistant { + assert_eq!(a.content.len(), 2); match &a.content[0] { AssistantPart::Text { text } => assert_eq!(text, "first"), - _ => panic!("expected text part"), + _ => panic!("expected first part to be text"), } - } else { - panic!("expected assistant preamble message"); - } - // Next assistant message should carry the tool call - let call_msg = &updated[2]; - if let ChatMessage::Assistant(a) = call_msg { - assert_eq!(a.content.len(), 1); - match &a.content[0] { + match &a.content[1] { AssistantPart::ToolCall(tc) => assert_eq!(tc.name, "test"), - _ => panic!("expected tool call part"), + _ => panic!("expected second part to be tool call"), } } else { - panic!("expected assistant message with tool call"); + panic!("expected combined assistant message"); } + // Final assistant message should be the follow-up content let final_msg = updated.last().unwrap(); if let ChatMessage::Assistant(a) = final_msg { assert_eq!(a.content.len(), 1);