615 lines
20 KiB
Go
615 lines
20 KiB
Go
package renderers
|
|
|
|
import (
|
|
"encoding/json"
|
|
"os"
|
|
"os/exec"
|
|
"path/filepath"
|
|
"runtime"
|
|
"strings"
|
|
"testing"
|
|
|
|
"github.com/google/go-cmp/cmp"
|
|
|
|
"github.com/ollama/ollama/api"
|
|
)
|
|
|
|
const nemotron3NanoTemplate = "testdata/nemotron3nano_chat_template.jinja2"
|
|
|
|
func TestNemotron3NanoRendererMatchesReference(t *testing.T) {
|
|
toolText := `<|im_start|>system
|
|
# Tools
|
|
|
|
You have access to the following functions:
|
|
|
|
<tools>
|
|
<function>
|
|
<name>search_docs</name>
|
|
<description>Search docs</description>
|
|
<parameters>
|
|
<parameter>
|
|
<name>query</name>
|
|
<type>string</type>
|
|
<description>Search query</description>
|
|
<enum>["api", "cli"]</enum>
|
|
</parameter>
|
|
<parameter>
|
|
<name>mode</name>
|
|
<type>['string', 'null']</type>
|
|
<description>Mode</description>
|
|
<anyOf>[{"type": "string"}, {"type": "number"}]</anyOf>
|
|
</parameter>
|
|
<parameter>
|
|
<name>payload</name>
|
|
<type>object</type>
|
|
<description>Payload</description>
|
|
<properties>{"enabled": {"type": "boolean"}}</properties>
|
|
<required>["enabled"]</required>
|
|
</parameter>
|
|
<parameter>
|
|
<name>tags</name>
|
|
<type>array</type>
|
|
<description>Tags</description>
|
|
<items>{"type": "string"}</items>
|
|
</parameter>
|
|
<$defs>{"shared": {"type": "string"}}</$defs>
|
|
<required>["query"]</required>
|
|
</parameters>
|
|
</function>
|
|
</tools>
|
|
|
|
If you choose to call a function ONLY reply in the following format with NO suffix:
|
|
|
|
<tool_call>
|
|
<function=example_function_name>
|
|
<parameter=example_parameter_1>
|
|
value_1
|
|
</parameter>
|
|
<parameter=example_parameter_2>
|
|
This is the value for the second parameter
|
|
that can span
|
|
multiple lines
|
|
</parameter>
|
|
</function>
|
|
</tool_call>
|
|
|
|
<IMPORTANT>
|
|
Reminder:
|
|
- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags
|
|
- Required parameters MUST be specified
|
|
- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after
|
|
- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls
|
|
</IMPORTANT><|im_end|>
|
|
`
|
|
toolTextWithSystem := `<|im_start|>system
|
|
Follow policy.
|
|
|
|
# Tools
|
|
|
|
You have access to the following functions:
|
|
|
|
<tools>
|
|
<function>
|
|
<name>search_docs</name>
|
|
<description>Search docs</description>
|
|
<parameters>
|
|
<parameter>
|
|
<name>query</name>
|
|
<type>string</type>
|
|
<description>Search query</description>
|
|
<enum>["api", "cli"]</enum>
|
|
</parameter>
|
|
<parameter>
|
|
<name>mode</name>
|
|
<type>['string', 'null']</type>
|
|
<description>Mode</description>
|
|
<anyOf>[{"type": "string"}, {"type": "number"}]</anyOf>
|
|
</parameter>
|
|
<parameter>
|
|
<name>payload</name>
|
|
<type>object</type>
|
|
<description>Payload</description>
|
|
<properties>{"enabled": {"type": "boolean"}}</properties>
|
|
<required>["enabled"]</required>
|
|
</parameter>
|
|
<parameter>
|
|
<name>tags</name>
|
|
<type>array</type>
|
|
<description>Tags</description>
|
|
<items>{"type": "string"}</items>
|
|
</parameter>
|
|
<$defs>{"shared": {"type": "string"}}</$defs>
|
|
<required>["query"]</required>
|
|
</parameters>
|
|
</function>
|
|
</tools>
|
|
|
|
If you choose to call a function ONLY reply in the following format with NO suffix:
|
|
|
|
<tool_call>
|
|
<function=example_function_name>
|
|
<parameter=example_parameter_1>
|
|
value_1
|
|
</parameter>
|
|
<parameter=example_parameter_2>
|
|
This is the value for the second parameter
|
|
that can span
|
|
multiple lines
|
|
</parameter>
|
|
</function>
|
|
</tool_call>
|
|
|
|
<IMPORTANT>
|
|
Reminder:
|
|
- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags
|
|
- Required parameters MUST be specified
|
|
- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after
|
|
- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls
|
|
</IMPORTANT><|im_end|>
|
|
`
|
|
|
|
tests := []struct {
|
|
name string
|
|
messages []api.Message
|
|
tools []api.Tool
|
|
think *api.ThinkValue
|
|
expected string
|
|
}{
|
|
{
|
|
name: "no system default thinking on",
|
|
messages: []api.Message{
|
|
{Role: "user", Content: "Hello"},
|
|
},
|
|
expected: "\n\n\n<|im_start|>system\n<|im_end|>\n\n<|im_start|>user\nHello<|im_end|>\n\n<|im_start|>assistant\n<think>\n",
|
|
},
|
|
{
|
|
name: "no system explicit thinking off",
|
|
messages: []api.Message{
|
|
{Role: "user", Content: "Hello"},
|
|
},
|
|
think: thinkFalse(),
|
|
expected: "\n\n\n<|im_start|>system\n<|im_end|>\n\n<|im_start|>user\nHello<|im_end|>\n\n<|im_start|>assistant\n<think></think>",
|
|
},
|
|
{
|
|
name: "literal endthink does not enable thinking",
|
|
messages: []api.Message{
|
|
{Role: "user", Content: "literal </think> only"},
|
|
},
|
|
think: thinkFalse(),
|
|
expected: "\n\n\n<|im_start|>system\n<|im_end|>\n\n<|im_start|>user\nliteral </think> only<|im_end|>\n\n<|im_start|>assistant\n<think></think>",
|
|
},
|
|
{
|
|
name: "user no think toggle",
|
|
messages: []api.Message{
|
|
{Role: "user", Content: "Hello /no_think"},
|
|
},
|
|
think: thinkTrue(),
|
|
expected: "\n\n\n<|im_start|>system\n<|im_end|>\n\n<|im_start|>user\nHello /no_think<|im_end|>\n\n<|im_start|>assistant\n<think></think>",
|
|
},
|
|
{
|
|
name: "system think toggle overrides false",
|
|
messages: []api.Message{
|
|
{Role: "system", Content: "Policy /think"},
|
|
{Role: "user", Content: "Hello"},
|
|
},
|
|
think: thinkFalse(),
|
|
expected: "\n\n\n<|im_start|>system\nPolicy <|im_end|>\n\n<|im_start|>user\nHello<|im_end|>\n\n<|im_start|>assistant\n<think>\n",
|
|
},
|
|
{
|
|
name: "later toggle wins",
|
|
messages: []api.Message{
|
|
{Role: "system", Content: "Policy /no_think"},
|
|
{Role: "user", Content: "Actually /think"},
|
|
},
|
|
think: thinkFalse(),
|
|
expected: "\n\n\n<|im_start|>system\nPolicy <|im_end|>\n\n<|im_start|>user\nActually /think<|im_end|>\n\n<|im_start|>assistant\n<think>\n",
|
|
},
|
|
{
|
|
name: "system sanitizes toggles but preserves closing tag",
|
|
messages: []api.Message{
|
|
{Role: "system", Content: "A /think B /no_think C </think>"},
|
|
{Role: "user", Content: "Hello"},
|
|
},
|
|
think: thinkFalse(),
|
|
expected: "\n\n\n<|im_start|>system\nA B C </think><|im_end|>\n\n<|im_start|>user\nHello<|im_end|>\n\n<|im_start|>assistant\n<think>\n",
|
|
},
|
|
{
|
|
name: "assistant plain content adds empty think block",
|
|
messages: []api.Message{
|
|
{Role: "user", Content: "Hi"},
|
|
{Role: "assistant", Content: "Hello there"},
|
|
},
|
|
think: thinkTrue(),
|
|
expected: "\n\n\n<|im_start|>system\n<|im_end|>\n\n<|im_start|>user\nHi<|im_end|>\n<|im_start|>assistant\n<think></think>Hello there<|im_end|>\n\n<|im_start|>assistant\n<think>\n",
|
|
},
|
|
{
|
|
name: "assistant reasoning content",
|
|
messages: []api.Message{
|
|
{Role: "user", Content: "Hi"},
|
|
{Role: "assistant", Content: "Answer", Thinking: "Need to think"},
|
|
},
|
|
think: thinkTrue(),
|
|
expected: "\n\n\n<|im_start|>system\n<|im_end|>\n\n<|im_start|>user\nHi<|im_end|>\n<|im_start|>assistant\n<think>\nNeed to think\n</think>\nAnswer<|im_end|>\n\n<|im_start|>assistant\n<think>\n",
|
|
},
|
|
{
|
|
name: "assistant preserves existing think tags",
|
|
messages: []api.Message{
|
|
{Role: "user", Content: "Hi"},
|
|
{Role: "assistant", Content: "<think>kept</think>Answer"},
|
|
},
|
|
think: thinkTrue(),
|
|
expected: "\n\n\n<|im_start|>system\n<|im_end|>\n\n<|im_start|>user\nHi<|im_end|>\n<|im_start|>assistant\n<think>kept</think>Answer<|im_end|>\n\n<|im_start|>assistant\n<think>\n",
|
|
},
|
|
{
|
|
name: "tools without system",
|
|
messages: []api.Message{
|
|
{Role: "user", Content: "Use a tool"},
|
|
},
|
|
tools: nemotron3NanoReferenceTools(),
|
|
think: thinkTrue(),
|
|
expected: "\n\n\n" + toolText + "\n<|im_start|>user\nUse a tool<|im_end|>\n\n<|im_start|>assistant\n<think>\n",
|
|
},
|
|
{
|
|
name: "system with tools",
|
|
messages: []api.Message{
|
|
{Role: "system", Content: "Follow policy."},
|
|
{Role: "user", Content: "Use a tool"},
|
|
},
|
|
tools: nemotron3NanoReferenceTools(),
|
|
think: thinkTrue(),
|
|
expected: "\n\n\n" + toolTextWithSystem + "\n<|im_start|>user\nUse a tool<|im_end|>\n\n<|im_start|>assistant\n<think>\n",
|
|
},
|
|
{
|
|
name: "assistant tool call with content",
|
|
messages: []api.Message{
|
|
{Role: "user", Content: "Weather?"},
|
|
{
|
|
Role: "assistant",
|
|
Content: "Checking now.",
|
|
ToolCalls: []api.ToolCall{{
|
|
Function: api.ToolCallFunction{
|
|
Name: "get_weather",
|
|
Arguments: testArgs(map[string]any{"city": "Paris"}),
|
|
},
|
|
}},
|
|
},
|
|
},
|
|
think: thinkTrue(),
|
|
expected: "\n\n\n<|im_start|>system\n<|im_end|>\n\n<|im_start|>user\nWeather?<|im_end|>\n<|im_start|>assistant\n<think></think>Checking now.\n<tool_call>\n<function=get_weather>\n<parameter=city>\nParis\n</parameter>\n</function>\n</tool_call>\n<|im_end|>\n\n<|im_start|>assistant\n<think>\n",
|
|
},
|
|
{
|
|
name: "assistant tool call with structured arguments",
|
|
messages: []api.Message{
|
|
{Role: "user", Content: "Create data"},
|
|
{
|
|
Role: "assistant",
|
|
ToolCalls: []api.ToolCall{{
|
|
Function: api.ToolCallFunction{
|
|
Name: "create",
|
|
Arguments: testArgsOrdered([]orderedArg{
|
|
{Key: "payload", Value: map[string]any{"count": 42, "nested": map[string]any{"value": "ok"}}},
|
|
{Key: "tags", Value: []any{"a", "b"}},
|
|
}),
|
|
},
|
|
}},
|
|
},
|
|
},
|
|
think: thinkTrue(),
|
|
expected: "\n\n\n<|im_start|>system\n<|im_end|>\n\n<|im_start|>user\nCreate data<|im_end|>\n<|im_start|>assistant\n<think></think>\n<tool_call>\n<function=create>\n<parameter=payload>\n{\"count\": 42, \"nested\": {\"value\": \"ok\"}}\n</parameter>\n<parameter=tags>\n[\"a\", \"b\"]\n</parameter>\n</function>\n</tool_call>\n<|im_end|>\n\n<|im_start|>assistant\n<think>\n",
|
|
},
|
|
{
|
|
name: "assistant tool call truncated with reasoning",
|
|
messages: []api.Message{
|
|
{Role: "user", Content: "Weather?"},
|
|
{
|
|
Role: "assistant",
|
|
Content: "Checking now.",
|
|
Thinking: "Need weather",
|
|
ToolCalls: []api.ToolCall{{
|
|
Function: api.ToolCallFunction{
|
|
Name: "get_weather",
|
|
Arguments: testArgs(map[string]any{"city": "Paris"}),
|
|
},
|
|
}},
|
|
},
|
|
{Role: "user", Content: "And tomorrow?"},
|
|
},
|
|
think: thinkTrue(),
|
|
expected: "\n\n\n<|im_start|>system\n<|im_end|>\n\n<|im_start|>user\nWeather?<|im_end|>\n<|im_start|>assistant\n<think></think>Checking now.\n<tool_call>\n<function=get_weather>\n<parameter=city>\nParis\n</parameter>\n</function>\n</tool_call>\n<|im_end|>\n<|im_start|>user\nAnd tomorrow?<|im_end|>\n\n<|im_start|>assistant\n<think>\n",
|
|
},
|
|
{
|
|
name: "assistant tool call truncated open think only",
|
|
messages: []api.Message{
|
|
{Role: "user", Content: "Weather?"},
|
|
{
|
|
Role: "assistant",
|
|
Content: "<think>draft",
|
|
ToolCalls: []api.ToolCall{{
|
|
Function: api.ToolCallFunction{
|
|
Name: "get_weather",
|
|
Arguments: testArgs(map[string]any{"city": "Paris"}),
|
|
},
|
|
}},
|
|
},
|
|
{Role: "user", Content: "And tomorrow?"},
|
|
},
|
|
think: thinkTrue(),
|
|
expected: "\n\n\n<|im_start|>system\n<|im_end|>\n\n<|im_start|>user\nWeather?<|im_end|>\n<|im_start|>assistant\n<think></think>\n<tool_call>\n<function=get_weather>\n<parameter=city>\nParis\n</parameter>\n</function>\n</tool_call>\n<|im_end|>\n<|im_start|>user\nAnd tomorrow?<|im_end|>\n\n<|im_start|>assistant\n<think>\n",
|
|
},
|
|
{
|
|
name: "assistant tool call empty content",
|
|
messages: []api.Message{
|
|
{Role: "user", Content: "Weather?"},
|
|
{
|
|
Role: "assistant",
|
|
Content: "",
|
|
ToolCalls: []api.ToolCall{{
|
|
Function: api.ToolCallFunction{
|
|
Name: "get_weather",
|
|
Arguments: testArgs(map[string]any{"city": "Paris"}),
|
|
},
|
|
}},
|
|
},
|
|
},
|
|
think: thinkTrue(),
|
|
expected: "\n\n\n<|im_start|>system\n<|im_end|>\n\n<|im_start|>user\nWeather?<|im_end|>\n<|im_start|>assistant\n<think></think>\n<tool_call>\n<function=get_weather>\n<parameter=city>\nParis\n</parameter>\n</function>\n</tool_call>\n<|im_end|>\n\n<|im_start|>assistant\n<think>\n",
|
|
},
|
|
{
|
|
name: "assistant truncated with think pair",
|
|
messages: []api.Message{
|
|
{Role: "user", Content: "Hi"},
|
|
{Role: "assistant", Content: "<think>hidden</think>Visible"},
|
|
{Role: "user", Content: "Next"},
|
|
},
|
|
think: thinkTrue(),
|
|
expected: "\n\n\n<|im_start|>system\n<|im_end|>\n\n<|im_start|>user\nHi<|im_end|>\n<|im_start|>assistant\n<think></think>Visible<|im_end|>\n<|im_start|>user\nNext<|im_end|>\n\n<|im_start|>assistant\n<think>\n",
|
|
},
|
|
{
|
|
name: "assistant truncated reasoning content",
|
|
messages: []api.Message{
|
|
{Role: "user", Content: "Hi"},
|
|
{Role: "assistant", Thinking: "hidden", Content: "Visible"},
|
|
{Role: "user", Content: "Next"},
|
|
},
|
|
think: thinkTrue(),
|
|
expected: "\n\n\n<|im_start|>system\n<|im_end|>\n\n<|im_start|>user\nHi<|im_end|>\n<|im_start|>assistant\n<think></think>\nVisible<|im_end|>\n<|im_start|>user\nNext<|im_end|>\n\n<|im_start|>assistant\n<think>\n",
|
|
},
|
|
{
|
|
name: "assistant truncated plain content",
|
|
messages: []api.Message{
|
|
{Role: "user", Content: "Hi"},
|
|
{Role: "assistant", Content: "Visible"},
|
|
{Role: "user", Content: "Next"},
|
|
},
|
|
think: thinkTrue(),
|
|
expected: "\n\n\n<|im_start|>system\n<|im_end|>\n\n<|im_start|>user\nHi<|im_end|>\n<|im_start|>assistant\n<think></think>Visible<|im_end|>\n<|im_start|>user\nNext<|im_end|>\n\n<|im_start|>assistant\n<think>\n",
|
|
},
|
|
{
|
|
name: "assistant truncated empty content",
|
|
messages: []api.Message{
|
|
{Role: "user", Content: "Hi"},
|
|
{Role: "assistant", Content: ""},
|
|
{Role: "user", Content: "Next"},
|
|
},
|
|
think: thinkTrue(),
|
|
expected: "\n\n\n<|im_start|>system\n<|im_end|>\n\n<|im_start|>user\nHi<|im_end|>\n<|im_start|>assistant\n<think></think><|im_end|>\n<|im_start|>user\nNext<|im_end|>\n\n<|im_start|>assistant\n<think>\n",
|
|
},
|
|
{
|
|
name: "consecutive tool messages grouped",
|
|
messages: []api.Message{
|
|
{Role: "user", Content: "Do work"},
|
|
{
|
|
Role: "assistant",
|
|
ToolCalls: []api.ToolCall{{
|
|
Function: api.ToolCallFunction{
|
|
Name: "step",
|
|
Arguments: testArgs(map[string]any{"value": 1}),
|
|
},
|
|
}},
|
|
},
|
|
{Role: "tool", Content: "one"},
|
|
{Role: "tool", Content: "two"},
|
|
},
|
|
think: thinkTrue(),
|
|
expected: "\n\n\n<|im_start|>system\n<|im_end|>\n\n<|im_start|>user\nDo work<|im_end|>\n<|im_start|>assistant\n<think></think>\n<tool_call>\n<function=step>\n<parameter=value>\n1\n</parameter>\n</function>\n</tool_call>\n<|im_end|>\n<|im_start|>user\n<tool_response>\none\n</tool_response>\n<tool_response>\ntwo\n</tool_response>\n<|im_end|>\n\n<|im_start|>assistant\n<think>\n",
|
|
},
|
|
{
|
|
name: "fallback role",
|
|
messages: []api.Message{
|
|
{Role: "developer", Content: "Custom role content"},
|
|
},
|
|
think: thinkTrue(),
|
|
expected: "\n\n\n<|im_start|>system\n<|im_end|>\n\n<|im_start|>developer\nCustom role content<|im_end|>\n\n<|im_start|>assistant\n<think>\n",
|
|
},
|
|
}
|
|
|
|
verifyJinja2 := os.Getenv("VERIFY_JINJA2") != ""
|
|
if verifyJinja2 {
|
|
if _, err := os.Stat(filepath.Join(nemotron3NanoRepoRoot(t), ".venv", "bin", "python3")); err != nil {
|
|
t.Fatal("VERIFY_JINJA2=1 requires .venv/bin/python3")
|
|
}
|
|
}
|
|
|
|
renderer := &Nemotron3NanoRenderer{}
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
got, err := renderer.Render(tt.messages, tt.tools, tt.think)
|
|
if err != nil {
|
|
t.Fatalf("Render() error = %v", err)
|
|
}
|
|
|
|
if diff := cmp.Diff(tt.expected, got); diff != "" {
|
|
t.Fatalf("renderer mismatch (-want +got):\n%s", diff)
|
|
}
|
|
|
|
if verifyJinja2 {
|
|
jinja2Output := renderNemotron3NanoWithJinja2(t, tt.messages, tt.tools, tt.think)
|
|
if diff := cmp.Diff(tt.expected, jinja2Output); diff != "" {
|
|
t.Fatalf("reference template mismatch (-want +got):\n%s", diff)
|
|
}
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func nemotron3NanoReferenceTools() []api.Tool {
|
|
return []api.Tool{{
|
|
Type: "function",
|
|
Function: api.ToolFunction{
|
|
Name: "search_docs",
|
|
Description: "Search docs",
|
|
Parameters: api.ToolFunctionParameters{
|
|
Type: "object",
|
|
Defs: map[string]any{"shared": map[string]any{"type": "string"}},
|
|
Required: []string{"query"},
|
|
Properties: testPropsOrdered([]orderedProp{
|
|
{
|
|
Key: "query",
|
|
Value: api.ToolProperty{
|
|
Type: api.PropertyType{"string"},
|
|
Description: "Search query",
|
|
Enum: []any{"api", "cli"},
|
|
},
|
|
},
|
|
{
|
|
Key: "mode",
|
|
Value: api.ToolProperty{
|
|
Type: api.PropertyType{"string", "null"},
|
|
Description: "Mode",
|
|
AnyOf: []api.ToolProperty{
|
|
{Type: api.PropertyType{"string"}},
|
|
{Type: api.PropertyType{"number"}},
|
|
},
|
|
},
|
|
},
|
|
{
|
|
Key: "payload",
|
|
Value: api.ToolProperty{
|
|
Type: api.PropertyType{"object"},
|
|
Description: "Payload",
|
|
Properties: testPropsOrdered([]orderedProp{{Key: "enabled", Value: api.ToolProperty{Type: api.PropertyType{"boolean"}}}}),
|
|
Required: []string{"enabled"},
|
|
},
|
|
},
|
|
{
|
|
Key: "tags",
|
|
Value: api.ToolProperty{
|
|
Type: api.PropertyType{"array"},
|
|
Description: "Tags",
|
|
Items: map[string]any{"type": "string"},
|
|
},
|
|
},
|
|
}),
|
|
},
|
|
},
|
|
}}
|
|
}
|
|
|
|
func renderNemotron3NanoWithJinja2(t *testing.T, messages []api.Message, tools []api.Tool, think *api.ThinkValue) string {
|
|
t.Helper()
|
|
|
|
type jinja2ToolCall struct {
|
|
ID string `json:"id,omitempty"`
|
|
Function struct {
|
|
Name string `json:"name"`
|
|
Arguments any `json:"arguments"`
|
|
} `json:"function"`
|
|
}
|
|
type jinja2Message struct {
|
|
Role string `json:"role"`
|
|
Content string `json:"content"`
|
|
ReasoningContent string `json:"reasoning_content,omitempty"`
|
|
ToolCalls []jinja2ToolCall `json:"tool_calls,omitempty"`
|
|
Name string `json:"name,omitempty"`
|
|
ToolCallID string `json:"tool_call_id,omitempty"`
|
|
}
|
|
|
|
var jMsgs []jinja2Message
|
|
for _, m := range messages {
|
|
jm := jinja2Message{
|
|
Role: m.Role,
|
|
Content: m.Content,
|
|
ReasoningContent: m.Thinking,
|
|
Name: m.ToolName,
|
|
ToolCallID: m.ToolCallID,
|
|
}
|
|
for _, tc := range m.ToolCalls {
|
|
jtc := jinja2ToolCall{ID: tc.ID}
|
|
jtc.Function.Name = tc.Function.Name
|
|
var args map[string]any
|
|
raw, _ := tc.Function.Arguments.MarshalJSON()
|
|
if err := json.Unmarshal(raw, &args); err != nil {
|
|
t.Fatalf("failed to unmarshal tool args: %v", err)
|
|
}
|
|
jtc.Function.Arguments = args
|
|
jm.ToolCalls = append(jm.ToolCalls, jtc)
|
|
}
|
|
jMsgs = append(jMsgs, jm)
|
|
}
|
|
|
|
msgsJSON, err := json.Marshal(jMsgs)
|
|
if err != nil {
|
|
t.Fatalf("failed to marshal messages: %v", err)
|
|
}
|
|
|
|
toolsJSON := "None"
|
|
if len(tools) > 0 {
|
|
b, err := json.Marshal(tools)
|
|
if err != nil {
|
|
t.Fatalf("failed to marshal tools: %v", err)
|
|
}
|
|
toolsJSON = string(b)
|
|
}
|
|
|
|
thinking := "unset"
|
|
if think != nil {
|
|
if think.Bool() {
|
|
thinking = "true"
|
|
} else {
|
|
thinking = "false"
|
|
}
|
|
}
|
|
|
|
repoRoot := nemotron3NanoRepoRoot(t)
|
|
templatePath := filepath.Join(repoRoot, "model", "renderers", nemotron3NanoTemplate)
|
|
pythonPath := filepath.Join(repoRoot, ".venv", "bin", "python3")
|
|
script := `
|
|
import json
|
|
import sys
|
|
from pathlib import Path
|
|
from transformers.utils.chat_template_utils import _compile_jinja_template
|
|
|
|
template_path, messages_json, tools_json, thinking = sys.argv[1:5]
|
|
tmpl = _compile_jinja_template(Path(template_path).read_text())
|
|
kwargs = {
|
|
"messages": json.loads(messages_json),
|
|
"add_generation_prompt": True,
|
|
}
|
|
if tools_json != "None":
|
|
kwargs["tools"] = json.loads(tools_json)
|
|
if thinking == "true":
|
|
kwargs["enable_thinking"] = True
|
|
elif thinking == "false":
|
|
kwargs["enable_thinking"] = False
|
|
print(tmpl.render(**kwargs), end="")
|
|
`
|
|
|
|
cmd := exec.Command(pythonPath, "-c", script, templatePath, string(msgsJSON), toolsJSON, thinking)
|
|
var stdout, stderr strings.Builder
|
|
cmd.Stdout = &stdout
|
|
cmd.Stderr = &stderr
|
|
if err := cmd.Run(); err != nil {
|
|
t.Fatalf("python render failed: %v\nstderr: %s", err, stderr.String())
|
|
}
|
|
return stdout.String()
|
|
}
|
|
|
|
func nemotron3NanoRepoRoot(t *testing.T) string {
|
|
t.Helper()
|
|
_, filename, _, ok := runtime.Caller(0)
|
|
if !ok {
|
|
t.Fatal("failed to locate test file")
|
|
}
|
|
return filepath.Dir(filepath.Dir(filepath.Dir(filename)))
|
|
}
|