Refactor llm package
This commit is contained in:
dwrz
@@ -2,96 +2,159 @@ package llm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"log/slog"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestCompletionsStreamStreamsReasoningDeltas(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path != "/chat/completions" {
|
||||
t.Fatalf("unexpected path: %s", r.URL.Path)
|
||||
}
|
||||
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"}}]}`,
|
||||
`{"choices":[{"index":0,"delta":{"content":"answer"},"finish_reason":"stop"}]}`,
|
||||
} {
|
||||
fmt.Fprintf(w, "data: %s\n\n", data)
|
||||
}
|
||||
fmt.Fprint(w, "data: [DONE]\n\n")
|
||||
}))
|
||||
defer srv.Close()
|
||||
})
|
||||
|
||||
client, err := NewClient(Config{URL: srv.URL}, nil, nil)
|
||||
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("NewClient: %v", err)
|
||||
}
|
||||
|
||||
events := make(chan StreamEvent)
|
||||
errCh := make(chan error, 1)
|
||||
go func() {
|
||||
errCh <- client.CompletionsStream(
|
||||
context.Background(), CompletionsParams{
|
||||
Model: "test-model",
|
||||
SystemMessage: "",
|
||||
MaxIterations: 0,
|
||||
Messages: []Message{{
|
||||
Role: RoleUser,
|
||||
ContentParts: []ContentPart{{
|
||||
Type: ContentTypeText,
|
||||
Text: "hello",
|
||||
}},
|
||||
}},
|
||||
}, events,
|
||||
)
|
||||
}()
|
||||
|
||||
var got []StreamEvent
|
||||
for evt := range events {
|
||||
got = append(got, evt)
|
||||
}
|
||||
if err := <-errCh; 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)
|
||||
}
|
||||
}
|
||||
|
||||
wantTypes := []string{
|
||||
streamEventReasoningDelta,
|
||||
streamEventReasoningDelta,
|
||||
streamEventDelta,
|
||||
streamEventDone,
|
||||
}
|
||||
if len(got) != len(wantTypes) {
|
||||
t.Fatalf("got %d events, want %d: %#v", len(got), len(wantTypes), got)
|
||||
}
|
||||
for i, want := range wantTypes {
|
||||
if got[i].Type != want {
|
||||
t.Fatalf("event %d type = %q, want %q", i, got[i].Type, want)
|
||||
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 got[0].ReasoningDelta != "thinking" {
|
||||
t.Errorf("first reasoning delta = %q", got[0].ReasoningDelta)
|
||||
if err := stream.Err(); err != nil {
|
||||
t.Fatalf("stream.Err: %v", err)
|
||||
}
|
||||
if got[1].ReasoningDelta != " more" {
|
||||
t.Errorf("second reasoning delta = %q", got[1].ReasoningDelta)
|
||||
if len(chunks) != 4 || len(chunks[1].Choices[0].Delta.ToolCalls) != 2 {
|
||||
t.Fatalf("tool-call chunks not exposed unchanged: %#v", chunks)
|
||||
}
|
||||
if got[2].Delta != "answer" {
|
||||
t.Errorf("content delta = %q", got[2].Delta)
|
||||
result := stream.Result()
|
||||
toolCalls := result.Message.ToolCalls
|
||||
if len(toolCalls) != 2 {
|
||||
t.Fatalf("tool calls = %d, want 2", len(toolCalls))
|
||||
}
|
||||
final := got[3].Message
|
||||
if final.ReasoningContent != "thinking more" {
|
||||
t.Errorf("final reasoning = %q", final.ReasoningContent)
|
||||
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 final.Text() != "answer" {
|
||||
t.Errorf("final text = %q", final.Text())
|
||||
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) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
client := newCompletionTestClient(t, func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path != "/chat/completions" {
|
||||
t.Fatalf("unexpected path: %s", r.URL.Path)
|
||||
}
|
||||
@@ -122,135 +185,127 @@ func TestCompletionsReturnsUsageAndTimings(t *testing.T) {
|
||||
"predicted_per_second": 50.0
|
||||
}
|
||||
}`)
|
||||
}))
|
||||
defer srv.Close()
|
||||
})
|
||||
|
||||
client, err := NewClient(Config{URL: srv.URL}, nil, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("NewClient: %v", err)
|
||||
}
|
||||
|
||||
res, err := client.Completions(
|
||||
context.Background(), CompletionsParams{
|
||||
Model: "test-model",
|
||||
SystemMessage: "",
|
||||
MaxIterations: 0,
|
||||
Messages: []Message{{
|
||||
Role: RoleUser,
|
||||
ContentParts: []ContentPart{{
|
||||
Type: ContentTypeText,
|
||||
Text: "hi",
|
||||
}},
|
||||
}},
|
||||
},
|
||||
)
|
||||
res, err := client.Completions(context.Background(), CompletionsParams{Model: "test-model"})
|
||||
if err != nil {
|
||||
t.Fatalf("Completions: %v", err)
|
||||
}
|
||||
if len(res.Messages) != 1 {
|
||||
t.Fatalf("got %d messages, want 1", len(res.Messages))
|
||||
if res.Message.Text() != "hello" || res.Usage == nil || res.Timings == nil {
|
||||
t.Fatalf("unexpected result: %#v", res)
|
||||
}
|
||||
if res.Usage == nil {
|
||||
t.Fatal("usage is nil")
|
||||
}
|
||||
if res.Usage.PromptTokens != 10 {
|
||||
t.Errorf("prompt_tokens = %d, want 10", res.Usage.PromptTokens)
|
||||
}
|
||||
if res.Usage.CompletionTokens != 5 {
|
||||
t.Errorf("completion_tokens = %d, want 5", res.Usage.CompletionTokens)
|
||||
}
|
||||
if res.Usage.TotalTokens != 15 {
|
||||
t.Errorf("total_tokens = %d, want 15", res.Usage.TotalTokens)
|
||||
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("cached_tokens unexpected: %#v", res.Usage.PromptTokensDetails)
|
||||
t.Errorf("prompt token details = %#v", res.Usage.PromptTokensDetails)
|
||||
}
|
||||
if res.Usage.CompletionTokensDetails == nil || res.Usage.CompletionTokensDetails.ReasoningTokens != 2 {
|
||||
t.Errorf("reasoning_tokens unexpected: %#v", res.Usage.CompletionTokensDetails)
|
||||
t.Errorf("completion token details = %#v", res.Usage.CompletionTokensDetails)
|
||||
}
|
||||
if res.Timings == nil {
|
||||
t.Fatal("timings is nil")
|
||||
}
|
||||
if res.Timings.PromptMS != 50.5 {
|
||||
t.Errorf("prompt_ms = %f, want 50.5", res.Timings.PromptMS)
|
||||
}
|
||||
if res.Timings.PredictedMS != 100.0 {
|
||||
t.Errorf("predicted_ms = %f, want 100.0", res.Timings.PredictedMS)
|
||||
if res.Timings.PromptMS != 50.5 || res.Timings.PredictedMS != 100 {
|
||||
t.Errorf("timings = %#v", res.Timings)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCompletionsStreamEmitsUsageEvent(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(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", "text/event-stream")
|
||||
// Stream some content
|
||||
fmt.Fprintf(w, "data: %s\n\n",
|
||||
`{"choices":[{"index":0,"delta":{"role":"assistant"}}]}`)
|
||||
fmt.Fprintf(w, "data: %s\n\n",
|
||||
`{"choices":[{"index":0,"delta":{"content":"hello"}}]}`)
|
||||
// Final chunk with usage
|
||||
fmt.Fprintf(w, "data: %s\n\n",
|
||||
`{"choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":10,"completion_tokens":5,"total_tokens":15}}`)
|
||||
fmt.Fprint(w, "data: [DONE]\n\n")
|
||||
}))
|
||||
defer srv.Close()
|
||||
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")
|
||||
})
|
||||
|
||||
client, err := NewClient(Config{URL: srv.URL}, nil, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("NewClient: %v", err)
|
||||
stream := openCompletionTestStream(t, client)
|
||||
var count int
|
||||
for stream.Next() {
|
||||
count++
|
||||
}
|
||||
|
||||
events := make(chan StreamEvent)
|
||||
errCh := make(chan error, 1)
|
||||
go func() {
|
||||
errCh <- client.CompletionsStream(
|
||||
context.Background(), CompletionsParams{
|
||||
Model: "test-model",
|
||||
SystemMessage: "",
|
||||
MaxIterations: 0,
|
||||
Messages: []Message{{
|
||||
Role: RoleUser,
|
||||
ContentParts: []ContentPart{{
|
||||
Type: ContentTypeText,
|
||||
Text: "hi",
|
||||
}},
|
||||
}},
|
||||
}, events,
|
||||
)
|
||||
}()
|
||||
|
||||
var got []StreamEvent
|
||||
for evt := range events {
|
||||
got = append(got, evt)
|
||||
if err := stream.Err(); err != nil {
|
||||
t.Fatalf("stream.Err: %v", err)
|
||||
}
|
||||
if err := <-errCh; err != nil {
|
||||
t.Fatalf("CompletionsStream: %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)
|
||||
}
|
||||
}
|
||||
|
||||
// Expect: delta, done events (usage attached to done event)
|
||||
wantTypes := []string{
|
||||
streamEventDelta,
|
||||
streamEventDone,
|
||||
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)
|
||||
}
|
||||
})
|
||||
}
|
||||
if len(got) != len(wantTypes) {
|
||||
t.Fatalf("got %d events, want %d: %#v", len(got), len(wantTypes), 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)
|
||||
}
|
||||
for i, want := range wantTypes {
|
||||
if got[i].Type != want {
|
||||
t.Fatalf("event %d type = %q, want %q", i, got[i].Type, want)
|
||||
}
|
||||
if err := stream.Close(); err != nil {
|
||||
t.Fatalf("second Close: %v", err)
|
||||
}
|
||||
// Verify usage is on the done event
|
||||
doneEvt := got[1]
|
||||
if doneEvt.Usage == nil {
|
||||
t.Fatal("done event has nil Usage")
|
||||
if stream.Err() != nil {
|
||||
t.Errorf("Err after Close = %v", stream.Err())
|
||||
}
|
||||
if doneEvt.Usage.PromptTokens != 10 {
|
||||
t.Errorf("prompt_tokens = %d, want 10", doneEvt.Usage.PromptTokens)
|
||||
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 doneEvt.Usage.TotalTokens != 15 {
|
||||
t.Errorf("total_tokens = %d, want 15", doneEvt.Usage.TotalTokens)
|
||||
if !errors.Is(stream.Err(), io.ErrUnexpectedEOF) {
|
||||
t.Fatalf("Err = %v, want io.ErrUnexpectedEOF", stream.Err())
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user