Add Scriptorium-backed LLM runtime
This commit is contained in:
299
internal/framework/llm/scriptorium_client_test.go
Normal file
299
internal/framework/llm/scriptorium_client_test.go
Normal file
@@ -0,0 +1,299 @@
|
||||
package llm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"testing/fstest"
|
||||
"time"
|
||||
|
||||
"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
|
||||
"gitea.maximumdirect.net/eric/scriptorium"
|
||||
)
|
||||
|
||||
func TestScriptoriumClientMapsPromptRequestAndUnmarshalsOutput(t *testing.T) {
|
||||
fake := &fakeScriptoriumLLM{content: `{"ok":true}`}
|
||||
client := newTestScriptoriumClient(t, fake)
|
||||
|
||||
var out struct {
|
||||
OK bool `json:"ok"`
|
||||
}
|
||||
resp, err := client.CompleteStructured(context.Background(), contracts.StructuredCompletionRequest{
|
||||
StageName: "test-stage",
|
||||
PromptID: "adapter.test",
|
||||
PromptVersion: "v1",
|
||||
ProfileID: "explicit-profile",
|
||||
SessionID: "session-123",
|
||||
Inputs: contracts.LLMInputSet{
|
||||
"transcript": contracts.NewLLMInputMaterial("transcript", "application/json", []byte(`{"source":true}`), "sha256:source", "file:///source.json"),
|
||||
},
|
||||
Vars: map[string]any{"custom": "value"},
|
||||
}, &out)
|
||||
if err != nil {
|
||||
t.Fatalf("CompleteStructured() error = %v, want nil", err)
|
||||
}
|
||||
if !out.OK {
|
||||
t.Fatalf("decoded output OK = false, want true")
|
||||
}
|
||||
if resp.Provider != scriptoriumProviderName || resp.Model != "explicit-model" || resp.ProfileID != "explicit-profile" {
|
||||
t.Fatalf("response metadata = %#v", resp)
|
||||
}
|
||||
if resp.PromptTokens != 11 || resp.CompletionTokens != 7 || resp.TotalTokens != 18 {
|
||||
t.Fatalf("usage = %#v, want mapped token counts", resp)
|
||||
}
|
||||
gotReq := fake.lastRequest()
|
||||
if gotReq.Prompt.SessionID != "session-123" {
|
||||
t.Fatalf("session id = %q, want session-123", gotReq.Prompt.SessionID)
|
||||
}
|
||||
if gotReq.Target.Model != "explicit-model" {
|
||||
t.Fatalf("model = %q, want explicit-model", gotReq.Target.Model)
|
||||
}
|
||||
if len(gotReq.Prompt.Messages) != 1 || !strings.Contains(gotReq.Prompt.Messages[0].Content, `{"source":true}`) {
|
||||
t.Fatalf("rendered messages = %#v, want transcript input content", gotReq.Prompt.Messages)
|
||||
}
|
||||
if gotReq.StructuredOutput == nil {
|
||||
t.Fatalf("structured output = nil, want JSON schema")
|
||||
}
|
||||
manifests := client.LLMProfileManifests()
|
||||
if len(manifests) != 1 || manifests[0].ID != "explicit-profile" || manifests[0].Model != "explicit-model" {
|
||||
t.Fatalf("profile manifests = %#v", manifests)
|
||||
}
|
||||
}
|
||||
|
||||
func TestScriptoriumClientUsesPromptDefaultProfileWhenRequestProfileEmpty(t *testing.T) {
|
||||
fake := &fakeScriptoriumLLM{content: `{"ok":true}`}
|
||||
client := newTestScriptoriumClient(t, fake)
|
||||
|
||||
var out map[string]any
|
||||
if _, err := client.CompleteStructured(context.Background(), contracts.StructuredCompletionRequest{
|
||||
PromptID: "adapter.test",
|
||||
SessionID: "session-123",
|
||||
Inputs: contracts.LLMInputSet{
|
||||
"transcript": contracts.NewLLMInputMaterial("transcript", "application/json", []byte(`{"source":true}`), "", ""),
|
||||
},
|
||||
}, &out); err != nil {
|
||||
t.Fatalf("CompleteStructured() error = %v, want nil", err)
|
||||
}
|
||||
if got := fake.lastRequest().Target.Model; got != "default-model" {
|
||||
t.Fatalf("model = %q, want prompt default profile model", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestScriptoriumClientValidationFailureReturnsError(t *testing.T) {
|
||||
client := newTestScriptoriumClient(t, &fakeScriptoriumLLM{content: `{"bad":true}`})
|
||||
|
||||
var out map[string]any
|
||||
_, err := client.CompleteStructured(context.Background(), contracts.StructuredCompletionRequest{
|
||||
PromptID: "adapter.test",
|
||||
SessionID: "session-123",
|
||||
Inputs: contracts.LLMInputSet{
|
||||
"transcript": contracts.NewLLMInputMaterial("transcript", "application/json", []byte(`{"source":true}`), "", ""),
|
||||
},
|
||||
}, &out)
|
||||
if err == nil || !strings.Contains(err.Error(), "validation failed") {
|
||||
t.Fatalf("CompleteStructured() error = %v, want validation failure", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestScriptoriumClientProviderFailureIncludesContextAndRedactsBearerToken(t *testing.T) {
|
||||
client := newTestScriptoriumClient(t, &fakeScriptoriumLLM{err: errors.New("provider failed with Bearer secret-token")})
|
||||
|
||||
var out map[string]any
|
||||
_, err := client.CompleteStructured(context.Background(), contracts.StructuredCompletionRequest{
|
||||
PromptID: "adapter.test",
|
||||
SessionID: "session-123",
|
||||
Inputs: contracts.LLMInputSet{
|
||||
"transcript": contracts.NewLLMInputMaterial("transcript", "application/json", []byte(`{"source":true}`), "", ""),
|
||||
},
|
||||
}, &out)
|
||||
if err == nil {
|
||||
t.Fatalf("CompleteStructured() error = nil, want provider error")
|
||||
}
|
||||
if !strings.Contains(err.Error(), `run Scriptorium prompt "adapter.test"`) {
|
||||
t.Fatalf("error = %q, want operation context", err.Error())
|
||||
}
|
||||
if strings.Contains(err.Error(), "secret-token") || !strings.Contains(err.Error(), "Bearer [REDACTED]") {
|
||||
t.Fatalf("error = %q, want redacted bearer token", err.Error())
|
||||
}
|
||||
}
|
||||
|
||||
func TestScriptoriumClientContextCancellationIsRespected(t *testing.T) {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
client := newTestScriptoriumClient(t, &fakeScriptoriumLLM{content: `{"ok":true}`})
|
||||
|
||||
var out map[string]any
|
||||
_, err := client.CompleteStructured(ctx, contracts.StructuredCompletionRequest{
|
||||
PromptID: "adapter.test",
|
||||
Inputs: contracts.LLMInputSet{
|
||||
"transcript": contracts.NewLLMInputMaterial("transcript", "application/json", []byte(`{"source":true}`), "", ""),
|
||||
},
|
||||
}, &out)
|
||||
if !errors.Is(err, context.Canceled) {
|
||||
t.Fatalf("CompleteStructured() error = %v, want context canceled", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestScheduledScriptoriumClientBoundsConcurrentCalls(t *testing.T) {
|
||||
fake := &fakeScriptoriumLLM{
|
||||
content: `{"ok":true}`,
|
||||
block: make(chan struct{}),
|
||||
}
|
||||
client := newTestScriptoriumClient(t, fake)
|
||||
scheduler, err := NewScheduler(1)
|
||||
if err != nil {
|
||||
t.Fatalf("NewScheduler() error = %v, want nil", err)
|
||||
}
|
||||
scheduled := NewScheduledClient(client, scheduler)
|
||||
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < 3; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
var out map[string]any
|
||||
_, callErr := scheduled.CompleteStructured(context.Background(), contracts.StructuredCompletionRequest{
|
||||
PromptID: "adapter.test",
|
||||
SessionID: "session-123",
|
||||
Inputs: contracts.LLMInputSet{
|
||||
"transcript": contracts.NewLLMInputMaterial("transcript", "application/json", []byte(`{"source":true}`), "", ""),
|
||||
},
|
||||
}, &out)
|
||||
if callErr != nil {
|
||||
t.Errorf("CompleteStructured() error = %v, want nil", callErr)
|
||||
}
|
||||
}()
|
||||
}
|
||||
waitForAtomicAtLeast(t, &fake.calls, 1)
|
||||
time.Sleep(20 * time.Millisecond)
|
||||
if got := atomic.LoadInt32(&fake.maxInFlight); got > 1 {
|
||||
t.Fatalf("max in-flight calls = %d, want <= 1", got)
|
||||
}
|
||||
close(fake.block)
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
func TestScriptoriumClientValidatesRequest(t *testing.T) {
|
||||
client := newTestScriptoriumClient(t, &fakeScriptoriumLLM{content: `{"ok":true}`})
|
||||
var out map[string]any
|
||||
if _, err := client.CompleteStructured(context.Background(), contracts.StructuredCompletionRequest{}, &out); err == nil || !strings.Contains(err.Error(), "prompt_id") {
|
||||
t.Fatalf("missing prompt id error = %v, want prompt_id validation", err)
|
||||
}
|
||||
if _, err := client.CompleteStructured(context.Background(), contracts.StructuredCompletionRequest{PromptID: "adapter.test"}, nil); err == nil || !strings.Contains(err.Error(), "non-nil pointer") {
|
||||
t.Fatalf("nil output error = %v, want output validation", err)
|
||||
}
|
||||
}
|
||||
|
||||
func newTestScriptoriumClient(t *testing.T, fake *fakeScriptoriumLLM) *ScriptoriumClient {
|
||||
t.Helper()
|
||||
registry := NewAssetRegistry()
|
||||
if err := registry.RegisterPromptFS(fstest.MapFS{
|
||||
"adapter.test.yaml": {Data: []byte(`id: adapter.test
|
||||
version: "v1"
|
||||
default_profile: default-profile
|
||||
session_id: "{{ .session_id }}"
|
||||
inputs:
|
||||
- name: transcript
|
||||
required: true
|
||||
content_type: application/json
|
||||
messages:
|
||||
- role: user
|
||||
content: "Transcript: {{ input \"transcript\" }}"
|
||||
output:
|
||||
format: json
|
||||
validation_mode: json_schema
|
||||
schema_path: adapter.schema.json
|
||||
repair_attempts: 0
|
||||
`)},
|
||||
}, "."); err != nil {
|
||||
t.Fatalf("RegisterPromptFS() error = %v", err)
|
||||
}
|
||||
if err := registry.RegisterSchemaFS(fstest.MapFS{
|
||||
"adapter.schema.json": {Data: []byte(`{"type":"object","required":["ok"],"properties":{"ok":{"type":"boolean"}}}`)},
|
||||
}, "."); err != nil {
|
||||
t.Fatalf("RegisterSchemaFS() error = %v", err)
|
||||
}
|
||||
client, err := NewScriptoriumClient(ScriptoriumClientConfig{
|
||||
Assets: registry,
|
||||
EngineOptions: []scriptorium.Option{
|
||||
scriptorium.WithProfiles(
|
||||
scriptorium.OpenAICompatibleProfile(scriptorium.OpenAICompatibleProfileConfig{
|
||||
ID: "default-profile",
|
||||
Endpoint: "http://127.0.0.1:1/v1",
|
||||
Model: "default-model",
|
||||
}),
|
||||
scriptorium.OpenAICompatibleProfile(scriptorium.OpenAICompatibleProfileConfig{
|
||||
ID: "explicit-profile",
|
||||
Endpoint: "http://127.0.0.1:1/v1",
|
||||
Model: "explicit-model",
|
||||
}),
|
||||
),
|
||||
scriptorium.WithLLMClient(fake),
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("NewScriptoriumClient() error = %v, want nil", err)
|
||||
}
|
||||
return client
|
||||
}
|
||||
|
||||
type fakeScriptoriumLLM struct {
|
||||
content string
|
||||
err error
|
||||
block chan struct{}
|
||||
mu sync.Mutex
|
||||
last scriptorium.GenerateRequest
|
||||
calls int32
|
||||
inFlight int32
|
||||
maxInFlight int32
|
||||
}
|
||||
|
||||
func (f *fakeScriptoriumLLM) Generate(ctx context.Context, req scriptorium.GenerateRequest) (*scriptorium.GenerateResponse, error) {
|
||||
f.mu.Lock()
|
||||
f.last = req
|
||||
f.mu.Unlock()
|
||||
atomic.AddInt32(&f.calls, 1)
|
||||
current := atomic.AddInt32(&f.inFlight, 1)
|
||||
for {
|
||||
seen := atomic.LoadInt32(&f.maxInFlight)
|
||||
if current <= seen || atomic.CompareAndSwapInt32(&f.maxInFlight, seen, current) {
|
||||
break
|
||||
}
|
||||
}
|
||||
defer atomic.AddInt32(&f.inFlight, -1)
|
||||
if f.block != nil {
|
||||
select {
|
||||
case <-f.block:
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
}
|
||||
if f.err != nil {
|
||||
return nil, f.err
|
||||
}
|
||||
content := f.content
|
||||
if content == "" {
|
||||
content = `{"ok":true}`
|
||||
}
|
||||
if !json.Valid([]byte(content)) {
|
||||
return nil, errors.New("test fake must return JSON content")
|
||||
}
|
||||
return &scriptorium.GenerateResponse{
|
||||
Content: content,
|
||||
Usage: scriptorium.TokenUsage{
|
||||
PromptTokens: 11,
|
||||
CompletionTokens: 7,
|
||||
TotalTokens: 18,
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
|
||||
func (f *fakeScriptoriumLLM) lastRequest() scriptorium.GenerateRequest {
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
return f.last
|
||||
}
|
||||
Reference in New Issue
Block a user