Refactor llm pkg

This commit is contained in:
dwrz
2026-10-09 15:29:37 +00:00
parent 30d33977fc
commit c1a64be31c
3 changed files with 958 additions and 153 deletions
+101 -33
View File
@@ -28,14 +28,11 @@ type CompletionsParams struct {
// CompletionResult contains one generated Message and its response metadata. // CompletionResult contains one generated Message and its response metadata.
type CompletionResult struct { type CompletionResult struct {
// FinishReason is the endpoint's reason for ending generation.
FinishReason string
// Message is the generated message. // Message is the generated message.
Message Message Message Message
// Timings contains llama.cpp performance data, when reported. // Metadata describes this inference's selected choice and endpoint,
Timings *Timings // separately from the model-visible message.
// Usage contains token counts, when reported. Metadata Metadata
Usage *Usage
} }
// Completions requests a non-streaming chat completion and returns the first // Completions requests a non-streaming chat completion and returns the first
@@ -52,10 +49,21 @@ func (c *Client) Completions(
} }
choice := cr.Choices[0] choice := cr.Choices[0]
return &CompletionResult{ return &CompletionResult{
Message: choice.Message, Message: choice.Message,
FinishReason: choice.FinishReason, Metadata: Metadata{
Usage: cr.Usage, ChoiceIndex: choice.Index,
Timings: cr.Timings, Created: cr.Created,
FinishReason: choice.FinishReason,
ID: cr.ID,
Logprobs: choice.Logprobs,
Model: cr.Model,
Provider: new(c.provider),
RequestedModel: &p.Model,
ServiceTier: cr.ServiceTier,
SystemFingerprint: cr.SystemFingerprint,
Timings: cr.Timings,
Usage: cr.Usage,
},
}, nil }, nil
} }
@@ -240,11 +248,13 @@ type CompletionStream struct {
body io.ReadCloser body io.ReadCloser
cancel context.CancelFunc cancel context.CancelFunc
acc StreamAccumulator acc StreamAccumulator
current StreamChunk current StreamChunk
finishReason string // metadata retains reported scalars and the latest usage/timing
usage *Usage // snapshots.
timings *Timings // Logprob sequences contain every reported token in arrival order.
// Nested metadata values are shared with Current and Result.
metadata Metadata
err error err error
cleanupErr error cleanupErr error
terminated bool terminated bool
@@ -352,10 +362,11 @@ func (c *Client) CompletionsStream(
} }
return &CompletionStream{ return &CompletionStream{
ctx: ctx, ctx: ctx,
reader: bufio.NewReader(res.Body), reader: bufio.NewReader(res.Body),
body: res.Body, body: res.Body,
cancel: cancel, cancel: cancel,
metadata: Metadata{Provider: new(c.provider), RequestedModel: &p.Model},
}, nil }, nil
} }
@@ -412,45 +423,102 @@ func (s *CompletionStream) Next() bool {
return false return false
} }
for _, choice := range chunk.Choices { for _, choice := range chunk.Choices {
if choice.Index != 0 { if choice.Index != nil && *choice.Index != 0 {
s.fail(fmt.Errorf( s.fail(fmt.Errorf(
"completions stream: unexpected choice index %d", "completions stream: unexpected choice index %d",
choice.Index, *choice.Index,
)) ))
return false return false
} }
if choice.FinishReason != "" { if choice.Index != nil {
s.finishReason = choice.FinishReason s.metadata.ChoiceIndex = choice.Index
}
if choice.FinishReason != nil {
s.metadata.FinishReason = choice.FinishReason
}
if choice.Logprobs != nil {
if s.metadata.Logprobs == nil {
s.metadata.Logprobs = &Logprobs{}
}
if choice.Logprobs.Content != nil &&
s.metadata.Logprobs.Content == nil {
s.metadata.Logprobs.Content = []TokenLogprob{}
}
if choice.Logprobs.Refusal != nil &&
s.metadata.Logprobs.Refusal == nil {
s.metadata.Logprobs.Refusal = []TokenLogprob{}
}
s.metadata.Logprobs.Content = append(
s.metadata.Logprobs.Content,
choice.Logprobs.Content...,
)
s.metadata.Logprobs.Refusal = append(
s.metadata.Logprobs.Refusal,
choice.Logprobs.Refusal...,
)
} }
s.acc.Add(choice.Delta) s.acc.Add(choice.Delta)
} }
if chunk.Created != nil {
s.metadata.Created = chunk.Created
}
if chunk.ID != nil {
s.metadata.ID = chunk.ID
}
if chunk.Model != nil {
s.metadata.Model = chunk.Model
}
if chunk.ServiceTier != nil {
s.metadata.ServiceTier = chunk.ServiceTier
}
if chunk.SystemFingerprint != nil {
s.metadata.SystemFingerprint = chunk.SystemFingerprint
}
if chunk.Usage != nil { if chunk.Usage != nil {
s.usage = chunk.Usage s.metadata.Usage = chunk.Usage
} }
if chunk.Timings != nil { if chunk.Timings != nil {
s.timings = chunk.Timings s.metadata.Timings = chunk.Timings
} }
s.current = chunk s.current = chunk
return true return true
} }
// Current returns the most recently decoded chunk. The value is replaced by // Current returns the most recently decoded chunk. The value is replaced by
// the next successful call to Next. // the next successful call to Next. Previously returned chunks remain valid.
// Reported metadata shares pointer targets and nested logprob values with
// the stream and Result. Mutations through those references affect both.
func (s *CompletionStream) Current() StreamChunk { func (s *CompletionStream) Current() StreamChunk {
return s.current return s.current
} }
// Result returns a snapshot of the state accumulated so far. Err reports // Result returns the message and metadata accumulated so far. Advancing the
// whether the snapshot is a complete response or a partial result. // stream does not change previously returned metadata values. Metadata pointer
// targets and nested logprob slices are shared with the stream and other
// results; mutations through those references are visible through the aliases.
// Err reports failures. Nil Err alone does not establish completion: the
// stream completes successfully when Next consumes the DONE marker.
func (s *CompletionStream) Result() CompletionResult { func (s *CompletionStream) Result() CompletionResult {
message := s.acc.Message() message := s.acc.Message()
message.ToolCalls = append([]ToolCall(nil), message.ToolCalls...) message.ToolCalls = append([]ToolCall(nil), message.ToolCalls...)
return CompletionResult{ metadata := s.metadata
Message: message, if metadata.Logprobs != nil {
FinishReason: s.finishReason, logprobs := &Logprobs{}
Usage: s.usage, if metadata.Logprobs.Content != nil {
Timings: s.timings, logprobs.Content = make(
[]TokenLogprob, len(metadata.Logprobs.Content),
)
copy(logprobs.Content, metadata.Logprobs.Content)
}
if metadata.Logprobs.Refusal != nil {
logprobs.Refusal = make(
[]TokenLogprob, len(metadata.Logprobs.Refusal),
)
copy(logprobs.Refusal, metadata.Logprobs.Refusal)
}
metadata.Logprobs = logprobs
} }
return CompletionResult{Message: message, Metadata: metadata}
} }
// Err returns the stream failure, if any. Closing a stream does not create or // Err returns the stream failure, if any. Closing a stream does not create or
+636 -74
View File
@@ -2,12 +2,15 @@ package llm
import ( import (
"context" "context"
"encoding/json"
jsonv2 "encoding/json/v2"
"errors" "errors"
"fmt" "fmt"
"io" "io"
"log/slog" "log/slog"
"net/http" "net/http"
"net/http/httptest" "net/http/httptest"
"reflect"
"strings" "strings"
"testing" "testing"
) )
@@ -70,8 +73,8 @@ func TestCompletionsStreamReconstructsReasoning(t *testing.T) {
if result.Message.Text() != "answer" { if result.Message.Text() != "answer" {
t.Errorf("text = %q", result.Message.Text()) t.Errorf("text = %q", result.Message.Text())
} }
if result.FinishReason != "stop" { if result.Metadata.FinishReason == nil || *result.Metadata.FinishReason != "stop" {
t.Errorf("finish reason = %q", result.FinishReason) t.Errorf("finish reason = %v", result.Metadata.FinishReason)
} }
if stream.Next() { if stream.Next() {
t.Error("Next returned true after termination") t.Error("Next returned true after termination")
@@ -148,90 +151,644 @@ func TestCompletionsStreamReconstructsFragmentedParallelToolCalls(t *testing.T)
toolCalls[1].Function.Arguments != `{"format":"%s"}` { toolCalls[1].Function.Arguments != `{"format":"%s"}` {
t.Errorf("tool call 1 = %#v", toolCalls[1]) t.Errorf("tool call 1 = %#v", toolCalls[1])
} }
if result.FinishReason != "tool_calls" { if result.Metadata.FinishReason == nil || *result.Metadata.FinishReason != "tool_calls" {
t.Errorf("finish reason = %q", result.FinishReason) t.Errorf("finish reason = %v", result.Metadata.FinishReason)
} }
} }
func TestCompletionsReturnsUsageAndTimings(t *testing.T) { func TestCompletionsReturnsUsageAndTimings(t *testing.T) {
client := newCompletionTestClient(t, func(w http.ResponseWriter, r *http.Request) { tests := []struct {
if r.URL.Path != "/chat/completions" { name string
t.Fatalf("unexpected path: %s", r.URL.Path) body string
} want string
w.Header().Set("Content-Type", "application/json") }{
fmt.Fprint(w, `{ {
"id": "test-123", name: "all known metadata",
"model": "test-model", body: `{
"choices": [{ "id":"body-id","created":123,"model":"reported-model",
"index": 0, "service_tier":"default","system_fingerprint":"backend",
"message": {"role": "assistant", "content": "hello"}, "choices":[{"index":3,"finish_reason":"stop",
"finish_reason": "stop" "message":{"role":"assistant","content":"hello"},
}], "logprobs":{"content":[{"token":"hello","logprob":0,
"usage": { "bytes":[104,195,169],"top_logprobs":[{
"prompt_tokens": 10, "token":"other","logprob":-2,"bytes":[111]}]}],
"completion_tokens": 5, "refusal":[{"token":"no","logprob":-9999}]}}],
"total_tokens": 15, "usage":{"prompt_tokens":10,"completion_tokens":5,
"prompt_tokens_details": {"cached_tokens": 3}, "total_tokens":15,"prompt_tokens_details":{
"completion_tokens_details": {"reasoning_tokens": 2} "audio_tokens":1,"cache_write_tokens":2,
}, "cached_tokens":3,"image_tokens":4,"text_tokens":5},
"timings": { "completion_tokens_details":{"accepted_prediction_tokens":1,
"prompt_n": 10, "audio_tokens":2,"reasoning_tokens":3,
"prompt_ms": 50.5, "rejected_prediction_tokens":4,"text_tokens":5}},
"prompt_per_token_ms": 5.05, "timings":{"cache_n":1,"draft_n":2,"draft_n_accepted":1,
"prompt_per_second": 198.0, "prompt_n":10,"prompt_ms":50.5,"prompt_per_token_ms":5.05,
"predicted_n": 5, "prompt_per_second":198,"predicted_n":5,"predicted_ms":100,
"predicted_ms": 100.0, "predicted_per_token_ms":20,"predicted_per_second":50}
"predicted_per_token_ms": 20.0, }`,
"predicted_per_second": 50.0 want: `{
"choice_index":3,
"created":123,
"finish_reason":"stop",
"id":"body-id",
"logprobs":{
"content":[
{
"bytes":[104,195,169],
"logprob":0,
"token":"hello",
"top_logprobs":[
{
"bytes":[111],
"logprob":-2,
"token":"other"
}
]
}
],
"refusal":[
{
"bytes":null,
"logprob":-9999,
"token":"no",
"top_logprobs":null
}
]
},
"model":"reported-model",
"provider":"test",
"requested_model":"test-model",
"service_tier":"default",
"system_fingerprint":"backend",
"timings":{
"cache_n":1,
"draft_n":2,
"draft_n_accepted":1,
"predicted_ms":100,
"predicted_n":5,
"predicted_per_second":50,
"predicted_per_token_ms":20,
"prompt_ms":50.5,
"prompt_n":10,
"prompt_per_second":198,
"prompt_per_token_ms":5.05
},
"usage":{
"completion_tokens":5,
"completion_tokens_details":{
"accepted_prediction_tokens":1,
"audio_tokens":2,
"reasoning_tokens":3,
"rejected_prediction_tokens":4,
"text_tokens":5
},
"prompt_tokens":10,
"prompt_tokens_details":{
"audio_tokens":1,
"cache_write_tokens":2,
"cached_tokens":3,
"image_tokens":4,
"text_tokens":5
},
"total_tokens":15
}
}`,
},
{
name: "omitted metadata",
body: `{"choices":[{"message":{"content":"hello"}}]}`,
want: `{
"choice_index":null,
"created":null,
"finish_reason":null,
"id":null,
"logprobs":null,
"model":null,
"provider":"test",
"requested_model":"test-model",
"service_tier":null,
"system_fingerprint":null,
"timings":null,
"usage":null
}`,
},
{
name: "reported zero and empty values",
body: `{"id":"","created":0,"model":"","service_tier":"",
"system_fingerprint":"","choices":[{"index":0,
"finish_reason":"","message":{"content":"hello"}}],
"usage":{"prompt_tokens":0,"completion_tokens":0,"total_tokens":0,
"prompt_tokens_details":{"cached_tokens":0},
"completion_tokens_details":{"reasoning_tokens":0}},
"timings":{"prompt_ms":0,"predicted_n":0}}`,
want: `{
"choice_index":0,
"created":0,
"finish_reason":"",
"id":"",
"logprobs":null,
"model":"",
"provider":"test",
"requested_model":"test-model",
"service_tier":"",
"system_fingerprint":"",
"timings":{
"cache_n":null,
"draft_n":null,
"draft_n_accepted":null,
"predicted_ms":null,
"predicted_n":0,
"predicted_per_second":null,
"predicted_per_token_ms":null,
"prompt_ms":0,
"prompt_n":null,
"prompt_per_second":null,
"prompt_per_token_ms":null
},
"usage":{
"completion_tokens":0,
"completion_tokens_details":{
"accepted_prediction_tokens":null,
"audio_tokens":null,
"reasoning_tokens":0,
"rejected_prediction_tokens":null,
"text_tokens":null
},
"prompt_tokens":0,
"prompt_tokens_details":{
"audio_tokens":null,
"cache_write_tokens":null,
"cached_tokens":0,
"image_tokens":null,
"text_tokens":null
},
"total_tokens":0
}
}`,
},
{
name: "empty and null logprob arrays",
body: `{"choices":[{"message":{"content":"hello"},
"logprobs":{"content":[{"bytes":[],"top_logprobs":[]},
{"bytes":null,"top_logprobs":null},
{"top_logprobs":[{"bytes":[]},{}]}],"refusal":[]}}]}`,
want: `{
"choice_index":null,
"created":null,
"finish_reason":null,
"id":null,
"logprobs":{
"content":[
{
"bytes":[],
"logprob":null,
"token":null,
"top_logprobs":[]
},
{
"bytes":null,
"logprob":null,
"token":null,
"top_logprobs":null
},
{
"bytes":null,
"logprob":null,
"token":null,
"top_logprobs":[
{
"bytes":[],
"logprob":null,
"token":null
},
{
"bytes":null,
"logprob":null,
"token":null
}
]
}
],
"refusal":[]
},
"model":null,
"provider":"test",
"requested_model":"test-model",
"service_tier":null,
"system_fingerprint":null,
"timings":null,
"usage":null
}`,
},
{
name: "empty nested metadata objects",
body: `{"choices":[{"message":{"content":"hello"},
"logprobs":{}}],"usage":{"prompt_tokens_details":{},
"completion_tokens_details":{}},"timings":{}}`,
want: `{
"choice_index":null,
"created":null,
"finish_reason":null,
"id":null,
"logprobs":{"content":null,"refusal":null},
"model":null,
"provider":"test",
"requested_model":"test-model",
"service_tier":null,
"system_fingerprint":null,
"timings":{
"cache_n":null,
"draft_n":null,
"draft_n_accepted":null,
"predicted_ms":null,
"predicted_n":null,
"predicted_per_second":null,
"predicted_per_token_ms":null,
"prompt_ms":null,
"prompt_n":null,
"prompt_per_second":null,
"prompt_per_token_ms":null
},
"usage":{
"completion_tokens":null,
"completion_tokens_details":{
"accepted_prediction_tokens":null,
"audio_tokens":null,
"reasoning_tokens":null,
"rejected_prediction_tokens":null,
"text_tokens":null
},
"prompt_tokens":null,
"prompt_tokens_details":{
"audio_tokens":null,
"cache_write_tokens":null,
"cached_tokens":null,
"image_tokens":null,
"text_tokens":null
},
"total_tokens":null
}
}`,
},
{
name: "null metadata",
body: `{"id":null,"usage":null,"timings":null,
"choices":[{"index":null,"logprobs":null,
"message":{"content":"hello"}}]}`,
want: `{
"choice_index":null,
"created":null,
"finish_reason":null,
"id":null,
"logprobs":null,
"model":null,
"provider":"test",
"requested_model":"test-model",
"service_tier":null,
"system_fingerprint":null,
"timings":null,
"usage":null
}`,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
client := newCompletionTestClient(t, func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/chat/completions" {
t.Errorf("unexpected path: %s", r.URL.Path)
}
var request ChatRequest
if err := json.NewDecoder(r.Body).Decode(&request); err != nil {
t.Error(err)
}
if request.LogProbs || request.TopLogProbs != 0 {
t.Error("log probabilities requested")
}
w.Header().Set("Request-ID", "not-the-body-id")
fmt.Fprint(w, tt.body)
})
res, err := client.Completions(t.Context(), CompletionsParams{Model: "test-model"})
if err != nil {
t.Fatalf("Completions: %v", err)
} }
}`) if res.Message.Text() != "hello" {
}) t.Errorf("text = %q", res.Message.Text())
}
res, err := client.Completions(context.Background(), CompletionsParams{Model: "test-model"}) var want Metadata
if err != nil { if err := json.Unmarshal([]byte(tt.want), &want); err != nil {
t.Fatalf("Completions: %v", err) t.Fatal(err)
} }
if res.Message.Text() != "hello" || res.Usage == nil || res.Timings == nil { if !reflect.DeepEqual(res.Metadata, want) {
t.Fatalf("unexpected result: %#v", res) t.Errorf("metadata = %#v, want %#v", res.Metadata, want)
} }
if res.Usage.PromptTokens != 10 || res.Usage.CompletionTokens != 5 || res.Usage.TotalTokens != 15 { encoded, err := jsonv2.Marshal(
t.Errorf("usage = %#v", res.Usage) res.Metadata, jsonv2.FormatNilSliceAsNull(true),
} )
if res.Usage.PromptTokensDetails == nil || res.Usage.PromptTokensDetails.CachedTokens != 3 { if err != nil {
t.Errorf("prompt token details = %#v", res.Usage.PromptTokensDetails) t.Fatal(err)
} }
if res.Usage.CompletionTokensDetails == nil || res.Usage.CompletionTokensDetails.ReasoningTokens != 2 { var gotJSON, wantJSON map[string]any
t.Errorf("completion token details = %#v", res.Usage.CompletionTokensDetails) if err := jsonv2.Unmarshal(encoded, &gotJSON); err != nil {
} t.Fatal(err)
if res.Timings.PromptMS != 50.5 || res.Timings.PredictedMS != 100 { }
t.Errorf("timings = %#v", res.Timings) if err := jsonv2.Unmarshal([]byte(tt.want), &wantJSON); err != nil {
t.Fatal(err)
}
if !reflect.DeepEqual(gotJSON, wantJSON) {
t.Errorf("JSON fields = %s, want %s", encoded, tt.want)
}
var roundtrip Metadata
if err := json.Unmarshal(encoded, &roundtrip); err != nil {
t.Fatal(err)
}
if !reflect.DeepEqual(roundtrip, want) {
t.Errorf("roundtrip = %s", encoded)
}
if strings.Contains(tt.body, `"bytes":[104,195,169]`) &&
!strings.Contains(string(encoded), `"bytes":[104,195,169]`) {
t.Errorf("token bytes are not integer arrays: %s", encoded)
}
})
} }
} }
func TestCompletionsStreamSSEAndMetadata(t *testing.T) { func TestCompletionsStreamSSEAndMetadata(t *testing.T) {
large := strings.Repeat("x", 70*1024) large := strings.Repeat("x", 70*1024)
client := newCompletionTestClient(t, func(w http.ResponseWriter, _ *http.Request) { tests := []struct {
fmt.Fprint(w, ": comment\n") name string
fmt.Fprintf(w, "data:{\"choices\":[{\"index\":0,\"delta\":{\"content\":%q}}],\n", large) response string
fmt.Fprint(w, "data: \"usage\":{\"prompt_tokens\":10,\"total_tokens\":15}}\n\n") wantText string
fmt.Fprint(w, "data: {\"choices\":[],\"timings\":{\"prompt_ms\":50.5}}\n\n") wantMetadata string
fmt.Fprint(w, "data:[DONE]\n\n") wantChunks int
}) }{
{
name: "large multiline event and metadata-only timing",
response: ": comment\n" +
fmt.Sprintf("data:{\"choices\":[{\"index\":0,\"delta\":{\"content\":%q}}],\n", large) +
"data: \"usage\":{\"prompt_tokens\":10,\"total_tokens\":15}}\n\n" +
"data: {\"choices\":[],\"timings\":{\"prompt_ms\":50.5}}\n\n",
wantText: large,
wantMetadata: `{
"choice_index":0,
"created":null,
"finish_reason":null,
"id":null,
"logprobs":null,
"model":null,
"provider":"test",
"requested_model":"test-model",
"service_tier":null,
"system_fingerprint":null,
"timings":{
"cache_n":null,
"draft_n":null,
"draft_n_accepted":null,
"predicted_ms":null,
"predicted_n":null,
"predicted_per_second":null,
"predicted_per_token_ms":null,
"prompt_ms":50.5,
"prompt_n":null,
"prompt_per_second":null,
"prompt_per_token_ms":null
},
"usage":{
"completion_tokens":null,
"completion_tokens_details":null,
"prompt_tokens":10,
"prompt_tokens_details":null,
"total_tokens":15
}
}`,
wantChunks: 2,
},
{
name: "last scalars and snapshots, ordered logprobs",
response: `data: {"id":"first","created":123,"model":"first-model","service_tier":"first-tier","system_fingerprint":"first-backend","choices":[{"index":0,"delta":{"content":"answer"},"finish_reason":"stop","logprobs":{"content":[{"token":"a","logprob":0,"bytes":[97],"top_logprobs":[{"token":"b","logprob":-1,"bytes":[98]}]}],"refusal":[{"token":"r1","logprob":-2}]}}],"usage":{"prompt_tokens":5,"total_tokens":9},"timings":{"prompt_ms":5,"predicted_n":3}}
stream := openCompletionTestStream(t, client) data: {"choices":[{"index":0,"delta":{},"finish_reason":"","logprobs":{"content":[{"token":"b","logprob":-1}],"refusal":[{"token":"r2","logprob":0}]}}]}
var count int
for stream.Next() { data: {"id":"","created":0,"model":"reported-model","service_tier":"last-tier","system_fingerprint":"last-backend","choices":[],"usage":{"prompt_tokens":0},"timings":{"predicted_n":0}}
count++
data: {"choices":[],"usage":null,"timings":null,"id":null}
`,
wantText: "answer",
wantMetadata: `{
"choice_index":0,
"created":0,
"finish_reason":"",
"id":"",
"logprobs":{
"content":[
{
"bytes":[97],
"logprob":0,
"token":"a",
"top_logprobs":[
{
"bytes":[98],
"logprob":-1,
"token":"b"
}
]
},
{
"bytes":null,
"logprob":-1,
"token":"b",
"top_logprobs":null
}
],
"refusal":[
{
"bytes":null,
"logprob":-2,
"token":"r1",
"top_logprobs":null
},
{
"bytes":null,
"logprob":0,
"token":"r2",
"top_logprobs":null
}
]
},
"model":"reported-model",
"provider":"test",
"requested_model":"test-model",
"service_tier":"last-tier",
"system_fingerprint":"last-backend",
"timings":{
"cache_n":null,
"draft_n":null,
"draft_n_accepted":null,
"predicted_ms":null,
"predicted_n":0,
"predicted_per_second":null,
"predicted_per_token_ms":null,
"prompt_ms":null,
"prompt_n":null,
"prompt_per_second":null,
"prompt_per_token_ms":null
},
"usage":{
"completion_tokens":null,
"completion_tokens_details":null,
"prompt_tokens":0,
"prompt_tokens_details":null,
"total_tokens":null
}
}`,
wantChunks: 4,
},
{
name: "empty and null logprob arrays",
response: `data: {"choices":[{"index":0,"delta":{},"logprobs":{"content":[],"refusal":null}}]}
data: {"choices":[{"index":0,"delta":{},"logprobs":{"content":null,"refusal":[{"bytes":[],"top_logprobs":[]},{"top_logprobs":[{"bytes":[]}]}]}}]}
`,
wantMetadata: `{
"choice_index":0,
"created":null,
"finish_reason":null,
"id":null,
"logprobs":{
"content":[],
"refusal":[
{
"bytes":[],
"logprob":null,
"token":null,
"top_logprobs":[]
},
{
"bytes":null,
"logprob":null,
"token":null,
"top_logprobs":[
{
"bytes":[],
"logprob":null,
"token":null
}
]
}
]
},
"model":null,
"provider":"test",
"requested_model":"test-model",
"service_tier":null,
"system_fingerprint":null,
"timings":null,
"usage":null
}`,
wantChunks: 2,
},
{
name: "metadata-only without choice",
response: "data: {\"created\":0,\"usage\":{\"total_tokens\":0}}\n\n",
wantMetadata: `{
"choice_index":null,
"created":0,
"finish_reason":null,
"id":null,
"logprobs":null,
"model":null,
"provider":"test",
"requested_model":"test-model",
"service_tier":null,
"system_fingerprint":null,
"timings":null,
"usage":{
"completion_tokens":null,
"completion_tokens_details":null,
"prompt_tokens":null,
"prompt_tokens_details":null,
"total_tokens":0
}
}`,
wantChunks: 1,
},
{
name: "omitted choice index and metadata",
response: "data: {\"choices\":[{\"delta\":{\"content\":\"answer\"}}]}\n\n",
wantText: "answer",
wantMetadata: `{
"choice_index":null,
"created":null,
"finish_reason":null,
"id":null,
"logprobs":null,
"model":null,
"provider":"test",
"requested_model":"test-model",
"service_tier":null,
"system_fingerprint":null,
"timings":null,
"usage":null
}`,
wantChunks: 1,
},
} }
if err := stream.Err(); err != nil { for _, tt := range tests {
t.Fatalf("stream.Err: %v", err) t.Run(tt.name, func(t *testing.T) {
} client := newCompletionTestClient(t, func(w http.ResponseWriter, r *http.Request) {
result := stream.Result() var request ChatRequest
if count != 2 || result.Message.Text() != large { if err := json.NewDecoder(r.Body).Decode(&request); err != nil {
t.Fatalf("count = %d, text length = %d", count, len(result.Message.Text())) t.Error(err)
} }
if result.Usage == nil || result.Usage.PromptTokens != 10 || result.Timings == nil || result.Timings.PromptMS != 50.5 { if request.LogProbs || request.TopLogProbs != 0 {
t.Errorf("metadata = usage %#v, timings %#v", result.Usage, result.Timings) t.Error("log probabilities requested")
}
fmt.Fprint(w, tt.response+"data:[DONE]\n\n")
})
stream := openCompletionTestStream(t, client)
var count int
for stream.Next() {
count++
}
if err := stream.Err(); err != nil {
t.Fatalf("stream.Err: %v", err)
}
result := stream.Result()
if count != tt.wantChunks || result.Message.Text() != tt.wantText {
t.Fatalf("count = %d, text length = %d", count, len(result.Message.Text()))
}
var want Metadata
if err := json.Unmarshal([]byte(tt.wantMetadata), &want); err != nil {
t.Fatal(err)
}
encoded, err := jsonv2.Marshal(
result.Metadata, jsonv2.FormatNilSliceAsNull(true),
)
if err != nil {
t.Fatal(err)
}
var gotJSON, wantJSON map[string]any
if err := jsonv2.Unmarshal(encoded, &gotJSON); err != nil {
t.Fatal(err)
}
if err := jsonv2.Unmarshal([]byte(tt.wantMetadata), &wantJSON); err != nil {
t.Fatal(err)
}
if !reflect.DeepEqual(gotJSON, wantJSON) {
t.Errorf("JSON fields = %s, want %s", encoded, tt.wantMetadata)
}
var roundtrip Metadata
if err := json.Unmarshal(encoded, &roundtrip); err != nil {
t.Fatal(err)
}
if !reflect.DeepEqual(roundtrip, want) {
t.Errorf("metadata roundtrip = %s", encoded)
}
if !reflect.DeepEqual(result.Metadata, want) {
encoded, _ := json.Marshal(result.Metadata)
t.Errorf("metadata = %s, want %s", encoded, tt.wantMetadata)
}
})
} }
} }
@@ -246,6 +803,11 @@ func TestCompletionsStreamFailuresPreservePartialResult(t *testing.T) {
response: "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"partial\"}}]}\n\n", response: "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"partial\"}}]}\n\n",
wantErr: io.ErrUnexpectedEOF, wantErr: io.ErrUnexpectedEOF,
}, },
{
name: "finish reason without DONE",
response: "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"partial\"},\"finish_reason\":\"stop\"}]}\n\n",
wantErr: io.ErrUnexpectedEOF,
},
{ {
name: "malformed JSON", name: "malformed JSON",
response: "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"partial\"}}]}\n\ndata: {bad}\n\n", response: "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"partial\"}}]}\n\ndata: {bad}\n\n",
+221 -46
View File
@@ -1,6 +1,5 @@
// Package llm implements an OpenAI-compatible inference client and its // Package llm implements an OpenAI-compatible inference client and its
// protocol types. // protocol types.
package llm package llm
import ( import (
@@ -114,69 +113,229 @@ type ChatRequest struct {
TopP *float32 `json:"top_p,omitempty"` TopP *float32 `json:"top_p,omitempty"`
} }
// Usage reports input, output, and total token counts for a completion. // Metadata describes one selected completion, separately from its message.
type Usage struct { // Nil optional fields mean the endpoint did not report a value. Reported zero
// CompletionTokens is the number of tokens generated by the model. // counts and empty strings remain distinct from absent values. Only known body
CompletionTokens int `json:"completion_tokens"` // fields are retained; response headers are not part of Metadata.
// CompletionTokensDetails breaks down CompletionTokens by category. //
CompletionTokensDetails *TokenDetails `json:"completion_tokens_details,omitempty"` // See the response and streaming definitions in
// PromptTokens is the number of tokens in the input. // [OpenAI completion response] and [OpenAI completion chunk].
PromptTokens int `json:"prompt_tokens"` //
// PromptTokensDetails breaks down PromptTokens by category. // [OpenAI completion response]: https://developers.openai.com/api/reference/resources/chat/subresources/completions
PromptTokensDetails *TokenDetails `json:"prompt_tokens_details,omitempty"` // [OpenAI completion chunk]: https://developers.openai.com/api/reference/resources/chat/subresources/completions/streaming-events
// TotalTokens is the sum of PromptTokens and CompletionTokens. type Metadata struct {
TotalTokens int `json:"total_tokens"` // ChoiceIndex is the endpoint's zero-based index of the selected
// choice.
ChoiceIndex *int `json:"choice_index"`
// Created is the endpoint's completion creation time in Unix seconds,
// not the database message's creation time.
Created *int64 `json:"created"`
// FinishReason is the endpoint's reason for ending the selected choice.
// A reported reason does not establish receipt of the stream's DONE
// marker.
FinishReason *string `json:"finish_reason"`
// ID is the completion ID from the response body, not a request-ID
// header or a database message ID.
ID *string `json:"id"`
// Logprobs contains token probabilities for the selected choice.
Logprobs *Logprobs `json:"logprobs"`
// Model is the model identifier reported in the response body. It can
// differ from RequestedModel, which was sent in the request.
Model *string `json:"model"`
// Provider is the application-defined label of the configured endpoint.
// It is not a provider-reported identity or a response header.
Provider *string `json:"provider"`
// RequestedModel is the model identifier sent to the endpoint for this
// inference, before any endpoint-side model alias resolution.
RequestedModel *string `json:"requested_model"`
// ServiceTier is the processing tier reported in the response body,
// which can differ from a requested tier.
ServiceTier *string `json:"service_tier"`
// SystemFingerprint identifies the backend configuration reported by
// the endpoint; it is not a model identifier.
SystemFingerprint *string `json:"system_fingerprint"`
// Timings contains performance measurements for the LLM request that
// produced this assistant message. During streaming, a newer timing
// report replaces the earlier report: the measurements cover the same
// request at a later point, not additional work to add together. Fields
// missing from the newer report remain nil.
Timings *Timings `json:"timings"`
// Usage contains token counts for the LLM request that produced this
// assistant message. During streaming, a newer usage report replaces
// the earlier report rather than adding to its counts. Fields missing
// from the newer report remain nil.
Usage *Usage `json:"usage"`
} }
// TokenDetails reports token categories within a Usage count. A zero value may // Usage reports input, output, and total token counts for one inference.
// mean either that no tokens belong to the category or that it was not // Nil counts mean unreported, not zero. Category counts are endpoint-reported
// reported. // breakdowns, not additional tokens to add to the totals.
type TokenDetails struct { //
// CachedTokens is the number of input tokens served from a cache. // See [OpenAI completion usage].
CachedTokens int `json:"cached_tokens,omitempty"` //
// [OpenAI completion usage]: https://developers.openai.com/api/reference/resources/chat/subresources/completions
type Usage struct {
// CompletionTokens is the number of tokens generated by the model.
CompletionTokens *int `json:"completion_tokens"`
// CompletionTokensDetails breaks down CompletionTokens by category.
CompletionTokensDetails *CompletionTokensDetails `json:"completion_tokens_details"`
// PromptTokens is the number of tokens in the input.
PromptTokens *int `json:"prompt_tokens"`
// PromptTokensDetails breaks down PromptTokens by category.
PromptTokensDetails *PromptTokensDetails `json:"prompt_tokens_details"`
// TotalTokens is the reported sum of prompt and completion token
// counts.
TotalTokens *int `json:"total_tokens"`
}
// CompletionTokensDetails reports output-token categories in [Usage].
// Nil counts mean the category was not reported. The categories follow
// [OpenAI completion usage].
//
// [OpenAI completion usage]: https://developers.openai.com/api/reference/resources/chat/subresources/completions
type CompletionTokensDetails struct {
// AcceptedPredictionTokens counts predicted-output tokens that appeared
// in the completion.
AcceptedPredictionTokens *int `json:"accepted_prediction_tokens"`
// AudioTokens counts audio tokens generated by the model.
AudioTokens *int `json:"audio_tokens"`
// ReasoningTokens is the number of output tokens used for reasoning. // ReasoningTokens is the number of output tokens used for reasoning.
ReasoningTokens int `json:"reasoning_tokens,omitempty"` ReasoningTokens *int `json:"reasoning_tokens"`
// RejectedPredictionTokens counts predicted-output tokens that did not
// appear in the completion. OpenAI includes these in completion counts
// for billing, output, and context-window limits.
RejectedPredictionTokens *int `json:"rejected_prediction_tokens"`
// TextTokens counts text tokens generated by the model.
TextTokens *int `json:"text_tokens"`
}
// PromptTokensDetails reports input-token categories in [Usage].
// Nil counts mean the category was not reported. The categories follow
// [OpenAI completion usage].
//
// [OpenAI completion usage]: https://developers.openai.com/api/reference/resources/chat/subresources/completions
type PromptTokensDetails struct {
// AudioTokens counts audio input tokens in the prompt.
AudioTokens *int `json:"audio_tokens"`
// CacheWriteTokens is the unadjusted count of prompt tokens written to
// cache, as reported by the endpoint.
CacheWriteTokens *int `json:"cache_write_tokens"`
// CachedTokens is the number of input tokens served from a cache.
CachedTokens *int `json:"cached_tokens"`
// ImageTokens counts image input tokens in the prompt.
ImageTokens *int `json:"image_tokens"`
// TextTokens counts text input tokens in the prompt.
TextTokens *int `json:"text_tokens"`
}
// Logprobs reports token log probabilities for one choice, when returned.
// Content and refusal tokens are separate ordered sequences. Receiving these
// values does not cause the client to request log probabilities. Nil slices
// represent absent or null arrays; non-nil empty slices retain reported [].
//
// See [OpenAI token log probabilities] and [OpenAI completion response].
//
// [OpenAI token log probabilities]: https://developers.openai.com/api/reference/resources/chat/subresources/completions
// [OpenAI completion response]: https://developers.openai.com/api/reference/resources/chat/subresources/completions
type Logprobs struct {
// Content contains log probabilities for message content tokens in
// order.
Content []TokenLogprob `json:"content"`
// Refusal contains log probabilities for refusal tokens in order.
Refusal []TokenLogprob `json:"refusal"`
}
// TokenLogprob describes a generated token and its alternatives at one
// position, using [OpenAI token log probabilities]. Nil optional scalars mean
// unreported; a zero log probability remains a reported value.
//
// [OpenAI token log probabilities]: https://developers.openai.com/api/reference/resources/chat/subresources/completions
type TokenLogprob struct {
// Bytes is the token's UTF-8 byte representation as integers, not
// base64.
// A token can contain only part of a multibyte character. Nil means no
// byte representation was reported.
Bytes []int `json:"bytes"`
// Logprob is the token's log probability. OpenAI uses -9999 for a token
// outside the top 20 most likely tokens.
Logprob *float64 `json:"logprob"`
// Token is the token text reported by the endpoint.
Token *string `json:"token"`
// TopLogprobs contains alternatives at this position; the endpoint can
// return fewer alternatives than requested.
TopLogprobs []TopLogprob `json:"top_logprobs"`
}
// TopLogprob describes one alternative token at a generated token's position.
// Alternatives have no nested alternatives in [OpenAI token log probabilities].
// Nil optional scalars mean unreported, rather than zero or empty text.
//
// [OpenAI token log probabilities]: https://developers.openai.com/api/reference/resources/chat/subresources/completions
type TopLogprob struct {
// Bytes is the alternative's UTF-8 byte representation as integers,
// not base64. Nil means no byte representation was reported.
Bytes []int `json:"bytes"`
// Logprob is the alternative's log probability. OpenAI uses -9999 for
// a token outside the top 20 most likely tokens.
Logprob *float64 `json:"logprob"`
// Token is the alternative token text reported by the endpoint.
Token *string `json:"token"`
} }
// Timings reports llama.cpp prompt-processing and token-generation performance. // Timings reports llama.cpp prompt-processing and token-generation performance.
// The N fields are token counts, MS fields are elapsed milliseconds, // The N fields are token counts, MS fields are elapsed milliseconds,
// PerTokenMS fields are milliseconds per token, and PerSecond fields are tokens // PerTokenMS fields are milliseconds per token, and PerSecond fields are tokens
// per second. Timings is a llama.cpp extension to the OpenAI API. // per second. Timings is a llama.cpp extension to the OpenAI API. Nil fields
// mean unreported. Streaming objects are snapshots, not increments to sum.
// //
// See [llama.cpp API documentation]. // See [llama.cpp API documentation].
// //
// [llama.cpp API documentation]: https://llama.app/docs/api // [llama.cpp API documentation]: https://llama.app/docs/api#timings-and-usage
type Timings struct { type Timings struct {
// CacheN is the number of prompt tokens reused from the cache.
CacheN *int `json:"cache_n"`
// DraftN is the number of tokens drafted during speculative decoding.
DraftN *int `json:"draft_n"`
// DraftNAccepted is the number of drafted tokens accepted during
// speculative decoding.
DraftNAccepted *int `json:"draft_n_accepted"`
// PredictedMS is the time spent generating output, in milliseconds. // PredictedMS is the time spent generating output, in milliseconds.
PredictedMS float64 `json:"predicted_ms"` PredictedMS *float64 `json:"predicted_ms"`
// PredictedN is the number of generated tokens. // PredictedN is the number of generated tokens.
PredictedN int `json:"predicted_n"` PredictedN *int `json:"predicted_n"`
// PredictedPerSecond is the output-generation rate in tokens per // PredictedPerSecond is the output-generation rate in tokens per
// second. // second.
PredictedPerSecond float64 `json:"predicted_per_second"` PredictedPerSecond *float64 `json:"predicted_per_second"`
// PredictedPerTokenMS is the mean output-generation time per token, in // PredictedPerTokenMS is the mean output-generation time per token, in
// milliseconds. // milliseconds.
PredictedPerTokenMS float64 `json:"predicted_per_token_ms"` PredictedPerTokenMS *float64 `json:"predicted_per_token_ms"`
// PromptMS is the time spent processing the prompt, in milliseconds. // PromptMS is the time spent processing the prompt, in milliseconds.
PromptMS float64 `json:"prompt_ms"` PromptMS *float64 `json:"prompt_ms"`
// PromptN is the number of prompt tokens processed. // PromptN is the number of prompt tokens processed.
PromptN int `json:"prompt_n"` PromptN *int `json:"prompt_n"`
// PromptPerSecond is the prompt-processing rate in tokens per second. // PromptPerSecond is the prompt-processing rate in tokens per second.
PromptPerSecond float64 `json:"prompt_per_second"` PromptPerSecond *float64 `json:"prompt_per_second"`
// PromptPerTokenMS is the mean prompt-processing time per token, in // PromptPerTokenMS is the mean prompt-processing time per token, in
// milliseconds. // milliseconds.
PromptPerTokenMS float64 `json:"prompt_per_token_ms"` PromptPerTokenMS *float64 `json:"prompt_per_token_ms"`
} }
// ChatResponse is the body of a non-streaming chat completion // ChatResponse is the body of a non-streaming chat completion
// response. // response. Nil metadata fields mean unreported values.
type ChatResponse struct { type ChatResponse struct {
// Choices contains the generated alternatives. // Choices contains the generated alternatives.
Choices []Choice `json:"choices"` Choices []Choice `json:"choices"`
// ID identifies the completion. // Created is the endpoint's completion creation time in Unix seconds.
ID string `json:"id"` Created *int64 `json:"created,omitempty"`
// ID identifies the completion in the response body, not a request
// header.
ID *string `json:"id,omitempty"`
// Model identifies the model that generated the response. // Model identifies the model that generated the response.
Model string `json:"model"` Model *string `json:"model,omitempty"`
// ServiceTier is the processing tier reported by the endpoint.
ServiceTier *string `json:"service_tier,omitempty"`
// SystemFingerprint identifies the reported backend configuration.
SystemFingerprint *string `json:"system_fingerprint,omitempty"`
// Timings contains llama.cpp performance data, when reported. // Timings contains llama.cpp performance data, when reported.
Timings *Timings `json:"timings,omitempty"` Timings *Timings `json:"timings,omitempty"`
// Usage contains token counts, when reported. // Usage contains token counts, when reported.
@@ -185,12 +344,15 @@ type ChatResponse struct {
// Choice is one generated alternative in a non-streaming response. Index is its // Choice is one generated alternative in a non-streaming response. Index is its
// position among the alternatives, and FinishReason reports why generation // position among the alternatives, and FinishReason reports why generation
// ended. // ended. Nil metadata fields mean the endpoint did not report those values.
type Choice struct { type Choice struct {
// FinishReason is the endpoint's reason for ending generation. // FinishReason is the endpoint's reason for ending generation.
FinishReason string `json:"finish_reason"` FinishReason *string `json:"finish_reason,omitempty"`
// Index is the position of this choice in the response. // Index is the zero-based position of this choice in the response.
Index int `json:"index"` Index *int `json:"index,omitempty"`
// Logprobs contains probabilities for this choice's content and
// refusal.
Logprobs *Logprobs `json:"logprobs,omitempty"`
// Message is the generated message. // Message is the generated message.
Message Message `json:"message"` Message Message `json:"message"`
} }
@@ -468,22 +630,35 @@ type ResponseFormat struct {
// StreamChunk is one data event in a streaming response. Delta fields are // StreamChunk is one data event in a streaming response. Delta fields are
// incremental and must be combined in order to reconstruct a message. A chunk // incremental and must be combined in order to reconstruct a message. A chunk
// may contain Usage or Timings without containing a choice. // may contain Usage or Timings without containing a choice. Nil metadata fields
// mean unreported values, not reported zero counts or empty strings.
type StreamChunk struct { type StreamChunk struct {
// Choices contains incremental updates for generated alternatives. // Choices contains incremental updates for generated alternatives.
// It may be empty in a metadata-only chunk. // It may be empty in a metadata-only chunk.
Choices []struct { Choices []struct {
// Index identifies the choice updated by this delta.
Index int `json:"index"`
// Delta contains the incremental message update. // Delta contains the incremental message update.
Delta Delta `json:"delta"` Delta Delta `json:"delta"`
// FinishReason is set when generation for the choice ends. // FinishReason is the reported reason for ending this choice;
FinishReason string `json:"finish_reason,omitempty"` // it does not establish receipt of DONE.
FinishReason *string `json:"finish_reason,omitempty"`
// Index identifies the choice updated by this delta. Nil means
// unreported and selects the first choice for accumulation.
Index *int `json:"index,omitempty"`
// Logprobs contains the next content and refusal token
// probabilities.
Logprobs *Logprobs `json:"logprobs,omitempty"`
} `json:"choices"` } `json:"choices"`
// ID identifies the completion to which the chunk belongs. // Created is the endpoint's completion creation time in Unix seconds.
ID string `json:"id"` Created *int64 `json:"created,omitempty"`
// ID identifies the completion in the response body, not a request
// header.
ID *string `json:"id,omitempty"`
// Model identifies the model generating the response. // Model identifies the model generating the response.
Model string `json:"model"` Model *string `json:"model,omitempty"`
// ServiceTier is the processing tier reported by the endpoint.
ServiceTier *string `json:"service_tier,omitempty"`
// SystemFingerprint identifies the reported backend configuration.
SystemFingerprint *string `json:"system_fingerprint,omitempty"`
// Timings contains llama.cpp performance data, when reported. // Timings contains llama.cpp performance data, when reported.
Timings *Timings `json:"timings,omitempty"` Timings *Timings `json:"timings,omitempty"`
// Usage contains aggregate token counts, when reported. // Usage contains aggregate token counts, when reported.