Make PromptKit profile handling safer and more consistent

This commit is contained in:
2026-08-03 18:35:40 +00:00
parent 12ac25bd63
commit 39388e96d4
19 changed files with 220 additions and 111 deletions

View File

@@ -143,17 +143,6 @@ func isEmptyRegistries(registries pipeline.Registries) bool {
registries.Outputs == nil
}
func productionLLMClientFactory(ctx context.Context, cfg config.Config, profileID string, overrides LLMRuntimeOverrides) (contracts.StructuredLLMClient, []artifacts.LLMProfileManifest, error) {
if err := ctx.Err(); err != nil {
return nil, nil, err
}
assets, err := productionPromptAssets()
if err != nil {
return nil, nil, err
}
return buildProductionLLMClient(ctx, cfg, profileID, overrides, assets)
}
func productionLLMClientFactoryWithAssets(assets *llm.AssetRegistry) LLMClientFactory {
return func(ctx context.Context, cfg config.Config, profileID string, overrides LLMRuntimeOverrides) (contracts.StructuredLLMClient, []artifacts.LLMProfileManifest, error) {
return buildProductionLLMClient(ctx, cfg, profileID, overrides, assets)

View File

@@ -474,39 +474,67 @@ func setNormalizeSpellCatalogSource(t *testing.T, resolved *pipeline.ResolvedPip
resolved.Steps[0].ArtifactLanes[0].NormalizeReferences.Bindings = bindings
}
func TestProductionLLMClientFactoriesBuildOfflineRuntime(t *testing.T) {
func TestProductionLLMClientFactoryBuildsOfflineRuntime(t *testing.T) {
components := productionTestComponents(t)
factories := []struct {
name string
factory LLMClientFactory
}{
{name: "default production assets", factory: productionLLMClientFactory},
{name: "provided production assets", factory: productionLLMClientFactoryWithAssets(components.assets)},
client, manifests, err := productionLLMClientFactoryWithAssets(components.assets)(context.Background(), config.Default(), "test-profile", LLMRuntimeOverrides{})
if err != nil {
t.Fatalf("build production LLM runtime: %v", err)
}
for _, tt := range factories {
t.Run(tt.name, func(t *testing.T) {
client, manifests, err := tt.factory(context.Background(), config.Default(), "test-profile", LLMRuntimeOverrides{})
if err != nil {
t.Fatalf("build production LLM runtime: %v", err)
}
if client == nil {
t.Fatal("production LLM runtime returned a nil client")
}
if len(manifests) != 0 {
t.Fatalf("eager profile manifests = %#v, want none", manifests)
}
fingerprintProvider, ok := client.(llm.CheckpointFingerprintProvider)
if !ok {
t.Fatalf("production LLM client %T does not provide one profile-source checkpoint fingerprint", client)
}
fingerprints, err := fingerprintProvider.LLMCheckpointFingerprints()
if err != nil || len(fingerprints) != 1 {
t.Fatalf("production LLM checkpoint fingerprints = %#v, error = %v, want one profile-source identity", fingerprints, err)
}
if _, ok := client.(contracts.LLMProfileManifestProvider); !ok {
t.Fatalf("production LLM client %T does not provide profile manifests", client)
}
})
if client == nil {
t.Fatal("production LLM runtime returned a nil client")
}
if len(manifests) != 0 {
t.Fatalf("eager profile manifests = %#v, want none", manifests)
}
fingerprintProvider, ok := client.(llm.CheckpointFingerprintProvider)
if !ok {
t.Fatalf("production LLM client %T does not provide one profile-source checkpoint fingerprint", client)
}
fingerprints, err := fingerprintProvider.LLMCheckpointFingerprints()
if err != nil || len(fingerprints) != 1 {
t.Fatalf("production LLM checkpoint fingerprints = %#v, error = %v, want one profile-source identity", fingerprints, err)
}
if _, ok := client.(contracts.LLMProfileManifestProvider); !ok {
t.Fatalf("production LLM client %T does not provide profile manifests", client)
}
}
func TestNormalizeOptionsSharesProductionProfileAssetsWithDefaultRuntime(t *testing.T) {
opts, err := normalizeOptions(Options{
Catalog: pipeline.ModuleCatalog{Inputs: pipeline.NewInputAdapterRegistry()},
})
if err != nil {
t.Fatal(err)
}
if opts.promptKitAssets == nil || opts.LLMClientFactory == nil {
t.Fatalf("normalized options = %#v, want shared profile assets and default runtime factory", opts)
}
if err := validateExplicitPromptKitProfiles(context.Background(), config.Default(), []string{"dnd-extraction"}, opts.promptKitAssets); err != nil {
t.Fatalf("inspect application fallback profile: %v", err)
}
client, _, err := opts.LLMClientFactory(context.Background(), config.Default(), "dnd-extraction", LLMRuntimeOverrides{})
if err != nil {
t.Fatalf("build default runtime: %v", err)
}
fingerprintProvider, ok := client.(llm.CheckpointFingerprintProvider)
if !ok {
t.Fatalf("default runtime client %T does not provide checkpoint fingerprints", client)
}
runtimeFingerprints, err := fingerprintProvider.LLMCheckpointFingerprints()
if err != nil {
t.Fatal(err)
}
directClient, err := llm.NewPromptKitClient(llm.PromptKitClientConfig{Assets: opts.promptKitAssets})
if err != nil {
t.Fatal(err)
}
inspectionFingerprints, err := directClient.LLMCheckpointFingerprints()
if err != nil {
t.Fatal(err)
}
if !reflect.DeepEqual(runtimeFingerprints, inspectionFingerprints) {
t.Fatalf("runtime profile fingerprints = %#v, inspection profile fingerprints = %#v", runtimeFingerprints, inspectionFingerprints)
}
}
@@ -597,7 +625,8 @@ func TestProductionLLMClientFactoriesRejectInvalidConstruction(t *testing.T) {
t.Run("canceled context", func(t *testing.T) {
ctx, cancel := context.WithCancel(context.Background())
cancel()
client, manifests, err := productionLLMClientFactory(ctx, config.Default(), "test-profile", LLMRuntimeOverrides{})
components := productionTestComponents(t)
client, manifests, err := productionLLMClientFactoryWithAssets(components.assets)(ctx, config.Default(), "test-profile", LLMRuntimeOverrides{})
if !errors.Is(err, context.Canceled) || client != nil || len(manifests) != 0 {
t.Fatalf("client=%T manifests=%#v error=%v, want canceled construction", client, manifests, err)
}

View File

@@ -47,6 +47,7 @@ api_key_env: NOTARIUS_PROMPTKIT_PROFILE_INSPECTION_TEST_KEY
localBackend bool
canceled bool
wantErr []string
rejectErr []string
}{
{
name: "configured local backend",
@@ -59,7 +60,7 @@ api_key_env: NOTARIUS_PROMPTKIT_PROFILE_INSPECTION_TEST_KEY
name: "missing local backend registration",
profilePath: writeProfile(t, "local-profile", localProfile),
profileID: "local-profile",
wantErr: []string{`inspect PromptKit profile "local-profile"`, `backend "local"`},
wantErr: []string{`PromptKit profile "local-profile" is invalid or unreadable`},
},
{
name: "absent profile",
@@ -72,13 +73,14 @@ api_key_env: NOTARIUS_PROMPTKIT_PROFILE_INSPECTION_TEST_KEY
name: "malformed profile",
profilePath: writeProfile(t, "malformed-profile", "id: malformed-profile\nbackend: [\n"),
profileID: "malformed-profile",
wantErr: []string{`inspect PromptKit profile "malformed-profile"`},
wantErr: []string{`PromptKit profile "malformed-profile" is invalid or unreadable`},
rejectErr: []string{"malformed-profile.yaml", "backend: ["},
},
{
name: "invalid profile source",
profilePath: filepath.Join(t.TempDir(), "missing-profile.yaml"),
profileID: "missing-profile",
wantErr: []string{"load PromptKit profiles", "failed to access source file"},
wantErr: []string{"load PromptKit profiles", "profile configuration is invalid or unreadable"},
},
{
name: "credential environment intentionally unset",
@@ -132,6 +134,11 @@ api_key_env: NOTARIUS_PROMPTKIT_PROFILE_INSPECTION_TEST_KEY
t.Fatalf("validation error = %q, want %q", err, want)
}
}
for _, rejected := range append(tt.rejectErr, tt.profilePath) {
if rejected != "" && strings.Contains(err.Error(), rejected) {
t.Fatalf("validation error = %q, must not expose %q", err, rejected)
}
}
})
}
if providerCalls.Load() != 0 {

View File

@@ -123,12 +123,16 @@ func normalizeOptions(opts Options) (Options, error) {
opts.Registries = components.registries
opts.Catalog = catalogFromRegistries(components.registries)
opts.promptKitAssets = components.assets
if opts.LLMClientFactory == nil {
opts.LLMClientFactory = productionLLMClientFactoryWithAssets(components.assets)
}
}
if opts.LLMClientFactory == nil {
opts.LLMClientFactory = productionLLMClientFactory
if opts.promptKitAssets == nil {
assets, err := productionPromptAssets()
if err != nil {
return Options{}, err
}
opts.promptKitAssets = assets
}
opts.LLMClientFactory = productionLLMClientFactoryWithAssets(opts.promptKitAssets)
}
return opts, nil
}