495 lines
14 KiB
Go
495 lines
14 KiB
Go
package llm
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strings"
|
|
"sync/atomic"
|
|
"testing"
|
|
|
|
"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
|
|
)
|
|
|
|
type testArtifact struct {
|
|
Value string `json:"value"`
|
|
}
|
|
|
|
func TestNewOpenAICompatibleClientValidation(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
cfg OpenAICompatibleClientConfig
|
|
want string
|
|
}{
|
|
{
|
|
name: "empty base URL",
|
|
cfg: OpenAICompatibleClientConfig{
|
|
BaseURL: " ",
|
|
Model: "model",
|
|
},
|
|
want: "base URL",
|
|
},
|
|
{
|
|
name: "invalid base URL",
|
|
cfg: OpenAICompatibleClientConfig{
|
|
BaseURL: "://bad",
|
|
Model: "model",
|
|
},
|
|
want: "base URL",
|
|
},
|
|
{
|
|
name: "empty model",
|
|
cfg: OpenAICompatibleClientConfig{
|
|
BaseURL: "https://example.test/v1",
|
|
Model: " ",
|
|
},
|
|
want: "model",
|
|
},
|
|
{
|
|
name: "negative retries",
|
|
cfg: OpenAICompatibleClientConfig{
|
|
BaseURL: "https://example.test/v1",
|
|
Model: "model",
|
|
MaxRetries: -1,
|
|
},
|
|
want: "max retries",
|
|
},
|
|
}
|
|
|
|
for _, tc := range tests {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
_, err := NewOpenAICompatibleClient(tc.cfg)
|
|
if err == nil || !strings.Contains(err.Error(), tc.want) {
|
|
t.Fatalf("expected error containing %q, got %v", tc.want, err)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestOpenAICompatibleClientSuccessfulStructuredCompletion(t *testing.T) {
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_, _ = io.WriteString(w, `{
|
|
"model":"provider-model",
|
|
"choices":[{"message":{"content":"{\"value\":\"ok\"}"}}],
|
|
"usage":{"prompt_tokens":11,"completion_tokens":7,"total_tokens":18}
|
|
}`)
|
|
}))
|
|
defer server.Close()
|
|
|
|
client := newTestClient(t, server.URL, "default-model", 0)
|
|
var out testArtifact
|
|
resp, err := client.CompleteStructured(context.Background(), validStructuredRequest(""), &out)
|
|
if err != nil {
|
|
t.Fatalf("CompleteStructured: %v", err)
|
|
}
|
|
|
|
if out.Value != "ok" {
|
|
t.Fatalf("unexpected decoded output: %+v", out)
|
|
}
|
|
if string(resp.Content) != `{"value":"ok"}` {
|
|
t.Fatalf("unexpected raw content: %s", resp.Content)
|
|
}
|
|
if resp.Provider != openAICompatibleProviderName {
|
|
t.Fatalf("unexpected provider: %q", resp.Provider)
|
|
}
|
|
if resp.Model != "provider-model" {
|
|
t.Fatalf("unexpected model: %q", resp.Model)
|
|
}
|
|
if resp.PromptTokens != 11 || resp.CompletionTokens != 7 || resp.TotalTokens != 18 {
|
|
t.Fatalf("unexpected token metadata: %+v", resp)
|
|
}
|
|
}
|
|
|
|
func TestOpenAICompatibleClientRequestBodyIncludesStructuredOutputShape(t *testing.T) {
|
|
var seenPath string
|
|
var seenAuthorization string
|
|
var seenReq map[string]any
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
seenPath = r.URL.Path
|
|
seenAuthorization = r.Header.Get("Authorization")
|
|
if err := json.NewDecoder(r.Body).Decode(&seenReq); err != nil {
|
|
t.Fatalf("decode request: %v", err)
|
|
}
|
|
_, _ = io.WriteString(w, `{"choices":[{"message":{"content":"{\"value\":\"ok\"}"}}]}`)
|
|
}))
|
|
defer server.Close()
|
|
|
|
client, err := NewOpenAICompatibleClient(OpenAICompatibleClientConfig{
|
|
BaseURL: server.URL + "/v1",
|
|
Model: "default-model",
|
|
APIKey: "secret-key",
|
|
MaxRetries: 0,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NewOpenAICompatibleClient: %v", err)
|
|
}
|
|
|
|
var out testArtifact
|
|
_, err = client.CompleteStructured(context.Background(), validStructuredRequest("request-model"), &out)
|
|
if err != nil {
|
|
t.Fatalf("CompleteStructured: %v", err)
|
|
}
|
|
|
|
if seenPath != "/v1/chat/completions" {
|
|
t.Fatalf("unexpected request path: %q", seenPath)
|
|
}
|
|
if seenAuthorization != "Bearer secret-key" {
|
|
t.Fatalf("unexpected authorization header: %q", seenAuthorization)
|
|
}
|
|
if seenReq["model"] != "request-model" {
|
|
t.Fatalf("unexpected model: %v", seenReq["model"])
|
|
}
|
|
|
|
messages, ok := seenReq["messages"].([]any)
|
|
if !ok || len(messages) != 1 {
|
|
t.Fatalf("unexpected messages: %#v", seenReq["messages"])
|
|
}
|
|
message, ok := messages[0].(map[string]any)
|
|
if !ok {
|
|
t.Fatalf("unexpected message shape: %#v", messages[0])
|
|
}
|
|
if message["role"] != "user" || message["content"] != "extract this" {
|
|
t.Fatalf("unexpected message: %#v", message)
|
|
}
|
|
|
|
responseFormat, ok := seenReq["response_format"].(map[string]any)
|
|
if !ok {
|
|
t.Fatalf("expected response_format object, got %T", seenReq["response_format"])
|
|
}
|
|
if responseFormat["type"] != "json_schema" {
|
|
t.Fatalf("unexpected response_format.type: %v", responseFormat["type"])
|
|
}
|
|
jsonSchema, ok := responseFormat["json_schema"].(map[string]any)
|
|
if !ok {
|
|
t.Fatalf("expected response_format.json_schema object, got %T", responseFormat["json_schema"])
|
|
}
|
|
if jsonSchema["name"] != "test_artifact" {
|
|
t.Fatalf("unexpected schema name: %v", jsonSchema["name"])
|
|
}
|
|
if jsonSchema["strict"] != true {
|
|
t.Fatalf("expected strict=true, got %v", jsonSchema["strict"])
|
|
}
|
|
schema, ok := jsonSchema["schema"].(map[string]any)
|
|
if !ok {
|
|
t.Fatalf("expected schema object, got %T", jsonSchema["schema"])
|
|
}
|
|
if schema["type"] != "object" {
|
|
t.Fatalf("unexpected schema: %#v", schema)
|
|
}
|
|
}
|
|
|
|
func TestOpenAICompatibleClientDefaultModelFallbackAndOverride(t *testing.T) {
|
|
var seenModels []string
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
var req map[string]any
|
|
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
|
|
t.Fatalf("decode request: %v", err)
|
|
}
|
|
seenModels = append(seenModels, fmt.Sprint(req["model"]))
|
|
_, _ = io.WriteString(w, `{"choices":[{"message":{"content":"{\"value\":\"ok\"}"}}]}`)
|
|
}))
|
|
defer server.Close()
|
|
|
|
client := newTestClient(t, server.URL, "default-model", 0)
|
|
var first testArtifact
|
|
if _, err := client.CompleteStructured(context.Background(), validStructuredRequest(""), &first); err != nil {
|
|
t.Fatalf("first CompleteStructured: %v", err)
|
|
}
|
|
var second testArtifact
|
|
if _, err := client.CompleteStructured(context.Background(), validStructuredRequest("override-model"), &second); err != nil {
|
|
t.Fatalf("second CompleteStructured: %v", err)
|
|
}
|
|
|
|
if len(seenModels) != 2 || seenModels[0] != "default-model" || seenModels[1] != "override-model" {
|
|
t.Fatalf("unexpected models: %v", seenModels)
|
|
}
|
|
}
|
|
|
|
func TestOpenAICompatibleClientInvalidOutputTarget(t *testing.T) {
|
|
client := newTestClient(t, "https://example.test/v1", "default-model", 0)
|
|
|
|
tests := []struct {
|
|
name string
|
|
out any
|
|
}{
|
|
{name: "nil", out: nil},
|
|
{name: "non-pointer", out: testArtifact{}},
|
|
{name: "nil pointer", out: (*testArtifact)(nil)},
|
|
}
|
|
|
|
for _, tc := range tests {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
_, err := client.CompleteStructured(context.Background(), validStructuredRequest(""), tc.out)
|
|
if err == nil || !strings.Contains(err.Error(), "output target") {
|
|
t.Fatalf("expected output target error, got %v", err)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestOpenAICompatibleClientMissingAndInvalidSchema(t *testing.T) {
|
|
client := newTestClient(t, "https://example.test/v1", "default-model", 0)
|
|
|
|
tests := []struct {
|
|
name string
|
|
mutate func(*contracts.StructuredCompletionRequest)
|
|
want string
|
|
}{
|
|
{
|
|
name: "missing schema name",
|
|
mutate: func(req *contracts.StructuredCompletionRequest) {
|
|
req.ResponseSchemaName = " "
|
|
},
|
|
want: "schema name",
|
|
},
|
|
{
|
|
name: "missing schema JSON",
|
|
mutate: func(req *contracts.StructuredCompletionRequest) {
|
|
req.ResponseSchema = nil
|
|
},
|
|
want: "schema JSON",
|
|
},
|
|
{
|
|
name: "invalid schema JSON",
|
|
mutate: func(req *contracts.StructuredCompletionRequest) {
|
|
req.ResponseSchema = json.RawMessage(`{"type":`)
|
|
},
|
|
want: "schema JSON",
|
|
},
|
|
}
|
|
|
|
for _, tc := range tests {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
req := validStructuredRequest("")
|
|
tc.mutate(&req)
|
|
var out testArtifact
|
|
_, err := client.CompleteStructured(context.Background(), req, &out)
|
|
if err == nil || !strings.Contains(err.Error(), tc.want) {
|
|
t.Fatalf("expected error containing %q, got %v", tc.want, err)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestOpenAICompatibleClientRejectsEmptyMessages(t *testing.T) {
|
|
client := newTestClient(t, "https://example.test/v1", "default-model", 0)
|
|
|
|
tests := []struct {
|
|
name string
|
|
mutate func(*contracts.StructuredCompletionRequest)
|
|
want string
|
|
}{
|
|
{
|
|
name: "no messages",
|
|
mutate: func(req *contracts.StructuredCompletionRequest) {
|
|
req.Messages = nil
|
|
},
|
|
want: "messages",
|
|
},
|
|
{
|
|
name: "empty role",
|
|
mutate: func(req *contracts.StructuredCompletionRequest) {
|
|
req.Messages[0].Role = " "
|
|
},
|
|
want: "role",
|
|
},
|
|
{
|
|
name: "empty content",
|
|
mutate: func(req *contracts.StructuredCompletionRequest) {
|
|
req.Messages[0].Content = " "
|
|
},
|
|
want: "content",
|
|
},
|
|
}
|
|
|
|
for _, tc := range tests {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
req := validStructuredRequest("")
|
|
tc.mutate(&req)
|
|
var out testArtifact
|
|
_, err := client.CompleteStructured(context.Background(), req, &out)
|
|
if err == nil || !strings.Contains(err.Error(), tc.want) {
|
|
t.Fatalf("expected error containing %q, got %v", tc.want, err)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestOpenAICompatibleClientProviderNon2xxBehavior(t *testing.T) {
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.WriteHeader(http.StatusBadRequest)
|
|
_, _ = io.WriteString(w, `{"error":{"message":"bad request"}}`)
|
|
}))
|
|
defer server.Close()
|
|
|
|
client := newTestClient(t, server.URL, "default-model", 0)
|
|
var out testArtifact
|
|
_, err := client.CompleteStructured(context.Background(), validStructuredRequest(""), &out)
|
|
if err == nil || !strings.Contains(err.Error(), "status 400: bad request") {
|
|
t.Fatalf("expected provider status error, got %v", err)
|
|
}
|
|
}
|
|
|
|
func TestOpenAICompatibleClientRetries429And5xx(t *testing.T) {
|
|
var attempts atomic.Int32
|
|
statuses := []int{http.StatusTooManyRequests, http.StatusInternalServerError, http.StatusOK}
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
attempt := int(attempts.Add(1)) - 1
|
|
if statuses[attempt] != http.StatusOK {
|
|
w.WriteHeader(statuses[attempt])
|
|
_, _ = io.WriteString(w, `{"error":{"message":"try again"}}`)
|
|
return
|
|
}
|
|
_, _ = io.WriteString(w, `{"choices":[{"message":{"content":"{\"value\":\"ok\"}"}}]}`)
|
|
}))
|
|
defer server.Close()
|
|
|
|
client := newTestClient(t, server.URL, "default-model", 2)
|
|
var out testArtifact
|
|
if _, err := client.CompleteStructured(context.Background(), validStructuredRequest(""), &out); err != nil {
|
|
t.Fatalf("CompleteStructured: %v", err)
|
|
}
|
|
if attempts.Load() != 3 {
|
|
t.Fatalf("expected 3 attempts, got %d", attempts.Load())
|
|
}
|
|
}
|
|
|
|
func TestOpenAICompatibleClientRetriesMalformedResponses(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
firstBody string
|
|
}{
|
|
{
|
|
name: "malformed provider envelope",
|
|
firstBody: `{"choices":[]}`,
|
|
},
|
|
{
|
|
name: "malformed assistant JSON",
|
|
firstBody: `{"choices":[{"message":{"content":"{"}}]}`,
|
|
},
|
|
}
|
|
|
|
for _, tc := range tests {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
var attempts atomic.Int32
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
if attempts.Add(1) == 1 {
|
|
_, _ = io.WriteString(w, tc.firstBody)
|
|
return
|
|
}
|
|
_, _ = io.WriteString(w, `{"choices":[{"message":{"content":"{\"value\":\"ok\"}"}}]}`)
|
|
}))
|
|
defer server.Close()
|
|
|
|
client := newTestClient(t, server.URL, "default-model", 1)
|
|
var out testArtifact
|
|
if _, err := client.CompleteStructured(context.Background(), validStructuredRequest(""), &out); err != nil {
|
|
t.Fatalf("CompleteStructured: %v", err)
|
|
}
|
|
if attempts.Load() != 2 {
|
|
t.Fatalf("expected 2 attempts, got %d", attempts.Load())
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestOpenAICompatibleClientNoRetryForNonRetryable4xx(t *testing.T) {
|
|
var attempts atomic.Int32
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
attempts.Add(1)
|
|
w.WriteHeader(http.StatusForbidden)
|
|
_, _ = io.WriteString(w, `{"error":{"message":"forbidden"}}`)
|
|
}))
|
|
defer server.Close()
|
|
|
|
client := newTestClient(t, server.URL, "default-model", 3)
|
|
var out testArtifact
|
|
_, err := client.CompleteStructured(context.Background(), validStructuredRequest(""), &out)
|
|
if err == nil || !strings.Contains(err.Error(), "status 403") {
|
|
t.Fatalf("expected forbidden error, got %v", err)
|
|
}
|
|
if attempts.Load() != 1 {
|
|
t.Fatalf("expected 1 attempt, got %d", attempts.Load())
|
|
}
|
|
}
|
|
|
|
func TestOpenAICompatibleClientProviderErrorRedactsAPIKey(t *testing.T) {
|
|
const apiKey = "secret-api-key"
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.WriteHeader(http.StatusInternalServerError)
|
|
_, _ = io.WriteString(w, `{"error":{"message":"Bearer secret-api-key failed for secret-api-key"}}`)
|
|
}))
|
|
defer server.Close()
|
|
|
|
client, err := NewOpenAICompatibleClient(OpenAICompatibleClientConfig{
|
|
BaseURL: server.URL,
|
|
Model: "default-model",
|
|
APIKey: apiKey,
|
|
MaxRetries: 0,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NewOpenAICompatibleClient: %v", err)
|
|
}
|
|
|
|
var out testArtifact
|
|
_, err = client.CompleteStructured(context.Background(), validStructuredRequest(""), &out)
|
|
if err == nil {
|
|
t.Fatalf("expected provider error")
|
|
}
|
|
if strings.Contains(err.Error(), apiKey) || strings.Contains(err.Error(), "Bearer "+apiKey) {
|
|
t.Fatalf("expected API key to be redacted, got %q", err.Error())
|
|
}
|
|
}
|
|
|
|
func TestOpenAICompatibleClientRespectsContextCancellation(t *testing.T) {
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
cancel()
|
|
|
|
client := newTestClient(t, "https://example.test/v1", "default-model", 1)
|
|
var out testArtifact
|
|
_, err := client.CompleteStructured(ctx, validStructuredRequest(""), &out)
|
|
if !errors.Is(err, context.Canceled) {
|
|
t.Fatalf("expected context canceled, got %v", err)
|
|
}
|
|
}
|
|
|
|
func newTestClient(t *testing.T, baseURL string, model string, maxRetries int) *OpenAICompatibleClient {
|
|
t.Helper()
|
|
client, err := NewOpenAICompatibleClient(OpenAICompatibleClientConfig{
|
|
BaseURL: baseURL,
|
|
Model: model,
|
|
MaxRetries: maxRetries,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NewOpenAICompatibleClient: %v", err)
|
|
}
|
|
return client
|
|
}
|
|
|
|
func validStructuredRequest(model string) contracts.StructuredCompletionRequest {
|
|
return contracts.StructuredCompletionRequest{
|
|
Messages: []contracts.LLMMessage{
|
|
{Role: " user ", Content: " extract this "},
|
|
},
|
|
Model: model,
|
|
ResponseSchemaName: " test_artifact ",
|
|
ResponseSchema: testResponseSchema(),
|
|
}
|
|
}
|
|
|
|
func testResponseSchema() json.RawMessage {
|
|
return json.RawMessage(`{
|
|
"type": "object",
|
|
"properties": {
|
|
"value": {"type": "string"}
|
|
},
|
|
"required": ["value"],
|
|
"additionalProperties": false
|
|
}`)
|
|
}
|