ollama source for Momentry Core verification

model/renderers/nemotron3nano_reference_test.go

@@ -0,0 +1,614 @@
+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)))
+}