makarna/pkg/chat/render_qwen3.go

package chat

import (
	"encoding/json"
	"fmt"
	"strings"
)

// Qwen3Renderer implements the Qwen3 chat template behavior commonly shipped
// in tokenizer_config.json (including tool-call wrappers).
// It is a deterministic renderer (not a general Jinja interpreter).
//
//nolint:revive // exported API
type Qwen3Renderer struct{}

func (r *Qwen3Renderer) Render(messages []Message, opts Options) (string, error) {
	var sb strings.Builder

	// Tools header block (only when tools are provided)
	if len(opts.Tools) > 0 {
		sb.WriteString("<|im_start|>system\n")
		if len(messages) > 0 && messages[0].Role == "system" {
			sb.WriteString(messages[0].Content)
			sb.WriteString("\n\n")
		}
		sb.WriteString("# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>")
		for _, tool := range opts.Tools {
			b, err := json.Marshal(tool)
			if err != nil {
				return "", fmt.Errorf("marshal tool: %w", err)
			}
			sb.WriteString("\n")
			sb.Write(b)
		}
		sb.WriteString("\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n")
	} else {
		// No-tools path: include system message only if first message is system.
		if len(messages) > 0 && messages[0].Role == "system" {
			sb.WriteString("<|im_start|>system\n")
			sb.WriteString(messages[0].Content)
			sb.WriteString("<|im_end|>\n")
		}
	}

	// Find last user query index for thinking rendering.
	lastQueryIdx := len(messages) - 1
	multiStepTool := true
	for i := len(messages) - 1; i >= 0; i-- {
		m := messages[i]
		if multiStepTool && m.Role == "user" && !strings.HasPrefix(m.Content, "<tool_response>") {
			multiStepTool = false
			lastQueryIdx = i
			break
		}
	}

	// Render message stream.
	inToolGroup := false
	for i := 0; i < len(messages); i++ {
		m := messages[i]

		switch m.Role {
		case "tool":
			// Tool responses are wrapped as a fake user message with <tool_response> blocks.
			if !inToolGroup {
				sb.WriteString("<|im_start|>user")
				inToolGroup = true
			}
			sb.WriteString("\n<tool_response>\n")
			sb.WriteString(m.Content)
			sb.WriteString("\n</tool_response>")

			// Close tool group if next isn't tool.
			if i == len(messages)-1 || messages[i+1].Role != "tool" {
				sb.WriteString("<|im_end|>\n")
				inToolGroup = false
			}
			continue
		default:
			if inToolGroup {
				// Safety: close an unterminated tool group.
				sb.WriteString("<|im_end|>\n")
				inToolGroup = false
			}
		}

		content := m.Content
		if m.Role == "user" || (m.Role == "system" && i != 0) {
			sb.WriteString("<|im_start|>")
			sb.WriteString(m.Role)
			sb.WriteString("\n")
			sb.WriteString(content)
			sb.WriteString("<|im_end|>\n")
			continue
		}

		if m.Role == "assistant" {
			sb.WriteString("<|im_start|>assistant\n")

			reasoning := m.ReasoningContent
			if opts.EnableThinking {
				// Render thinking only after the last user query index.
				if i > lastQueryIdx {
					sb.WriteString("<think>\n")
					sb.WriteString(strings.Trim(reasoning, "\n"))
					sb.WriteString("\n</think>\n\n")
				}
			}

			sb.WriteString(strings.TrimLeft(content, "\n"))

			// Render tool calls, if any.
			if len(m.ToolCalls) > 0 {
				for j, tc := range m.ToolCalls {
					if (j == 0 && content != "") || j > 0 {
						sb.WriteString("\n")
					}
					sb.WriteString("<tool_call>\n{\"name\": \"")
					sb.WriteString(tc.Name)
					sb.WriteString("\", \"arguments\": ")
					if len(tc.Arguments) == 0 {
						sb.WriteString("{}")
					} else {
						sb.Write(tc.Arguments)
					}
					sb.WriteString("}\n</tool_call>")
				}
			}

			sb.WriteString("<|im_end|>\n")
			continue
		}
	}

	if opts.AddGenerationPrompt {
		sb.WriteString("<|im_start|>assistant\n")
	}

	return sb.String(), nil
}