server: apply structured outputs in a single pass on thinking models

A format on a thinking model ran two generations: an unconstrained one,
cancelled once the parser reported content, then a re-rendered prompt with
the parsed thinking under the grammar. The restart cost a second prefill,
dropped the chunk that crossed the boundary, needed a harmony prompt hack,
stitched metrics across the two requests, and on MLX could leak a stray
first token into the JSON. The generate endpoint never deferred at all, so
its JSON was forced inside the thinking.

Both handlers now make one completion request that names the strings
ending the response's thinking, from the builtin parser or the generic
thinking parser, and the runner constrains only the content after them in
a single generation. The prompt is evaluated once and metrics pass straight
through. A raw generate prompt names no strings, since nothing says where
its response starts, and its format applies from the first token as
before. A format now applies to whatever follows the thinking, so a tool
call can no longer take the place of formatted content, which was already
the case with thinking off; harmony is the exception, since its tool calls
precede the final message.

The per-token metrics flag both runners carried for the cancelled first
pass has no caller left and goes with the two-pass code and its tests.

Fixes #18441
Fixes #17544
Fixes #14196
Fixes #10929
This commit is contained in:
Jesse Gross
2026-09-17 13:40:55 -07:00
parent 1042fead49
commit 152f1df7b2
9 changed files with 128 additions and 674 deletions
+4 -11
View File
@@ -1669,7 +1669,6 @@ func (s *llamaServerRunner) Completion(ctx context.Context, req CompletionReques
TypicalP: req.Options.TypicalP,
Seed: req.Options.Seed,
PreservedTokens: llamaServerPreservedTokens(req.PreservedTokens, req.ToolCallTag),
TimingsPerToken: req.IncludeIntermediateMetrics,
}
if req.Logprobs {
@@ -1810,16 +1809,10 @@ func (s *llamaServerRunner) Completion(ctx context.Context, req CompletionReques
}
if lsResp.Content != "" && !lsResp.Stop {
resp := CompletionResponse{Content: lsResp.Content}
if req.IncludeIntermediateMetrics {
resp.PromptEvalCount = lsResp.Timings.promptEvalCount()
resp.PromptEvalCachedCount = lsResp.Timings.CacheN
resp.PromptEvalDuration = time.Duration(lsResp.Timings.PromptMS * float64(time.Millisecond))
resp.EvalCount = lsResp.Timings.PredictN
resp.EvalDuration = time.Duration(lsResp.Timings.PredictMS * float64(time.Millisecond))
}
resp.Logprobs = convertLogprobs(lsResp.CompletionProbabilities, req.TopLogprobs > 0)
fn(resp)
fn(CompletionResponse{
Content: lsResp.Content,
Logprobs: convertLogprobs(lsResp.CompletionProbabilities, req.TopLogprobs > 0),
})
}
if lsResp.Stop {
+2 -24
View File
@@ -187,9 +187,6 @@ func TestLlamaServerCompletionSSEParsing(t *testing.T) {
if !reqBody.Stream {
t.Error("stream should be true")
}
if !reqBody.TimingsPerToken {
t.Error("timings_per_token should be true")
}
w.Header().Set("Content-Type", "text/event-stream")
for _, line := range sseLines {
@@ -212,9 +209,8 @@ func TestLlamaServerCompletionSSEParsing(t *testing.T) {
var responses []CompletionResponse
opts := api.DefaultOptions()
err := runner.Completion(t.Context(), CompletionRequest{
Prompt: "test prompt",
Options: &opts,
IncludeIntermediateMetrics: true,
Prompt: "test prompt",
Options: &opts,
}, func(cr CompletionResponse) {
responses = append(responses, cr)
})
@@ -233,29 +229,11 @@ func TestLlamaServerCompletionSSEParsing(t *testing.T) {
if responses[0].Done {
t.Error("response[0] should not be done")
}
if responses[0].PromptEvalCount != 5 || responses[0].EvalCount != 1 {
t.Errorf("response[0] counts = (%d, %d), want (5, 1)", responses[0].PromptEvalCount, responses[0].EvalCount)
}
if got := responses[0].PromptEvalCachedCount; got == nil || *got != 2 {
t.Errorf("response[0] cached prompt count = %v, want 2", got)
}
if responses[0].PromptEvalDuration != 10500*time.Microsecond || responses[0].EvalDuration != 9100*time.Microsecond {
t.Errorf("response[0] durations = (%s, %s), want (10.5ms, 9.1ms)", responses[0].PromptEvalDuration, responses[0].EvalDuration)
}
// Second token
if responses[1].Content != " world" {
t.Errorf("response[1].Content = %q, want %q", responses[1].Content, " world")
}
if responses[1].PromptEvalCount != 5 || responses[1].EvalCount != 2 {
t.Errorf("response[1] counts = (%d, %d), want (5, 2)", responses[1].PromptEvalCount, responses[1].EvalCount)
}
if got := responses[1].PromptEvalCachedCount; got == nil || *got != 2 {
t.Errorf("response[1] cached prompt count = %v, want 2", got)
}
if responses[1].PromptEvalDuration != 10500*time.Microsecond || responses[1].EvalDuration != 20300*time.Microsecond {
t.Errorf("response[1] durations = (%s, %s), want (10.5ms, 20.3ms)", responses[1].PromptEvalDuration, responses[1].EvalDuration)
}
// Final response
if !responses[2].Done {
-2
View File
@@ -209,8 +209,6 @@ type CompletionRequest struct {
// 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
// Logprobs specifies whether to include log probabilities in the response
Logprobs bool
+11 -13
View File
@@ -114,13 +114,12 @@ func (c *Client) WaitUntilRunning(ctx context.Context) error {
}
type CompletionRequest struct {
Prompt string
Media []llm.MediaData
Format json.RawMessage
Options api.Options
Logprobs bool
TopLogprobs int
IncludeIntermediateMetrics bool
Prompt string
Media []llm.MediaData
Format json.RawMessage
Options api.Options
Logprobs bool
TopLogprobs int
}
type CompletionResponse struct {
@@ -197,12 +196,11 @@ func jsonString(s string) string {
// Completion implements llm.LlamaServer.
func (c *Client) Completion(ctx context.Context, req llm.CompletionRequest, fn func(llm.CompletionResponse)) error {
creq := CompletionRequest{
Prompt: req.Prompt,
Media: req.Media,
Format: requestGrammar(req),
Logprobs: req.Logprobs,
TopLogprobs: req.TopLogprobs,
IncludeIntermediateMetrics: req.IncludeIntermediateMetrics,
Prompt: req.Prompt,
Media: req.Media,
Format: requestGrammar(req),
Logprobs: req.Logprobs,
TopLogprobs: req.TopLogprobs,
}
if req.Options != nil {
creq.Options = *req.Options
-52
View File
@@ -2,20 +2,11 @@ package mlxrunner
import (
"encoding/json"
"net"
"net/http"
"net/http/httptest"
"strconv"
"testing"
"github.com/ollama/ollama/api"
"github.com/ollama/ollama/llm"
)
func testIntPtr(v int) *int {
return &v
}
func TestRequestGrammar(t *testing.T) {
schema := `{"type":"object","properties":{"answer":{"type":"string"}}}`
tag := `{"type":"structural_tag","format":{"type":"json_schema","json_schema":` + schema + `}}`
@@ -56,46 +47,3 @@ func TestRequestGrammar(t *testing.T) {
})
}
}
func TestClientCompletionRequestsIntermediateMetrics(t *testing.T) {
var request CompletionRequest
want := CompletionResponse{
Done: true,
PromptEvalCount: 10,
PromptEvalCachedCount: testIntPtr(4),
}
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if err := json.NewDecoder(r.Body).Decode(&request); err != nil {
t.Errorf("decode request: %v", err)
return
}
if err := json.NewEncoder(w).Encode(want); err != nil {
t.Errorf("encode response: %v", err)
}
}))
t.Cleanup(srv.Close)
_, portString, err := net.SplitHostPort(srv.Listener.Addr().String())
if err != nil {
t.Fatalf("parse server port: %v", err)
}
port, err := strconv.Atoi(portString)
if err != nil {
t.Fatalf("parse server port: %v", err)
}
client := &Client{port: port, client: srv.Client()}
opts := api.DefaultOptions()
var got llm.CompletionResponse
if err := client.Completion(t.Context(), llm.CompletionRequest{
Options: &opts,
IncludeIntermediateMetrics: true,
}, func(response llm.CompletionResponse) { got = response }); err != nil {
t.Fatalf("Completion: %v", err)
}
if !request.IncludeIntermediateMetrics {
t.Fatal("metrics per token was not forwarded to the MLX runner")
}
if got.PromptEvalCount != want.PromptEvalCount || got.PromptEvalCachedCount == nil || *got.PromptEvalCachedCount != *want.PromptEvalCachedCount {
t.Errorf("prompt counts = (%d, %v), want (%d, %d)", got.PromptEvalCount, got.PromptEvalCachedCount, want.PromptEvalCount, *want.PromptEvalCachedCount)
}
}
-51
View File
@@ -9,7 +9,6 @@ import (
"slices"
"strings"
"testing"
"time"
"github.com/ollama/ollama/api"
"github.com/ollama/ollama/mlx"
@@ -433,56 +432,6 @@ func TestRunMTPDecodeGreedy(t *testing.T) {
})
}
func TestDecodeIntermediateMetrics(t *testing.T) {
mlxtest.Run(t, func(t *mlxtest.T) {
const eos int32 = 7
predict := map[int32]int32{1: 2, 2: 3, 3: eos, eos: 0}
r := mtpTestRunner(t, predict, []int32{eos}, sampler.Options{})
caches, _ := newMTPTestCaches(1)
session, ch := newMTPTestSession(caches)
session.inputs = []int32{0, 1}
session.remaining = []int32{1}
req := Request{
Responses: ch,
Tokens: session.inputs,
CompletionRequest: CompletionRequest{
Options: api.Options{NumPredict: 20},
IncludeIntermediateMetrics: true,
},
}
d := testDecoder(t, r, req, caches, []int32{1}, 1)
promptEval := 5 * time.Millisecond
if err := r.decode(context.Background(), req, session, d, promptEval); err != nil {
t.Fatalf("decode: %v", err)
}
d.close()
var responses []CompletionResponse
for len(ch) > 0 {
resp := <-ch
responses = append(responses, resp)
}
if len(responses) != 3 {
t.Fatalf("got %d responses, want 3", len(responses))
}
for i, resp := range responses[:2] {
if resp.PromptEvalCount != 2 || resp.PromptEvalCachedCount == nil || *resp.PromptEvalCachedCount != 1 || resp.PromptEvalDuration != promptEval {
t.Errorf("response[%d] prompt metrics = (%d, %v, %s), want (2, 1, %s)", i, resp.PromptEvalCount, resp.PromptEvalCachedCount, resp.PromptEvalDuration, promptEval)
}
if resp.EvalCount != i+1 || resp.EvalDuration <= 0 {
t.Errorf("response[%d] eval metrics = (%d, %s), want count %d and positive duration", i, resp.EvalCount, resp.EvalDuration, i+1)
}
}
final := responses[2]
if final.PromptEvalCount != 2 || final.PromptEvalCachedCount == nil || *final.PromptEvalCachedCount != 1 || final.PromptEvalDuration != promptEval {
t.Errorf("final prompt metrics = (%d, %v, %s), want (2, 1, %s)", final.PromptEvalCount, final.PromptEvalCachedCount, final.PromptEvalDuration, promptEval)
}
if final.EvalCount != 2 || final.EvalDuration <= 0 {
t.Errorf("final eval metrics = (%d, %s), want count 2 and positive duration", final.EvalCount, final.EvalDuration)
}
})
}
func TestRunMTPDecodeSampled(t *testing.T) {
mlxtest.Run(t, func(t *mlxtest.T) {
// The same chain at temperature 1: because oneHotLogits uses a large gap,
-8
View File
@@ -327,14 +327,6 @@ func (r *Runner) decode(ctx context.Context, request Request, session *cacheSess
if !ok {
continue
}
// Two-pass structured output cancels the first pass before its final response.
if request.IncludeIntermediateMetrics {
resp.PromptEvalCount = len(request.Tokens)
resp.PromptEvalCachedCount = final.PromptEvalCachedCount
resp.PromptEvalDuration = promptEval
resp.EvalCount = generated
resp.EvalDuration = time.Since(now)
}
select {
case <-ctx.Done():
err = ctx.Err()
+111 -188
View File
@@ -664,6 +664,13 @@ func (s *Server) GenerateHandler(c *gin.Context) {
defer cancel()
var parserErr error
// A raw prompt gives no way to tell where the response starts, so the
// format applies from its first token.
var thinkingClose []string
if !req.Raw {
thinkingClose = thinkingCloseForCompletion(builtinParser, thinkingState)
}
if err := r.Completion(ctx, llm.CompletionRequest{
Prompt: prompt,
Media: media,
@@ -675,6 +682,7 @@ func (s *Server) GenerateHandler(c *gin.Context) {
TopLogprobs: req.TopLogprobs,
PreservedTokens: preservedTokensForCompletion(builtinParser),
LeadingBOS: leadingBOS,
ThinkingClose: thinkingClose,
}, func(cr llm.CompletionResponse) {
res := api.GenerateResponse{
Model: req.Model,
@@ -2317,6 +2325,16 @@ func toolCallTagForCompletion(toolParser *tools.Parser) string {
return toolParser.Tag()
}
func thinkingCloseForCompletion(builtinParser parsers.Parser, thinkingState *thinking.Parser) []string {
if builtinParser != nil {
return builtinParser.ThinkingClose()
}
if thinkingState != nil {
return []string{thinkingState.ClosingTag}
}
return nil
}
func leadingBOSForModel(m *Model) string {
if m == nil || m.Config.Renderer == "" {
return ""
@@ -2729,224 +2747,129 @@ func (s *Server) ChatHandler(c *gin.Context) {
toolParser = tools.NewParser(m.Template.Template, req.Tools)
}
type structuredOutputsState int
const (
structuredOutputsState_None structuredOutputsState = iota
structuredOutputsState_ReadyToApply
structuredOutputsState_Applying
)
ch := make(chan any)
go func() {
defer close(ch)
structuredOutputsState := structuredOutputsState_None
var firstPassMetrics api.Metrics
ctx, cancel := context.WithCancel(c.Request.Context())
defer cancel()
for {
var tb strings.Builder
var parserErr error
currentFormat := req.Format
// structured outputs via double request is enabled when:
// 1. the model supports the thinking capability and
// 2. it uses a built-in parser or our generic thinking parser
// Note that the current approach does not work for (potential future)
// non-thinking models that emit anything before actual content. This
// current approach uses the transition from parsed thinking content to
// parsed non-thinking content as the signal to turn constraining on
forceImmediate := builtinParser != nil && builtinParser.HasThinkingSupport() && req.Think != nil && !req.Think.Bool()
if req.Format != nil && structuredOutputsState == structuredOutputsState_None && !forceImmediate && ((builtinParser != nil || thinkingState != nil) && slices.Contains(m.Capabilities(), model.CapabilityThinking)) {
currentFormat = nil
}
includeIntermediateMetrics := req.Format != nil && currentFormat == nil
// sets up new context given parent context per request
ctx, cancel := context.WithCancel(c.Request.Context())
var parserErr error
err := r.Completion(ctx, llm.CompletionRequest{
Prompt: prompt,
Media: media,
Format: currentFormat,
Options: opts,
Shift: req.Shift == nil || *req.Shift,
Truncate: truncate,
Logprobs: req.Logprobs,
TopLogprobs: req.TopLogprobs,
PreservedTokens: preservedTokensForCompletion(builtinParser),
ToolCallTag: toolCallTagForCompletion(toolParser),
LeadingBOS: leadingBOSForModel(m),
IncludeIntermediateMetrics: includeIntermediateMetrics,
}, func(r llm.CompletionResponse) {
metrics := api.Metrics{
err := r.Completion(ctx, llm.CompletionRequest{
Prompt: prompt,
Media: media,
Format: req.Format,
Options: opts,
Shift: req.Shift == nil || *req.Shift,
Truncate: truncate,
Logprobs: req.Logprobs,
TopLogprobs: req.TopLogprobs,
PreservedTokens: preservedTokensForCompletion(builtinParser),
ToolCallTag: toolCallTagForCompletion(toolParser),
LeadingBOS: leadingBOSForModel(m),
ThinkingClose: thinkingCloseForCompletion(builtinParser, thinkingState),
}, func(r llm.CompletionResponse) {
res := api.ChatResponse{
Model: req.Model,
CreatedAt: time.Now().UTC(),
Message: api.Message{Role: "assistant", Content: r.Content},
Done: r.Done,
Metrics: api.Metrics{
PromptEvalCount: r.PromptEvalCount,
PromptEvalCachedCount: r.PromptEvalCachedCount,
PromptEvalDuration: r.PromptEvalDuration,
EvalCount: r.EvalCount,
EvalDuration: r.EvalDuration,
}
if includeIntermediateMetrics {
firstPassMetrics = metrics
if !r.Done {
metrics = api.Metrics{}
}
} else if structuredOutputsState == structuredOutputsState_Applying && r.Done {
// Treat the restart as generation work: retain the original prompt metrics and fold in the second prefill.
metrics.PromptEvalCount = firstPassMetrics.PromptEvalCount
metrics.PromptEvalCachedCount = firstPassMetrics.PromptEvalCachedCount
metrics.PromptEvalDuration = firstPassMetrics.PromptEvalDuration
metrics.EvalCount += firstPassMetrics.EvalCount
metrics.EvalDuration += firstPassMetrics.EvalDuration + r.PromptEvalDuration
},
Logprobs: toAPILogprobs(r.Logprobs),
}
if r.Done {
res.DoneReason = r.DoneReason.String()
res.TotalDuration = time.Since(checkpointStart)
res.LoadDuration = checkpointLoaded.Sub(checkpointStart)
}
if builtinParser != nil {
slog.Log(context.TODO(), logutil.LevelTrace, "builtin parser input", "parser", m.Config.Parser, "content", r.Content)
content, thinking, toolCalls, err := builtinParser.Add(r.Content, r.Done)
if err != nil {
parserErr = err
cancel()
return
}
res := api.ChatResponse{
Model: req.Model,
CreatedAt: time.Now().UTC(),
Message: api.Message{Role: "assistant", Content: r.Content},
Done: r.Done,
Metrics: metrics,
Logprobs: toAPILogprobs(r.Logprobs),
res.Message.Content = content
res.Message.Thinking = thinking
for i := range toolCalls {
toolCalls[i].ID = toolCallId()
}
res.Message.ToolCalls = toolCalls
if r.Done {
res.DoneReason = r.DoneReason.String()
res.TotalDuration = time.Since(checkpointStart)
res.LoadDuration = checkpointLoaded.Sub(checkpointStart)
if res.Message.Content != "" || res.Message.Thinking != "" || len(res.Message.ToolCalls) > 0 || r.Done || len(res.Logprobs) > 0 {
slog.Log(context.TODO(), logutil.LevelTrace, "builtin parser output", "parser", m.Config.Parser, "content", content, "thinking", thinking, "toolCalls", toolCalls, "done", r.Done)
ch <- res
} else {
slog.Log(context.TODO(), logutil.LevelTrace, "builtin parser empty output", "parser", m.Config.Parser)
}
return
}
if builtinParser != nil {
slog.Log(context.TODO(), logutil.LevelTrace, "builtin parser input", "parser", m.Config.Parser, "content", r.Content)
content, thinking, toolCalls, err := builtinParser.Add(r.Content, r.Done)
if err != nil {
parserErr = err
cancel()
return
}
if thinkingState != nil {
thinkingContent, remainingContent := thinkingState.AddContent(res.Message.Content)
if thinkingContent == "" && remainingContent == "" && !r.Done {
// need to accumulate more to decide what to send
return
}
res.Message.Thinking = thinkingContent
res.Message.Content = remainingContent
}
if len(req.Tools) > 0 {
toolCalls, content := toolParser.Add(res.Message.Content)
if len(content) > 0 {
res.Message.Content = content
res.Message.Thinking = thinking
} else if len(toolCalls) > 0 {
for i := range toolCalls {
toolCalls[i].ID = toolCallId()
}
res.Message.ToolCalls = toolCalls
tb.WriteString(thinking)
// we are now receiving content from the model - we should start applying structured outputs
if structuredOutputsState == structuredOutputsState_None && req.Format != nil && tb.String() != "" && res.Message.Content != "" {
structuredOutputsState = structuredOutputsState_ReadyToApply
cancel()
return
}
if res.Message.Content != "" || res.Message.Thinking != "" || len(res.Message.ToolCalls) > 0 || r.Done || len(res.Logprobs) > 0 {
slog.Log(context.TODO(), logutil.LevelTrace, "builtin parser output", "parser", m.Config.Parser, "content", content, "thinking", thinking, "toolCalls", toolCalls, "done", r.Done)
ch <- res
} else {
slog.Log(context.TODO(), logutil.LevelTrace, "builtin parser empty output", "parser", m.Config.Parser)
}
return
}
if thinkingState != nil {
thinkingContent, remainingContent := thinkingState.AddContent(res.Message.Content)
if thinkingContent == "" && remainingContent == "" && !r.Done {
// need to accumulate more to decide what to send
return
}
res.Message.Thinking = thinkingContent
tb.WriteString(thinkingContent)
// emit the collected thinking text before restarting with structured outputs and clear unstructured content
// to avoid leaking mixed tokens like "</think>Hello"
if structuredOutputsState == structuredOutputsState_None && req.Format != nil && tb.String() != "" && remainingContent != "" {
structuredOutputsState = structuredOutputsState_ReadyToApply
res.Message.Content = ""
ch <- res
cancel()
return
}
res.Message.Content = remainingContent
}
if len(req.Tools) > 0 {
toolCalls, content := toolParser.Add(res.Message.Content)
if len(content) > 0 {
res.Message.Content = content
} else if len(toolCalls) > 0 {
for i := range toolCalls {
toolCalls[i].ID = toolCallId()
}
res.Message.ToolCalls = toolCalls
res.Message.Content = ""
} else if res.Message.Thinking != "" {
// don't return, fall through to send
} else {
// Send logprobs while content is being buffered by the parser for tool calls
if len(res.Logprobs) > 0 && !r.Done {
logprobRes := res
logprobRes.Message.Content = ""
logprobRes.Message.ToolCalls = nil
ch <- logprobRes
}
if r.Done {
res.Message.Content = toolParser.Content()
ch <- res
}
return
}
}
ch <- res
})
if parserErr != nil {
ch <- gin.H{"error": parserErr.Error()}
return
}
if err != nil {
if structuredOutputsState == structuredOutputsState_ReadyToApply && strings.Contains(err.Error(), "context canceled") && c.Request.Context().Err() == nil {
// only ignores error if it's a context cancellation due to setting structured outputs
res.Message.Content = ""
} else if res.Message.Thinking != "" {
// don't return, fall through to send
} else {
s.sched.expireRunnersForRuntimeOOM(m, err)
var serr api.StatusError
if errors.As(err, &serr) {
ch <- gin.H{"error": serr.ErrorMessage, "status": serr.StatusCode}
} else {
ch <- gin.H{"error": err.Error()}
// Send logprobs while content is being buffered by the parser for tool calls
if len(res.Logprobs) > 0 && !r.Done {
logprobRes := res
logprobRes.Message.Content = ""
logprobRes.Message.ToolCalls = nil
ch <- logprobRes
}
if r.Done {
res.Message.Content = toolParser.Content()
ch <- res
}
return
}
}
// ignored structured outputs cancellation falls through to here, start a new request with the structured outputs and updated prompt. use the
if structuredOutputsState == structuredOutputsState_ReadyToApply {
structuredOutputsState = structuredOutputsState_Applying
msg := api.Message{
Role: "assistant",
Thinking: tb.String(),
}
msgs = append(msgs, msg)
prompt, _, err = chatPrompt(c.Request.Context(), m, r.Tokenize, promptOpts, msgs, processedTools, req.Think, truncate)
if err != nil {
slog.Error("chat prompt error applying structured outputs", "error", err)
ch <- gin.H{"error": err.Error()}
return
}
// force constraining by terminating thinking header, the parser is already at this state
// when the last message is thinking, the rendered for gpt-oss cannot disambiguate between having the
// model continue thinking or ending thinking and outputting the final message.
// TODO(parthsareen): consider adding prefill disambiguation logic to the renderer for structured outputs.
if shouldUseHarmony(m) || (builtinParser != nil && m.Config.Parser == "harmony") {
prompt += "<|end|><|start|>assistant<|channel|>final<|message|>"
}
continue
ch <- res
})
if parserErr != nil {
ch <- gin.H{"error": parserErr.Error()}
return
}
if err != nil {
s.sched.expireRunnersForRuntimeOOM(m, err)
var serr api.StatusError
if errors.As(err, &serr) {
ch <- gin.H{"error": serr.ErrorMessage, "status": serr.StatusCode}
} else {
ch <- gin.H{"error": err.Error()}
}
break
}
}()
-325
View File
@@ -2646,331 +2646,6 @@ func TestChatWithPromptEndingInThinkTag(t *testing.T) {
t.Errorf("expected content %q, got %q", "Based on my analysis, the solution is straightforward.", got)
}
})
earlyFirstPassMetrics := api.Metrics{
PromptEvalCount: 4,
PromptEvalCachedCount: testIntPtr(1),
PromptEvalDuration: 5 * time.Millisecond,
EvalCount: 6,
EvalDuration: 7 * time.Millisecond,
}
firstPassMetrics := api.Metrics{
PromptEvalCount: 10,
PromptEvalCachedCount: testIntPtr(4),
PromptEvalDuration: 11 * time.Millisecond,
EvalCount: 12,
EvalDuration: 13 * time.Millisecond,
}
secondPassMetrics := api.Metrics{
PromptEvalCount: 20_000,
PromptEvalCachedCount: testIntPtr(19_000),
PromptEvalDuration: 21 * time.Millisecond,
EvalCount: 22,
EvalDuration: 23 * time.Millisecond,
}
wantMetrics := api.Metrics{
PromptEvalCount: firstPassMetrics.PromptEvalCount,
PromptEvalCachedCount: firstPassMetrics.PromptEvalCachedCount,
PromptEvalDuration: firstPassMetrics.PromptEvalDuration,
EvalCount: firstPassMetrics.EvalCount + secondPassMetrics.EvalCount,
EvalDuration: firstPassMetrics.EvalDuration + secondPassMetrics.PromptEvalDuration + secondPassMetrics.EvalDuration,
}
t.Run("structured outputs restart non-stream", func(t *testing.T) {
var (
requestsMu sync.Mutex
requests []llm.CompletionRequest
wg sync.WaitGroup
)
wg.Add(2)
format := json.RawMessage(`{"type":"object","properties":{"answer":{"type":"string"}}}`)
mock.CompletionFn = func(ctx context.Context, r llm.CompletionRequest, fn func(r llm.CompletionResponse)) error {
defer wg.Done()
requestsMu.Lock()
requests = append(requests, r)
callNum := len(requests)
requestsMu.Unlock()
switch callNum {
case 1:
fn(llm.CompletionResponse{
Content: " I am thinking through this problem.",
Done: false,
PromptEvalCount: earlyFirstPassMetrics.PromptEvalCount,
PromptEvalCachedCount: earlyFirstPassMetrics.PromptEvalCachedCount,
PromptEvalDuration: earlyFirstPassMetrics.PromptEvalDuration,
EvalCount: earlyFirstPassMetrics.EvalCount,
EvalDuration: earlyFirstPassMetrics.EvalDuration,
})
fn(llm.CompletionResponse{
Content: " </think> {\"answer\":\"42\"}",
Done: false,
PromptEvalCount: firstPassMetrics.PromptEvalCount,
PromptEvalCachedCount: firstPassMetrics.PromptEvalCachedCount,
PromptEvalDuration: firstPassMetrics.PromptEvalDuration,
EvalCount: firstPassMetrics.EvalCount,
EvalDuration: firstPassMetrics.EvalDuration,
})
select {
case <-ctx.Done():
return ctx.Err()
case <-time.After(time.Second):
t.Fatalf("timeout waiting for structured outputs cancellation")
return nil
}
case 2:
fn(llm.CompletionResponse{
Content: `{"answer":"42"}`,
Done: true,
DoneReason: llm.DoneReasonStop,
PromptEvalCount: secondPassMetrics.PromptEvalCount,
PromptEvalCachedCount: secondPassMetrics.PromptEvalCachedCount,
PromptEvalDuration: secondPassMetrics.PromptEvalDuration,
EvalCount: secondPassMetrics.EvalCount,
EvalDuration: secondPassMetrics.EvalDuration,
})
return nil
default:
t.Fatalf("unexpected number of completion calls: %d", callNum)
return nil
}
}
think := true
streamRequest := false
w := createRequest(t, s.ChatHandler, api.ChatRequest{
Model: "test-thinking",
Messages: []api.Message{{Role: "user", Content: "Please respond in JSON."}},
Think: &api.ThinkValue{Value: think},
Stream: &streamRequest,
Format: format,
})
wg.Wait()
mock.CompletionFn = nil
if w.Code != http.StatusOK {
t.Fatalf("expected status 200, got %d", w.Code)
}
if len(requests) != 2 {
t.Fatalf("expected two completion calls, got %d", len(requests))
}
if requests[0].Format != nil {
t.Errorf("expected first completion format to be nil, got %q", requests[0].Format)
}
if !requests[0].IncludeIntermediateMetrics {
t.Error("expected first completion to request per-token metrics")
}
if !bytes.Equal([]byte(format), []byte(requests[1].Format)) {
t.Errorf("expected second completion format to match original format")
}
if requests[1].IncludeIntermediateMetrics {
t.Error("expected second completion to use terminal metrics")
}
var resp api.ChatResponse
if err := json.NewDecoder(w.Body).Decode(&resp); err != nil {
t.Fatal(err)
}
if resp.Message.Thinking != "I am thinking through this problem. " {
t.Errorf("expected thinking %q, got %q", "I am thinking through this problem. ", resp.Message.Thinking)
}
if resp.Message.Content != `{"answer":"42"}` {
t.Errorf("expected content %q, got %q", `{"answer":"42"}`, resp.Message.Content)
}
if !resp.Done {
t.Errorf("expected response to be done")
}
if resp.DoneReason != "stop" {
t.Errorf("expected done reason stop, got %s", resp.DoneReason)
}
if resp.PromptEvalCount != wantMetrics.PromptEvalCount || resp.EvalCount != wantMetrics.EvalCount {
t.Errorf("response counts = (%d, %d), want (%d, %d)", resp.PromptEvalCount, resp.EvalCount, wantMetrics.PromptEvalCount, wantMetrics.EvalCount)
}
if diff := cmp.Diff(wantMetrics.PromptEvalCachedCount, resp.PromptEvalCachedCount); diff != "" {
t.Errorf("response cached prompt count mismatch (-want +got):\n%s", diff)
}
if resp.PromptEvalDuration != wantMetrics.PromptEvalDuration || resp.EvalDuration != wantMetrics.EvalDuration {
t.Errorf("response durations = (%s, %s), want (%s, %s)", resp.PromptEvalDuration, resp.EvalDuration, wantMetrics.PromptEvalDuration, wantMetrics.EvalDuration)
}
})
t.Run("structured outputs restart streaming", func(t *testing.T) {
var (
requestsMu sync.Mutex
requests []llm.CompletionRequest
wg sync.WaitGroup
)
wg.Add(2)
format := json.RawMessage(`{"type":"object","properties":{"answer":{"type":"string"}}}`)
mock.CompletionFn = func(ctx context.Context, r llm.CompletionRequest, fn func(r llm.CompletionResponse)) error {
defer wg.Done()
requestsMu.Lock()
requests = append(requests, r)
callNum := len(requests)
requestsMu.Unlock()
switch callNum {
case 1:
fn(llm.CompletionResponse{
Content: " I am thinking through this problem.",
Done: false,
PromptEvalCount: earlyFirstPassMetrics.PromptEvalCount,
PromptEvalCachedCount: earlyFirstPassMetrics.PromptEvalCachedCount,
PromptEvalDuration: earlyFirstPassMetrics.PromptEvalDuration,
EvalCount: earlyFirstPassMetrics.EvalCount,
EvalDuration: earlyFirstPassMetrics.EvalDuration,
})
fn(llm.CompletionResponse{
Content: " </think> {\"answer\":\"42\"}",
Done: false,
PromptEvalCount: firstPassMetrics.PromptEvalCount,
PromptEvalCachedCount: firstPassMetrics.PromptEvalCachedCount,
PromptEvalDuration: firstPassMetrics.PromptEvalDuration,
EvalCount: firstPassMetrics.EvalCount,
EvalDuration: firstPassMetrics.EvalDuration,
})
select {
case <-ctx.Done():
return ctx.Err()
case <-time.After(time.Second):
t.Fatalf("timeout waiting for structured outputs cancellation")
return nil
}
case 2:
fn(llm.CompletionResponse{
Content: `{"answer":"42"}`,
Done: true,
DoneReason: llm.DoneReasonStop,
PromptEvalCount: secondPassMetrics.PromptEvalCount,
PromptEvalCachedCount: secondPassMetrics.PromptEvalCachedCount,
PromptEvalDuration: secondPassMetrics.PromptEvalDuration,
EvalCount: secondPassMetrics.EvalCount,
EvalDuration: secondPassMetrics.EvalDuration,
})
return nil
default:
t.Fatalf("unexpected number of completion calls: %d", callNum)
return nil
}
}
think := true
streamRequest := true
w := createRequest(t, s.ChatHandler, api.ChatRequest{
Model: "test-thinking",
Messages: []api.Message{{Role: "user", Content: "Please respond in JSON."}},
Think: &api.ThinkValue{Value: think},
Stream: &streamRequest,
Format: format,
})
wg.Wait()
mock.CompletionFn = nil
if w.Code != http.StatusOK {
t.Fatalf("expected status 200, got %d", w.Code)
}
if len(requests) != 2 {
t.Fatalf("expected two completion calls, got %d", len(requests))
}
if requests[0].Format != nil {
t.Errorf("expected first completion format to be nil, got %q", requests[0].Format)
}
if !requests[0].IncludeIntermediateMetrics {
t.Error("expected first completion to request per-token metrics")
}
if !bytes.Equal([]byte(format), []byte(requests[1].Format)) {
t.Errorf("expected second completion format to match original format")
}
if requests[1].IncludeIntermediateMetrics {
t.Error("expected second completion to use terminal metrics")
}
decoder := json.NewDecoder(w.Body)
var events []api.ChatResponse
for {
var event api.ChatResponse
if err := decoder.Decode(&event); err == io.EOF {
break
} else if err != nil {
t.Fatal(err)
}
events = append(events, event)
if event.Done {
break
}
}
if len(events) < 2 {
t.Fatalf("expected at least two streaming events, got %d", len(events))
}
first := events[0]
var thinking strings.Builder
for _, event := range events {
thinking.WriteString(event.Message.Thinking)
if !event.Done && (event.PromptEvalCount != 0 || event.PromptEvalCachedCount != nil || event.PromptEvalDuration != 0 || event.EvalCount != 0 || event.EvalDuration != 0) {
t.Errorf("non-terminal event unexpectedly exposed metrics: %+v", event.Metrics)
}
}
if got := thinking.String(); got != "I am thinking through this problem. " {
t.Errorf("thinking = %q, want %q", got, "I am thinking through this problem. ")
}
if first.Message.Content != "" {
t.Errorf("expected first event content to be empty, got %q", first.Message.Content)
}
if first.Done {
t.Error("expected first event to be non-terminal")
}
last := events[len(events)-1]
if last.Message.Thinking != "" {
t.Errorf("expected final event thinking to be empty, got %q", last.Message.Thinking)
}
if last.Message.Content != `{"answer":"42"}` {
t.Errorf("expected final event content %q, got %q", `{"answer":"42"}`, last.Message.Content)
}
if !last.Done {
t.Error("expected final event to be done")
}
if last.DoneReason != "stop" {
t.Errorf("expected final done reason stop, got %s", last.DoneReason)
}
if last.PromptEvalCount != wantMetrics.PromptEvalCount || last.EvalCount != wantMetrics.EvalCount {
t.Errorf("final counts = (%d, %d), want (%d, %d)", last.PromptEvalCount, last.EvalCount, wantMetrics.PromptEvalCount, wantMetrics.EvalCount)
}
if diff := cmp.Diff(wantMetrics.PromptEvalCachedCount, last.PromptEvalCachedCount); diff != "" {
t.Errorf("final cached prompt count mismatch (-want +got):\n%s", diff)
}
if last.PromptEvalDuration != wantMetrics.PromptEvalDuration || last.EvalDuration != wantMetrics.EvalDuration {
t.Errorf("final durations = (%s, %s), want (%s, %s)", last.PromptEvalDuration, last.EvalDuration, wantMetrics.PromptEvalDuration, wantMetrics.EvalDuration)
}
})
}
// TestChatFormatWithThinkFalse verifies that when a model uses a builtin