From 54b9851f65e5ee82a5cbef661eba1caeeb1627ab Mon Sep 17 00:00:00 2001 From: Eric Rakestraw Date: Fri, 3 Jul 2026 16:40:56 +0000 Subject: [PATCH] Add embedded response schema registry --- .../llm/assets/schemas/test_artifact.v1.json | 31 ++++ .../schemas/test_validator_decision.v1.json | 21 +++ internal/framework/llm/schema_registry.go | 150 ++++++++++++++++++ .../framework/llm/schema_registry_test.go | 113 +++++++++++++ 4 files changed, 315 insertions(+) create mode 100644 internal/framework/llm/assets/schemas/test_artifact.v1.json create mode 100644 internal/framework/llm/assets/schemas/test_validator_decision.v1.json create mode 100644 internal/framework/llm/schema_registry.go create mode 100644 internal/framework/llm/schema_registry_test.go diff --git a/internal/framework/llm/assets/schemas/test_artifact.v1.json b/internal/framework/llm/assets/schemas/test_artifact.v1.json new file mode 100644 index 0000000..60190bb --- /dev/null +++ b/internal/framework/llm/assets/schemas/test_artifact.v1.json @@ -0,0 +1,31 @@ +{ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "$id": "notarius.test_artifact", + "type": "object", + "additionalProperties": false, + "required": ["value"], + "properties": { + "value": { + "type": "string", + "minLength": 1 + }, + "source_refs": { + "type": "array", + "items": { + "type": "object", + "additionalProperties": false, + "required": ["source_id", "unit_id"], + "properties": { + "source_id": { + "type": "string", + "minLength": 1 + }, + "unit_id": { + "type": "string", + "minLength": 1 + } + } + } + } + } +} diff --git a/internal/framework/llm/assets/schemas/test_validator_decision.v1.json b/internal/framework/llm/assets/schemas/test_validator_decision.v1.json new file mode 100644 index 0000000..8c706b3 --- /dev/null +++ b/internal/framework/llm/assets/schemas/test_validator_decision.v1.json @@ -0,0 +1,21 @@ +{ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "$id": "notarius.test_validator_decision", + "type": "object", + "additionalProperties": false, + "required": ["accepted", "reason"], + "properties": { + "accepted": { + "type": "boolean" + }, + "reason": { + "type": "string", + "minLength": 1 + }, + "confidence": { + "type": "number", + "minimum": 0, + "maximum": 1 + } + } +} diff --git a/internal/framework/llm/schema_registry.go b/internal/framework/llm/schema_registry.go new file mode 100644 index 0000000..678bd00 --- /dev/null +++ b/internal/framework/llm/schema_registry.go @@ -0,0 +1,150 @@ +package llm + +import ( + "crypto/sha256" + "embed" + "encoding/hex" + "encoding/json" + "fmt" + "sort" + "strings" +) + +//go:embed assets/schemas/*.json +var schemaAssets embed.FS + +// ResponseSchemaKey identifies one structured response schema. +type ResponseSchemaKey string + +const ( + TestArtifactSchemaKey ResponseSchemaKey = "test_artifact" + TestValidatorDecisionSchemaKey ResponseSchemaKey = "test_validator_decision" + + schemaVersionV1 = "v1" +) + +// ResponseSchema describes one registered structured response schema. +type ResponseSchema struct { + Key ResponseSchemaKey `json:"key"` + ID string `json:"id"` + Version string `json:"version"` + Name string `json:"name"` + JSONSchema json.RawMessage `json:"json_schema"` + SHA256 string `json:"sha256"` +} + +var responseSchemaRegistry = map[ResponseSchemaKey]ResponseSchema{ + TestArtifactSchemaKey: mustLoadResponseSchema( + TestArtifactSchemaKey, + "notarius.test_artifact", + schemaVersionV1, + "notarius_test_artifact_v1", + "assets/schemas/test_artifact.v1.json", + ), + TestValidatorDecisionSchemaKey: mustLoadResponseSchema( + TestValidatorDecisionSchemaKey, + "notarius.test_validator_decision", + schemaVersionV1, + "notarius_test_validator_decision_v1", + "assets/schemas/test_validator_decision.v1.json", + ), +} + +// RegisteredResponseSchemas returns all registered response schemas sorted by key. +func RegisteredResponseSchemas() []ResponseSchema { + keys := make([]string, 0, len(responseSchemaRegistry)) + for key := range responseSchemaRegistry { + keys = append(keys, string(key)) + } + sort.Strings(keys) + + out := make([]ResponseSchema, 0, len(keys)) + for _, key := range keys { + out = append(out, cloneResponseSchema(responseSchemaRegistry[ResponseSchemaKey(key)])) + } + return out +} + +// LookupResponseSchema returns a copy of the schema for key. +func LookupResponseSchema(key ResponseSchemaKey) (ResponseSchema, bool) { + schema, ok := responseSchemaRegistry[key] + if !ok { + return ResponseSchema{}, false + } + return cloneResponseSchema(schema), true +} + +// MustLookupResponseSchema returns a copy of the schema for key and panics when missing. +func MustLookupResponseSchema(key ResponseSchemaKey) ResponseSchema { + schema, ok := LookupResponseSchema(key) + if !ok { + panic(fmt.Sprintf("unknown structured response schema key %q", key)) + } + return schema +} + +// DiagnosticsMap returns schema metadata without raw schema content. +func (s ResponseSchema) DiagnosticsMap() map[string]any { + return map[string]any{ + "key": s.Key, + "id": s.ID, + "version": s.Version, + "name": s.Name, + "sha256": s.SHA256, + } +} + +func mustLoadResponseSchema( + key ResponseSchemaKey, + id string, + version string, + name string, + path string, +) ResponseSchema { + key = ResponseSchemaKey(strings.TrimSpace(string(key))) + id = strings.TrimSpace(id) + version = strings.TrimSpace(version) + name = strings.TrimSpace(name) + path = strings.TrimSpace(path) + if key == "" { + panic("response schema key must not be empty") + } + if id == "" { + panic("response schema id must not be empty") + } + if version == "" { + panic("response schema version must not be empty") + } + if name == "" { + panic("response schema name must not be empty") + } + if path == "" { + panic("response schema asset path must not be empty") + } + + rawSchema, err := schemaAssets.ReadFile(path) + if err != nil { + panic(fmt.Sprintf("read response schema %s: %v", path, err)) + } + if !json.Valid(rawSchema) { + panic(fmt.Sprintf("response schema %s is not valid JSON", path)) + } + + hash := sha256.Sum256(rawSchema) + return ResponseSchema{ + Key: key, + ID: id, + Version: version, + Name: name, + JSONSchema: append(json.RawMessage(nil), rawSchema...), + SHA256: "sha256:" + hex.EncodeToString(hash[:]), + } +} + +func cloneResponseSchema(in ResponseSchema) ResponseSchema { + out := in + if in.JSONSchema != nil { + out.JSONSchema = append(json.RawMessage(nil), in.JSONSchema...) + } + return out +} diff --git a/internal/framework/llm/schema_registry_test.go b/internal/framework/llm/schema_registry_test.go new file mode 100644 index 0000000..662660c --- /dev/null +++ b/internal/framework/llm/schema_registry_test.go @@ -0,0 +1,113 @@ +package llm + +import ( + "encoding/json" + "sort" + "strings" + "testing" +) + +func TestLookupResponseSchemaSucceedsForTestSchemas(t *testing.T) { + tests := []ResponseSchemaKey{ + TestArtifactSchemaKey, + TestValidatorDecisionSchemaKey, + } + + for _, key := range tests { + t.Run(string(key), func(t *testing.T) { + schema, ok := LookupResponseSchema(key) + if !ok { + t.Fatalf("expected schema for key %q", key) + } + if schema.Key != key { + t.Fatalf("unexpected key: got %q want %q", schema.Key, key) + } + if schema.ID == "" || schema.Version == "" || schema.Name == "" { + t.Fatalf("expected schema metadata, got %+v", schema) + } + if !strings.HasPrefix(schema.SHA256, "sha256:") { + t.Fatalf("expected prefixed hash, got %q", schema.SHA256) + } + }) + } +} + +func TestLookupResponseSchemaUnknownReturnsFalse(t *testing.T) { + if schema, ok := LookupResponseSchema("unknown"); ok { + t.Fatalf("expected unknown schema lookup to fail, got %+v", schema) + } +} + +func TestMustLookupResponseSchemaPanicsForUnknownKey(t *testing.T) { + defer func() { + if recover() == nil { + t.Fatalf("expected panic") + } + }() + + _ = MustLookupResponseSchema("unknown") +} + +func TestRegisteredResponseSchemasSortedByKey(t *testing.T) { + schemas := RegisteredResponseSchemas() + if len(schemas) != 2 { + t.Fatalf("expected two schemas, got %d", len(schemas)) + } + + keys := make([]string, len(schemas)) + for i, schema := range schemas { + keys[i] = string(schema.Key) + } + if !sort.StringsAreSorted(keys) { + t.Fatalf("expected sorted keys, got %v", keys) + } +} + +func TestResponseSchemaContentIsValidJSON(t *testing.T) { + for _, schema := range RegisteredResponseSchemas() { + if !json.Valid(schema.JSONSchema) { + t.Fatalf("schema %q has invalid JSON: %s", schema.Key, schema.JSONSchema) + } + } +} + +func TestResponseSchemaJSONIsMutationSafe(t *testing.T) { + first := MustLookupResponseSchema(TestArtifactSchemaKey) + first.JSONSchema[0] = '[' + + second := MustLookupResponseSchema(TestArtifactSchemaKey) + if !json.Valid(second.JSONSchema) { + t.Fatalf("schema JSON was mutated: %s", second.JSONSchema) + } + if len(second.JSONSchema) > 0 && second.JSONSchema[0] == '[' { + t.Fatalf("schema JSON did not use defensive copy") + } + + registered := RegisteredResponseSchemas() + for i := range registered { + if registered[i].Key == TestArtifactSchemaKey { + registered[i].JSONSchema[0] = '[' + } + } + again := MustLookupResponseSchema(TestArtifactSchemaKey) + if !json.Valid(again.JSONSchema) || again.JSONSchema[0] == '[' { + t.Fatalf("registered schema JSON did not use defensive copy") + } +} + +func TestResponseSchemaDiagnosticsMapOmitsRawSchemaContent(t *testing.T) { + schema := MustLookupResponseSchema(TestValidatorDecisionSchemaKey) + diagnostics := schema.DiagnosticsMap() + + for _, key := range []string{"id", "version", "name", "sha256"} { + if diagnostics[key] == "" { + t.Fatalf("expected diagnostics key %q, got %#v", key, diagnostics) + } + } + if _, ok := diagnostics["json_schema"]; ok { + t.Fatalf("diagnostics should omit raw schema content: %#v", diagnostics) + } + if _, ok := diagnostics["JSONSchema"]; ok { + t.Fatalf("diagnostics should omit raw schema content: %#v", diagnostics) + } +}