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
+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",