alex/g3
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

mild json filtering improvement

Browse Source
This commit is contained in:
Dhanji Prasanna
2025-10-31 15:19:33 +11:00
parent e59ce2f93f
commit 57d473c19d
4 changed files with 248 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files

223
crates/g3-core/src/fixed_filter_json.rs
View File

@@ -4,6 +4,11 @@
// 3. Only elide JSON content between first '{' and last '}' (inclusive)
// 4. Return everything else as the final filtered string
//! JSON tool call filtering for streaming LLM responses.
//!
//! This module filters out JSON tool calls from LLM output streams while preserving
//! regular text content. It uses a state machine to handle streaming chunks.
use regex::Regex;
use std::cell::RefCell;
use tracing::debug;
@@ -13,37 +18,51 @@ thread_local! {
static FIXED_JSON_TOOL_STATE: RefCell<FixedJsonToolState> = RefCell::new(FixedJsonToolState::new());
}
/// Internal state for tracking JSON tool call filtering across streaming chunks.
#[derive(Debug, Clone)]
struct FixedJsonToolState {
/// True when actively suppressing a confirmed tool call
suppression_mode: bool,
/// True when buffering potential JSON (saw { but not yet confirmed as tool call)
potential_json_mode: bool,
/// Tracks nesting depth of braces within JSON
brace_depth: i32,
buffer: String,
json_start_in_buffer: Option<usize>,
json_start_in_buffer: Option<usize>, // Position where confirmed JSON tool call starts
content_returned_up_to: usize, // Track how much content we've already returned
potential_json_start: Option<usize>, // Where the potential JSON started
}
impl FixedJsonToolState {
fn new() -> Self {
Self {
suppression_mode: false,
potential_json_mode: false,
brace_depth: 0,
buffer: String::new(),
json_start_in_buffer: None,
content_returned_up_to: 0,
potential_json_start: None,
}
}
fn reset(&mut self) {
self.suppression_mode = false;
self.potential_json_mode = false;
self.brace_depth = 0;
self.buffer.clear();
self.json_start_in_buffer = None;
self.content_returned_up_to = 0;
self.potential_json_start = None;
}
}
// FINAL CORRECTED implementation according to specification
/// Filters JSON tool calls from streaming LLM content.
///
/// Processes content chunks and removes JSON tool calls while preserving regular text.
/// Maintains state across calls to handle tool calls spanning multiple chunks.
pub fn fixed_filter_json_tool_calls(content: &str) -> String {
if content.is_empty() {
return String::new();
@@ -91,9 +110,187 @@ pub fn fixed_filter_json_tool_calls(content: &str) -> String {
return String::new();
}
// Check if we're in potential JSON mode (saw { but waiting to confirm it's a tool call)
if state.potential_json_mode {
// Check if the buffer contains a confirmed tool call pattern
let tool_call_regex = Regex::new(r#"(?m)^\s*\{\s*"tool"\s*:\s*""#).unwrap();
if let Some(captures) = tool_call_regex.find(&state.buffer) {
// Confirmed! This is a tool call - enter suppression mode
let match_text = captures.as_str();
if let Some(brace_offset) = match_text.find('{') {
let json_start = captures.start() + brace_offset;
debug!("Confirmed JSON tool call at position {} - entering suppression mode", json_start);
state.potential_json_mode = false;
state.suppression_mode = true;
state.brace_depth = 0;
state.json_start_in_buffer = Some(json_start);
// Count braces from json_start to see if JSON is complete
let buffer_slice = state.buffer[json_start..].to_string();
for ch in buffer_slice.chars() {
match ch {
'{' => state.brace_depth += 1,
'}' => {
state.brace_depth -= 1;
if state.brace_depth <= 0 {
debug!("JSON tool call completed immediately");
let result = extract_fixed_content(&state.buffer, json_start);
let new_content = if result.len() > state.content_returned_up_to {
result[state.content_returned_up_to..].to_string()
} else {
String::new()
};
state.reset();
return new_content;
}
}
_ => {}
}
}
// JSON incomplete, stay in suppression mode, return nothing
return String::new();
}
}
// Check if we can rule out this being a tool call
// If we have enough content after the { and it doesn't match the pattern, release it
if let Some(potential_start) = state.potential_json_start {
let content_after_brace = &state.buffer[potential_start..];
// Rule out as a tool call if:
// 1. Closing } appears before we see the full pattern
// 2. Content clearly doesn't match the tool call pattern
// 3. Newline appears after the opening brace (tool calls should be compact)
let has_closing_brace = content_after_brace.contains('}');
let has_newline = content_after_brace[1..].contains('\n'); // Skip first char which is {
let long_enough = content_after_brace.len() >= 10;
// Detect non-tool JSON patterns:
// - { followed by " and a key that doesn't start with "tool"
// - { followed by "t" but not "to"
// - { followed by "to" but not "too", etc.
let not_tool_pattern = Regex::new(r#"^\{\s*"(?:[^t]|t(?:[^o]|o(?:[^o]|o(?:[^l]|l[^"\s:]))))"#).unwrap();
let definitely_not_tool = not_tool_pattern.is_match(content_after_brace);
if has_closing_brace || has_newline || (long_enough && definitely_not_tool) {
debug!("Potential JSON ruled out - not a tool call");
state.potential_json_mode = false;
state.potential_json_start = None;
// Return the buffered content we've been holding
let new_content = if state.buffer.len() > state.content_returned_up_to {
state.buffer[state.content_returned_up_to..].to_string()
} else {
String::new()
};
state.content_returned_up_to = state.buffer.len();
return new_content;
}
}
// Still in potential mode, keep buffering
return String::new();
}
// Detect potential JSON start: { at the beginning of a line
let potential_json_regex = Regex::new(r"(?m)^\s*\{\s*").unwrap();
if let Some(captures) = potential_json_regex.find(&state.buffer[state.content_returned_up_to..]) {
let match_start = state.content_returned_up_to + captures.start();
let brace_pos = match_start + captures.as_str().find('{').unwrap();
debug!("Potential JSON detected at position {} - entering buffering mode", brace_pos);
// Fast path: check if this is already a confirmed tool call
let tool_call_regex = Regex::new(r#"(?m)^\s*\{\s*"tool"\s*:\s*""#).unwrap();
if tool_call_regex.is_match(&state.buffer[brace_pos..]) {
// This is a confirmed tool call! Process it immediately
let json_start = brace_pos;
debug!("Immediately confirmed tool call at position {}", json_start);
// Return content before JSON
let content_before = if json_start > state.content_returned_up_to {
state.buffer[state.content_returned_up_to..json_start].to_string()
} else {
String::new()
};
state.content_returned_up_to = json_start;
state.suppression_mode = true;
state.brace_depth = 0;
state.json_start_in_buffer = Some(json_start);
// Count braces to see if JSON is complete
let buffer_slice = state.buffer[json_start..].to_string();
for ch in buffer_slice.chars() {
match ch {
'{' => state.brace_depth += 1,
'}' => {
state.brace_depth -= 1;
if state.brace_depth <= 0 {
debug!("JSON tool call completed in same chunk");
let result = extract_fixed_content(&state.buffer, json_start);
let content_after = if result.len() > json_start {
&result[json_start..]
} else {
""
};
let final_result = format!("{}{}", content_before, content_after);
state.reset();
return final_result;
}
}
_ => {}
}
}
// JSON incomplete, return content before and stay in suppression mode
return content_before;
}
// Return content before the potential JSON
let content_before = if brace_pos > state.content_returned_up_to {
state.buffer[state.content_returned_up_to..brace_pos].to_string()
} else {
String::new()
};
state.content_returned_up_to = brace_pos;
state.potential_json_mode = true;
state.potential_json_start = Some(brace_pos);
// Optimization: immediately check if we can rule this out for single-chunk processing
let content_after_brace = &state.buffer[brace_pos..];
let has_closing_brace = content_after_brace.contains('}');
let has_newline = content_after_brace.len() > 1 && content_after_brace[1..].contains('\n');
let long_enough = content_after_brace.len() >= 10;
let not_tool_pattern = Regex::new(r#"^\{\s*"(?:[^t]|t(?:[^o]|o(?:[^o]|o(?:[^l]|l[^"\s:]))))"#).unwrap();
let definitely_not_tool = not_tool_pattern.is_match(content_after_brace);
if has_closing_brace || has_newline || (long_enough && definitely_not_tool) {
debug!("Immediately ruled out as not a tool call");
state.potential_json_mode = false;
state.potential_json_start = None;
// Return all the buffered content
let new_content = if state.buffer.len() > state.content_returned_up_to {
state.buffer[state.content_returned_up_to..].to_string()
} else {
String::new()
};
state.content_returned_up_to = state.buffer.len();
return format!("{}{}", content_before, new_content);
}
return content_before;
}
// Check for tool call pattern using corrected regex
// More flexible than the strict specification to handle real-world JSON
let tool_call_regex = Regex::new(r#"(?m)^\s*\{\s*"tool"\s*:\s*""#).unwrap();
let tool_call_regex = Regex::new(r#"(?m)^\s*\{\s*"tool"\s*:\s*"[^"]*""#).unwrap();
if let Some(captures) = tool_call_regex.find(&state.buffer) {
let match_text = captures.as_str();
@@ -168,9 +365,17 @@ pub fn fixed_filter_json_tool_calls(content: &str) -> String {
})
}
// Helper function to extract content with JSON tool call filtered out
// Returns everything except the JSON between the first '{' and last '}' (inclusive)
/// Extracts content from buffer, removing the JSON tool call.
///
/// Given a buffer and the start position of a JSON tool call, this function:
/// 1. Extracts all content before the JSON
/// 2. Finds the end of the JSON (matching closing brace)
/// 3. Extracts all content after the JSON
/// 4. Returns the concatenation of before + after (JSON removed)
///
/// # Arguments
/// * `full_content` - The full content buffer
/// * `json_start` - Position where the JSON tool call begins
fn extract_fixed_content(full_content: &str, json_start: usize) -> String {
// Find the end of the JSON using proper brace counting with string handling
let mut brace_depth = 0;
@@ -212,8 +417,10 @@ fn extract_fixed_content(full_content: &str, json_start: usize) -> String {
format!("{}{}", before, after)
}
// Reset function for testing
/// Resets the global JSON filtering state.
///
/// Call this between independent filtering sessions to ensure clean state.
/// This is particularly important in tests and when starting new conversations.
pub fn reset_fixed_json_tool_state() {
FIXED_JSON_TOOL_STATE.with(|state| {
let mut state = state.borrow_mut();

42
crates/g3-core/src/fixed_filter_tests.rs
View File

@@ -1,8 +1,14 @@
//! Tests for JSON tool call filtering.
//!
//! These tests verify that the filter correctly identifies and removes JSON tool calls
//! from LLM output streams while preserving all other content.
#[cfg(test)]
mod fixed_filter_tests {
use crate::fixed_filter_json::{fixed_filter_json_tool_calls, reset_fixed_json_tool_state};
use regex::Regex;
/// Test that regular text without tool calls passes through unchanged.
#[test]
fn test_no_tool_call_passthrough() {
reset_fixed_json_tool_state();
@@ -11,6 +17,7 @@ mod fixed_filter_tests {
assert_eq!(result, input);
}
/// Test detection and removal of a complete tool call in a single chunk.
#[test]
fn test_simple_tool_call_detection() {
reset_fixed_json_tool_state();
@@ -23,6 +30,7 @@ Some text after"#;
assert_eq!(result, expected);
}
/// Test handling of tool calls that arrive across multiple streaming chunks.
#[test]
fn test_streaming_chunks() {
reset_fixed_json_tool_state();
@@ -48,6 +56,7 @@ Some text after"#;
assert_eq!(final_result, expected);
}
/// Test correct handling of nested braces within JSON strings.
#[test]
fn test_nested_braces_in_tool_call() {
reset_fixed_json_tool_state();
@@ -61,6 +70,7 @@ Text after"#;
assert_eq!(result, expected);
}
/// Verify the regex pattern matches the specification with flexible whitespace.
#[test]
fn test_regex_pattern_specification() {
// Test the corrected regex pattern that's more flexible with whitespace
@@ -84,11 +94,6 @@ Text after"#;
), // Space after { DOES match with \s*
(
r#"line
abc{"tool":"#,
true,
),
(
r#"line
{"tool123":"#,
false,
), // "tool123" is not exactly "tool"
@@ -109,6 +114,7 @@ abc{"tool":"#,
}
}
/// Test that tool calls must appear at the start of a line (after newline).
#[test]
fn test_newline_requirement() {
reset_fixed_json_tool_state();
@@ -122,13 +128,14 @@ abc{"tool":"#,
reset_fixed_json_tool_state();
let result2 = fixed_filter_json_tool_calls(input_without_newline);
// Both cases currently trigger suppression due to regex pattern
// TODO: Fix regex to only match after actual newlines
// With the new aggressive filtering, only the newline case should trigger suppression
// The pattern requires { to be at the start of a line (after ^)
assert_eq!(result1, "Text\n");
// This currently fails because our regex matches both cases
assert_eq!(result2, "Text ");
// Without newline before {, it should pass through unchanged
assert_eq!(result2, input_without_newline);
}
/// Test handling of escaped quotes within JSON strings.
#[test]
fn test_json_with_escaped_quotes() {
reset_fixed_json_tool_state();
@@ -142,6 +149,7 @@ More text"#;
assert_eq!(result, expected);
}
/// Test graceful handling of incomplete/malformed JSON.
#[test]
fn test_edge_case_malformed_json() {
reset_fixed_json_tool_state();
@@ -157,6 +165,7 @@ More text"#;
assert_eq!(result, expected);
}
/// Test processing multiple independent tool calls sequentially.
#[test]
fn test_multiple_tool_calls_sequential() {
reset_fixed_json_tool_state();
@@ -179,6 +188,7 @@ Final text"#;
assert_eq!(result2, expected2);
}
/// Test tool calls with complex multi-line arguments.
#[test]
fn test_tool_call_with_complex_args() {
reset_fixed_json_tool_state();
@@ -192,6 +202,7 @@ After"#;
assert_eq!(result, expected);
}
/// Test input containing only a tool call with no surrounding text.
#[test]
fn test_tool_call_only() {
reset_fixed_json_tool_state();
@@ -204,6 +215,7 @@ After"#;
assert_eq!(result, expected);
}
/// Test accurate brace counting with deeply nested structures.
#[test]
fn test_brace_counting_accuracy() {
reset_fixed_json_tool_state();
@@ -218,6 +230,7 @@ End"#;
assert_eq!(result, expected);
}
/// Test that braces within strings don't affect brace counting.
#[test]
fn test_string_escaping_in_json() {
reset_fixed_json_tool_state();
@@ -232,6 +245,7 @@ More"#;
assert_eq!(result, expected);
}
/// Verify compliance with the exact specification requirements.
#[test]
fn test_specification_compliance() {
reset_fixed_json_tool_state();
@@ -248,6 +262,7 @@ More"#;
assert_eq!(result, expected);
}
/// Test that non-tool JSON objects are not filtered.
#[test]
fn test_no_false_positives() {
reset_fixed_json_tool_state();
@@ -261,6 +276,7 @@ More text"#;
assert_eq!(result, input);
}
/// Test patterns that look similar to tool calls but aren't exact matches.
#[test]
fn test_partial_tool_patterns() {
reset_fixed_json_tool_state();
@@ -280,6 +296,7 @@ More text"#;
}
}
/// Test streaming with very small chunks (character-by-character).
#[test]
fn test_streaming_edge_cases() {
reset_fixed_json_tool_state();
@@ -296,12 +313,13 @@ More text"#;
}
let final_result: String = results.join("");
// This test currently fails because the JSON is incomplete across chunks
// The function doesn't handle this edge case properly yet
let expected = "Text\n{\"tool\": \nAfter";
// With the new aggressive filtering, the JSON should be completely filtered out
// even when it arrives in very small chunks
let expected = "Text\n\nAfter";
assert_eq!(final_result, expected);
}
/// Debug test with detailed logging for streaming behavior.
#[test]
fn test_streaming_debug() {
reset_fixed_json_tool_state();

21
crates/g3-core/src/lib.rs
View File

@@ -786,7 +786,6 @@ impl<W: UiWriter> Agent<W> {
// Register embedded provider if configured AND it's the default provider
if let Some(embedded_config) = &config.providers.embedded {
if providers_to_register.contains(&"embedded".to_string()) {
info!("Initializing embedded provider");
let embedded_provider = g3_providers::EmbeddedProvider::new(
embedded_config.model_path.clone(),
embedded_config.model_type.clone(),
@@ -797,15 +796,12 @@ impl<W: UiWriter> Agent<W> {
embedded_config.threads,
)?;
providers.register(embedded_provider);
} else {
info!("Embedded provider configured but not needed, skipping initialization");
}
}
// Register OpenAI provider if configured AND it's the default provider
if let Some(openai_config) = &config.providers.openai {
if providers_to_register.contains(&"openai".to_string()) {
info!("Initializing OpenAI provider");
let openai_provider = g3_providers::OpenAIProvider::new(
openai_config.api_key.clone(),
Some(openai_config.model.clone()),
@@ -814,15 +810,12 @@ impl<W: UiWriter> Agent<W> {
openai_config.temperature,
)?;
providers.register(openai_provider);
} else {
info!("OpenAI provider configured but not needed, skipping initialization");
}
}
// Register Anthropic provider if configured AND it's the default provider
if let Some(anthropic_config) = &config.providers.anthropic {
if providers_to_register.contains(&"anthropic".to_string()) {
info!("Initializing Anthropic provider");
let anthropic_provider = g3_providers::AnthropicProvider::new(
anthropic_config.api_key.clone(),
Some(anthropic_config.model.clone()),
@@ -830,15 +823,12 @@ impl<W: UiWriter> Agent<W> {
anthropic_config.temperature,
)?;
providers.register(anthropic_provider);
} else {
info!("Anthropic provider configured but not needed, skipping initialization");
}
}
// Register Databricks provider if configured AND it's the default provider
if let Some(databricks_config) = &config.providers.databricks {
if providers_to_register.contains(&"databricks".to_string()) {
info!("Initializing Databricks provider");
let databricks_provider = if let Some(token) = &databricks_config.token {
// Use token-based authentication
@@ -861,8 +851,6 @@ impl<W: UiWriter> Agent<W> {
};
providers.register(databricks_provider);
} else {
info!("Databricks provider configured but not needed, skipping initialization");
}
}
@@ -885,16 +873,12 @@ impl<W: UiWriter> Agent<W> {
content: readme,
};
context_window.add_message(readme_message);
info!("Added project README to context window");
}
// Initialize computer controller if enabled
let computer_controller = if config.computer_control.enabled {
match g3_computer_control::create_controller() {
Ok(controller) => {
info!("Computer control enabled");
Some(controller)
}
Ok(controller) => Some(controller),
Err(e) => {
warn!("Failed to initialize computer control: {}", e);
None
@@ -2991,7 +2975,8 @@ Template:
"Using filtered parser text as last resort: {} chars",
filtered_text.len()
);
current_response = filtered_text;
// Note: This assignment is currently unused but kept for potential future use
let _ = filtered_text;
}
}