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.
type CompletionResult struct {
// FinishReason is the endpoint's reason for ending generation.
FinishReason string
// Message is the generated message.
Message Message
// Timings contains llama.cpp performance data, when reported.
Timings *Timings
// Usage contains token counts, when reported.
Usage *Usage
// Metadata describes this inference's selected choice and endpoint,
// separately from the model-visible message.
Metadata Metadata
}
// Completions requests a non-streaming chat completion and returns the first
@@ -52,10 +49,21 @@ func (c *Client) Completions(
}
choice := cr.Choices[0]
return &CompletionResult{
Message: choice.Message,
FinishReason: choice.FinishReason,
Usage: cr.Usage,
Timings: cr.Timings,
Message: choice.Message,
Metadata: Metadata{
ChoiceIndex: choice.Index,
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
}
@@ -240,11 +248,13 @@ type CompletionStream struct {
body io.ReadCloser
cancel context.CancelFunc
acc StreamAccumulator
current StreamChunk
finishReason string
usage *Usage
timings *Timings
acc StreamAccumulator
current StreamChunk
// metadata retains reported scalars and the latest usage/timing
// snapshots.
// Logprob sequences contain every reported token in arrival order.
// Nested metadata values are shared with Current and Result.
metadata Metadata
err error
cleanupErr error
terminated bool
@@ -352,10 +362,11 @@ func (c *Client) CompletionsStream(
}
return &CompletionStream{
ctx: ctx,
reader: bufio.NewReader(res.Body),
body: res.Body,
cancel: cancel,
ctx: ctx,
reader: bufio.NewReader(res.Body),
body: res.Body,
cancel: cancel,
metadata: Metadata{Provider: new(c.provider), RequestedModel: &p.Model},
}, nil
}
@@ -412,45 +423,102 @@ func (s *CompletionStream) Next() bool {
return false
}
for _, choice := range chunk.Choices {
if choice.Index != 0 {
if choice.Index != nil && *choice.Index != 0 {
s.fail(fmt.Errorf(
"completions stream: unexpected choice index %d",
choice.Index,
*choice.Index,
))
return false
}
if choice.FinishReason != "" {
s.finishReason = choice.FinishReason
if choice.Index != nil {
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)
}
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 {
s.usage = chunk.Usage
s.metadata.Usage = chunk.Usage
}
if chunk.Timings != nil {
s.timings = chunk.Timings
s.metadata.Timings = chunk.Timings
}
s.current = chunk
return true
}
// 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 {
return s.current
}
// Result returns a snapshot of the state accumulated so far. Err reports
// whether the snapshot is a complete response or a partial result.
// Result returns the message and metadata accumulated so far. Advancing the
// 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 {
message := s.acc.Message()
message.ToolCalls = append([]ToolCall(nil), message.ToolCalls...)
return CompletionResult{
Message: message,
FinishReason: s.finishReason,
Usage: s.usage,
Timings: s.timings,
metadata := s.metadata
if metadata.Logprobs != nil {
logprobs := &Logprobs{}
if metadata.Logprobs.Content != nil {
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
+636 -74
View File
@@ -2,12 +2,15 @@ package llm
import (
"context"
"encoding/json"
jsonv2 "encoding/json/v2"
"errors"
"fmt"
"io"
"log/slog"
"net/http"
"net/http/httptest"
"reflect"
"strings"
"testing"
)
@@ -70,8 +73,8 @@ func TestCompletionsStreamReconstructsReasoning(t *testing.T) {
if result.Message.Text() != "answer" {
t.Errorf("text = %q", result.Message.Text())
}
if result.FinishReason != "stop" {
t.Errorf("finish reason = %q", result.FinishReason)
if result.Metadata.FinishReason == nil || *result.Metadata.FinishReason != "stop" {
t.Errorf("finish reason = %v", result.Metadata.FinishReason)
}
if stream.Next() {
t.Error("Next returned true after termination")
@@ -148,90 +151,644 @@ func TestCompletionsStreamReconstructsFragmentedParallelToolCalls(t *testing.T)
toolCalls[1].Function.Arguments != `{"format":"%s"}` {
t.Errorf("tool call 1 = %#v", toolCalls[1])
}
if result.FinishReason != "tool_calls" {
t.Errorf("finish reason = %q", result.FinishReason)
if result.Metadata.FinishReason == nil || *result.Metadata.FinishReason != "tool_calls" {
t.Errorf("finish reason = %v", result.Metadata.FinishReason)
}
}
func TestCompletionsReturnsUsageAndTimings(t *testing.T) {
client := newCompletionTestClient(t, func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/chat/completions" {
t.Fatalf("unexpected path: %s", r.URL.Path)
}
w.Header().Set("Content-Type", "application/json")
fmt.Fprint(w, `{
"id": "test-123",
"model": "test-model",
"choices": [{
"index": 0,
"message": {"role": "assistant", "content": "hello"},
"finish_reason": "stop"
}],
"usage": {
"prompt_tokens": 10,
"completion_tokens": 5,
"total_tokens": 15,
"prompt_tokens_details": {"cached_tokens": 3},
"completion_tokens_details": {"reasoning_tokens": 2}
},
"timings": {
"prompt_n": 10,
"prompt_ms": 50.5,
"prompt_per_token_ms": 5.05,
"prompt_per_second": 198.0,
"predicted_n": 5,
"predicted_ms": 100.0,
"predicted_per_token_ms": 20.0,
"predicted_per_second": 50.0
tests := []struct {
name string
body string
want string
}{
{
name: "all known metadata",
body: `{
"id":"body-id","created":123,"model":"reported-model",
"service_tier":"default","system_fingerprint":"backend",
"choices":[{"index":3,"finish_reason":"stop",
"message":{"role":"assistant","content":"hello"},
"logprobs":{"content":[{"token":"hello","logprob":0,
"bytes":[104,195,169],"top_logprobs":[{
"token":"other","logprob":-2,"bytes":[111]}]}],
"refusal":[{"token":"no","logprob":-9999}]}}],
"usage":{"prompt_tokens":10,"completion_tokens":5,
"total_tokens":15,"prompt_tokens_details":{
"audio_tokens":1,"cache_write_tokens":2,
"cached_tokens":3,"image_tokens":4,"text_tokens":5},
"completion_tokens_details":{"accepted_prediction_tokens":1,
"audio_tokens":2,"reasoning_tokens":3,
"rejected_prediction_tokens":4,"text_tokens":5}},
"timings":{"cache_n":1,"draft_n":2,"draft_n_accepted":1,
"prompt_n":10,"prompt_ms":50.5,"prompt_per_token_ms":5.05,
"prompt_per_second":198,"predicted_n":5,"predicted_ms":100,
"predicted_per_token_ms":20,"predicted_per_second":50}
}`,
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)
}
}`)
})
res, err := client.Completions(context.Background(), CompletionsParams{Model: "test-model"})
if err != nil {
t.Fatalf("Completions: %v", err)
}
if res.Message.Text() != "hello" || res.Usage == nil || res.Timings == nil {
t.Fatalf("unexpected result: %#v", res)
}
if res.Usage.PromptTokens != 10 || res.Usage.CompletionTokens != 5 || res.Usage.TotalTokens != 15 {
t.Errorf("usage = %#v", res.Usage)
}
if res.Usage.PromptTokensDetails == nil || res.Usage.PromptTokensDetails.CachedTokens != 3 {
t.Errorf("prompt token details = %#v", res.Usage.PromptTokensDetails)
}
if res.Usage.CompletionTokensDetails == nil || res.Usage.CompletionTokensDetails.ReasoningTokens != 2 {
t.Errorf("completion token details = %#v", res.Usage.CompletionTokensDetails)
}
if res.Timings.PromptMS != 50.5 || res.Timings.PredictedMS != 100 {
t.Errorf("timings = %#v", res.Timings)
if res.Message.Text() != "hello" {
t.Errorf("text = %q", res.Message.Text())
}
var want Metadata
if err := json.Unmarshal([]byte(tt.want), &want); err != nil {
t.Fatal(err)
}
if !reflect.DeepEqual(res.Metadata, want) {
t.Errorf("metadata = %#v, want %#v", res.Metadata, want)
}
encoded, err := jsonv2.Marshal(
res.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.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) {
large := strings.Repeat("x", 70*1024)
client := newCompletionTestClient(t, func(w http.ResponseWriter, _ *http.Request) {
fmt.Fprint(w, ": comment\n")
fmt.Fprintf(w, "data:{\"choices\":[{\"index\":0,\"delta\":{\"content\":%q}}],\n", large)
fmt.Fprint(w, "data: \"usage\":{\"prompt_tokens\":10,\"total_tokens\":15}}\n\n")
fmt.Fprint(w, "data: {\"choices\":[],\"timings\":{\"prompt_ms\":50.5}}\n\n")
fmt.Fprint(w, "data:[DONE]\n\n")
})
tests := []struct {
name string
response string
wantText string
wantMetadata string
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)
var count int
for stream.Next() {
count++
data: {"choices":[{"index":0,"delta":{},"finish_reason":"","logprobs":{"content":[{"token":"b","logprob":-1}],"refusal":[{"token":"r2","logprob":0}]}}]}
data: {"id":"","created":0,"model":"reported-model","service_tier":"last-tier","system_fingerprint":"last-backend","choices":[],"usage":{"prompt_tokens":0},"timings":{"predicted_n":0}}
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 {
t.Fatalf("stream.Err: %v", err)
}
result := stream.Result()
if count != 2 || result.Message.Text() != large {
t.Fatalf("count = %d, text length = %d", count, len(result.Message.Text()))
}
if result.Usage == nil || result.Usage.PromptTokens != 10 || result.Timings == nil || result.Timings.PromptMS != 50.5 {
t.Errorf("metadata = usage %#v, timings %#v", result.Usage, result.Timings)
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
client := newCompletionTestClient(t, func(w http.ResponseWriter, r *http.Request) {
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")
}
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",
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",
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
// protocol types.
package llm
import (
@@ -114,69 +113,229 @@ type ChatRequest struct {
TopP *float32 `json:"top_p,omitempty"`
}
// Usage reports input, output, and total token counts for a completion.
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 *TokenDetails `json:"completion_tokens_details,omitempty"`
// PromptTokens is the number of tokens in the input.
PromptTokens int `json:"prompt_tokens"`
// PromptTokensDetails breaks down PromptTokens by category.
PromptTokensDetails *TokenDetails `json:"prompt_tokens_details,omitempty"`
// TotalTokens is the sum of PromptTokens and CompletionTokens.
TotalTokens int `json:"total_tokens"`
// Metadata describes one selected completion, separately from its message.
// Nil optional fields mean the endpoint did not report a value. Reported zero
// counts and empty strings remain distinct from absent values. Only known body
// fields are retained; response headers are not part of Metadata.
//
// See the response and streaming definitions in
// [OpenAI completion response] and [OpenAI completion chunk].
//
// [OpenAI completion response]: https://developers.openai.com/api/reference/resources/chat/subresources/completions
// [OpenAI completion chunk]: https://developers.openai.com/api/reference/resources/chat/subresources/completions/streaming-events
type Metadata struct {
// 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
// mean either that no tokens belong to the category or that it was not
// reported.
type TokenDetails struct {
// CachedTokens is the number of input tokens served from a cache.
CachedTokens int `json:"cached_tokens,omitempty"`
// Usage reports input, output, and total token counts for one inference.
// Nil counts mean unreported, not zero. Category counts are endpoint-reported
// breakdowns, not additional tokens to add to the totals.
//
// See [OpenAI completion usage].
//
// [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 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.
// The N fields are token counts, MS fields are elapsed milliseconds,
// 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].
//
// [llama.cpp API documentation]: https://llama.app/docs/api
// [llama.cpp API documentation]: https://llama.app/docs/api#timings-and-usage
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 float64 `json:"predicted_ms"`
PredictedMS *float64 `json:"predicted_ms"`
// 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
// second.
PredictedPerSecond float64 `json:"predicted_per_second"`
PredictedPerSecond *float64 `json:"predicted_per_second"`
// PredictedPerTokenMS is the mean output-generation time per token, in
// 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 float64 `json:"prompt_ms"`
PromptMS *float64 `json:"prompt_ms"`
// 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 float64 `json:"prompt_per_second"`
PromptPerSecond *float64 `json:"prompt_per_second"`
// PromptPerTokenMS is the mean prompt-processing time per token, in
// 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
// response.
// response. Nil metadata fields mean unreported values.
type ChatResponse struct {
// Choices contains the generated alternatives.
Choices []Choice `json:"choices"`
// ID identifies the completion.
ID string `json:"id"`
// Created is the endpoint's completion creation time in Unix seconds.
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 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 *Timings `json:"timings,omitempty"`
// 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
// 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 {
// FinishReason is the endpoint's reason for ending generation.
FinishReason string `json:"finish_reason"`
// Index is the position of this choice in the response.
Index int `json:"index"`
FinishReason *string `json:"finish_reason,omitempty"`
// Index is the zero-based position of this choice in the response.
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 Message `json:"message"`
}
@@ -468,22 +630,35 @@ type ResponseFormat struct {
// 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
// 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 {
// Choices contains incremental updates for generated alternatives.
// It may be empty in a metadata-only chunk.
Choices []struct {
// Index identifies the choice updated by this delta.
Index int `json:"index"`
// Delta contains the incremental message update.
Delta Delta `json:"delta"`
// FinishReason is set when generation for the choice ends.
FinishReason string `json:"finish_reason,omitempty"`
// FinishReason is the reported reason for ending this choice;
// 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"`
// ID identifies the completion to which the chunk belongs.
ID string `json:"id"`
// Created is the endpoint's completion creation time in Unix seconds.
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 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 *Timings `json:"timings,omitempty"`
// Usage contains aggregate token counts, when reported.