notarius/internal/modules/dnd/extract/scenedescriptions/extractor.go

package scenedescriptions

import (
	"context"
	"fmt"
	"strings"

	"gitea.maximumdirect.net/eric/notarius/internal/core/source"
	"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
	"gitea.maximumdirect.net/eric/notarius/internal/framework/pipeline"
	"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd"
	"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/shared"
)

const Key = "dnd/scene-descriptions"

var requiredCapabilities = []string{
	"chunks",
	"source.transcript",
}

var providedCapabilities = []string{
	"dnd.scene_descriptions",
}

var referenceSlotDescriptions = shared.ReferenceSlotDescriptions{
	Glossary: "Optional campaign glossary reference material used only to disambiguate scene descriptions.",
	Party:    "Optional party roster reference material used only to disambiguate scene descriptions.",
	Players:  "Optional player list reference material used only to disambiguate scene descriptions.",
}

func referenceSlots() []contracts.ReferenceSlot {
	all := shared.ReferenceSlots(referenceSlotDescriptions)
	slots := make([]contracts.ReferenceSlot, 0, 3)
	for _, slot := range all {
		if slot.Name != "roster" {
			slots = append(slots, slot)
		}
	}
	return slots
}

var _ contracts.Extractor[dnd.SceneDescriptionList] = (*Extractor)(nil)
var _ contracts.ManifestMetadataProvider = (*Extractor)(nil)
var _ pipeline.CheckpointFingerprintProvider = (*Extractor)(nil)

type Options struct{}

type Extractor struct {
	llm               contracts.StructuredLLMClient
	promptSHA         string
	responseSchemaSHA string
}

func New(llmClient contracts.StructuredLLMClient, _ Options, references ...contracts.ReferenceSet) (*Extractor, error) {
	if llmClient == nil {
		return nil, extractorErrorf("LLM client must not be nil")
	}
	if len(references) > 1 {
		return nil, extractorErrorf("at most one reference set may be supplied")
	}
	promptSHA, err := promptAssetMetadata()
	if err != nil {
		return nil, extractorErrorf("load prompt metadata: %w", err)
	}
	responseSchema, err := loadResponseSchema()
	if err != nil {
		return nil, extractorErrorf("load response schema: %w", err)
	}
	return &Extractor{
		llm:               llmClient,
		promptSHA:         promptSHA,
		responseSchemaSHA: responseSchema.SHA256,
	}, nil
}

func (e *Extractor) Key() string { return Key }

func (e *Extractor) ReferenceSlots() []contracts.ReferenceSlot { return referenceSlots() }

func (e *Extractor) ManifestMetadata() map[string]any {
	if e == nil {
		return nil
	}
	return map[string]any{
		"prompt_id":               PromptID,
		"prompt_version":          SchemaVersion,
		"prompt_sha256":           e.promptSHA,
		"response_schema_key":     string(ResponseSchemaKey),
		"response_schema_id":      ResponseSchemaID,
		"response_schema_name":    ResponseSchemaName,
		"response_schema_version": SchemaVersion,
		"response_schema_sha256":  e.responseSchemaSHA,
		"mapping_policy":          mappingPolicy,
	}
}

func (e *Extractor) CheckpointFingerprints() []pipeline.CheckpointFingerprint {
	if e == nil {
		return nil
	}
	return []pipeline.CheckpointFingerprint{
		{Name: "prompt", Value: e.promptSHA},
		{Name: "response_schema", Value: e.responseSchemaSHA},
		{Name: "mapping_policy", Value: mappingPolicy},
	}
}

func (e *Extractor) Extract(ctx context.Context, req contracts.TypedExtractionRequest) (contracts.TypedExtractionResult[dnd.SceneDescriptionList], error) {
	if e == nil {
		return contracts.TypedExtractionResult[dnd.SceneDescriptionList]{}, extractorErrorf("extractor must not be nil")
	}
	if e.llm == nil {
		return contracts.TypedExtractionResult[dnd.SceneDescriptionList]{}, extractorErrorf("LLM client must not be nil")
	}
	sourceInput, err := shared.PrepareChunkExtraction(ctx, req)
	if err != nil {
		return contracts.TypedExtractionResult[dnd.SceneDescriptionList]{}, extractorErrorf("%w", err)
	}

	var response extractionResponse
	if _, err := e.llm.CompleteStructured(ctx, contracts.StructuredCompletionRequest{
		StageName:     Key,
		PromptID:      PromptID,
		PromptVersion: SchemaVersion,
		ProfileID:     req.LLMProfile,
		SessionID:     req.SessionID,
		Inputs:        shared.PromptInputs(sourceInput, req.References),
	}, &response); err != nil {
		return contracts.TypedExtractionResult[dnd.SceneDescriptionList]{}, extractorErrorf("complete structured output: %w", err)
	}
	return contracts.TypedExtractionResult[dnd.SceneDescriptionList]{Value: mapResponse(response, req.Chunk)}, nil
}

func mapResponse(response extractionResponse, chunk *source.Chunk) dnd.SceneDescriptionList {
	return dnd.SceneDescriptionList{Scenes: []dnd.SceneDescription{{
		ID:        chunk.ID,
		SourceRef: chunk.Ref,
		Kind:      response.Kind,
		Title:     strings.TrimSpace(response.Title),
		Summary:   strings.TrimSpace(response.Summary),
	}}}
}

func ModuleSpec() pipeline.ModuleSpec {
	return pipeline.ModuleSpec{
		Key:            Key,
		Stage:          pipeline.StageExtract,
		ExecutionClass: contracts.ExecutionClassLLMBacked,
		Requires:       append([]string(nil), requiredCapabilities...),
		Provides:       append([]string(nil), providedCapabilities...),
		ArtifactKind:   dnd.SceneDescriptionListKind,
		ReferenceSlots: referenceSlots(),
	}
}

func Register(registry *pipeline.ExtractorRegistry) error {
	return pipeline.RegisterExtractorBuilder(registry, ModuleSpec(), validateOptions, func(request pipeline.BuildRequest) (contracts.Extractor[dnd.SceneDescriptionList], error) {
		options, err := DecodeOptions(request.Options)
		if err != nil {
			return nil, err
		}
		return New(request.Dependencies.LLM, options, request.References)
	})
}

func validateOptions(options map[string]any) error {
	_, err := DecodeOptions(options)
	return err
}

func DecodeOptions(options map[string]any) (Options, error) {
	if err := pipeline.RejectUnknownOptions(options); err != nil {
		return Options{}, extractorErrorf("%w", err)
	}
	return Options{}, nil
}

func extractorErrorf(format string, args ...any) error {
	return fmt.Errorf("dnd scene descriptions extractor: "+format, args...)
}