From 71fbe88a300319351c5591acb2e99b13b7dd714b Mon Sep 17 00:00:00 2001 From: drbh Date: Mon, 7 Jul 2025 14:35:04 +0000 Subject: [PATCH 1/2] fix: enable defs references in tool calls --- .../test_tool_def/test_flash_gemma3_defs.json | 36 +++++++++++++ integration-tests/models/test_tool_def.py | 51 +++++++++++++++++++ router/src/infer/tool_grammar.rs | 28 +++++++++- 3 files changed, 114 insertions(+), 1 deletion(-) create mode 100644 integration-tests/models/__snapshots__/test_tool_def/test_flash_gemma3_defs.json create mode 100644 integration-tests/models/test_tool_def.py diff --git a/integration-tests/models/__snapshots__/test_tool_def/test_flash_gemma3_defs.json b/integration-tests/models/__snapshots__/test_tool_def/test_flash_gemma3_defs.json new file mode 100644 index 00000000000..827730dd57f --- /dev/null +++ b/integration-tests/models/__snapshots__/test_tool_def/test_flash_gemma3_defs.json @@ -0,0 +1,36 @@ +{ + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "logprobs": null, + "message": { + "content": null, + "name": null, + "role": "assistant", + "tool_calls": [ + { + "function": { + "arguments": "{\"weather\":\"sunny\"}", + "description": null, + "name": "classify_weather" + }, + "id": "0", + "type": "function" + } + ] + }, + "usage": null + } + ], + "created": 1751896977, + "id": "", + "model": "google/gemma-3-4b-it", + "object": "chat.completion", + "system_fingerprint": "3.3.4-dev0-native", + "usage": { + "completion_tokens": 20, + "prompt_tokens": 196, + "total_tokens": 216 + } +} diff --git a/integration-tests/models/test_tool_def.py b/integration-tests/models/test_tool_def.py new file mode 100644 index 00000000000..e8fee2580c4 --- /dev/null +++ b/integration-tests/models/test_tool_def.py @@ -0,0 +1,51 @@ +import pytest + + +@pytest.fixture(scope="module") +def flash_gemma3_handle(launcher): + with launcher("google/gemma-3-4b-it", num_shard=2) as handle: + yield handle + + +@pytest.fixture(scope="module") +async def flash_gemma3(flash_gemma3_handle): + await flash_gemma3_handle.health(300) + return flash_gemma3_handle.client + + +async def test_flash_gemma3_defs(flash_gemma3, response_snapshot): + response = await flash_gemma3.chat( + messages=[ + { + "content": "Classify the weather: It's sunny outside with clear skies", + "role": "user", + } + ], + tools=[ + { + "type": "function", + "function": { + "name": "classify_weather", + "description": "Classify weather conditions", + "parameters": { + "$defs": { + "WeatherType": { + "enum": ["sunny", "cloudy", "rainy", "snowy"], + "type": "string", + } + }, + "properties": {"weather": {"$ref": "#/$defs/WeatherType"}}, + "required": ["weather"], + "type": "object", + }, + }, + } + ], + tool_choice="auto", + max_tokens=100, + seed=42, + ) + + assert response.choices[0].message.tool_calls[0]["function"]["name"] == "classify_weather" + assert response.choices[0].message.tool_calls[0]["function"]["arguments"] == '{"weather":"sunny"}' + assert response == response_snapshot diff --git a/router/src/infer/tool_grammar.rs b/router/src/infer/tool_grammar.rs index e4e2085983f..a9072030723 100644 --- a/router/src/infer/tool_grammar.rs +++ b/router/src/infer/tool_grammar.rs @@ -72,6 +72,7 @@ impl ToolGrammar { Value::String(func.description.unwrap_or_default()), ); + let mut defs = Map::new(); let mut properties = Map::new(); let mut required = vec![Value::String("_name".to_string())]; @@ -85,11 +86,35 @@ impl ToolGrammar { if let Value::Object(args) = func.arguments { if let Some(Value::Object(props)) = args.get("properties") { - properties.extend(props.clone()); + let mut updated_props = Map::new(); + // Update $ref paths in properties by iterating through + for (key, value) in props.iter() { + let updated_value = match value { + Value::Object(obj) if obj.contains_key("$ref") => { + let mut new_obj = obj.clone(); + if let Some(Value::String(ref_str)) = new_obj.get("$ref") { + if ref_str.starts_with("#/$defs/") { + // Replace $defs with $functions/{func.name}/$defs to handle + // function-specific definitions + new_obj.insert("$ref".to_string(), Value::String( + ref_str.replace("#/$defs/", &format!("#/$functions/{}/$defs/", func.name)) + )); + } + } + Value::Object(new_obj) + } + _ => value.clone(), + }; + updated_props.insert(key.clone(), updated_value); + } + properties.extend(updated_props); } if let Some(Value::Array(reqs)) = args.get("required") { required.extend(reqs.clone()); } + if let Some(Value::Object(definitions)) = args.get("$defs") { + defs.extend(definitions.clone()); + } params.insert( "additionalProperties".to_string(), Value::Bool( @@ -101,6 +126,7 @@ impl ToolGrammar { params.insert("properties".to_string(), Value::Object(properties)); params.insert("required".to_string(), Value::Array(required)); + params.insert("$defs".to_string(), Value::Object(defs)); (func.name, Value::Object(params)) }) From b6540cea50dd076f8f378feddeee93044655bf34 Mon Sep 17 00:00:00 2001 From: drbh Date: Mon, 7 Jul 2025 16:14:00 +0000 Subject: [PATCH 2/2] fix: lint and format --- integration-tests/models/test_tool_def.py | 10 ++++++++-- router/src/infer/tool_grammar.rs | 10 +++++++--- 2 files changed, 15 insertions(+), 5 deletions(-) diff --git a/integration-tests/models/test_tool_def.py b/integration-tests/models/test_tool_def.py index e8fee2580c4..2aa6c82e4da 100644 --- a/integration-tests/models/test_tool_def.py +++ b/integration-tests/models/test_tool_def.py @@ -46,6 +46,12 @@ async def test_flash_gemma3_defs(flash_gemma3, response_snapshot): seed=42, ) - assert response.choices[0].message.tool_calls[0]["function"]["name"] == "classify_weather" - assert response.choices[0].message.tool_calls[0]["function"]["arguments"] == '{"weather":"sunny"}' + assert ( + response.choices[0].message.tool_calls[0]["function"]["name"] + == "classify_weather" + ) + assert ( + response.choices[0].message.tool_calls[0]["function"]["arguments"] + == '{"weather":"sunny"}' + ) assert response == response_snapshot diff --git a/router/src/infer/tool_grammar.rs b/router/src/infer/tool_grammar.rs index a9072030723..eea2f55ec2e 100644 --- a/router/src/infer/tool_grammar.rs +++ b/router/src/infer/tool_grammar.rs @@ -96,9 +96,13 @@ impl ToolGrammar { if ref_str.starts_with("#/$defs/") { // Replace $defs with $functions/{func.name}/$defs to handle // function-specific definitions - new_obj.insert("$ref".to_string(), Value::String( - ref_str.replace("#/$defs/", &format!("#/$functions/{}/$defs/", func.name)) - )); + new_obj.insert( + "$ref".to_string(), + Value::String(ref_str.replace( + "#/$defs/", + &format!("#/$functions/{}/$defs/", func.name), + )), + ); } } Value::Object(new_obj)