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

package spells

import (
	"context"
	"encoding/json"
	"fmt"
	"strings"

	"gitea.maximumdirect.net/eric/notarius/internal/core/artifacts"
	"gitea.maximumdirect.net/eric/notarius/internal/core/source"
	"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
	"gitea.maximumdirect.net/eric/notarius/internal/framework/pipeline"
)

const Key = "dnd/spells"
const ArtifactType = "dnd.spell_cast"
const SchemaVersion = "v1"

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

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

var referenceSlots = []contracts.ReferenceSlot{
	{
		Name:               "glossary",
		Description:        "Optional campaign glossary reference material used only for disambiguation.",
		AcceptedMediaTypes: []string{"text/plain; charset=utf-8"},
	},
	{
		Name:               "roster",
		Description:        "Optional campaign roster or player-character reference material used only for disambiguation.",
		AcceptedMediaTypes: []string{"text/plain; charset=utf-8"},
	},
}

var _ contracts.Extractor = (*Extractor)(nil)

type Extractor struct{}

func New() *Extractor {
	return &Extractor{}
}

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

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

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

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

func (e *Extractor) ManifestMetadata() map[string]any {
	promptMetadata := spellsPromptBundle.Metadata()
	metadata := map[string]any{
		"prompt_id":            PromptID,
		"prompt_version":       promptMetadata.PromptVersion,
		"prompt_sha256":        promptMetadata.SHA256,
		"response_schema_key":  string(ResponseSchemaKey),
		"response_schema_id":   ResponseSchemaID,
		"response_schema_name": ResponseSchemaName,
	}
	if schema, err := loadResponseSchema(); err == nil {
		metadata["response_schema_version"] = schema.Version
		metadata["response_schema_sha256"] = schema.SHA256
	}
	return metadata
}

func (e *Extractor) Validators() []contracts.Validator {
	return []contracts.Validator{
		ShapeValidator{},
		SourceRefValidator{},
	}
}

func (e *Extractor) Extract(ctx context.Context, req contracts.ExtractionRequest) (contracts.ExtractionResult, error) {
	if e == nil {
		return contracts.ExtractionResult{}, extractorErrorf("extractor must not be nil")
	}
	if ctx == nil {
		return contracts.ExtractionResult{}, extractorErrorf("context must not be nil")
	}
	if err := ctx.Err(); err != nil {
		return contracts.ExtractionResult{}, extractorErrorf("context error before extraction: %w", err)
	}
	if req.Source == nil {
		return contracts.ExtractionResult{}, extractorErrorf("source must not be nil")
	}
	if req.Chunk == nil {
		return contracts.ExtractionResult{}, extractorErrorf("chunk must not be nil")
	}
	if len(req.Chunk.Units) == 0 {
		return contracts.ExtractionResult{}, extractorErrorf("chunk %q units must not be empty", req.Chunk.ID)
	}
	if req.LLMClient == nil {
		return contracts.ExtractionResult{}, extractorErrorf("LLM client must not be nil")
	}

	system, user, _, err := renderPrompt(req)
	if err != nil {
		return contracts.ExtractionResult{}, extractorErrorf("render prompt: %w", err)
	}
	schema, err := loadResponseSchema()
	if err != nil {
		return contracts.ExtractionResult{}, extractorErrorf("load response schema %q: %w", ResponseSchemaKey, err)
	}

	var response extractionResponse
	if _, err := req.LLMClient.CompleteStructured(ctx, contracts.StructuredCompletionRequest{
		StageName: Key,
		Messages: []contracts.LLMMessage{
			{Role: "system", Content: system},
			{Role: "user", Content: user},
		},
		ResponseSchemaName: schema.Name,
		ResponseSchema:     schema.JSONSchema,
	}, &response); err != nil {
		return contracts.ExtractionResult{}, extractorErrorf("complete structured output: %w", err)
	}
	if response.SpellCasts == nil {
		return contracts.ExtractionResult{}, extractorErrorf("malformed structured output: spell_casts must be present")
	}
	if len(response.SpellCasts) == 0 {
		return contracts.ExtractionResult{}, nil
	}

	candidates := make([]artifacts.ArtifactCandidate, 0, len(response.SpellCasts))
	for i, spellCast := range response.SpellCasts {
		payload, err := spellCastPayload(spellCast)
		if err != nil {
			return contracts.ExtractionResult{}, extractorErrorf("marshal spell cast[%d]: %w", i, err)
		}
		candidates = append(candidates, artifacts.ArtifactCandidate{
			Payload:    payload,
			SourceRefs: append([]source.SourceRef(nil), spellCast.SourceRefs...),
		})
	}
	return contracts.ExtractionResult{Candidates: candidates}, nil
}

func ModuleSpec() pipeline.ModuleSpec {
	return pipeline.ModuleSpec{
		Key:            Key,
		Stage:          pipeline.StageExtract,
		Requires:       append([]string(nil), requiredCapabilities...),
		Provides:       append([]string(nil), providedCapabilities...),
		ReferenceSlots: cloneReferenceSlots(referenceSlots),
	}
}

func Register(registry *pipeline.ExtractorRegistry) error {
	return registry.RegisterWithSpec(ModuleSpec(), func() (contracts.Extractor, error) {
		return New(), nil
	})
}

func spellCastPayload(spellCast spellCastResponse) (json.RawMessage, error) {
	return json.Marshal(SpellCast{
		Caster:               strings.TrimSpace(spellCast.Caster),
		Spell:                strings.TrimSpace(spellCast.Spell),
		Effect:               strings.TrimSpace(spellCast.Effect),
		NarrativeDescription: strings.TrimSpace(spellCast.NarrativeDescription),
	})
}

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

func cloneReferenceSlots(slots []contracts.ReferenceSlot) []contracts.ReferenceSlot {
	if len(slots) == 0 {
		return nil
	}
	out := make([]contracts.ReferenceSlot, len(slots))
	for i, slot := range slots {
		slot.AcceptedMediaTypes = append([]string(nil), slot.AcceptedMediaTypes...)
		out[i] = slot
	}
	return out
}