notarius/internal/framework/llm/schema_registry.go

package llm

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io/fs"
	"sort"
	"strings"

	rootassets "gitea.maximumdirect.net/eric/notarius/assets"
)

var schemaAssets = mustSubFS(rootassets.FS(), "generic")

// ResponseSchemaKey identifies one structured response schema.
type ResponseSchemaKey string

const (
	TestArtifactSchemaKey          ResponseSchemaKey = "test_artifact"
	TestValidatorDecisionSchemaKey ResponseSchemaKey = "test_validator_decision"

	schemaVersionV1 = "v1"
)

// ResponseSchemaDefinition identifies a caller-owned structured response schema asset.
type ResponseSchemaDefinition struct {
	Key       ResponseSchemaKey
	ID        string
	Version   string
	Name      string
	AssetPath string
}

// 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(schemaAssets, ResponseSchemaDefinition{
		Key:       TestArtifactSchemaKey,
		ID:        "notarius.test_artifact",
		Version:   schemaVersionV1,
		Name:      "notarius_test_artifact_v1",
		AssetPath: "schemas/test_artifact.v1.json",
	}),
	TestValidatorDecisionSchemaKey: mustLoadResponseSchema(schemaAssets, ResponseSchemaDefinition{
		Key:       TestValidatorDecisionSchemaKey,
		ID:        "notarius.test_validator_decision",
		Version:   schemaVersionV1,
		Name:      "notarius_test_validator_decision_v1",
		AssetPath: "schemas/test_validator_decision.v1.json",
	}),
}

func mustSubFS(fsys fs.FS, dir string) fs.FS {
	sub, err := fs.Sub(fsys, dir)
	if err != nil {
		panic(fmt.Errorf("scope embedded assets to %s: %w", dir, err))
	}
	return sub
}

// 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,
	}
}

// LoadResponseSchema loads a structured response schema from a caller-owned filesystem.
func LoadResponseSchema(fsys fs.FS, def ResponseSchemaDefinition) (ResponseSchema, error) {
	key := ResponseSchemaKey(strings.TrimSpace(string(def.Key)))
	id := strings.TrimSpace(def.ID)
	version := strings.TrimSpace(def.Version)
	name := strings.TrimSpace(def.Name)
	path := strings.TrimSpace(def.AssetPath)
	if key == "" {
		return ResponseSchema{}, fmt.Errorf("response schema key must not be empty")
	}
	if id == "" {
		return ResponseSchema{}, fmt.Errorf("response schema id must not be empty")
	}
	if version == "" {
		return ResponseSchema{}, fmt.Errorf("response schema version must not be empty")
	}
	if name == "" {
		return ResponseSchema{}, fmt.Errorf("response schema name must not be empty")
	}
	if path == "" {
		return ResponseSchema{}, fmt.Errorf("response schema asset path must not be empty")
	}

	rawSchema, err := fs.ReadFile(fsys, path)
	if err != nil {
		return ResponseSchema{}, fmt.Errorf("read response schema %s: %w", path, err)
	}
	if !json.Valid(rawSchema) {
		return ResponseSchema{}, fmt.Errorf("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[:]),
	}, nil
}

func mustLoadResponseSchema(fsys fs.FS, def ResponseSchemaDefinition) ResponseSchema {
	schema, err := LoadResponseSchema(fsys, def)
	if err != nil {
		panic(err)
	}
	return schema
}

func cloneResponseSchema(in ResponseSchema) ResponseSchema {
	out := in
	if in.JSONSchema != nil {
		out.JSONSchema = append(json.RawMessage(nil), in.JSONSchema...)
	}
	return out
}