873 lines
25 KiB
Go
873 lines
25 KiB
Go
package llm
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json/v2"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"log/slog"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"reflect"
|
|
"strings"
|
|
"testing"
|
|
)
|
|
|
|
func newCompletionTestClient(t *testing.T, handler http.HandlerFunc) *Client {
|
|
t.Helper()
|
|
srv := httptest.NewServer(handler)
|
|
t.Cleanup(srv.Close)
|
|
client, err := NewClient(Config{Provider: "test", URL: srv.URL}, slog.Default())
|
|
if err != nil {
|
|
t.Fatalf("NewClient: %v", err)
|
|
}
|
|
return client
|
|
}
|
|
|
|
func openCompletionTestStream(t *testing.T, client *Client) *CompletionStream {
|
|
t.Helper()
|
|
stream, err := client.CompletionsStream(
|
|
t.Context(), CompletionsParams{Model: "test-model"},
|
|
)
|
|
if err != nil {
|
|
t.Fatalf("CompletionsStream: %v", err)
|
|
}
|
|
t.Cleanup(func() { _ = stream.Close() })
|
|
return stream
|
|
}
|
|
|
|
func TestCompletionsStreamReconstructsReasoning(t *testing.T) {
|
|
client := newCompletionTestClient(t, func(w http.ResponseWriter, _ *http.Request) {
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
for _, data := range []string{
|
|
`{"choices":[{"index":0,"delta":{"role":"assistant"}}]}`,
|
|
`{"choices":[{"index":0,"delta":{"reasoning_content":"thinking"}}]}`,
|
|
`{"choices":[{"index":0,"delta":{"reasoning_content":" more"}}]}`,
|
|
`{"choices":[{"index":0,"delta":{"content":"answer"},"finish_reason":"stop"}]}`,
|
|
} {
|
|
fmt.Fprintf(w, "data: %s\n\n", data)
|
|
}
|
|
fmt.Fprint(w, "data: [DONE]\n\n")
|
|
})
|
|
|
|
stream := openCompletionTestStream(t, client)
|
|
if got := stream.Result(); got.Message.Text() != "" {
|
|
t.Fatalf("initial result text = %q, want empty", got.Message.Text())
|
|
}
|
|
var chunks []StreamChunk
|
|
for stream.Next() {
|
|
chunks = append(chunks, stream.Current())
|
|
}
|
|
if err := stream.Err(); err != nil {
|
|
t.Fatalf("stream.Err: %v", err)
|
|
}
|
|
if len(chunks) != 4 {
|
|
t.Fatalf("chunks = %d, want 4", len(chunks))
|
|
}
|
|
result := stream.Result()
|
|
if result.Message.ReasoningContent != "thinking more" {
|
|
t.Errorf("reasoning = %q", result.Message.ReasoningContent)
|
|
}
|
|
if result.Message.Text() != "answer" {
|
|
t.Errorf("text = %q", result.Message.Text())
|
|
}
|
|
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")
|
|
}
|
|
}
|
|
|
|
func TestCompletionsStreamReportsCancellationWithPartialResult(t *testing.T) {
|
|
firstSent := make(chan struct{})
|
|
client := newCompletionTestClient(t, func(w http.ResponseWriter, r *http.Request) {
|
|
flusher := w.(http.Flusher)
|
|
fmt.Fprint(w, "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"first\"}}]}\n\n")
|
|
flusher.Flush()
|
|
close(firstSent)
|
|
<-r.Context().Done()
|
|
})
|
|
|
|
ctx, cancel := context.WithCancel(t.Context())
|
|
stream, err := client.CompletionsStream(ctx, CompletionsParams{Model: "test-model"})
|
|
if err != nil {
|
|
t.Fatalf("CompletionsStream: %v", err)
|
|
}
|
|
defer stream.Close()
|
|
<-firstSent
|
|
if !stream.Next() {
|
|
t.Fatalf("first Next = false: %v", stream.Err())
|
|
}
|
|
cancel()
|
|
if stream.Next() {
|
|
t.Fatal("Next returned true after cancellation")
|
|
}
|
|
if !errors.Is(stream.Err(), context.Canceled) {
|
|
t.Fatalf("Err = %v, want context.Canceled", stream.Err())
|
|
}
|
|
result := stream.Result()
|
|
if got := result.Message.Text(); got != "first" {
|
|
t.Errorf("partial text = %q, want first", got)
|
|
}
|
|
}
|
|
|
|
func TestCompletionsStreamReconstructsFragmentedParallelToolCalls(t *testing.T) {
|
|
client := newCompletionTestClient(t, func(w http.ResponseWriter, _ *http.Request) {
|
|
for _, data := range []string{
|
|
`{"choices":[{"index":0,"delta":{"role":"assistant"}}]}`,
|
|
`{"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"id":"call_0","function":{"name":"ec"}},{"index":1,"id":"call_1","function":{"name":"pr"}}]}}]}`,
|
|
`{"choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"name":"ho","arguments":"{\"value\":"}},{"index":1,"function":{"name":"intf","arguments":"{\"format\":"}}]}}]}`,
|
|
`{"choices":[{"index":0,"delta":{"tool_calls":[{"index":1,"function":{"arguments":"\"%s\"}"}},{"index":0,"function":{"arguments":"\"one\"}"}}]},"finish_reason":"tool_calls"}]}`,
|
|
} {
|
|
fmt.Fprintf(w, "data: %s\n\n", data)
|
|
}
|
|
fmt.Fprint(w, "data: [DONE]\n\n")
|
|
})
|
|
|
|
stream := openCompletionTestStream(t, client)
|
|
var chunks []StreamChunk
|
|
for stream.Next() {
|
|
chunks = append(chunks, stream.Current())
|
|
}
|
|
if err := stream.Err(); err != nil {
|
|
t.Fatalf("stream.Err: %v", err)
|
|
}
|
|
if len(chunks) != 4 || len(chunks[1].Choices[0].Delta.ToolCalls) != 2 {
|
|
t.Fatalf("tool-call chunks not exposed unchanged: %#v", chunks)
|
|
}
|
|
result := stream.Result()
|
|
toolCalls := result.Message.ToolCalls
|
|
if len(toolCalls) != 2 {
|
|
t.Fatalf("tool calls = %d, want 2", len(toolCalls))
|
|
}
|
|
if toolCalls[0].ID != "call_0" || toolCalls[0].Function.Name != "echo" ||
|
|
toolCalls[0].Function.Arguments != `{"value":"one"}` {
|
|
t.Errorf("tool call 0 = %#v", toolCalls[0])
|
|
}
|
|
if toolCalls[1].ID != "call_1" || toolCalls[1].Function.Name != "printf" ||
|
|
toolCalls[1].Function.Arguments != `{"format":"%s"}` {
|
|
t.Errorf("tool call 1 = %#v", toolCalls[1])
|
|
}
|
|
if result.Metadata.FinishReason == nil || *result.Metadata.FinishReason != "tool_calls" {
|
|
t.Errorf("finish reason = %v", result.Metadata.FinishReason)
|
|
}
|
|
}
|
|
|
|
func TestCompletionsReturnsUsageAndTimings(t *testing.T) {
|
|
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.UnmarshalRead(r.Body, &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())
|
|
}
|
|
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 := json.Marshal(
|
|
res.Metadata, json.FormatNilSliceAsNull(true),
|
|
)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
var gotJSON, wantJSON map[string]any
|
|
if err := json.Unmarshal(encoded, &gotJSON); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := json.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)
|
|
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}}
|
|
|
|
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,
|
|
},
|
|
}
|
|
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.UnmarshalRead(r.Body, &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 := json.Marshal(
|
|
result.Metadata, json.FormatNilSliceAsNull(true),
|
|
)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
var gotJSON, wantJSON map[string]any
|
|
if err := json.Unmarshal(encoded, &gotJSON); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if err := json.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)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestCompletionsStreamFailuresPreservePartialResult(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
response string
|
|
wantErr error
|
|
}{
|
|
{
|
|
name: "EOF before done",
|
|
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",
|
|
},
|
|
{
|
|
name: "unexpected choice",
|
|
response: "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"partial\"}}]}\n\ndata: {\"choices\":[{\"index\":1,\"delta\":{}}]}\n\n",
|
|
},
|
|
}
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
client := newCompletionTestClient(t, func(w http.ResponseWriter, _ *http.Request) {
|
|
fmt.Fprint(w, tt.response)
|
|
})
|
|
stream := openCompletionTestStream(t, client)
|
|
for stream.Next() {
|
|
}
|
|
if stream.Err() == nil {
|
|
t.Fatal("Err = nil")
|
|
}
|
|
if tt.wantErr != nil && !errors.Is(stream.Err(), tt.wantErr) {
|
|
t.Errorf("Err = %v, want %v", stream.Err(), tt.wantErr)
|
|
}
|
|
result := stream.Result()
|
|
if got := result.Message.Text(); got != "partial" {
|
|
t.Errorf("partial text = %q", got)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestCompletionStreamCloseIsIdempotent(t *testing.T) {
|
|
client := newCompletionTestClient(t, func(w http.ResponseWriter, r *http.Request) {
|
|
w.(http.Flusher).Flush()
|
|
<-r.Context().Done()
|
|
})
|
|
stream := openCompletionTestStream(t, client)
|
|
if err := stream.Close(); err != nil {
|
|
t.Fatalf("first Close: %v", err)
|
|
}
|
|
if err := stream.Close(); err != nil {
|
|
t.Fatalf("second Close: %v", err)
|
|
}
|
|
if stream.Err() != nil {
|
|
t.Errorf("Err after Close = %v", stream.Err())
|
|
}
|
|
if stream.Next() {
|
|
t.Error("Next after Close returned true")
|
|
}
|
|
}
|
|
|
|
func TestCompletionsStreamRequiresCompleteDoneEvent(t *testing.T) {
|
|
client := newCompletionTestClient(t, func(w http.ResponseWriter, _ *http.Request) {
|
|
fmt.Fprint(w, "data: [DONE]")
|
|
})
|
|
stream := openCompletionTestStream(t, client)
|
|
if stream.Next() {
|
|
t.Fatal("Next returned true")
|
|
}
|
|
if !errors.Is(stream.Err(), io.ErrUnexpectedEOF) {
|
|
t.Fatalf("Err = %v, want io.ErrUnexpectedEOF", stream.Err())
|
|
}
|
|
}
|