287 lines
12 KiB
Go
287 lines
12 KiB
Go
package llm
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"os"
|
|
"path/filepath"
|
|
"slices"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"gitea.maximumdirect.net/eric/audita/v3/internal/core/config"
|
|
"gitea.maximumdirect.net/eric/audita/v3/internal/framework/outcomes"
|
|
"gitea.maximumdirect.net/eric/audita/v3/internal/framework/responseschema"
|
|
"gitea.maximumdirect.net/eric/audita/v3/internal/framework/structuredoutput"
|
|
"gitea.maximumdirect.net/eric/audita/v3/internal/prompts"
|
|
ck "go.maximumdirect.net/configkit"
|
|
)
|
|
|
|
type transportFunc func(*http.Request) (*http.Response, error)
|
|
|
|
func (f transportFunc) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) }
|
|
|
|
func TestRoleClientUsesNativeTargetAndRedactsCredential(t *testing.T) {
|
|
t.Setenv("AUDITA_TEST_KEY", "synthetic-secret-value")
|
|
dir := t.TempDir()
|
|
if err := os.WriteFile(filepath.Join(dir, "config.yml"), []byte("version: 3\n"), 0644); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
document := "profiles:\n - id: fixture\n endpoint: https://example.invalid/v1\n model: test-model\n api_key_env: AUDITA_TEST_KEY\n"
|
|
if err := os.WriteFile(filepath.Join(dir, "prompts.yml"), []byte(document), 0644); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
loaded, err := config.Load(ck.Input{Source: ck.File(filepath.Join(dir, "config.yml"))}, "")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
calls := 0
|
|
providerErr := errors.New("provider echoed synthetic-secret-value")
|
|
httpClient := &http.Client{Transport: transportFunc(func(r *http.Request) (*http.Response, error) {
|
|
calls++
|
|
if r.Header.Get("Authorization") != "Bearer synthetic-secret-value" {
|
|
t.Errorf("authorization missing")
|
|
}
|
|
return nil, providerErr
|
|
})}
|
|
runtime, err := NewRuntime(context.Background(), loaded, httpClient)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
client, err := runtime.NewRoleClient("proposal", "fixture", config.LLM{OperationTimeout: time.Second}, CallHooks{})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
var out map[string]any
|
|
_, err = client.CompleteStructured(context.Background(), structuredoutput.Request{PromptID: prompts.PromptIDModuleGrammarProposal, Schema: responseschema.CorrectionSetKey, Variables: map[string]string{"TranscriptDescriptionBlock": "", "GlossaryJSON": "[]", "SectionJSON": "[]"}}, &out)
|
|
if err == nil || calls != 1 {
|
|
t.Fatalf("error = %v, calls = %d", err, calls)
|
|
}
|
|
masked := runtime.RedactError(context.Background(), providerErr)
|
|
if strings.Contains(masked, "synthetic-secret-value") {
|
|
t.Fatal("credential not redacted")
|
|
}
|
|
}
|
|
|
|
func TestRoleClientExecutionBudgetsAndFailureClasses(t *testing.T) {
|
|
dir := t.TempDir()
|
|
for name, body := range map[string]string{
|
|
"config.yml": "version: 3\n",
|
|
"prompts.yml": "profiles:\n - id: fixture\n endpoint: https://example.invalid/v1\n model: test-model\n",
|
|
} {
|
|
if err := os.WriteFile(filepath.Join(dir, name), []byte(body), 0644); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
loaded, err := config.Load(ck.Input{Source: ck.File(filepath.Join(dir, "config.yml"))}, "")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
request := structuredoutput.Request{PromptID: prompts.PromptIDModuleGrammarProposal, Schema: responseschema.CorrectionSetKey, Variables: map[string]string{"TranscriptDescriptionBlock": "", "GlossaryJSON": "[]", "SectionJSON": "[]"}}
|
|
for _, tc := range []struct {
|
|
name string
|
|
repair int
|
|
outputs []string
|
|
wantCalls int
|
|
success bool
|
|
}{
|
|
{"success", 0, []string{`{"corrections":[]}`}, 1, true},
|
|
{"repair", 1, []string{`{"wrong":true}`, `{"corrections":[]}`}, 2, true},
|
|
{"no repair", 0, []string{`{"wrong":true}`}, 1, false},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
calls := 0
|
|
httpClient := &http.Client{Transport: transportFunc(func(r *http.Request) (*http.Response, error) {
|
|
item := tc.outputs[min(calls, len(tc.outputs)-1)]
|
|
calls++
|
|
payload, _ := json.Marshal(item)
|
|
body := fmt.Sprintf(`{"model":"test-model","choices":[{"message":{"content":%s}}]}`, payload)
|
|
return &http.Response{StatusCode: 200, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(body)), Request: r}, nil
|
|
})}
|
|
runtime, err := NewRuntime(context.Background(), loaded, httpClient)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
client, err := runtime.NewRoleClient("proposal", "fixture", config.LLM{OperationTimeout: 10 * time.Second, RepairAttempts: tc.repair}, CallHooks{})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
var result map[string]any
|
|
response, err := client.CompleteStructured(context.Background(), request, &result)
|
|
if (err == nil) != tc.success || calls != tc.wantCalls {
|
|
t.Fatalf("error=%v calls=%d", err, calls)
|
|
}
|
|
if response.Prompt.ProfileID != "fixture" || response.Prompt.Model != "test-model" {
|
|
t.Fatalf("prepared target = %+v", response.Prompt)
|
|
}
|
|
if !tc.success && !structuredoutput.IsRecoverable(err) {
|
|
t.Fatalf("expected recoverable failure: %v", err)
|
|
}
|
|
})
|
|
}
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
cancel()
|
|
runtime, err := NewRuntime(context.Background(), loaded, nil)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
client, err := runtime.NewRoleClient("proposal", "fixture", config.LLM{OperationTimeout: time.Second}, CallHooks{})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
_, err = client.CompleteStructured(ctx, request, &map[string]any{})
|
|
if !errors.Is(err, context.Canceled) {
|
|
t.Fatalf("cancellation = %v", err)
|
|
}
|
|
}
|
|
|
|
func TestRoleClientsShareBackendCapacity(t *testing.T) {
|
|
dir := t.TempDir()
|
|
for name, body := range map[string]string{
|
|
"config.yml": "version: 3\n",
|
|
"prompts.yml": "backends:\n - id: shared\n endpoint: https://example.invalid/v1\n concurrency_limit: 1\n queue_capacity: 0\nprofiles:\n - id: first\n backend: shared\n model: one\n - id: second\n backend: shared\n model: two\n",
|
|
} {
|
|
if err := os.WriteFile(filepath.Join(dir, name), []byte(body), 0644); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
loaded, err := config.Load(ck.Input{Source: ck.File(filepath.Join(dir, "config.yml"))}, "")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
entered, release := make(chan struct{}), make(chan struct{})
|
|
transport := transportFunc(func(r *http.Request) (*http.Response, error) {
|
|
close(entered)
|
|
<-release
|
|
body := `{"model":"one","choices":[{"message":{"content":"{\"corrections\":[]}"}}]}`
|
|
return &http.Response{StatusCode: 200, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(body)), Request: r}, nil
|
|
})
|
|
runtime, err := NewRuntime(context.Background(), loaded, &http.Client{Transport: transport})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
settings := config.LLM{OperationTimeout: 10 * time.Second}
|
|
first, err := runtime.NewRoleClient("proposal", "first", settings, CallHooks{})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
second, err := runtime.NewRoleClient("validation", "second", settings, CallHooks{})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
request := structuredoutput.Request{PromptID: prompts.PromptIDModuleGrammarProposal, Schema: responseschema.CorrectionSetKey, Variables: map[string]string{"TranscriptDescriptionBlock": "", "GlossaryJSON": "[]", "SectionJSON": "[]"}}
|
|
done := make(chan error, 1)
|
|
go func() {
|
|
_, err := first.CompleteStructured(context.Background(), request, &map[string]any{})
|
|
done <- err
|
|
}()
|
|
<-entered
|
|
_, secondErr := second.CompleteStructured(context.Background(), request, &map[string]any{})
|
|
close(release)
|
|
if err := <-done; err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
var failure *structuredoutput.Failure
|
|
if !errors.As(secondErr, &failure) || failure.Reason != outcomes.ReasonCapacityExceeded {
|
|
t.Fatalf("second call = %v", secondErr)
|
|
}
|
|
}
|
|
|
|
func TestRoleClientRetriesFollowRequestBudget(t *testing.T) {
|
|
dir := t.TempDir()
|
|
for name, body := range map[string]string{"config.yml": "version: 3\n", "prompts.yml": "profiles:\n - id: fixture\n endpoint: https://example.invalid/v1\n model: test-model\n"} {
|
|
if err := os.WriteFile(filepath.Join(dir, name), []byte(body), 0644); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
loaded, err := config.Load(ck.Input{Source: ck.File(filepath.Join(dir, "config.yml"))}, "")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
calls := 0
|
|
transport := transportFunc(func(r *http.Request) (*http.Response, error) {
|
|
calls++
|
|
status, body := 503, `{"error":{"message":"unavailable"}}`
|
|
if calls > 1 {
|
|
status, body = 200, `{"model":"test-model","choices":[{"message":{"content":"{\"corrections\":[]}"}}]}`
|
|
}
|
|
return &http.Response{StatusCode: status, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(body)), Request: r}, nil
|
|
})
|
|
runtime, err := NewRuntime(context.Background(), loaded, &http.Client{Transport: transport})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
client, err := runtime.NewRoleClient("proposal", "fixture", config.LLM{OperationTimeout: 10 * time.Second, TransientRetries: 1}, CallHooks{})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
req := structuredoutput.Request{PromptID: prompts.PromptIDModuleGrammarProposal, Schema: responseschema.CorrectionSetKey, Variables: map[string]string{"TranscriptDescriptionBlock": "", "GlossaryJSON": "[]", "SectionJSON": "[]"}}
|
|
response, err := client.CompleteStructured(context.Background(), req, &map[string]any{})
|
|
if err != nil || calls != 2 || response.Execution.TransientRetries != 1 {
|
|
t.Fatalf("error=%v calls=%d execution=%+v", err, calls, response.Execution)
|
|
}
|
|
}
|
|
|
|
func TestNativeCredentialSelectionAndOptionalAbsence(t *testing.T) {
|
|
t.Setenv("AUDITA_BACKEND_KEY", "backend-secret-value")
|
|
t.Setenv("AUDITA_PROFILE_KEY", "profile-secret-value")
|
|
dir := t.TempDir()
|
|
files := map[string]string{
|
|
"config.yml": "version: 3\n",
|
|
"prompts.yml": "backends:\n - id: service\n endpoint: https://example.invalid/v1\n api_key_env: AUDITA_BACKEND_KEY\nprofiles:\n - id: inherited\n backend: service\n model: fixture\n api_key_env: null\n - id: explicit\n backend: service\n model: fixture\n api_key_env: AUDITA_PROFILE_KEY\n - id: cleared\n backend: service\n model: fixture\n api_key_env: ''\n",
|
|
}
|
|
for name, body := range files {
|
|
if err := os.WriteFile(filepath.Join(dir, name), []byte(body), 0644); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
loaded, err := config.Load(ck.Input{Source: ck.File(filepath.Join(dir, "config.yml"))}, "")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
authorization := ""
|
|
transport := transportFunc(func(r *http.Request) (*http.Response, error) {
|
|
authorization = r.Header.Get("Authorization")
|
|
body := `{"model":"fixture","choices":[{"message":{"content":"{\"corrections\":[]}"}}]}`
|
|
return &http.Response{StatusCode: 200, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(body)), Request: r}, nil
|
|
})
|
|
runtime, err := NewRuntime(context.Background(), loaded, &http.Client{Transport: transport})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
req := structuredoutput.Request{PromptID: prompts.PromptIDModuleGrammarProposal, Schema: responseschema.CorrectionSetKey, Variables: map[string]string{"TranscriptDescriptionBlock": "", "GlossaryJSON": "[]", "SectionJSON": "[]"}}
|
|
for _, tc := range []struct{ profile, auth string }{{"inherited", "Bearer backend-secret-value"}, {"explicit", "Bearer profile-secret-value"}, {"cleared", ""}} {
|
|
t.Run(tc.profile, func(t *testing.T) {
|
|
client, err := runtime.NewRoleClient("proposal", tc.profile, config.LLM{OperationTimeout: 5 * time.Second}, CallHooks{})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
response, err := client.CompleteStructured(context.Background(), req, &map[string]any{})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if authorization != tc.auth {
|
|
t.Fatalf("authorization = %q", authorization)
|
|
}
|
|
snapshot, err := json.Marshal(response.Prepared)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if strings.Contains(string(snapshot), "secret-value") {
|
|
t.Fatal("prepared snapshot contains credential")
|
|
}
|
|
})
|
|
}
|
|
secrets := runtime.ConfiguredSecrets(context.Background())
|
|
for _, value := range []string{"backend-secret-value", "profile-secret-value"} {
|
|
if !slices.Contains(secrets, value) {
|
|
t.Fatalf("missing credential masking value")
|
|
}
|
|
}
|
|
}
|