312 lines
9.9 KiB
Go
312 lines
9.9 KiB
Go
package llm
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"log/slog"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"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.FinishReason != "stop" {
|
|
t.Errorf("finish reason = %q", result.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.FinishReason != "tool_calls" {
|
|
t.Errorf("finish reason = %q", result.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
|
|
}
|
|
}`)
|
|
})
|
|
|
|
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)
|
|
}
|
|
}
|
|
|
|
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")
|
|
})
|
|
|
|
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 != 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)
|
|
}
|
|
}
|
|
|
|
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: "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())
|
|
}
|
|
}
|