diff --git a/llm/server.go b/llm/server.go index 7d8db0b1c..bc118c02f 100644 --- a/llm/server.go +++ b/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 diff --git a/mlxrunner/client.go b/mlxrunner/client.go index 8c0288ec0..69d0c7c09 100644 --- a/mlxrunner/client.go +++ b/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. diff --git a/mlxrunner/client_test.go b/mlxrunner/client_test.go index 908374fd8..aa7da1473 100644 --- a/mlxrunner/client_test.go +++ b/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{""}}, + want: `{"type":"structural_tag","format":{"type":"sequence","elements":[{"type":"any_text","excludes":[""]},` + + `{"type":"optional","content":{"type":"sequence","elements":[{"type":"const_string","value":""},` + + `{"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{""}}}, } { t.Run(tt.name, func(t *testing.T) { if got := string(requestGrammar(tt.req)); got != tt.want {