Configure PromptKit transient retries

This commit is contained in:
2026-09-25 21:56:29 +00:00
parent d29a5cecd3
commit 10ac29d149
19 changed files with 214 additions and 37 deletions

View File

@@ -157,13 +157,15 @@ func buildProductionLLMClient(ctx context.Context, cfg config.Config, profileID
return nil, nil, fmt.Errorf("production asset registry must not be nil")
}
recorder := llm.NewLLMProfileRecorder()
transientRetries := cfg.PromptKit.TransientRetries
client, err := llm.NewPromptKitClient(llm.PromptKitClientConfig{
ProfileDir: cfg.PromptKit.ProfileDir,
ProfileFile: cfg.PromptKit.ProfileFile,
LocalBackend: mapPromptKitLocalBackend(cfg.PromptKit.LocalBackend),
Assets: assets,
Recorder: recorder,
ReasoningEffort: overrides.ReasoningEffort,
ProfileDir: cfg.PromptKit.ProfileDir,
ProfileFile: cfg.PromptKit.ProfileFile,
TransientRetries: &transientRetries,
LocalBackend: mapPromptKitLocalBackend(cfg.PromptKit.LocalBackend),
Assets: assets,
Recorder: recorder,
ReasoningEffort: overrides.ReasoningEffort,
})
if err != nil {
return nil, nil, fmt.Errorf("create PromptKit-backed LLM client: %w", err)

View File

@@ -557,11 +557,16 @@ func TestNormalizeOptionsSharesProductionProfileAssetsWithDefaultRuntime(t *test
func TestProductionLLMClientFactoryUsesConfiguredLocalBackend(t *testing.T) {
var providerCalls atomic.Int32
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
providerCalls.Add(1)
call := providerCalls.Add(1)
if r.URL.Path != "/v1/chat/completions" {
t.Errorf("provider path = %q, want /v1/chat/completions", r.URL.Path)
}
w.Header().Set("Content-Type", "application/json")
if call == 1 {
w.WriteHeader(http.StatusServiceUnavailable)
_, _ = w.Write([]byte(`{"error":{"message":"transient"}}`))
return
}
_, _ = w.Write([]byte(`{
"choices": [{"message": {"role": "assistant", "content": "{\"ok\":true}"}}],
"usage": {"prompt_tokens": 3, "completion_tokens": 4, "total_tokens": 7}
@@ -596,6 +601,7 @@ output:
cfg := config.Default()
cfg.PromptKit.ProfileFile = profilePath
cfg.PromptKit.TransientRetries = 1
cfg.PromptKit.LocalBackend = &config.PromptKitLocalBackendConfig{
Endpoint: server.URL + "/v1",
ConcurrencyLimit: 2,
@@ -624,8 +630,8 @@ output:
if err != nil {
t.Fatalf("CompleteStructured() error = %v, want nil", err)
}
if providerCalls.Load() != 1 {
t.Fatalf("provider calls = %d, want 1", providerCalls.Load())
if providerCalls.Load() != 2 {
t.Fatalf("provider calls = %d, want 2", providerCalls.Load())
}
provider, ok := client.(contracts.LLMProfileManifestProvider)
if !ok {

View File

@@ -14,9 +14,10 @@ type Config struct {
}
type PromptKitConfig struct {
ProfileDir string `json:"profile_dir,omitempty"`
ProfileFile string `json:"profile_file,omitempty"`
LocalBackend *PromptKitLocalBackendConfig `json:"local_backend,omitempty"`
ProfileDir string `json:"profile_dir,omitempty"`
ProfileFile string `json:"profile_file,omitempty"`
TransientRetries int `json:"transient_retries"`
LocalBackend *PromptKitLocalBackendConfig `json:"local_backend,omitempty"`
}
type PromptKitLocalBackendConfig struct {
@@ -29,7 +30,10 @@ type ConcurrencyConfig struct {
StageWorkers map[string]int `json:"stage_workers"`
}
const defaultLLMConcurrency = 16
const (
defaultLLMConcurrency = 16
defaultTransientRetries = 2
)
type OutputConfig struct {
Directory string `json:"directory"`
@@ -55,6 +59,7 @@ type DebugConfig struct {
func Default() Config {
return Config{
PromptKit: PromptKitConfig{TransientRetries: defaultTransientRetries},
Pipelines: map[string]pipeline.PipelineProfile{},
Concurrency: ConcurrencyConfig{
TotalLLM: defaultLLMConcurrency,

View File

@@ -16,9 +16,10 @@ type Document struct {
}
type PromptKitDocument struct {
ProfileDir *string `yaml:"profile_dir,omitempty"`
ProfileFile *string `yaml:"profile_file,omitempty"`
LocalBackend *PromptKitLocalBackendDocument `yaml:"local_backend,omitempty"`
ProfileDir *string `yaml:"profile_dir,omitempty"`
ProfileFile *string `yaml:"profile_file,omitempty"`
TransientRetries *int `yaml:"transient_retries,omitempty"`
LocalBackend *PromptKitLocalBackendDocument `yaml:"local_backend,omitempty"`
}
type PromptKitLocalBackendDocument struct {

View File

@@ -63,6 +63,9 @@ func documentDefinition(state *documentLoadState) configkit.Definition[Document]
func documentDefaults() Document {
return Document{
PromptKit: &PromptKitDocument{
TransientRetries: documentValue(defaultTransientRetries),
},
Pipelines: map[string]PipelineDocument{},
Concurrency: &ConcurrencyDocument{
TotalLLM: documentValue(defaultLLMConcurrency),

View File

@@ -94,6 +94,9 @@ func documentToConfig(document Document) (Config, error) {
if document.PromptKit.ProfileFile != nil {
config.PromptKit.ProfileFile = *document.PromptKit.ProfileFile
}
if document.PromptKit.TransientRetries != nil {
config.PromptKit.TransientRetries = *document.PromptKit.TransientRetries
}
if document.PromptKit.LocalBackend != nil {
backend := PromptKitLocalBackendConfig{}
if document.PromptKit.LocalBackend.Endpoint != nil {

View File

@@ -2,6 +2,7 @@ package config
import (
"bytes"
"strconv"
"strings"
"testing"
@@ -58,6 +59,44 @@ pipelines:
}
}
func TestLoadedConfigResolvesPromptKitTransientRetriesAndRetainsInspection(t *testing.T) {
t.Parallel()
for _, test := range []struct {
name string
document string
want int
}{
{name: "omitted", document: "version: 5\n", want: defaultTransientRetries},
{name: "disabled", document: "version: 5\npromptkit:\n transient_retries: 0\n", want: 0},
{name: "configured", document: "version: 5\npromptkit:\n transient_retries: 5\n", want: 5},
} {
t.Run(test.name, func(t *testing.T) {
loaded, err := loadConfigSource(configkit.Bytes("document.yml", []byte(test.document)), LoadOptions{})
if err != nil {
t.Fatal(err)
}
if got := loaded.Config.PromptKit.TransientRetries; got != test.want {
t.Fatalf("transient retries = %d, want %d", got, test.want)
}
var effective bytes.Buffer
if err := loaded.WriteEffective(&effective); err != nil {
t.Fatal(err)
}
if !strings.Contains(effective.String(), "transient_retries: "+strconv.Itoa(test.want)) {
t.Fatalf("effective configuration = %q", effective.String())
}
var explanation bytes.Buffer
if err := loaded.WriteExplanation(&explanation, "/promptkit/transient_retries"); err != nil {
t.Fatal(err)
}
if !strings.Contains(explanation.String(), strconv.Itoa(test.want)) {
t.Fatalf("explanation = %q", explanation.String())
}
})
}
}
func TestLoadedConfigOwnsRuntimeValuesAndKeepsFrozenInspection(t *testing.T) {
t.Parallel()

View File

@@ -167,6 +167,11 @@ func validateDocumentPromptKit(promptKit *PromptKitDocument, validation *configk
if err := validatePromptKitProfileSources(profileDir, profileFile); err != nil {
validation.Error(configkit.Keys("promptkit"), err.Error())
}
if promptKit.TransientRetries == nil {
validation.Error(configkit.Keys("promptkit", "transient_retries"), "must be an integer")
} else if err := validatePromptKitTransientRetries(*promptKit.TransientRetries); err != nil {
validation.Error(configkit.Keys("promptkit", "transient_retries"), err.Error())
}
if promptKit.LocalBackend == nil {
return
}
@@ -283,6 +288,13 @@ func validatePromptKitLocalBackendConcurrencyLimit(concurrencyLimit int) error {
return nil
}
func validatePromptKitTransientRetries(transientRetries int) error {
if transientRetries < 0 {
return fmt.Errorf("promptkit.transient_retries must not be negative")
}
return nil
}
func validateRequiredPath(name, value string) error {
if strings.TrimSpace(value) == "" {
return fmt.Errorf("%s must not be empty", name)

View File

@@ -223,6 +223,13 @@ cache:
assertDocumentDiagnostic(t, err, configkit.CodeValidateError, "/promptkit/local_backend/endpoint")
}
func TestDocumentLoadRejectsNegativePromptKitTransientRetries(t *testing.T) {
t.Parallel()
_, err := loadDocumentSource(configkit.Bytes("promptkit.yml", []byte("version: 5\npromptkit:\n transient_retries: -1\n")), LoadOptions{})
assertDocumentDiagnostic(t, err, configkit.CodeValidateError, "/promptkit/transient_retries")
}
func assertDocumentDiagnostic(t *testing.T, err error, code configkit.Code, path string) {
t.Helper()
var loadError *configkit.LoadError

View File

@@ -46,6 +46,9 @@ func validatePromptKit(cfg PromptKitConfig) error {
if err := validatePromptKitProfileSources(cfg.ProfileDir, cfg.ProfileFile); err != nil {
return err
}
if err := validatePromptKitTransientRetries(cfg.TransientRetries); err != nil {
return err
}
if cfg.LocalBackend == nil {
return nil
}

View File

@@ -89,6 +89,22 @@ func TestValidatePromptKitSourcesAreMutuallyExclusive(t *testing.T) {
assertValidationContains(t, cfg, "promptkit profile_dir and profile_file are mutually exclusive")
}
func TestValidatePromptKitTransientRetries(t *testing.T) {
if got := Default().PromptKit.TransientRetries; got != defaultTransientRetries {
t.Fatalf("default transient retries = %d, want %d", got, defaultTransientRetries)
}
for _, retries := range []int{0, 4} {
cfg := Default()
cfg.PromptKit.TransientRetries = retries
if err := cfg.Validate(); err != nil {
t.Fatalf("Validate() with transient retries %d error = %v", retries, err)
}
}
negative := Default()
negative.PromptKit.TransientRetries = -1
assertValidationContains(t, negative, "promptkit.transient_retries must not be negative")
}
func TestValidatePromptKitLocalBackendEndpoints(t *testing.T) {
tests := []struct {
name string

View File

@@ -26,15 +26,16 @@ type PromptKitLocalBackendConfig struct {
}
type PromptKitClientConfig struct {
ProfileDir string
ProfileFile string
LocalBackend *PromptKitLocalBackendConfig
Assets *AssetRegistry
Timeout time.Duration
HTTPClient *http.Client
EngineOptions []promptkit.Option
Recorder *LLMProfileRecorder
ReasoningEffort *string
ProfileDir string
ProfileFile string
TransientRetries *int
LocalBackend *PromptKitLocalBackendConfig
Assets *AssetRegistry
Timeout time.Duration
HTTPClient *http.Client
EngineOptions []promptkit.Option
Recorder *LLMProfileRecorder
ReasoningEffort *string
}
type PromptKitClient struct {
@@ -79,9 +80,10 @@ func NewPromptKitClient(cfg PromptKitClientConfig) (*PromptKitClient, error) {
options = append(options, cfg.EngineOptions...)
engine, err := promptkit.NewEngine(promptkit.Config{
ProfileDir: profileSource.ProfileDir,
Timeout: cfg.Timeout,
HTTPClient: cfg.HTTPClient,
ProfileDir: profileSource.ProfileDir,
Timeout: cfg.Timeout,
HTTPClient: cfg.HTTPClient,
TransientRetries: cfg.TransientRetries,
}, options...)
if err != nil {
return nil, fmt.Errorf("create PromptKit engine: %w", err)

View File

@@ -971,6 +971,70 @@ model: local-model
}
}
func TestPromptKitClientAppliesConfiguredBuiltInTransientRetries(t *testing.T) {
for _, test := range []struct {
name string
retries int
statuses []int
wantCalls int32
wantErr bool
}{
{name: "eligible failure then success", retries: 1, statuses: []int{http.StatusServiceUnavailable, http.StatusOK}, wantCalls: 2},
{name: "disabled", retries: 0, statuses: []int{http.StatusServiceUnavailable}, wantCalls: 1, wantErr: true},
{name: "exhausted", retries: 1, statuses: []int{http.StatusServiceUnavailable, http.StatusServiceUnavailable}, wantCalls: 2, wantErr: true},
{name: "noneligible client error", retries: 2, statuses: []int{http.StatusBadRequest}, wantCalls: 1, wantErr: true},
} {
t.Run(test.name, func(t *testing.T) {
var calls atomic.Int32
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
index := calls.Add(1) - 1
if int(index) >= len(test.statuses) {
t.Errorf("unexpected HTTP attempt %d", index+1)
w.WriteHeader(http.StatusInternalServerError)
return
}
status := test.statuses[index]
w.Header().Set("Content-Type", "application/json")
if status != http.StatusOK {
w.WriteHeader(status)
_, _ = w.Write([]byte(`{"error":{"message":"provider failure"}}`))
return
}
_, _ = w.Write([]byte(`{"choices":[{"message":{"content":"{\"ok\":true}"}}]}`))
}))
defer server.Close()
profilePath := filepath.Join(t.TempDir(), "profile.yml")
if err := os.WriteFile(profilePath, []byte("id: retry-profile\nendpoint: "+server.URL+"/v1\nmodel: retry-model\n"), 0o600); err != nil {
t.Fatal(err)
}
client, err := NewPromptKitClient(PromptKitClientConfig{
Assets: newTestPromptKitAssets(t),
ProfileFile: profilePath,
TransientRetries: &test.retries,
})
if err != nil {
t.Fatal(err)
}
var out map[string]any
_, err = client.CompleteStructured(context.Background(), contracts.StructuredCompletionRequest{
PromptID: "adapter.test",
ProfileID: "retry-profile",
SessionID: "retry-test",
Inputs: contracts.LLMInputSet{
"transcript": contracts.NewLLMInputMaterial("transcript", "application/json", []byte(`{"source":true}`), "", ""),
},
}, &out)
if (err != nil) != test.wantErr {
t.Fatalf("CompleteStructured() error = %v, want error=%t", err, test.wantErr)
}
if got := calls.Load(); got != test.wantCalls {
t.Fatalf("HTTP calls = %d, want %d", got, test.wantCalls)
}
})
}
}
func TestLLMProfileRecorderDistinguishesEffectiveTargets(t *testing.T) {
recorder := NewLLMProfileRecorder()
for _, profile := range []artifacts.LLMProfileManifest{