mirror of
https://github.com/ollama/ollama.git
synced 2026-09-21 13:38:14 -05:00
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:
+4
-11
@@ -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 {
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
@@ -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
|
||||
}
|
||||
}()
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user