mlxrunner: constrain a format after the thinking in one generation

The MLX runner applies a format's grammar from the first sampled token,
so a thinking model asked for a format cannot think first, and the server
has to run two generations to get both the thinking and the formatted
content.

A completion request now carries the strings that end the thinking its
response begins with, and the MLX client builds from them a structural
tag: free text that cannot contain any of them, then one of them, then
the schema. The tail is optional so a response may still end inside its
thinking, as an unconstrained one can. Without a closing string the tag
is the plain schema, as before. The server does not send the strings yet.

llm/server.go

@@ -205,6 +205,10 @@ type CompletionRequest struct {
 	PreservedTokens []string // parser tokens to render as text; ignored by non-llama-server runners
 	ToolCallTag     string   // raw generic tool parser tag, if any
 	LeadingBOS      string   // textual BOS emitted by Go rendering, if any
+	// ThinkingClose holds the strings any of which ends the thinking the
+	// response begins with, which Format leaves free; none when the response
+	// starts in content.
+	ThinkingClose []string
 	// IncludeIntermediateMetrics adds cumulative metrics to non-final responses; final responses always include metrics.
 	IncludeIntermediateMetrics bool
 

mlxrunner/client.go

@@ -155,7 +155,8 @@ func (c *Client) Close() error {
 }
 
 // requestGrammar returns the structural tag the runner decodes under: the
-// API's format wrapped into a json_schema tag.
+// API's format as a json_schema tag, behind the free thinking the response
+// begins with when there is any.
 func requestGrammar(req llm.CompletionRequest) json.RawMessage {
 	schema := req.Format
 	switch string(schema) {
@@ -165,10 +166,32 @@ func requestGrammar(req llm.CompletionRequest) json.RawMessage {
 		// The API documents "json" as producing a JSON object.
 		schema = json.RawMessage(`{"type":"object"}`)
 	}
-	tag := make(json.RawMessage, 0, len(schema)+64)
-	tag = append(tag, `{"type":"structural_tag","format":{"type":"json_schema","json_schema":`...)
-	tag = append(tag, schema...)
-	return append(tag, `}}`...)
+	format := `{"type":"json_schema","json_schema":` + string(schema) + `}`
+	if len(req.ThinkingClose) > 0 {
+		excludes := make([]string, len(req.ThinkingClose))
+		closings := make([]string, len(req.ThinkingClose))
+		for i, closing := range req.ThinkingClose {
+			excludes[i] = jsonString(closing)
+			closings[i] = `{"type":"const_string","value":` + excludes[i] + `}`
+		}
+		closing := closings[0]
+		if len(closings) > 1 {
+			closing = `{"type":"or","elements":[` + strings.Join(closings, ",") + `]}`
+		}
+		// The tail is optional so EOS stays legal mid-thinking.
+		format = `{"type":"sequence","elements":[{"type":"any_text","excludes":[` + strings.Join(excludes, ",") + `]},` +
+			`{"type":"optional","content":{"type":"sequence","elements":[` + closing + `,` + format + `]}}]}`
+	}
+	return json.RawMessage(`{"type":"structural_tag","format":` + format + `}`)
+}
+
+// jsonString quotes s without escaping the HTML characters tags carry.
+func jsonString(s string) string {
+	var b strings.Builder
+	enc := json.NewEncoder(&b)
+	enc.SetEscapeHTML(false)
+	_ = enc.Encode(s)
+	return strings.TrimSuffix(b.String(), "\n")
 }
 
 // Completion implements llm.LlamaServer.

mlxrunner/client_test.go

@@ -33,6 +33,21 @@ func TestRequestGrammar(t *testing.T) {
 			want: `{"type":"structural_tag","format":{"type":"json_schema","json_schema":{"type":"object"}}}`,
 		},
 		{name: "schema", req: llm.CompletionRequest{Format: json.RawMessage(schema)}, want: tag},
+		{
+			name: "schema after thinking",
+			req:  llm.CompletionRequest{Format: json.RawMessage(schema), ThinkingClose: []string{"</think>"}},
+			want: `{"type":"structural_tag","format":{"type":"sequence","elements":[{"type":"any_text","excludes":["</think>"]},` +
+				`{"type":"optional","content":{"type":"sequence","elements":[{"type":"const_string","value":"</think>"},` +
+				`{"type":"json_schema","json_schema":` + schema + `}]}}]}}`,
+		},
+		{
+			name: "schema after thinking with two closings",
+			req:  llm.CompletionRequest{Format: json.RawMessage(schema), ThinkingClose: []string{"<|final|>", "<|final|>json"}},
+			want: `{"type":"structural_tag","format":{"type":"sequence","elements":[{"type":"any_text","excludes":["<|final|>","<|final|>json"]},` +
+				`{"type":"optional","content":{"type":"sequence","elements":[{"type":"or","elements":[{"type":"const_string","value":"<|final|>"},{"type":"const_string","value":"<|final|>json"}]},` +
+				`{"type":"json_schema","json_schema":` + schema + `}]}}]}}`,
+		},
+		{name: "thinking without a format", req: llm.CompletionRequest{ThinkingClose: []string{"</think>"}}},
 	} {
 		t.Run(tt.name, func(t *testing.T) {
 			if got := string(requestGrammar(tt.req)); got != tt.want {