215 lines
7.1 KiB
Go
215 lines
7.1 KiB
Go
package llm
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
ck "gitea.maximumdirect.net/eric/ConfigKit"
|
|
"gitea.maximumdirect.net/eric/audita/internal/core/config"
|
|
"gitea.maximumdirect.net/eric/audita/internal/framework/responseschema"
|
|
"gitea.maximumdirect.net/eric/audita/internal/framework/structuredoutput"
|
|
"gitea.maximumdirect.net/eric/audita/internal/prompts"
|
|
pk "gitea.maximumdirect.net/eric/promptkit"
|
|
)
|
|
|
|
func TestRolesShareBackendCapacityAndReleaseAfterCancellation(t *testing.T) {
|
|
entered := make(chan struct{}, 3)
|
|
release := make(chan struct{})
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
entered <- struct{}{}
|
|
<-release
|
|
content, _ := json.Marshal(`{"corrections":[]}`)
|
|
_, _ = fmt.Fprintf(w, `{"choices":[{"message":{"content":%s}}]}`, content)
|
|
}))
|
|
defer server.Close()
|
|
dir := t.TempDir()
|
|
if err := os.Mkdir(filepath.Join(dir, "profiles"), 0700); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
main := filepath.Join(dir, "config.yml")
|
|
body := fmt.Sprintf("version: 2\nbackends:\n shared:\n endpoint: %s/v1\n concurrency: 1\n queue_capacity: 0\nllm:\n proposal_profile: proposal\n validation_profile: validation\n transient_retries: 0\n repair_attempts: 0\n", server.URL)
|
|
if err := os.WriteFile(main, []byte(body), 0600); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
for _, name := range []string{"proposal", "validation"} {
|
|
profile := "id: " + name + "\nbackend: shared\nmodel: local\n"
|
|
if err := os.WriteFile(filepath.Join(dir, "profiles", name+".yaml"), []byte(profile), 0600); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
loaded, err := config.Load(ck.File(main), "")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
runtime, err := NewRuntime(context.Background(), loaded.Config, main, server.Client())
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
proposal, err := runtime.NewRoleClient("proposal", CallHooks{})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
validation, err := runtime.NewRoleClient("validation", CallHooks{})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
request := structuredoutput.Request{PromptID: prompts.PromptIDModuleGlossaryProposal, PromptVersion: prompts.VersionV1, Variables: map[string]string{"TranscriptDescriptionBlock": "", "GlossaryJSON": "[]", "SectionJSON": "[]"}, Schema: responseschema.CorrectionSetKey}
|
|
firstDone := make(chan error, 1)
|
|
go func() {
|
|
var output map[string]any
|
|
_, callErr := proposal.CompleteStructured(context.Background(), request, &output)
|
|
firstDone <- callErr
|
|
}()
|
|
select {
|
|
case <-entered:
|
|
case <-time.After(2 * time.Second):
|
|
t.Fatal("first call did not enter provider")
|
|
}
|
|
var output map[string]any
|
|
_, err = validation.CompleteStructured(context.Background(), request, &output)
|
|
if !errors.Is(err, pk.ErrCapacityExceeded) {
|
|
t.Fatalf("other role did not share capacity: %v", err)
|
|
}
|
|
select {
|
|
case <-entered:
|
|
t.Fatal("second role reached provider despite full backend")
|
|
default:
|
|
}
|
|
close(release)
|
|
if err := <-firstDone; err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
_, err = validation.CompleteStructured(context.Background(), request, &output)
|
|
if err != nil {
|
|
t.Fatalf("capacity was not released: %v", err)
|
|
}
|
|
for len(entered) > 0 {
|
|
<-entered
|
|
}
|
|
release = make(chan struct{})
|
|
if err := os.WriteFile(main, []byte(strings.Replace(body, "queue_capacity: 0", "queue_capacity: 1", 1)), 0600); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
loaded, err = config.Load(ck.File(main), "")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
runtime, err = NewRuntime(context.Background(), loaded.Config, main, server.Client())
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
proposal, err = runtime.NewRoleClient("proposal", CallHooks{})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
validation, err = runtime.NewRoleClient("validation", CallHooks{})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
go func() {
|
|
var out map[string]any
|
|
_, callErr := proposal.CompleteStructured(context.Background(), request, &out)
|
|
firstDone <- callErr
|
|
}()
|
|
select {
|
|
case <-entered:
|
|
case <-time.After(2 * time.Second):
|
|
t.Fatal("first queued-capacity call did not dispatch")
|
|
}
|
|
queuedCtx, cancelQueued := context.WithTimeout(context.Background(), 30*time.Millisecond)
|
|
defer cancelQueued()
|
|
_, err = validation.CompleteStructured(queuedCtx, request, &output)
|
|
if !errors.Is(err, context.DeadlineExceeded) {
|
|
t.Fatalf("queued call did not honor deadline: %v", err)
|
|
}
|
|
close(release)
|
|
if err := <-firstDone; err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
_, err = validation.CompleteStructured(context.Background(), request, &output)
|
|
if err != nil {
|
|
t.Fatalf("queued cancellation leaked capacity: %v", err)
|
|
}
|
|
}
|
|
|
|
func TestDistinctBackendsHaveIndependentCapacity(t *testing.T) {
|
|
entered := make(chan struct{}, 2)
|
|
release := make(chan struct{})
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
entered <- struct{}{}
|
|
<-release
|
|
schema := map[string]any{}
|
|
_ = json.NewDecoder(r.Body).Decode(&schema)
|
|
content, _ := json.Marshal(`{"corrections":[]}`)
|
|
_, _ = fmt.Fprintf(w, `{"choices":[{"message":{"content":%s}}]}`, content)
|
|
}))
|
|
defer server.Close()
|
|
dir := t.TempDir()
|
|
if err := os.Mkdir(filepath.Join(dir, "profiles"), 0700); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
main := filepath.Join(dir, "config.yml")
|
|
body := fmt.Sprintf("version: 2\nbackends:\n first:\n endpoint: %s/v1\n concurrency: 1\n queue_capacity: 0\n second:\n endpoint: %s/v1\n concurrency: 1\n queue_capacity: 0\nllm:\n proposal_profile: proposal\n validation_profile: validation\n transient_retries: 0\n repair_attempts: 0\n", server.URL, server.URL)
|
|
if err := os.WriteFile(main, []byte(body), 0600); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
for name, backend := range map[string]string{"proposal": "first", "validation": "second"} {
|
|
if err := os.WriteFile(filepath.Join(dir, "profiles", name+".yaml"), []byte("id: "+name+"\nbackend: "+backend+"\nmodel: local\n"), 0600); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
loaded, err := config.Load(ck.File(main), "")
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
runtime, err := NewRuntime(context.Background(), loaded.Config, main, server.Client())
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
proposal, err := runtime.NewRoleClient("proposal", CallHooks{})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
validation, err := runtime.NewRoleClient("validation", CallHooks{})
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
request := structuredoutput.Request{PromptID: prompts.PromptIDModuleGlossaryProposal, PromptVersion: prompts.VersionV1, Variables: map[string]string{"TranscriptDescriptionBlock": "", "GlossaryJSON": "[]", "SectionJSON": "[]"}, Schema: responseschema.CorrectionSetKey}
|
|
done := make(chan error, 2)
|
|
go func() {
|
|
var out map[string]any
|
|
_, e := proposal.CompleteStructured(context.Background(), request, &out)
|
|
done <- e
|
|
}()
|
|
select {
|
|
case <-entered:
|
|
case <-time.After(2 * time.Second):
|
|
t.Fatal("first backend did not dispatch")
|
|
}
|
|
go func() {
|
|
var out map[string]any
|
|
_, e := validation.CompleteStructured(context.Background(), request, &out)
|
|
done <- e
|
|
}()
|
|
select {
|
|
case <-entered:
|
|
case <-time.After(2 * time.Second):
|
|
t.Fatal("second backend was blocked by first")
|
|
}
|
|
close(release)
|
|
for range 2 {
|
|
if err := <-done; err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
}
|
|
}
|