283 lines
7.8 KiB
Go
283 lines
7.8 KiB
Go
package llm
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"strings"
|
|
"sync/atomic"
|
|
"testing"
|
|
"time"
|
|
|
|
"gitea.maximumdirect.net/eric/audita/internal/framework/contracts"
|
|
|
|
openai "github.com/sashabaranov/go-openai"
|
|
)
|
|
|
|
func TestNewInstructorClientValidation(t *testing.T) {
|
|
_, err := NewInstructorClient(InstructorClientConfig{
|
|
BaseURL: " ",
|
|
Model: "x",
|
|
})
|
|
if err == nil || !strings.Contains(err.Error(), "base URL") {
|
|
t.Fatalf("expected base URL validation error, got %v", err)
|
|
}
|
|
|
|
_, err = NewInstructorClient(InstructorClientConfig{
|
|
BaseURL: "http://localhost:1234/v1",
|
|
Model: " ",
|
|
})
|
|
if err == nil || !strings.Contains(err.Error(), "model") {
|
|
t.Fatalf("expected model validation error, got %v", err)
|
|
}
|
|
|
|
_, err = NewInstructorClient(InstructorClientConfig{
|
|
BaseURL: "http://localhost:1234/v1",
|
|
Model: "test-model",
|
|
MaxRetries: -1,
|
|
})
|
|
if err == nil || !strings.Contains(err.Error(), "max retries") {
|
|
t.Fatalf("expected retries validation error, got %v", err)
|
|
}
|
|
|
|
_, err = NewInstructorClient(InstructorClientConfig{
|
|
BaseURL: "http://localhost:1234/v1",
|
|
Model: "test-model",
|
|
Mode: "unsupported",
|
|
})
|
|
if err == nil || !strings.Contains(err.Error(), "unsupported LLM mode") {
|
|
t.Fatalf("expected mode validation error, got %v", err)
|
|
}
|
|
}
|
|
|
|
func TestInstructorClientCompleteStructuredSuccessNoAPIKey(t *testing.T) {
|
|
var seenPath string
|
|
var seenHost string
|
|
var seenAuth string
|
|
var seenModel string
|
|
client, err := NewInstructorClient(InstructorClientConfig{
|
|
BaseURL: "https://local-compat.example/v1",
|
|
Model: "test-model",
|
|
HTTPClient: &http.Client{
|
|
Transport: roundTripFunc(func(r *http.Request) (*http.Response, error) {
|
|
seenPath = r.URL.Path
|
|
seenHost = r.URL.Host
|
|
seenAuth = r.Header.Get("Authorization")
|
|
|
|
var req map[string]any
|
|
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
|
t.Fatalf("decode request: %v", err)
|
|
}
|
|
model, _ := req["model"].(string)
|
|
seenModel = model
|
|
|
|
return newJSONHTTPResponse(http.StatusOK, chatCompletionResponseBody(`{"name":"Robby","age":22}`)), nil
|
|
}),
|
|
},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NewInstructorClient: %v", err)
|
|
}
|
|
|
|
type person struct {
|
|
Name string `json:"name"`
|
|
Age int `json:"age"`
|
|
}
|
|
var out person
|
|
resp, err := client.CompleteStructured(context.Background(), contracts.StructuredCompletionRequest{
|
|
StageName: "proposal:test",
|
|
Messages: []contracts.LLMMessage{
|
|
{Role: openai.ChatMessageRoleUser, Content: "extract person"},
|
|
},
|
|
}, &out)
|
|
if err != nil {
|
|
t.Fatalf("CompleteStructured: %v", err)
|
|
}
|
|
|
|
if seenPath != "/v1/chat/completions" {
|
|
t.Fatalf("unexpected request path: %q", seenPath)
|
|
}
|
|
if seenHost != "local-compat.example" {
|
|
t.Fatalf("unexpected request host: %q", seenHost)
|
|
}
|
|
if seenModel != "test-model" {
|
|
t.Fatalf("unexpected model: %q", seenModel)
|
|
}
|
|
if seenAuth != "" {
|
|
t.Fatalf("expected empty Authorization header for empty API key, got %q", seenAuth)
|
|
}
|
|
if out.Name != "Robby" || out.Age != 22 {
|
|
t.Fatalf("unexpected output: %+v", out)
|
|
}
|
|
if resp.Provider != "openai-compatible" || resp.Model != "test-model" {
|
|
t.Fatalf("unexpected response metadata: %+v", resp)
|
|
}
|
|
if resp.TotalTokens != 18 || resp.PromptTokens != 11 || resp.CompletionTokens != 7 {
|
|
t.Fatalf("unexpected usage metadata: %+v", resp)
|
|
}
|
|
}
|
|
|
|
func TestInstructorClientRetriesOnMalformedJSON(t *testing.T) {
|
|
var attempts int32
|
|
client, err := NewInstructorClient(InstructorClientConfig{
|
|
BaseURL: "https://retry.example/v1",
|
|
Model: "test-model",
|
|
MaxRetries: 1,
|
|
HTTPClient: &http.Client{
|
|
Transport: roundTripFunc(func(r *http.Request) (*http.Response, error) {
|
|
current := atomic.AddInt32(&attempts, 1)
|
|
if current == 1 {
|
|
return newJSONHTTPResponse(http.StatusOK, chatCompletionResponseBody(`{"name":"broken"`)), nil
|
|
}
|
|
return newJSONHTTPResponse(http.StatusOK, chatCompletionResponseBody(`{"name":"Recovered","age":30}`)), nil
|
|
}),
|
|
},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NewInstructorClient: %v", err)
|
|
}
|
|
|
|
type person struct {
|
|
Name string `json:"name"`
|
|
Age int `json:"age"`
|
|
}
|
|
var out person
|
|
_, err = client.CompleteStructured(context.Background(), contracts.StructuredCompletionRequest{
|
|
Messages: []contracts.LLMMessage{
|
|
{Role: openai.ChatMessageRoleUser, Content: "extract person"},
|
|
},
|
|
}, &out)
|
|
if err != nil {
|
|
t.Fatalf("CompleteStructured: %v", err)
|
|
}
|
|
if out.Name != "Recovered" || out.Age != 30 {
|
|
t.Fatalf("unexpected output after retry: %+v", out)
|
|
}
|
|
if got := atomic.LoadInt32(&attempts); got != 2 {
|
|
t.Fatalf("expected 2 attempts, got %d", got)
|
|
}
|
|
}
|
|
|
|
func TestInstructorClientContextDeadlinePropagates(t *testing.T) {
|
|
client, err := NewInstructorClient(InstructorClientConfig{
|
|
BaseURL: "https://slow.example/v1",
|
|
Model: "test-model",
|
|
HTTPClient: &http.Client{
|
|
Transport: roundTripFunc(func(r *http.Request) (*http.Response, error) {
|
|
select {
|
|
case <-r.Context().Done():
|
|
return nil, r.Context().Err()
|
|
case <-time.After(250 * time.Millisecond):
|
|
return newJSONHTTPResponse(http.StatusOK, chatCompletionResponseBody(`{"name":"slow","age":1}`)), nil
|
|
}
|
|
}),
|
|
},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NewInstructorClient: %v", err)
|
|
}
|
|
|
|
type person struct {
|
|
Name string `json:"name"`
|
|
}
|
|
var out person
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
|
|
defer cancel()
|
|
|
|
_, err = client.CompleteStructured(ctx, contracts.StructuredCompletionRequest{
|
|
Messages: []contracts.LLMMessage{
|
|
{Role: openai.ChatMessageRoleUser, Content: "extract person"},
|
|
},
|
|
}, &out)
|
|
if err == nil {
|
|
t.Fatalf("expected context deadline error")
|
|
}
|
|
if !strings.Contains(strings.ToLower(err.Error()), "context deadline") {
|
|
t.Fatalf("expected deadline-related error, got %v", err)
|
|
}
|
|
}
|
|
|
|
func TestInstructorClientSanitizesAPIKeyInErrors(t *testing.T) {
|
|
apiKey := "super-secret-key"
|
|
client := &InstructorClient{
|
|
cfg: InstructorClientConfig{
|
|
BaseURL: "http://localhost:1234/v1",
|
|
Model: "test-model",
|
|
APIKey: apiKey,
|
|
Mode: ModeJSON,
|
|
},
|
|
client: fakeChatCompletionClient{
|
|
err: fmt.Errorf("provider failed with Authorization: Bearer %s", apiKey),
|
|
},
|
|
}
|
|
|
|
var out map[string]any
|
|
_, err := client.CompleteStructured(context.Background(), contracts.StructuredCompletionRequest{
|
|
Messages: []contracts.LLMMessage{
|
|
{Role: openai.ChatMessageRoleUser, Content: "extract"},
|
|
},
|
|
}, &out)
|
|
if err == nil {
|
|
t.Fatalf("expected error")
|
|
}
|
|
if strings.Contains(err.Error(), apiKey) {
|
|
t.Fatalf("error leaked API key: %v", err)
|
|
}
|
|
if !strings.Contains(err.Error(), "[REDACTED]") {
|
|
t.Fatalf("expected redaction marker in error: %v", err)
|
|
}
|
|
}
|
|
|
|
type fakeChatCompletionClient struct {
|
|
err error
|
|
}
|
|
|
|
func (f fakeChatCompletionClient) CreateChatCompletion(ctx context.Context, request openai.ChatCompletionRequest, responseType any) (response openai.ChatCompletionResponse, err error) {
|
|
_ = ctx
|
|
_ = request
|
|
_ = responseType
|
|
return openai.ChatCompletionResponse{}, f.err
|
|
}
|
|
|
|
type roundTripFunc func(*http.Request) (*http.Response, error)
|
|
|
|
func (f roundTripFunc) RoundTrip(r *http.Request) (*http.Response, error) {
|
|
return f(r)
|
|
}
|
|
|
|
func newJSONHTTPResponse(status int, body string) *http.Response {
|
|
return &http.Response{
|
|
StatusCode: status,
|
|
Header: http.Header{"Content-Type": []string{"application/json"}},
|
|
Body: io.NopCloser(strings.NewReader(body)),
|
|
}
|
|
}
|
|
|
|
func chatCompletionResponseBody(content string) string {
|
|
body, _ := json.Marshal(map[string]any{
|
|
"id": "chatcmpl-test",
|
|
"object": "chat.completion",
|
|
"created": 12345,
|
|
"model": "test-model",
|
|
"choices": []map[string]any{
|
|
{
|
|
"index": 0,
|
|
"message": map[string]any{
|
|
"role": "assistant",
|
|
"content": content,
|
|
},
|
|
"finish_reason": "stop",
|
|
},
|
|
},
|
|
"usage": map[string]any{
|
|
"prompt_tokens": 11,
|
|
"completion_tokens": 7,
|
|
"total_tokens": 18,
|
|
},
|
|
})
|
|
return string(body)
|
|
}
|