Refactor llm pkg
This commit is contained in:
@@ -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",
|
||||
|
||||
Reference in New Issue
Block a user