Unverified Commit 9768c50d authored by Simo Lin's avatar Simo Lin Committed by GitHub
Browse files

[router] restructure tool parser module folder (#9693)

parent fd71b11b
/// Tool parser module for handling function/tool calls in model outputs /// Tool parser module for handling function/tool calls in model outputs
/// ///
/// This module provides infrastructure for parsing tool calls from various model formats. /// This module provides infrastructure for parsing tool calls from various model formats.
// Core modules
pub mod errors; pub mod errors;
pub mod json_parser;
pub mod llama_parser;
pub mod mistral_parser;
pub mod partial_json; pub mod partial_json;
pub mod python_literal_parser; pub mod python_literal_parser;
pub mod pythonic_parser;
pub mod qwen_parser;
pub mod registry; pub mod registry;
pub mod state; pub mod state;
pub mod traits; pub mod traits;
pub mod types; pub mod types;
// Parser implementations
pub mod parsers;
#[cfg(test)] #[cfg(test)]
mod tests; mod tests;
// Re-export commonly used types // Re-export commonly used types
pub use errors::{ToolParserError, ToolParserResult}; pub use errors::{ToolParserError, ToolParserResult};
pub use json_parser::JsonParser;
pub use llama_parser::LlamaParser;
pub use mistral_parser::MistralParser;
pub use pythonic_parser::PythonicParser;
pub use qwen_parser::QwenParser;
pub use registry::ParserRegistry; pub use registry::ParserRegistry;
pub use state::{ParsePhase, ParseState}; pub use state::{ParsePhase, ParseState};
pub use traits::{PartialJsonParser, ToolParser}; pub use traits::{PartialJsonParser, ToolParser};
pub use types::{FunctionCall, PartialToolCall, StreamResult, TokenConfig, ToolCall}; pub use types::{FunctionCall, PartialToolCall, StreamResult, TokenConfig, ToolCall};
// Re-export parsers for convenience
pub use parsers::{JsonParser, LlamaParser, MistralParser, PythonicParser, QwenParser};
...@@ -291,7 +291,8 @@ impl ToolParser for JsonParser { ...@@ -291,7 +291,8 @@ impl ToolParser for JsonParser {
// Clear buffer since we consumed everything // Clear buffer since we consumed everything
state.buffer.clear(); state.buffer.clear();
// Return the first tool as complete (simplified for Phase 2) // Return the first tool as complete
// TODO simplified version, address more complex version
if let Some(tool) = tools.into_iter().next() { if let Some(tool) = tools.into_iter().next() {
return Ok(StreamResult::ToolComplete(tool)); return Ok(StreamResult::ToolComplete(tool));
} }
...@@ -299,7 +300,7 @@ impl ToolParser for JsonParser { ...@@ -299,7 +300,7 @@ impl ToolParser for JsonParser {
} else { } else {
// Partial JSON, try to extract tool name // Partial JSON, try to extract tool name
if let Some(name) = value.get("name").and_then(|v| v.as_str()) { if let Some(name) = value.get("name").and_then(|v| v.as_str()) {
// Simple implementation for Phase 2 // TODO simplified version, address more complex version
// Just return the tool name once we see it // Just return the tool name once we see it
if !state.in_string { if !state.in_string {
state.in_string = true; // Use as a flag for "name sent" state.in_string = true; // Use as a flag for "name sent"
...@@ -430,12 +431,12 @@ mod tests { ...@@ -430,12 +431,12 @@ mod tests {
#[tokio::test] #[tokio::test]
async fn test_streaming_parse() { async fn test_streaming_parse() {
// Phase 2 simplified streaming test
// Just verify that streaming eventually produces a complete tool call // Just verify that streaming eventually produces a complete tool call
let parser = JsonParser::new(); let parser = JsonParser::new();
let mut state = ParseState::new(); let mut state = ParseState::new();
// Send complete JSON in one go (simplified for Phase 2) // Send complete JSON in one go
// TODO simplified version, address more complex version
let full_json = r#"{"name": "get_weather", "arguments": {"location": "SF"}}"#; let full_json = r#"{"name": "get_weather", "arguments": {"location": "SF"}}"#;
let result = parser let result = parser
......
use async_trait::async_trait; use async_trait::async_trait;
use super::json_parser::JsonParser;
use crate::tool_parser::{ use crate::tool_parser::{
errors::ToolParserResult, errors::ToolParserResult,
json_parser::JsonParser,
state::ParseState, state::ParseState,
traits::ToolParser, traits::ToolParser,
types::{StreamResult, TokenConfig, ToolCall}, types::{StreamResult, TokenConfig, ToolCall},
......
/// Parser implementations for different model formats
///
/// This module contains concrete parser implementations for various model-specific
/// tool/function call formats.
// Individual parser modules
pub mod json_parser;
pub mod llama_parser;
pub mod mistral_parser;
pub mod pythonic_parser;
pub mod qwen_parser;
pub use json_parser::JsonParser;
pub use llama_parser::LlamaParser;
pub use mistral_parser::MistralParser;
pub use pythonic_parser::PythonicParser;
pub use qwen_parser::QwenParser;
use crate::tool_parser::json_parser::JsonParser; use crate::tool_parser::parsers::{
use crate::tool_parser::llama_parser::LlamaParser; JsonParser, LlamaParser, MistralParser, PythonicParser, QwenParser,
use crate::tool_parser::mistral_parser::MistralParser; };
use crate::tool_parser::pythonic_parser::PythonicParser;
use crate::tool_parser::qwen_parser::QwenParser;
use crate::tool_parser::traits::ToolParser; use crate::tool_parser::traits::ToolParser;
use std::collections::HashMap; use std::collections::HashMap;
use std::sync::Arc; use std::sync::Arc;
......
use super::*; use super::*;
use crate::tool_parser::json_parser::JsonParser; use crate::tool_parser::parsers::JsonParser;
use crate::tool_parser::partial_json::{ use crate::tool_parser::partial_json::{
compute_diff, find_common_prefix, is_complete_json, PartialJson, compute_diff, find_common_prefix, is_complete_json, PartialJson,
}; };
...@@ -387,11 +387,10 @@ fn test_json_parser_format_detection() { ...@@ -387,11 +387,10 @@ fn test_json_parser_format_detection() {
#[tokio::test] #[tokio::test]
async fn test_json_parser_streaming() { async fn test_json_parser_streaming() {
// Phase 2 simplified streaming test
let parser = JsonParser::new(); let parser = JsonParser::new();
let mut state = ParseState::new(); let mut state = ParseState::new();
// Test with complete JSON (simplified for Phase 2) // Test with complete JSON
let full_json = r#"{"name": "get_weather", "arguments": {"location": "San Francisco"}}"#; let full_json = r#"{"name": "get_weather", "arguments": {"location": "San Francisco"}}"#;
let result = parser let result = parser
...@@ -739,7 +738,7 @@ mod edge_cases { ...@@ -739,7 +738,7 @@ mod edge_cases {
_ => panic!("Expected ToolComplete for complete JSON"), _ => panic!("Expected ToolComplete for complete JSON"),
} }
// Test 3: Partial JSON with name - Phase 2 behavior // Test 3: Partial JSON with name
// The PartialJson parser can complete partial JSON by filling in missing values // The PartialJson parser can complete partial JSON by filling in missing values
let mut state3 = ParseState::new(); let mut state3 = ParseState::new();
let partial_with_name = r#"{"name": "test", "argum"#; let partial_with_name = r#"{"name": "test", "argum"#;
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment